Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/database/_migrations/versions/postgresql/2023_01_26_045500_2882cd2df463_implement_artifact_table.py: 67%

33 statements  

« prev     ^ index     » next       coverage.py v7.15.2, created at 2026-10-07 02:04 +0000

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")