Skip to content

Commit 7782fc0

Browse files
tcjenningstcjennings
committed
feat(db): Add migrations and models for v2 tables
feat(db): Extend v2 model fields with pydantic aliases
1 parent 9a1cef7 commit 7782fc0

5 files changed

Lines changed: 589 additions & 17 deletions

File tree

‎alembic/versions/1da92a1c740f_create_v2_tables.py‎

Lines changed: 247 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,247 @@
1+
"""create v2 tables
2+
3+
Revision ID: 1da92a1c740f
4+
Revises: acf951c80750
5+
Create Date: 2025-06-13 14:56:31.238050+00:00
6+
7+
"""
8+
9+
from collections.abc import Sequence
10+
from enum import Enum
11+
from uuid import NAMESPACE_DNS
12+
13+
import sqlalchemy as sa
14+
from sqlalchemy.dialects import postgresql
15+
16+
from alembic import op
17+
18+
# revision identifiers, used by Alembic.
19+
revision: str = "1da92a1c740f"
20+
down_revision: str | None = "acf951c80750"
21+
branch_labels: str | Sequence[str] | None = None
22+
depends_on: str | Sequence[str] | None = None
23+
24+
25+
DEFAULT_CAMPAIGN_NAMESPACE = "dda54a0c-6878-5c95-ac4f-007f6808049e"
26+
"""UUID5 of name 'io.lsst.cmservice' in `uuid.NAMESPACE_DNS`."""
27+
28+
# DB model uses mapped columns with Python Enum types, but we do not care
29+
# to use native enums in the database, so when we have such a column, this
30+
# definition will produce a VARCHAR instead.
31+
ENUM_COLUMN_AS_VARCHAR = sa.Enum(Enum, length=20, native_enum=False, check_constraint=False)
32+
33+
34+
def upgrade() -> None:
35+
# Create table for machines v2
36+
machines_v2 = op.create_table(
37+
"machines_v2",
38+
sa.Column("id", postgresql.UUID(), nullable=False),
39+
sa.Column("state", sa.PickleType, nullable=False),
40+
sa.PrimaryKeyConstraint("id"),
41+
if_not_exists=True,
42+
)
43+
44+
# Create table for campaigns v2
45+
campaigns_v2 = op.create_table(
46+
"campaigns_v2",
47+
sa.Column("id", postgresql.UUID(), nullable=False),
48+
sa.Column("name", postgresql.VARCHAR(), nullable=False),
49+
sa.Column("namespace", postgresql.UUID(), nullable=False, default=DEFAULT_CAMPAIGN_NAMESPACE),
50+
sa.Column("owner", postgresql.VARCHAR(), nullable=True),
51+
sa.Column(
52+
"metadata",
53+
postgresql.JSONB(),
54+
nullable=False,
55+
default=dict,
56+
server_default=sa.text("'{}'::json"),
57+
),
58+
sa.Column(
59+
"configuration",
60+
postgresql.JSONB(),
61+
nullable=False,
62+
default=dict,
63+
server_default=sa.text("'{}'::json"),
64+
),
65+
sa.Column("status", ENUM_COLUMN_AS_VARCHAR, nullable=False, default="waiting"),
66+
sa.Column(
67+
"machine", postgresql.UUID(), sa.ForeignKey(machines_v2.c.id, ondelete="CASCADE"), nullable=True
68+
),
69+
sa.PrimaryKeyConstraint("id"),
70+
sa.UniqueConstraint("name", "namespace"),
71+
if_not_exists=True,
72+
)
73+
74+
# Create node and edges tables for campaign digraph
75+
nodes_v2 = op.create_table(
76+
"nodes_v2",
77+
sa.Column("id", postgresql.UUID(), nullable=False),
78+
sa.Column(
79+
"namespace",
80+
postgresql.UUID(),
81+
sa.ForeignKey(campaigns_v2.c.id, ondelete="CASCADE"),
82+
nullable=False,
83+
),
84+
sa.Column("name", postgresql.VARCHAR(), nullable=False),
85+
sa.Column("version", postgresql.INTEGER(), nullable=False, default=1),
86+
sa.Column("kind", ENUM_COLUMN_AS_VARCHAR, nullable=False, default="node"),
87+
sa.Column(
88+
"metadata",
89+
postgresql.JSONB(),
90+
nullable=False,
91+
default=dict,
92+
server_default=sa.text("'{}'::json"),
93+
),
94+
sa.Column(
95+
"configuration",
96+
postgresql.JSONB(),
97+
nullable=False,
98+
default=dict,
99+
server_default=sa.text("'{}'::json"),
100+
),
101+
sa.Column("status", ENUM_COLUMN_AS_VARCHAR, nullable=False, default="waiting"),
102+
sa.Column(
103+
"machine", postgresql.UUID(), sa.ForeignKey(machines_v2.c.id, ondelete="CASCADE"), nullable=True
104+
),
105+
sa.PrimaryKeyConstraint("id"),
106+
sa.UniqueConstraint("name", "version", "namespace"),
107+
if_not_exists=True,
108+
)
109+
110+
_ = op.create_table(
111+
"edges_v2",
112+
sa.Column("id", postgresql.UUID(), nullable=False),
113+
sa.Column("name", postgresql.VARCHAR(), nullable=False),
114+
sa.Column(
115+
"namespace",
116+
postgresql.UUID(),
117+
sa.ForeignKey(campaigns_v2.c.id, ondelete="CASCADE"),
118+
nullable=False,
119+
),
120+
sa.Column("source", postgresql.UUID(), sa.ForeignKey(nodes_v2.c.id), nullable=False),
121+
sa.Column("target", postgresql.UUID(), sa.ForeignKey(nodes_v2.c.id), nullable=False),
122+
sa.Column(
123+
"metadata",
124+
postgresql.JSONB(),
125+
nullable=False,
126+
default=dict,
127+
server_default=sa.text("'{}'::json"),
128+
),
129+
sa.Column(
130+
"configuration",
131+
postgresql.JSONB(),
132+
nullable=False,
133+
default=dict,
134+
server_default=sa.text("'{}'::json"),
135+
),
136+
sa.PrimaryKeyConstraint("id"),
137+
sa.UniqueConstraint("source", "target", "namespace"),
138+
if_not_exists=True,
139+
)
140+
141+
# Create table for spec blocks v2 ("manifests")
142+
_ = op.create_table(
143+
"manifests_v2",
144+
sa.Column("id", postgresql.UUID(), nullable=False),
145+
sa.Column("name", postgresql.VARCHAR(), nullable=False),
146+
sa.Column(
147+
"namespace",
148+
postgresql.UUID(),
149+
sa.ForeignKey(campaigns_v2.c.id, ondelete="CASCADE"),
150+
nullable=False,
151+
),
152+
sa.Column("version", postgresql.INTEGER(), nullable=False, default=1),
153+
sa.Column("kind", ENUM_COLUMN_AS_VARCHAR, nullable=False, default="other"),
154+
sa.Column(
155+
"metadata",
156+
postgresql.JSONB(),
157+
nullable=False,
158+
default=dict,
159+
server_default=sa.text("'{}'::json"),
160+
),
161+
sa.Column(
162+
"spec",
163+
postgresql.JSONB(),
164+
nullable=False,
165+
default=dict,
166+
server_default=sa.text("'{}'::json"),
167+
),
168+
sa.PrimaryKeyConstraint("id"),
169+
sa.UniqueConstraint("name", "version", "namespace"),
170+
if_not_exists=True,
171+
)
172+
173+
# Create table for tasks v2
174+
_ = op.create_table(
175+
"tasks_v2",
176+
sa.Column("id", postgresql.UUID(), nullable=False),
177+
sa.Column("namespace", postgresql.UUID(), nullable=False),
178+
sa.Column("node", postgresql.UUID(), nullable=False),
179+
sa.Column("priority", postgresql.INTEGER(), nullable=True),
180+
sa.Column("created_at", postgresql.TIMESTAMP(timezone=True), nullable=False),
181+
sa.Column("last_processed_at", postgresql.TIMESTAMP(timezone=True), nullable=True),
182+
sa.Column("finished_at", postgresql.TIMESTAMP(timezone=True), nullable=True),
183+
sa.Column("wms_id", postgresql.VARCHAR(), nullable=True),
184+
sa.Column("site_affinity", postgresql.ARRAY(postgresql.VARCHAR()), nullable=True),
185+
sa.Column("status", ENUM_COLUMN_AS_VARCHAR, nullable=False),
186+
sa.Column("previous_status", ENUM_COLUMN_AS_VARCHAR, nullable=True),
187+
sa.Column(
188+
"metadata",
189+
postgresql.JSONB(),
190+
nullable=False,
191+
default=dict,
192+
server_default=sa.text("'{}'::json"),
193+
),
194+
sa.PrimaryKeyConstraint("id"),
195+
sa.ForeignKeyConstraint(["node"], ["nodes_v2.id"]),
196+
sa.ForeignKeyConstraint(["namespace"], ["campaigns_v2.id"]),
197+
if_not_exists=True,
198+
)
199+
200+
_ = op.create_table(
201+
"activity_log_v2",
202+
sa.Column("id", postgresql.UUID(), nullable=False),
203+
sa.Column("namespace", postgresql.UUID(), nullable=False),
204+
sa.Column("node", postgresql.UUID(), sa.ForeignKey(nodes_v2.c.id), nullable=False),
205+
sa.Column("operator", postgresql.VARCHAR(), nullable=False, default="root"),
206+
sa.Column("from_status", ENUM_COLUMN_AS_VARCHAR, nullable=False),
207+
sa.Column("to_status", ENUM_COLUMN_AS_VARCHAR, nullable=False),
208+
sa.Column(
209+
"detail",
210+
postgresql.JSONB(),
211+
nullable=False,
212+
default=dict,
213+
server_default=sa.text("'{}'::json"),
214+
),
215+
sa.Column(
216+
"metadata",
217+
postgresql.JSONB(),
218+
nullable=False,
219+
default=dict,
220+
server_default=sa.text("'{}'::json"),
221+
),
222+
if_not_exists=True,
223+
)
224+
225+
# Insert default campaign (namespace) record
226+
op.bulk_insert(
227+
campaigns_v2,
228+
[
229+
{
230+
"id": DEFAULT_CAMPAIGN_NAMESPACE,
231+
"namespace": str(NAMESPACE_DNS),
232+
"name": "DEFAULT",
233+
"owner": "root",
234+
}
235+
],
236+
)
237+
238+
239+
def downgrade() -> None:
240+
"""Drop tables in the reverse order in which they were created."""
241+
op.drop_table("activity_log_v2", if_exists=True)
242+
op.drop_table("tasks_v2", if_exists=True)
243+
op.drop_table("manifests_v2", if_exists=True)
244+
op.drop_table("edges_v2", if_exists=True)
245+
op.drop_table("nodes_v2", if_exists=True)
246+
op.drop_table("campaigns_v2", if_exists=True)
247+
op.drop_table("machines_v2", if_exists=True)

‎src/lsst/cmservice/common/enums.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -286,3 +286,20 @@ class WmsComputeSite(enum.Enum):
286286
lanc = 2
287287
ral = 3
288288
in2p3 = 4
289+
290+
291+
class ManifestKind(enum.Enum):
292+
"""Define a manifest kind"""
293+
294+
campaign = enum.auto()
295+
node = enum.auto()
296+
edge = enum.auto()
297+
# Legacy kinds
298+
specification = enum.auto()
299+
spec_block = enum.auto()
300+
step = enum.auto()
301+
group = enum.auto()
302+
job = enum.auto()
303+
script = enum.auto()
304+
# Fallback kind
305+
other = enum.auto()

0 commit comments

Comments
 (0)