Add database schema update and database migration logic

conftest.py

@@ -2,7 +2,7 @@
 
 import pytest
 from sqlalchemy import create_engine
-from sqlalchemy.orm import Session, sessionmaker
+from sqlalchemy.orm import sessionmaker
 
 from jupyter_scheduler.orm import Base
 from jupyter_scheduler.scheduler import Scheduler

jupyter_scheduler/orm.py

@@ -1,11 +1,11 @@
 import json
-import os
 from sqlite3 import OperationalError
 from uuid import uuid4
 
 import sqlalchemy.types as types
-from sqlalchemy import Boolean, Column, Integer, String, create_engine
+from sqlalchemy import Boolean, Column, Integer, String, create_engine, inspect
 from sqlalchemy.orm import declarative_base, declarative_mixin, registry, sessionmaker
+from sqlalchemy.sql import text
 
 from jupyter_scheduler.models import EmailNotifications, Status
 from jupyter_scheduler.utils import get_utc_timestamp
@@ -91,6 +91,7 @@ class CommonColumns:
 
 class Job(CommonColumns, Base):
     __tablename__ = "jobs"
+    __table_args__ = {"extend_existing": True}
     job_id = Column(String(36), primary_key=True, default=generate_uuid)
     job_definition_id = Column(String(36))
     status = Column(String(64), default=Status.STOPPED)
@@ -104,6 +105,7 @@ class Job(CommonColumns, Base):
 
 class JobDefinition(CommonColumns, Base):
     __tablename__ = "job_definitions"
+    __table_args__ = {"extend_existing": True}
     job_definition_id = Column(String(36), primary_key=True, default=generate_uuid)
     schedule = Column(String(256))
     timezone = Column(String(36))
@@ -112,8 +114,29 @@ class JobDefinition(CommonColumns, Base):
     active = Column(Boolean, default=True)
 
 
-def create_tables(db_url, drop_tables=False):
+def update_db_schema(engine, Base):
+    inspector = inspect(engine)
+
+    with engine.connect() as connection:
+        for table_name, model in Base.metadata.tables.items():
+            if inspector.has_table(table_name):
+                columns_db = inspector.get_columns(table_name)
+                columns_db_names = {col["name"] for col in columns_db}
+
+                for column_model_name, column_model in model.c.items():
+                    if column_model_name not in columns_db_names:
+                        column_type = str(column_model.type.compile(dialect=engine.dialect))
+                        nullable = "NULL" if column_model.nullable else "NOT NULL"
+                        alter_statement = text(
+                            f"ALTER TABLE {table_name} ADD COLUMN {column_model_name} {column_type} {nullable}"
+                        )
+                        connection.execute(alter_statement)
+
+
+def create_tables(db_url, drop_tables=False, Base=Base):
     engine = create_engine(db_url)
+    update_db_schema(engine, Base)
+
     try:
         if drop_tables:
             Base.metadata.drop_all(engine)

jupyter_scheduler/tests/test_orm.py

@@ -0,0 +1,73 @@
+from typing import Type
+
+import pytest
+from sqlalchemy import Column, Integer, String, inspect
+from sqlalchemy.orm import DeclarativeMeta, sessionmaker
+
+from jupyter_scheduler.orm import (
+    create_session,
+    create_tables,
+    declarative_base,
+    generate_uuid,
+)
+
+
+@pytest.fixture
+def initial_db(jp_scheduler_db_url) -> tuple[Type[DeclarativeMeta], sessionmaker, str]:
+    TestBase = declarative_base()
+
+    class InitialJob(TestBase):
+        __tablename__ = "jobs"
+        job_id = Column(String(36), primary_key=True, default=generate_uuid)
+        runtime_environment_name = Column(String(256), nullable=False)
+        input_filename = Column(String(256), nullable=False)
+
+    initial_job = InitialJob(runtime_environment_name="abc", input_filename="input.ipynb")
+
+    create_tables(db_url=jp_scheduler_db_url, Base=TestBase)
+
+    Session = create_session(jp_scheduler_db_url)
+    session = Session()
+
+    session.add(initial_job)
+    session.commit()
+    job_id = initial_job.job_id
+    session.close()
+
+    return (TestBase, Session, job_id)
+
+
+@pytest.fixture
+def updated_job_model(initial_db) -> Type[DeclarativeMeta]:
+    Base = initial_db[0]
+
+    class UpdatedJob(Base):
+        __tablename__ = "jobs"
+        __table_args__ = {"extend_existing": True}
+        job_id = Column(String(36), primary_key=True, default=generate_uuid)
+        runtime_environment_name = Column(String(256), nullable=False)
+        input_filename = Column(String(256), nullable=False)
+        new_column = Column("new_column", Integer)
+
+    return UpdatedJob
+
+
+def test_create_tables_with_new_column(jp_scheduler_db_url, initial_db, updated_job_model):
+    Base, Session, initial_job_id = initial_db
+
+    session = Session()
+    initial_columns = {col["name"] for col in inspect(session.bind).get_columns("jobs")}
+    assert "new_column" not in initial_columns
+    session.close()
+
+    JobModel = updated_job_model
+    create_tables(db_url=jp_scheduler_db_url, Base=Base)
+
+    session = Session()
+    updated_columns = {col["name"] for col in inspect(session.bind).get_columns("jobs")}
+    assert "new_column" in updated_columns
+
+    updated_job = session.query(JobModel).filter(JobModel.job_id == initial_job_id).one()
+    assert hasattr(updated_job, "new_column")
+    assert updated_job.runtime_environment_name == "abc"
+    assert updated_job.input_filename == "input.ipynb"