Skip to content

Commit f011891

Browse files
committed
feat(db): Add migrations and models for v2 tables
1 parent 2e21c32 commit f011891

4 files changed

Lines changed: 499 additions & 17 deletions

File tree

Lines changed: 231 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,231 @@
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 campaigns v2
36+
campaigns_v2 = op.create_table(
37+
"campaigns_v2",
38+
sa.Column("id", postgresql.UUID(), nullable=False),
39+
sa.Column("name", postgresql.VARCHAR(), nullable=False),
40+
sa.Column("namespace", postgresql.UUID(), nullable=False, default=DEFAULT_CAMPAIGN_NAMESPACE),
41+
sa.Column("owner", postgresql.VARCHAR(), nullable=True),
42+
sa.Column(
43+
"metadata",
44+
postgresql.JSONB(),
45+
nullable=False,
46+
default=dict,
47+
server_default=sa.text("'{}'::json"),
48+
),
49+
sa.Column(
50+
"configuration",
51+
postgresql.JSONB(),
52+
nullable=False,
53+
default=dict,
54+
server_default=sa.text("'{}'::json"),
55+
),
56+
sa.Column("status", ENUM_COLUMN_AS_VARCHAR, nullable=False, default="waiting"),
57+
sa.PrimaryKeyConstraint("id"),
58+
sa.UniqueConstraint("name", "namespace"),
59+
if_not_exists=True,
60+
)
61+
62+
# Create node and edges tables for campaign digraph
63+
nodes_v2 = op.create_table(
64+
"nodes_v2",
65+
sa.Column("id", postgresql.UUID(), nullable=False),
66+
sa.Column(
67+
"namespace",
68+
postgresql.UUID(),
69+
sa.ForeignKey(campaigns_v2.c.id, ondelete="CASCADE"),
70+
nullable=False,
71+
),
72+
sa.Column("name", postgresql.VARCHAR(), nullable=False),
73+
sa.Column("version", postgresql.INTEGER(), nullable=False, default=1),
74+
sa.Column("kind", ENUM_COLUMN_AS_VARCHAR, nullable=False, default="node"),
75+
sa.Column(
76+
"metadata",
77+
postgresql.JSONB(),
78+
nullable=False,
79+
default=dict,
80+
server_default=sa.text("'{}'::json"),
81+
),
82+
sa.Column(
83+
"configuration",
84+
postgresql.JSONB(),
85+
nullable=False,
86+
default=dict,
87+
server_default=sa.text("'{}'::json"),
88+
),
89+
sa.Column("status", ENUM_COLUMN_AS_VARCHAR, nullable=False, default="waiting"),
90+
sa.PrimaryKeyConstraint("id"),
91+
sa.UniqueConstraint("name", "version", "namespace"),
92+
if_not_exists=True,
93+
)
94+
95+
_ = op.create_table(
96+
"edges_v2",
97+
sa.Column("id", postgresql.UUID(), nullable=False),
98+
sa.Column("name", postgresql.VARCHAR(), nullable=False),
99+
sa.Column(
100+
"namespace",
101+
postgresql.UUID(),
102+
sa.ForeignKey(campaigns_v2.c.id, ondelete="CASCADE"),
103+
nullable=False,
104+
),
105+
sa.Column("source", postgresql.UUID(), sa.ForeignKey(nodes_v2.c.id), nullable=False),
106+
sa.Column("target", postgresql.UUID(), sa.ForeignKey(nodes_v2.c.id), nullable=False),
107+
sa.Column(
108+
"metadata",
109+
postgresql.JSONB(),
110+
nullable=False,
111+
default=dict,
112+
server_default=sa.text("'{}'::json"),
113+
),
114+
sa.Column(
115+
"configuration",
116+
postgresql.JSONB(),
117+
nullable=False,
118+
default=dict,
119+
server_default=sa.text("'{}'::json"),
120+
),
121+
sa.PrimaryKeyConstraint("id"),
122+
sa.UniqueConstraint("source", "target", "namespace"),
123+
if_not_exists=True,
124+
)
125+
126+
# Create table for spec blocks v2 ("manifests")
127+
_ = op.create_table(
128+
"manifests_v2",
129+
sa.Column("id", postgresql.UUID(), nullable=False),
130+
sa.Column("name", postgresql.VARCHAR(), nullable=False),
131+
sa.Column(
132+
"namespace",
133+
postgresql.UUID(),
134+
sa.ForeignKey(campaigns_v2.c.id, ondelete="CASCADE"),
135+
nullable=False,
136+
),
137+
sa.Column("version", postgresql.INTEGER(), nullable=False, default=1),
138+
sa.Column("kind", ENUM_COLUMN_AS_VARCHAR, nullable=False, default="other"),
139+
sa.Column(
140+
"metadata",
141+
postgresql.JSONB(),
142+
nullable=False,
143+
default=dict,
144+
server_default=sa.text("'{}'::json"),
145+
),
146+
sa.Column(
147+
"spec",
148+
postgresql.JSONB(),
149+
nullable=False,
150+
default=dict,
151+
server_default=sa.text("'{}'::json"),
152+
),
153+
sa.PrimaryKeyConstraint("id"),
154+
sa.UniqueConstraint("name", "version", "namespace"),
155+
if_not_exists=True,
156+
)
157+
158+
# Create table for tasks v2
159+
_ = op.create_table(
160+
"tasks_v2",
161+
sa.Column("id", postgresql.UUID(), nullable=False),
162+
sa.Column("namespace", postgresql.UUID(), nullable=False),
163+
sa.Column("node", postgresql.UUID(), nullable=False),
164+
sa.Column("priority", postgresql.INTEGER(), nullable=True),
165+
sa.Column("created_at", postgresql.TIMESTAMP(timezone=True), nullable=False),
166+
sa.Column("last_processed_at", postgresql.TIMESTAMP(timezone=True), nullable=True),
167+
sa.Column("finished_at", postgresql.TIMESTAMP(timezone=True), nullable=True),
168+
sa.Column("wms_id", postgresql.VARCHAR(), nullable=True),
169+
sa.Column("site_affinity", postgresql.ARRAY(postgresql.VARCHAR()), nullable=True),
170+
sa.Column("status", ENUM_COLUMN_AS_VARCHAR, nullable=False),
171+
sa.Column("previous_status", ENUM_COLUMN_AS_VARCHAR, nullable=True),
172+
sa.Column(
173+
"metadata",
174+
postgresql.JSONB(),
175+
nullable=False,
176+
default=dict,
177+
server_default=sa.text("'{}'::json"),
178+
),
179+
sa.PrimaryKeyConstraint("id"),
180+
sa.ForeignKeyConstraint(["node"], ["nodes_v2.id"]),
181+
sa.ForeignKeyConstraint(["namespace"], ["campaigns_v2.id"]),
182+
if_not_exists=True,
183+
)
184+
185+
_ = op.create_table(
186+
"activity_log_v2",
187+
sa.Column("id", postgresql.UUID(), nullable=False),
188+
sa.Column("namespace", postgresql.UUID(), nullable=False),
189+
sa.Column("node", postgresql.UUID(), sa.ForeignKey(nodes_v2.c.id), nullable=False),
190+
sa.Column("operator", postgresql.VARCHAR(), nullable=False, default="root"),
191+
sa.Column("from_status", ENUM_COLUMN_AS_VARCHAR, nullable=False),
192+
sa.Column("to_status", ENUM_COLUMN_AS_VARCHAR, nullable=False),
193+
sa.Column(
194+
"detail",
195+
postgresql.JSONB(),
196+
nullable=False,
197+
default=dict,
198+
server_default=sa.text("'{}'::json"),
199+
),
200+
sa.Column(
201+
"metadata",
202+
postgresql.JSONB(),
203+
nullable=False,
204+
default=dict,
205+
server_default=sa.text("'{}'::json"),
206+
),
207+
if_not_exists=True,
208+
)
209+
210+
# Insert default campaign (namespace) record
211+
op.bulk_insert(
212+
campaigns_v2,
213+
[
214+
{
215+
"id": DEFAULT_CAMPAIGN_NAMESPACE,
216+
"namespace": str(NAMESPACE_DNS),
217+
"name": "DEFAULT",
218+
"owner": "root",
219+
}
220+
],
221+
)
222+
223+
224+
def downgrade() -> None:
225+
"""Drop tables in the reverse order in which they were created."""
226+
op.drop_table("activity_log_v2", if_exists=True)
227+
op.drop_table("tasks_v2", if_exists=True)
228+
op.drop_table("manifests_v2", if_exists=True)
229+
op.drop_table("edges_v2", if_exists=True)
230+
op.drop_table("nodes_v2", if_exists=True)
231+
op.drop_table("campaigns_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)