add DatabaseManager class

Author: andrii-i

jupyter_scheduler/executors.py

@@ -29,11 +29,12 @@ class ExecutionManager(ABC):
     _model = None
     _db_session = None
 
-    def __init__(self, job_id: str, root_dir: str, db_url: str, staging_paths: Dict[str, str]):
+    def __init__(self, job_id: str, root_dir: str, db_url: str, staging_paths: Dict[str, str], database_manager=None):
         self.job_id = job_id
         self.staging_paths = staging_paths
         self.root_dir = root_dir
         self.db_url = db_url
+        self.database_manager = database_manager
 
     @property
     def model(self):
@@ -46,7 +47,7 @@ def model(self):
     @property
     def db_session(self):
         if self._db_session is None:
-            self._db_session = create_session(self.db_url)
+            self._db_session = create_session(self.db_url, self.database_manager)
 
         return self._db_session
 

jupyter_scheduler/extension.py

@@ -45,6 +45,13 @@ class SchedulerApp(ExtensionApp):
     def _db_url_default(self):
         return f"sqlite:///{jupyter_data_dir()}/scheduler.sqlite"
 
+    database_manager_class = Type(
+        default_value=None,
+        klass="jupyter_scheduler.managers.DatabaseManager",
+        config=True,
+        help=_i18n("Database manager class. If None, uses SQLAlchemy with db_url."),
+    )
+
     environment_manager_class = Type(
         default_value="jupyter_scheduler.environments.CondaEnvironmentManager",
         klass="jupyter_scheduler.environments.EnvironmentManager",
@@ -69,7 +76,8 @@ def _db_url_default(self):
     def initialize_settings(self):
         super().initialize_settings()
 
-        create_tables(self.db_url, self.drop_tables)
+        database_manager = self.database_manager_class() if self.database_manager_class else None
+        create_tables(self.db_url, self.drop_tables, database_manager=database_manager)
 
         environments_manager = self.environment_manager_class()
 
@@ -78,6 +86,7 @@ def initialize_settings(self):
             environments_manager=environments_manager,
             db_url=self.db_url,
             config=self.config,
+            database_manager=database_manager,
         )
 
         job_files_manager = self.job_files_manager_class(scheduler=scheduler)

jupyter_scheduler/managers.py

@@ -0,0 +1,34 @@
+"""Database managers for custom database backends."""
+
+from abc import ABC, abstractmethod
+
+
+class DatabaseManager(ABC):
+    """Base class for database managers.
+    
+    Database managers handle database operations for jupyter-scheduler.
+    Subclasses can implement custom storage backends (K8s, Redis, etc.)
+    while maintaining compatibility with the scheduler's session interface.
+    """
+    
+    @abstractmethod
+    def create_session(self, db_url: str):
+        """Create a database session.
+        
+        Args:
+            db_url: Database URL (e.g., "k8s://namespace", "redis://localhost")
+            
+        Returns:
+            Session object compatible with SQLAlchemy session interface
+        """
+        pass
+    
+    @abstractmethod  
+    def create_tables(self, db_url: str, drop_tables: bool = False):
+        """Create database tables/schema.
+        
+        Args:
+            db_url: Database URL
+            drop_tables: Whether to drop existing tables first
+        """
+        pass

jupyter_scheduler/orm.py

@@ -146,7 +146,11 @@ def update_db_schema(engine, Base):
             connection.execute(alter_statement)
 
 
-def create_tables(db_url, drop_tables=False, Base=Base):
+def create_tables(db_url, drop_tables=False, Base=Base, database_manager=None):
+    if database_manager:
+        database_manager.create_tables(db_url, drop_tables)
+        return
+        
     engine = create_engine(db_url)
     update_db_schema(engine, Base)
 
@@ -159,7 +163,10 @@ def create_tables(db_url, drop_tables=False, Base=Base):
         Base.metadata.create_all(engine)
 
 
-def create_session(db_url):
+def create_session(db_url, database_manager=None):
+    if database_manager:
+        return database_manager.create_session(db_url)
+        
     engine = create_engine(db_url, echo=False)
     Session = sessionmaker(bind=engine)
 

jupyter_scheduler/scheduler.py

@@ -405,19 +405,21 @@ def __init__(
         environments_manager: Type[EnvironmentManager],
         db_url: str,
         config=None,
+        database_manager=None,
         **kwargs,
     ):
         super().__init__(
             root_dir=root_dir, environments_manager=environments_manager, config=config, **kwargs
         )
         self.db_url = db_url
+        self.database_manager = database_manager
         if self.task_runner_class:
             self.task_runner = self.task_runner_class(scheduler=self, config=config)
 
     @property
     def db_session(self):
         if not self._db_session:
-            self._db_session = create_session(self.db_url)
+            self._db_session = create_session(self.db_url, self.database_manager)
 
         return self._db_session
 
@@ -492,6 +494,7 @@ def create_job(self, model: CreateJob) -> str:
                     staging_paths=staging_paths,
                     root_dir=self.root_dir,
                     db_url=self.db_url,
+                    database_manager=self.database_manager,
                 ).process
             )
             p.start()