add DatabaseManager class

Author: andrii-i

conftest.py

@@ -4,6 +4,7 @@
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
 
+from jupyter_scheduler.managers import SQLAlchemyDatabaseManager
 from jupyter_scheduler.orm import Base
 from jupyter_scheduler.scheduler import Scheduler
 from jupyter_scheduler.tests.mocks import MockEnvironmentManager
@@ -59,6 +60,8 @@ def jp_scheduler(jp_scheduler_db_url, jp_scheduler_root_dir, jp_scheduler_db):
         db_url=jp_scheduler_db_url,
         root_dir=str(jp_scheduler_root_dir),
         environments_manager=MockEnvironmentManager(),
+        database_manager=SQLAlchemyDatabaseManager(),
+        database_manager_class="jupyter_scheduler.managers.SQLAlchemyDatabaseManager",
     )
 
 

jupyter_scheduler/executors.py

@@ -1,3 +1,4 @@
+import importlib
 import io
 import os
 import shutil
@@ -29,12 +30,24 @@ 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_class):
         self.job_id = job_id
         self.staging_paths = staging_paths
         self.root_dir = root_dir
         self.db_url = db_url
+        
+        self.database_manager = self._create_database_manager(database_manager_class)
 
+    
+    def _create_database_manager(self, database_manager_class):
+        try:
+            module_name, class_name = database_manager_class.rsplit('.', 1)
+            module = importlib.import_module(module_name)
+            DatabaseManagerClass = getattr(module, class_name)
+            return DatabaseManagerClass()
+        except (ValueError, ImportError, AttributeError) as e:
+            raise ValueError(f"Invalid database_manager_class '{database_manager_class}': {e}")
+    
     @property
     def model(self):
         if self._model is None:
@@ -46,7 +59,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="jupyter_scheduler.managers.SQLAlchemyDatabaseManager",
+        klass="jupyter_scheduler.managers.DatabaseManager",
+        config=True,
+        help=_i18n("Database manager class for custom database backends."),
+    )
+
     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()
+        create_tables(self.db_url, self.drop_tables, database_manager=database_manager)
 
         environments_manager = self.environment_manager_class()
 
@@ -78,6 +86,8 @@ def initialize_settings(self):
             environments_manager=environments_manager,
             db_url=self.db_url,
             config=self.config,
+            database_manager=database_manager,
+            database_manager_class=self.database_manager_class,
         )
 
         job_files_manager = self.job_files_manager_class(scheduler=scheduler)

jupyter_scheduler/managers.py

@@ -0,0 +1,64 @@
+from abc import ABC, abstractmethod
+from sqlite3 import OperationalError
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
+
+from jupyter_scheduler.orm import Base as DefaultBase, update_db_schema
+
+
+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, Base=None):
+        """Create database tables/schema.
+        
+        Args:
+            db_url: Database URL
+            drop_tables: Whether to drop existing tables first
+            Base: SQLAlchemy Base for custom schemas (tests)
+        """
+        pass
+
+
+class SQLAlchemyDatabaseManager(DatabaseManager):
+    """Default database manager using SQLAlchemy."""
+    
+    def create_session(self, db_url: str):
+        """Create SQLAlchemy session factory."""
+        engine = create_engine(db_url, echo=False)
+        Session = sessionmaker(bind=engine)
+        return Session
+    
+    def create_tables(self, db_url: str, drop_tables: bool = False, Base=None):
+        """Create database tables using SQLAlchemy."""
+        if Base is None:
+            Base = DefaultBase
+        
+        engine = create_engine(db_url)
+        update_db_schema(engine, Base)
+        
+        try:
+            if drop_tables:
+                Base.metadata.drop_all(engine)
+        except OperationalError:
+            pass
+        finally:
+            Base.metadata.create_all(engine)

jupyter_scheduler/orm.py

@@ -146,21 +146,9 @@ def update_db_schema(engine, Base):
             connection.execute(alter_statement)
 
 
-def create_tables(db_url, drop_tables=False, Base=Base):
-    engine = create_engine(db_url)
-    update_db_schema(engine, Base)
+def create_tables(db_url, drop_tables=False, Base=Base, *, database_manager):
+    database_manager.create_tables(db_url, drop_tables, Base)
 
-    try:
-        if drop_tables:
-            Base.metadata.drop_all(engine)
-    except OperationalError:
-        pass
-    finally:
-        Base.metadata.create_all(engine)
 
-
-def create_session(db_url):
-    engine = create_engine(db_url, echo=False)
-    Session = sessionmaker(bind=engine)
-
-    return Session
+def create_session(db_url, database_manager):
+    return database_manager.create_session(db_url)

jupyter_scheduler/scheduler.py

@@ -405,19 +405,23 @@ def __init__(
         environments_manager: Type[EnvironmentManager],
         db_url: str,
         config=None,
+        database_manager=None,
+        database_manager_class=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
+        self.database_manager_class = database_manager_class
         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 +496,7 @@ def create_job(self, model: CreateJob) -> str:
                     staging_paths=staging_paths,
                     root_dir=self.root_dir,
                     db_url=self.db_url,
+                    database_manager_class=self.database_manager_class,
                 ).process
             )
             p.start()

jupyter_scheduler/tests/test_execution_manager.py

@@ -53,6 +53,7 @@ def test_add_side_effects_files(
         root_dir=jp_scheduler_root_dir,
         db_url=jp_scheduler_db_url,
         staging_paths={"input": staged_notebook_file_path},
+        database_manager_class="jupyter_scheduler.managers.SQLAlchemyDatabaseManager",
     )
     manager.add_side_effects_files(staged_notebook_dir)
 

jupyter_scheduler/tests/test_orm.py

@@ -13,7 +13,13 @@
 
 
 @pytest.fixture
-def initial_db(jp_scheduler_db_url) -> tuple[Type[DeclarativeMeta], sessionmaker, str]:
+def database_manager():
+    from jupyter_scheduler.managers import SQLAlchemyDatabaseManager
+    return SQLAlchemyDatabaseManager()
+
+
+@pytest.fixture
+def initial_db(jp_scheduler_db_url, database_manager) -> tuple[Type[DeclarativeMeta], sessionmaker, str]:
     TestBase = declarative_base()
 
     class MockInitialJob(TestBase):
@@ -24,9 +30,9 @@ class MockInitialJob(TestBase):
 
     initial_job = MockInitialJob(runtime_environment_name="abc", input_filename="input.ipynb")
 
-    create_tables(db_url=jp_scheduler_db_url, Base=TestBase)
+    create_tables(db_url=jp_scheduler_db_url, Base=TestBase, database_manager=database_manager)
 
-    Session = create_session(jp_scheduler_db_url)
+    Session = create_session(jp_scheduler_db_url, database_manager)
     session = Session()
 
     session.add(initial_job)
@@ -52,7 +58,7 @@ class MockUpdatedJob(TestBase):
     return MockUpdatedJob
 
 
-def test_create_tables_with_new_column(jp_scheduler_db_url, initial_db, updated_job_model):
+def test_create_tables_with_new_column(jp_scheduler_db_url, initial_db, updated_job_model, database_manager):
     TestBase, Session, initial_job_id = initial_db
 
     session = Session()
@@ -61,7 +67,7 @@ def test_create_tables_with_new_column(jp_scheduler_db_url, initial_db, updated_
     session.close()
 
     JobModel = updated_job_model
-    create_tables(db_url=jp_scheduler_db_url, Base=TestBase)
+    create_tables(db_url=jp_scheduler_db_url, Base=TestBase, database_manager=database_manager)
 
     session = Session()
     updated_columns = {col["name"] for col in inspect(session.bind).get_columns("jobs")}