feat(conductor): update database models for new flow

conductor/src/conductor/database.py

@@ -13,7 +13,6 @@
 from conductor.models.job import Job  # noqa: F401
 from conductor.models.project import Project  # noqa: F401
 from conductor.models.run import Run  # noqa: F401
-from conductor.models.step import Step  # noqa: F401
 
 
 def get_db_session(db_uri: str) -> Session:
@@ -27,7 +26,7 @@ def add_db_session_to_req(db_uri: str, request: Request) -> None:
     Parameters
     ----------
     db_uri : str
-        Databse URI
+        Database URI
     request : Request
         FastAPI request.
 
@@ -41,7 +40,7 @@ def add_db_session_to_ws(db_uri: str, websocket: WebSocket) -> None:
     Parameters
     ----------
     db_uri : str
-        Databse URI
+        Database URI
     websocket : WebSocket
         FastAPI websocket.
 

conductor/src/conductor/models/__init__.py

@@ -1 +1 @@
-"""Databse models."""
+"""Database models."""

conductor/src/conductor/models/job.py

@@ -15,14 +15,14 @@ class Job(BaseSQLModel):
 
     __tablename__ = "job"
     id: Mapped[int] = mapped_column(primary_key=True, index=True)
-    module_id: Mapped[str] = mapped_column(nullable=False, index=True)
+    uid: Mapped[str] = mapped_column(nullable=False, index=True)
     name: Mapped[str] = mapped_column(String(100))
     description: Mapped[str] = mapped_column()
-    spec: Mapped[dict] = mapped_column(Container, nullable=False)
+    spec: Mapped[Container] = mapped_column(JSON, nullable=False)
     outputs: Mapped[dict] = mapped_column(JSON)
     inputs: Mapped[dict] = mapped_column(JSON)
-    step_id: Mapped[int] = mapped_column(ForeignKey("step.id"), nullable=False)
-    step = relationship("Step", back_populates="jobs")
+    project_id: Mapped[int] = mapped_column(ForeignKey("project.id"), nullable=False)
+    project = relationship("Project", back_populates="jobs")
     runs: Mapped[list[Run]] = relationship(
         back_populates="job", cascade="all, delete-orphan"
     )

conductor/src/conductor/models/project.py

@@ -3,9 +3,9 @@
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 from sqlalchemy.sql.sqltypes import Enum, String
 
-from conductor.constants import ParserType, ProjectType
+from conductor.constants import ParserType
 from conductor.models.base import BaseSQLModel
-from conductor.models.step import Step
+from conductor.models.job import Job
 
 
 class Project(BaseSQLModel):
@@ -16,12 +16,14 @@ class Project(BaseSQLModel):
     name: Mapped[str] = mapped_column(String(100), unique=True, index=True)
     dataset_uri: Mapped[str] = mapped_column()
     workspace_uri: Mapped[str] = mapped_column()
+    storage_uri: Mapped[str] = mapped_column()
     img_uri: Mapped[str | None] = mapped_column()
+    is_configured: Mapped[bool] = mapped_column(default=False)
     description: Mapped[str] = mapped_column()
-    type = mapped_column(Enum(ProjectType, create_constraint=True), nullable=False)
+    type: Mapped[str] = mapped_column()
     parser_type = mapped_column(
         Enum(ParserType, create_constraint=True), nullable=False
     )
-    steps: Mapped[list[Step]] = relationship(
+    jobs: Mapped[list[Job]] = relationship(
         back_populates="project", cascade="all, delete-orphan"
     )

conductor/src/conductor/models/run.py

@@ -3,6 +3,7 @@
 from sqlalchemy import JSON, Enum, ForeignKey
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 from sqlalchemy.sql.sqltypes import String
+from starrynight.modules.schema import Container
 
 from conductor.constants import ExecutorType, RunStatus
 from conductor.models.base import BaseSQLModel
@@ -18,6 +19,7 @@ class Run(BaseSQLModel):
     run_status = mapped_column(Enum(RunStatus, create_constraint=True), nullable=False)
     outputs: Mapped[dict] = mapped_column(JSON)
     inputs: Mapped[dict] = mapped_column(JSON)
+    spec: Mapped[Container] = mapped_column(JSON, nullable=False)
     executor_type = mapped_column(
         Enum(ExecutorType, create_constraint=True), nullable=False
     )

conductor/tests/conftest.py

@@ -3,19 +3,16 @@
 from collections.abc import Generator
 
 import pytest
+from sqlalchemy import Engine, create_engine, event
+from sqlalchemy.orm import Session, sessionmaker
+
 from conductor.constants import (
-    JobType,
     ParserType,
     ProjectType,
-    StepType,
-    job_output_dict,
 )
 from conductor.models.base import BaseSQLModel
 from conductor.models.job import Job
 from conductor.models.project import Project
-from conductor.models.step import Step
-from sqlalchemy import Engine, create_engine, event
-from sqlalchemy.orm import Session, sessionmaker
 
 DATABASE_URL = (
     "sqlite+pysqlite:///:memory:"  # Use an in-memory SQLite database for testing
@@ -48,8 +45,9 @@ def sample_project(db: Session) -> Project:
         name="test-project",
         dataset_uri="s3://test-project",
         workspace_uri="test-uri/workspace",
+        storage_uri="test-uri/workspace",
         description="This is a test project",
-        type=ProjectType.OPS_GENERIC,
+        type="randomproject",
         parser_type=ParserType.OPS_VINCENT,
     )
     db.add(project)
@@ -58,27 +56,15 @@ def sample_project(db: Session) -> Project:
 
 
 @pytest.fixture(scope="function")
-def sample_step(db: Session, sample_project: Project) -> Step:
-    step = Step(
-        name="test-step",
-        description="This is a test step",
-        project_id=sample_project.id,
-        type=StepType.CP_ILLUM_CALC,
-    )
-    db.add(step)
-    db.commit()
-    return step
-
-
-@pytest.fixture(scope="function")
-def sample_job(db: Session, sample_step: Step) -> Job:
+def sample_job(db: Session, sample_project: Project) -> Job:
     job = Job(
         name="test-job",
+        uid="Unique module name",
         description="This is a test job",
-        step_id=sample_step.id,
-        type=JobType.GEN_LOADDATA,
+        project_id=sample_project.id,
+        spec={},
         inputs={},
-        outputs=job_output_dict[JobType.GEN_LOADDATA],
+        outputs={},
     )
     db.add(job)
     db.commit()

conductor/tests/unit/models/test_job_model.py

@@ -1,21 +1,23 @@
 """Job model test suite."""
 
 import pytest
-from conductor.constants import JobType
-from conductor.models.job import Job
-from conductor.models.step import Step
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import Session
+from starrynight.modules.schema import Container
+
+from conductor.models.job import Job
+from conductor.models.project import Project
 
 
-def test_create_job(db: Session, sample_step: Step) -> None:
+def test_create_job(db: Session, sample_project: Project) -> None:
     job = Job(
         name="TestJob",
+        uid="module_id",
         description="A test job",
-        type=JobType.GEN_LOADDATA,
         outputs={},
         inputs={},
-        step_id=sample_step.id,
+        spec={},
+        project_id=sample_project.id,
     )
     db.add(job)
     db.commit()
@@ -47,46 +49,34 @@ def test_create_job(db: Session, sample_step: Step) -> None:
 #         db.commit()
 
 
-def test_relationship_step(db: Session, sample_step: Step) -> None:
+def test_relationship_project(db: Session, sample_project: Project) -> None:
     job = Job(
         name="TestJob",
+        uid="module_id",
         description="A test job",
-        type=JobType.GEN_LOADDATA,
         outputs={},
         inputs={},
-        step_id=sample_step.id,
+        spec={},
+        project_id=sample_project.id,
     )
 
     db.add(job)
     db.commit()
-    db.refresh(sample_step)
-    assert len(sample_step.jobs) == 1
-    assert sample_step.jobs[0].id == job.id
-
-
-def test_invalid_type(db: Session, sample_step: Step) -> None:
-    with pytest.raises(IntegrityError, match="CHECK constraint failed"):
-        job = Job(
-            name="TestJob",
-            description="A test job",
-            type="InvalidType",
-            outputs={},
-            inputs={},
-            step_id=sample_step.id,
-        )
-        db.add(job)
-        db.commit()
+    db.refresh(sample_project)
+    assert len(sample_project.jobs) == 1
+    assert sample_project.jobs[0].id == job.id
 
 
-def test_invalid_step(db: Session) -> None:
+def test_invalid_project(db: Session) -> None:
     with pytest.raises(IntegrityError, match="FOREIGN KEY constraint failed"):
         job = Job(
             name="TestJob",
+            uid="module_id",
             description="A test job",
-            type=JobType.GEN_LOADDATA,
             outputs={},
             inputs={},
-            step_id=2093482903,
+            spec={},
+            project_id=2093482903,
         )
         db.add(job)
         db.commit()