feat(db): Update Models

- add state machine db model
- rearrange shared types and serde helpers
- support arbitrary tbale schema in new db models and tests

Author: tcjennings

src/lsst/cmservice/common/types.py

@@ -1,12 +1,28 @@
+from typing import Annotated
+
 from sqlalchemy.ext.asyncio import AsyncSession as AsyncSessionSA
 from sqlalchemy.ext.asyncio import async_scoped_session
 from sqlmodel.ext.asyncio.session import AsyncSession
 
 from .. import models
+from ..models.serde import EnumSerializer, ManifestKindEnumValidator, StatusEnumValidator
+from .enums import ManifestKind, StatusEnum
 
 type AnyAsyncSession = AsyncSession | AsyncSessionSA | async_scoped_session
 """A type union of async database sessions the application may use"""
 
 
 type AnyCampaignElement = models.Group | models.Campaign | models.Step | models.Job
 """A type union of Campaign elements"""
+
+
+type StatusField = Annotated[StatusEnum, StatusEnumValidator, EnumSerializer]
+"""A type for fields representing a Status with a custom validator tuned for
+enums operations.
+"""
+
+
+type KindField = Annotated[ManifestKind, ManifestKindEnumValidator, EnumSerializer]
+"""A type for fields representing a Kind with a custom validator tuned for
+enums operations.
+"""

src/lsst/cmservice/config.py

@@ -546,8 +546,8 @@ class DatabaseConfiguration(BaseModel):
         description="The password for the cm-service database",
     )
 
-    table_schema: str | None = Field(
-        default=None,
+    table_schema: str = Field(
+        default="public",
         description="Schema to use for cm-service database",
     )
 

src/lsst/cmservice/db/campaigns_v2.py

@@ -1,19 +1,23 @@
 from datetime import datetime
-from typing import Annotated, Any
+from typing import Any
 from uuid import NAMESPACE_DNS, UUID, uuid5
 
-from pydantic import AliasChoices, PlainSerializer, PlainValidator, ValidationInfo, model_validator
+from pydantic import AliasChoices, ValidationInfo, model_validator
 from sqlalchemy.dialects import postgresql
-from sqlalchemy.ext.mutable import MutableDict
+from sqlalchemy.ext.mutable import MutableDict, MutableList
 from sqlalchemy.types import PickleType
-from sqlmodel import Column, Enum, Field, SQLModel, String
+from sqlmodel import Column, Enum, Field, MetaData, SQLModel, String
 
 from ..common.enums import ManifestKind, StatusEnum
+from ..common.types import KindField, StatusField
 from ..config import config
 
 _default_campaign_namespace = uuid5(namespace=NAMESPACE_DNS, name="io.lsst.cmservice")
 """Default UUID5 namespace for campaigns"""
 
+metadata: MetaData = MetaData(schema=config.db.table_schema)
+"""SQLModel metadata for table models"""
+
 
 def jsonb_column(name: str, aliases: list[str] | None = None) -> Any:
     """Constructor for a Field based on a JSONB database column.
@@ -47,34 +51,10 @@ def jsonb_column(name: str, aliases: list[str] | None = None) -> Any:
 #   3. the model of the manifest when updating an object
 #   4. a response model for APIs related to the object
 
-EnumSerializer = PlainSerializer(
-    lambda x: x.name,
-    return_type="str",
-    when_used="always",
-)
-"""A serializer for enums that produces its name, not the value."""
-
-
-StatusEnumValidator = PlainValidator(lambda x: StatusEnum[x] if isinstance(x, str) else StatusEnum(x))
-"""A validator for the StatusEnum that can parse the enum from either a name
-or a value.
-"""
-
-
-ManifestKindEnumValidator = PlainValidator(
-    lambda x: ManifestKind[x] if isinstance(x, str) else ManifestKind(x)
-)
-"""A validator for the ManifestKindEnum that can parse the enum from a name
-or a value.
-"""
-
-
-type StatusField = Annotated[StatusEnum, StatusEnumValidator, EnumSerializer]
-type KindField = Annotated[ManifestKind, ManifestKindEnumValidator, EnumSerializer]
-
 
 class BaseSQLModel(SQLModel):
     __table_args__ = {"schema": config.db.table_schema}
+    metadata = metadata
 
 
 class CampaignBase(BaseSQLModel):
@@ -116,10 +96,10 @@ class Campaign(CampaignModel, table=True):
 
     __tablename__: str = "campaigns_v2"  # type: ignore[misc]
 
-    machine: UUID | None
+    machine: UUID | None = None
 
 
-class CampaignUpdate(SQLModel):
+class CampaignUpdate(BaseSQLModel):
     """Model representing updatable fields for a PATCH operation on a Campaign
     using RFC7396.
     """
@@ -215,6 +195,12 @@ class MachineBase(BaseSQLModel):
     state: Any | None = Field(sa_column=Column("state", PickleType))
 
 
+class Machine(MachineBase, table=True):
+    """machines_v2 db table."""
+
+    __tablename__: str = "machiness_v2"  # type: ignore[misc]
+
+
 class ManifestBase(BaseSQLModel):
     """manifests_v2 db table"""
 
@@ -252,7 +238,7 @@ class Manifest(ManifestBase, table=True):
     __tablename__: str = "manifests_v2"  # type: ignore[misc]
 
 
-class Task(SQLModel, table=True):
+class Task(BaseSQLModel, table=True):
     """tasks_v2 db table"""
 
     __tablename__: str = "tasks_v2"  # type: ignore[misc]

src/lsst/cmservice/db/manifests_v2.py

@@ -4,7 +4,7 @@
 from sqlmodel import Field, SQLModel
 
 from ..common.enums import ManifestKind
-from .campaigns_v2 import EnumSerializer, ManifestKindEnumValidator
+from ..models.serde import EnumSerializer, ManifestKindEnumValidator
 
 
 # this can probably be a BaseModel since this is not a db relation, but the

src/lsst/cmservice/models/serde.py

@@ -0,0 +1,28 @@
+"""Module for serialization and deserialization support for pydantic and
+other derivative models.
+"""
+
+from pydantic import PlainSerializer, PlainValidator
+
+from ..common.enums import ManifestKind, StatusEnum
+
+EnumSerializer = PlainSerializer(
+    lambda x: x.name,
+    return_type="str",
+    when_used="always",
+)
+"""A serializer for enums that produces its name, not the value."""
+
+
+StatusEnumValidator = PlainValidator(lambda x: StatusEnum[x] if isinstance(x, str) else StatusEnum(x))
+"""A validator for the StatusEnum that can parse the enum from either a name
+or a value.
+"""
+
+
+ManifestKindEnumValidator = PlainValidator(
+    lambda x: ManifestKind[x] if isinstance(x, str) else ManifestKind(x)
+)
+"""A validator for the ManifestKindEnum that can parse the enum from a name
+or a value.
+"""

tests/v2/conftest.py

@@ -4,17 +4,19 @@
 import os
 from collections.abc import AsyncGenerator, Generator
 from typing import TYPE_CHECKING
+from uuid import uuid4
 
 import pytest
 import pytest_asyncio
 from fastapi.testclient import TestClient
 from httpx import ASGITransport, AsyncClient
 from sqlalchemy.pool import NullPool
-from sqlmodel import SQLModel
+from sqlalchemy.schema import CreateSchema, DropSchema
 from testcontainers.postgres import PostgresContainer
 
 from lsst.cmservice.common.types import AnyAsyncSession
 from lsst.cmservice.config import config
+from lsst.cmservice.db.campaigns_v2 import metadata
 from lsst.cmservice.db.session import DatabaseSessionDependency, db_session_dependency
 
 if TYPE_CHECKING:
@@ -37,9 +39,13 @@ async def rawdb(monkeypatch_module: pytest.MonkeyPatch) -> AsyncGenerator[Databa
     `TEST__LOCAL_DB` is not set; otherwise the fixture will assume that the
     correct database is available to the test environment through ordinary
     configuration parameters.
+
+    The tests are performed within a random temporary schema that is created
+    and dropped along with the tables.
     """
 
     monkeypatch_module.setattr(target=config.asgi, name="enable_frontend", value=False)
+    monkeypatch_module.setattr(target=config.db, name="table_schema", value=uuid4().hex[:8])
 
     if os.getenv("TEST__LOCAL_DB") is not None:
         db_session_dependency.pool_class = NullPool
@@ -77,11 +83,15 @@ async def testdb(rawdb: DatabaseSessionDependency) -> AsyncGenerator[DatabaseSes
     # v2 objects are created from the SQLModel metadata.
     assert rawdb.engine is not None
     async with rawdb.engine.begin() as aconn:
-        await aconn.run_sync(SQLModel.metadata.drop_all)
-        await aconn.run_sync(SQLModel.metadata.create_all)
+        await aconn.run_sync(metadata.drop_all)
+        await aconn.execute(CreateSchema(config.db.table_schema, if_not_exists=True))
+        await aconn.run_sync(metadata.create_all)
+        await aconn.commit()
     yield rawdb
     async with rawdb.engine.begin() as aconn:
-        await aconn.run_sync(SQLModel.metadata.drop_all)
+        await aconn.run_sync(metadata.drop_all)
+        await aconn.execute(DropSchema(config.db.table_schema, if_exists=True))
+        await aconn.commit()
 
 
 @pytest_asyncio.fixture(name="session", scope="module", loop_scope="module")

tests/v2/test_db.py

@@ -1,16 +1,18 @@
 """Tests v2 database operations"""
 
-from uuid import uuid5
+from uuid import uuid4, uuid5
 
 import pytest
 from sqlmodel import select
 
-from lsst.cmservice.db.campaigns_v2 import Campaign, _default_campaign_namespace
+from lsst.cmservice.db.campaigns_v2 import Campaign, Machine, _default_campaign_namespace
 from lsst.cmservice.db.session import DatabaseSessionDependency
 
 
 @pytest.mark.asyncio
 async def test_create_campaigns_v2(testdb: DatabaseSessionDependency) -> None:
+    """Tests the campaigns_v2 table by creating and updating a Campaign."""
+
     assert testdb.sessionmaker is not None
 
     campaign_name = "test_campaign"
@@ -50,3 +52,26 @@ async def test_create_campaigns_v2(testdb: DatabaseSessionDependency) -> None:
         assert campaign.configuration["crtime"] == 0
         assert campaign.metadata_["crtime"] == 1750107719
         assert campaign.metadata_["mtime"] == 0
+
+
+@pytest.mark.asyncio
+async def test_create_machines_v2(testdb: DatabaseSessionDependency) -> None:
+    """Tests the machines_v2 table by storing + retrieving a pickled object."""
+
+    assert testdb.sessionmaker is not None
+
+    # the machines table is a PickleType so it doesn't really matter for this
+    # test what kind of object is being pickled.
+    o = {"a": [1, 2, 3, 4, {"aa": [[0, 1], [2, 3]]}]}
+
+    machine_id = uuid4()
+    machine = Machine(id=machine_id, state=o)
+    async with testdb.sessionmaker() as session:
+        session.add(machine)
+        await session.commit()
+
+    async with testdb.sessionmaker() as session:
+        s = select(Machine).where(Machine.id == machine_id).limit(1)
+        unpickled = (await session.exec(s)).one()
+
+    assert unpickled.state == o