feat(db): Update Models

- add state machine db model
- rearrange shared types and serde helpers
- support arbitrary table schema in new db models and tests
- Refactor validation function for enum fields
- Add FK relations to model fields

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

@@ -551,8 +551,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, 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 = Field(foreign_key="machines_v2.id", default=None, ondelete="CASCADE")
 
 
-class CampaignUpdate(SQLModel):
+class CampaignUpdate(BaseSQLModel):
     """Model representing updatable fields for a PATCH operation on a Campaign
     using RFC7396.
     """
@@ -168,7 +148,7 @@ def custom_model_validator(cls, data: Any, info: ValidationInfo) -> Any:
 class Node(NodeModel, table=True):
     __tablename__: str = "nodes_v2"  # type: ignore[misc]
 
-    machine: UUID | None
+    machine: UUID | None = Field(foreign_key="machines_v2.id", default=None, ondelete="CASCADE")
 
 
 class EdgeBase(BaseSQLModel):
@@ -215,13 +195,19 @@ class MachineBase(BaseSQLModel):
     state: Any | None = Field(sa_column=Column("state", PickleType))
 
 
+class Machine(MachineBase, table=True):
+    """machines_v2 db table."""
+
+    __tablename__: str = "machines_v2"  # type: ignore[misc]
+
+
 class ManifestBase(BaseSQLModel):
     """manifests_v2 db table"""
 
     id: UUID = Field(primary_key=True)
     name: str
     version: int
-    namespace: UUID
+    namespace: UUID = Field(foreign_key="campaigns_v2.id")
     kind: KindField = Field(
         default=ManifestKind.other,
         sa_column=Column("kind", Enum(ManifestKind, length=20, native_enum=False, create_constraint=False)),
@@ -252,14 +238,14 @@ 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]
 
     id: UUID = Field(primary_key=True)
-    namespace: UUID
-    node: UUID
+    namespace: UUID = Field(foreign_key="campaigns_v2.id")
+    node: UUID = Field(foreign_key="nodes_v2.id")
     priority: int
     created_at: datetime
     last_processed_at: datetime
@@ -280,8 +266,8 @@ class Task(SQLModel, table=True):
 
 class ActivityLogBase(BaseSQLModel):
     id: UUID = Field(primary_key=True)
-    namespace: UUID
-    node: UUID
+    namespace: UUID = Field(foreign_key="campaigns_v2.id")
+    node: UUID = Field(foreign_key="nodes_v2.id")
     operator: str
     to_status: StatusField = Field(
         sa_column=Column(

src/lsst/cmservice/db/manifests_v2.py

@@ -1,23 +1,19 @@
-from typing import Annotated
+"""Module for models representing generic CM Service manifests."""
 
 from pydantic import AliasChoices
 from sqlmodel import Field, SQLModel
 
 from ..common.enums import ManifestKind
-from .campaigns_v2 import EnumSerializer, ManifestKindEnumValidator
+from ..common.types import KindField
 
 
-# this can probably be a BaseModel since this is not a db relation, but the
-# distinction probably doesn't matter
 class ManifestWrapper(SQLModel):
-    """a model for an object's Manifest wrapper, used by APIs where the `spec`
+    """A model for an object's Manifest wrapper, used by APIs where the `spec`
     should be the kind's table model, more or less.
     """
 
     apiversion: str = Field(default="io.lsst.cmservice/v1")
-    kind: Annotated[ManifestKind, ManifestKindEnumValidator, EnumSerializer] = Field(
-        default=ManifestKind.other,
-    )
+    kind: KindField = Field(default=ManifestKind.other)
     metadata_: dict = Field(
         default_factory=dict,
         schema_extra={

src/lsst/cmservice/models/serde.py

@@ -0,0 +1,44 @@
+"""Module for serialization and deserialization support for pydantic and
+other derivative models.
+"""
+
+from enum import EnumType
+from functools import partial
+from typing import Any
+
+from pydantic import PlainSerializer, PlainValidator
+
+from ..common.enums import ManifestKind, StatusEnum
+
+
+def EnumValidator[T: EnumType](value: Any, enum_: T) -> T:
+    """Create an enum from the input value. The input can be either the
+    enum name or its value.
+
+    Used as a Validator for a pydantic field.
+    """
+    try:
+        new_enum: T = enum_[value] if value in enum_.__members__ else enum_(value)
+    except (KeyError, ValueError):
+        raise ValueError(f"Value must be a member of {enum_.__qualname__}")
+    return new_enum
+
+
+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(partial(EnumValidator, enum_=StatusEnum))
+"""A validator for the StatusEnum that can parse the enum from either a name
+or a value.
+"""
+
+
+ManifestKindEnumValidator = PlainValidator(partial(EnumValidator, enum_=ManifestKind))
+"""A validator for the ManifestKindEnum that can parse the enum from a name
+or a value.
+"""

tests/models/test_serde.py

@@ -0,0 +1,33 @@
+import pytest
+from pydantic import BaseModel, ValidationError
+
+from lsst.cmservice.common.enums import ManifestKind, StatusEnum
+from lsst.cmservice.common.types import KindField, StatusField
+
+
+class TestModel(BaseModel):
+    status: StatusField
+    kind: KindField
+
+
+def test_validators() -> None:
+    """Test model field enum validators."""
+    # test enum validation by name and value
+    x = TestModel(status=0, kind="campaign")
+    assert x.status is StatusEnum.waiting
+    assert x.kind is ManifestKind.campaign
+
+    # test bad input (wrong name)
+    with pytest.raises(ValidationError):
+        x = TestModel(status="bad", kind="edge")
+
+    # test bad input (bad value)
+    with pytest.raises(ValidationError):
+        x = TestModel(status="waiting", kind=99)
+
+
+def test_serializers() -> None:
+    x = TestModel(status="accepted", kind="node")
+    y = x.model_dump()
+    assert y["status"] == "accepted"
+    assert y["kind"] == "node"

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