feat(db): Extend v2 model fields with pydantic aliases

Author: tcjennings

src/lsst/cmservice/db/campaigns_v2.py

@@ -2,7 +2,7 @@
 from typing import Annotated, Any
 from uuid import NAMESPACE_DNS, UUID, uuid5
 
-from pydantic import PlainSerializer, PlainValidator, ValidationInfo, model_validator
+from pydantic import AliasChoices, PlainSerializer, PlainValidator, ValidationInfo, model_validator
 from sqlalchemy.dialects import postgresql
 from sqlalchemy.ext.mutable import MutableDict
 from sqlalchemy.types import PickleType
@@ -15,11 +15,27 @@
 """Default UUID5 namespace for campaigns"""
 
 
-def jsonb_column(name: str) -> Any:
-    """Constructor for a Field based on a JSONB database column."""
+def jsonb_column(name: str, aliases: list[str] | None = None) -> Any:
+    """Constructor for a Field based on a JSONB database column.
+
+    If provided, a list of aliases will be used to construct a pydantic
+    ``AliasChoices`` object for the field's validation alias, which improves
+    usability by making model validation more flexible (e.g., having "metadata"
+    and "metadata_" refer to the same field).
+
+    Additionally, the first alias in the list will be used for the model's
+    serialization alias.
+    """
+    schema_extra = {}
+    if aliases:
+        schema_extra = {
+            "validation_alias": AliasChoices(*aliases),
+            "serialization_alias": aliases[0],
+        }
     return Field(
         sa_column=Column(name, MutableDict.as_mutable(postgresql.JSONB)),
         default_factory=dict,
+        schema_extra=schema_extra,
     )
 
 
@@ -68,8 +84,8 @@ class CampaignBase(BaseSQLModel):
         default=StatusEnum.waiting,
         sa_column=Column("status", Enum(StatusEnum, length=20, native_enum=False, create_constraint=False)),
     )
-    metadata_: dict = jsonb_column("metadata")
-    configuration: dict = jsonb_column("configuration")
+    metadata_: dict = jsonb_column("metadata", aliases=["metadata", "metadata_"])
+    configuration: dict = jsonb_column("configuration", aliases=["configuration", "data", "spec"])
 
 
 class CampaignModel(CampaignBase):
@@ -114,8 +130,8 @@ class NodeBase(BaseSQLModel):
         default=StatusEnum.waiting,
         sa_column=Column("status", Enum(StatusEnum, length=20, native_enum=False, create_constraint=False)),
     )
-    metadata_: dict = jsonb_column("metadata")
-    configuration: dict = jsonb_column("configuration")
+    metadata_: dict = jsonb_column("metadata", aliases=["metadata", "metadata_"])
+    configuration: dict = jsonb_column("configuration", aliases=["configuration", "data", "spec"])
 
 
 class NodeModel(NodeBase):
@@ -150,8 +166,8 @@ class EdgeBase(BaseSQLModel):
     namespace: UUID = Field(foreign_key="campaigns_v2.id")
     source: UUID = Field(foreign_key="nodes_v2.id")
     target: UUID = Field(foreign_key="nodes_v2.id")
-    metadata_: dict = jsonb_column("metadata")
-    configuration: dict = jsonb_column("configuration")
+    metadata_: dict = jsonb_column("metadata", aliases=["metadata", "metadata_"])
+    configuration: dict = jsonb_column("configuration", aliases=["configuration", "data", "spec"])
 
 
 class EdgeModel(EdgeBase):
@@ -197,8 +213,8 @@ class ManifestBase(BaseSQLModel):
         default=ManifestKind.other,
         sa_column=Column("kind", Enum(ManifestKind, length=20, native_enum=False, create_constraint=False)),
     )
-    metadata_: dict = jsonb_column("metadata")
-    spec: dict = jsonb_column("spec")
+    metadata_: dict = jsonb_column("metadata", aliases=["metadata", "metadata_"])
+    spec: dict = jsonb_column("spec", aliases=["spec", "configuration", "data"])
 
 
 class ManifestModel(ManifestBase):
@@ -236,10 +252,17 @@ class Task(SQLModel, table=True):
     last_processed_at: datetime
     finished_at: datetime
     wms_id: str
-    # TODO site_affinity should probably be a mutable array
-    site_affinity: list[str] = Field(sa_column=Column(postgresql.ARRAY(String())))
-    status: int  # TODO make enum
-    previous_status: int  # TODO make enum
+    site_affinity: list[str] = Field(
+        sa_column=Column("site_affinity", MutableList.as_mutable(postgresql.ARRAY(String())))
+    )
+    status: StatusField = Field(
+        sa_column=Column("status", Enum(StatusEnum, length=20, native_enum=False, create_constraint=False)),
+    )
+    previous_status: StatusField = Field(
+        sa_column=Column(
+            "previous_status", Enum(StatusEnum, length=20, native_enum=False, create_constraint=False)
+        ),
+    )
 
 
 class ActivityLogBase(BaseSQLModel):

src/lsst/cmservice/db/manifests_v2.py

@@ -0,0 +1,34 @@
+from typing import Annotated
+
+from pydantic import AliasChoices
+from sqlmodel import Field, SQLModel
+
+from ..common.enums import ManifestKind
+from .campaigns_v2 import EnumSerializer, ManifestKindEnumValidator
+
+
+# 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`
+    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,
+    )
+    metadata_: dict = Field(
+        default_factory=dict,
+        schema_extra={
+            "validation_alias": AliasChoices("metadata", "metadata_"),
+            "serialization_alias": "metadata",
+        },
+    )
+    spec: dict = Field(
+        default_factory=dict,
+        schema_extra={
+            "validation_alias": AliasChoices("spec", "configuration", "data"),
+            "serialization_alias": "spec",
+        },
+    )