1"""Implements the artifact table and migrates run results
2
3Revision ID: 2882cd2df463
4Revises: d481d5058a19
5Create Date: 2023-01-26 04:55:00.358638
6
7"""
8
9import sqlalchemy as sa
10from alembic import op
11
12import prefect
13
14# revision identifiers, used by Alembic.
15revision = "2882cd2df463"
16down_revision = "9326a6aee18b"
17branch_labels = None
18depends_on = None
19
20
21def upgrade():
22 op.create_table(
23 "artifact",
24 sa.Column(
25 "id",
26 prefect.server.utilities.database.UUID(),
27 server_default=sa.text("(GEN_RANDOM_UUID())"),
28 nullable=False,
29 ),
30 sa.Column(
31 "created",
32 prefect.server.utilities.database.Timestamp(timezone=True),
33 server_default=sa.text("CURRENT_TIMESTAMP"),
34 nullable=False,
35 ),
36 sa.Column(
37 "updated",
38 prefect.server.utilities.database.Timestamp(timezone=True),
39 server_default=sa.text("CURRENT_TIMESTAMP"),
40 nullable=False,
41 ),
42 sa.Column("key", sa.String(), nullable=True),
43 sa.Column("type", sa.String(), nullable=True),
44 sa.Column("data", sa.JSON(), nullable=True),
45 sa.Column("metadata_", sa.JSON(), nullable=True),
46 sa.Column(
47 "task_run_id", prefect.server.utilities.database.UUID(), nullable=True
48 ),
49 sa.Column(
50 "task_run_state_id", prefect.server.utilities.database.UUID(), nullable=True
51 ),
52 sa.Column(
53 "flow_run_id", prefect.server.utilities.database.UUID(), nullable=True
54 ),
55 sa.Column(
56 "flow_run_state_id", prefect.server.utilities.database.UUID(), nullable=True
57 ),
58 sa.ForeignKeyConstraint(
59 ["flow_run_id"],
60 ["flow_run.id"],
61 name=op.f("fk_artifact__flow_run_id__flow_run"),
62 ),
63 sa.ForeignKeyConstraint(
64 ["flow_run_state_id"],
65 ["flow_run_state.id"],
66 name=op.f("fk_artifact__flow_run_state_id__flow_run_state"),
67 ),
68 sa.ForeignKeyConstraint(
69 ["task_run_id"],
70 ["task_run.id"],
71 name=op.f("fk_artifact__task_run_id__task_run"),
72 ),
73 sa.ForeignKeyConstraint(
74 ["task_run_state_id"],
75 ["task_run_state.id"],
76 name=op.f("fk_artifact__task_run_state_id__task_run_state"),
77 ),
78 sa.PrimaryKeyConstraint("id", name=op.f("pk_artifact")),
79 )
80 op.create_index(
81 op.f("ix_artifact__flow_run_id"), "artifact", ["flow_run_id"], unique=False
82 )
83 op.create_index(
84 op.f("ix_artifact__flow_run_state_id"),
85 "artifact",
86 ["flow_run_state_id"],
87 unique=False,
88 )
89 op.create_index(op.f("ix_artifact__key"), "artifact", ["key"], unique=True)
90 op.create_index(
91 op.f("ix_artifact__task_run_id"), "artifact", ["task_run_id"], unique=False
92 )
93 op.create_index(
94 op.f("ix_artifact__task_run_state_id"),
95 "artifact",
96 ["task_run_state_id"],
97 unique=False,
98 )
99 op.create_index(op.f("ix_artifact__updated"), "artifact", ["updated"], unique=False)
100
101 op.add_column(
102 "flow_run_state",
103 sa.Column(
104 "result_artifact_id",
105 prefect.server.utilities.database.UUID(),
106 nullable=True,
107 ),
108 )
109 op.create_index(
110 op.f("ix_flow_run_state__result_artifact_id"),
111 "flow_run_state",
112 ["result_artifact_id"],
113 unique=False,
114 )
115 op.create_foreign_key(
116 op.f("fk_flow_run_state__result_artifact_id__artifact"),
117 "flow_run_state",
118 "artifact",
119 ["result_artifact_id"],
120 ["id"],
121 ondelete="SET NULL",
122 use_alter=True,
123 )
124 op.add_column(
125 "task_run_state",
126 sa.Column(
127 "result_artifact_id",
128 prefect.server.utilities.database.UUID(),
129 nullable=True,
130 ),
131 )
132 op.create_index(
133 op.f("ix_task_run_state__result_artifact_id"),
134 "task_run_state",
135 ["result_artifact_id"],
136 unique=False,
137 )
138 op.create_foreign_key(
139 op.f("fk_task_run_state__result_artifact_id__artifact"),
140 "task_run_state",
141 "artifact",
142 ["result_artifact_id"],
143 ["id"],
144 ondelete="SET NULL",
145 use_alter=True,
146 )
147
148
149def downgrade():
150 op.drop_constraint(
151 op.f("fk_task_run_state__result_artifact_id__artifact"),
152 "task_run_state",
153 type_="foreignkey",
154 )
155 op.drop_index(
156 op.f("ix_task_run_state__result_artifact_id"), table_name="task_run_state"
157 )
158 op.drop_column("task_run_state", "result_artifact_id")
159 op.drop_constraint(
160 op.f("fk_flow_run_state__result_artifact_id__artifact"),
161 "flow_run_state",
162 type_="foreignkey",
163 )
164 op.drop_index(
165 op.f("ix_flow_run_state__result_artifact_id"), table_name="flow_run_state"
166 )
167 op.drop_column("flow_run_state", "result_artifact_id")
168
169 op.drop_index(op.f("ix_artifact__updated"), table_name="artifact")
170 op.drop_index(op.f("ix_artifact__task_run_id"), table_name="artifact")
171 op.drop_index(op.f("ix_artifact__key"), table_name="artifact")
172 op.drop_index(op.f("ix_artifact__flow_run_id"), table_name="artifact")
173 op.drop_table("artifact")