Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/database/_migrations/versions/sqlite/2023_01_12_000042_f92143d30c24_implement_artifact_table.py: 63%

41 statements  

« prev     ^ index     » next       coverage.py v7.10.6, created at 2025-12-05 13:38 +0000

1"""Implements the artifact table and migrates run results 

2 

3Revision ID: f92143d30c24 

4Revises: bb38729c471a 

5Create Date: 2023-01-12 00:00:42.488367 

6 

7""" 

8 

9import sqlalchemy as sa 1a

10from alembic import op 1a

11 

12import prefect 1a

13 

14# revision identifiers, used by Alembic. 

15revision = "f92143d30c24" 1a

16down_revision = "bb38729c471a" 1a

17branch_labels = None 1a

18depends_on = None 1a

19 

20 

21def upgrade(): 1a

22 op.execute("PRAGMA foreign_keys=OFF") 1a

23 

24 op.create_table( 1a

25 "artifact", 

26 sa.Column( 

27 "id", 

28 prefect.server.utilities.database.UUID(), 

29 server_default=sa.text( 

30 "(\n (\n lower(hex(randomblob(4)))\n || '-'\n " 

31 " || lower(hex(randomblob(2)))\n || '-4'\n ||" 

32 " substr(lower(hex(randomblob(2))),2)\n || '-'\n ||" 

33 " substr('89ab',abs(random()) % 4 + 1, 1)\n ||" 

34 " substr(lower(hex(randomblob(2))),2)\n || '-'\n ||" 

35 " lower(hex(randomblob(6)))\n )\n )" 

36 ), 

37 nullable=False, 

38 ), 

39 sa.Column( 

40 "created", 

41 prefect.server.utilities.database.Timestamp(timezone=True), 

42 server_default=sa.text("(strftime('%Y-%m-%d %H:%M:%f000', 'now'))"), 

43 nullable=False, 

44 ), 

45 sa.Column( 

46 "updated", 

47 prefect.server.utilities.database.Timestamp(timezone=True), 

48 server_default=sa.text("(strftime('%Y-%m-%d %H:%M:%f000', 'now'))"), 

49 nullable=False, 

50 ), 

51 sa.Column("key", sa.String(), nullable=True), 

52 sa.Column("type", sa.String(), nullable=True), 

53 sa.Column("data", sa.JSON(), nullable=True), 

54 sa.Column("metadata_", sa.JSON(), nullable=True), 

55 sa.Column( 

56 "task_run_id", prefect.server.utilities.database.UUID(), nullable=True 

57 ), 

58 sa.Column( 

59 "task_run_state_id", prefect.server.utilities.database.UUID(), nullable=True 

60 ), 

61 sa.Column( 

62 "flow_run_id", prefect.server.utilities.database.UUID(), nullable=True 

63 ), 

64 sa.Column( 

65 "flow_run_state_id", prefect.server.utilities.database.UUID(), nullable=True 

66 ), 

67 sa.ForeignKeyConstraint( 

68 ["flow_run_state_id"], 

69 ["flow_run_state.id"], 

70 name=op.f("fk_artifact__flow_run_state_id__flow_run_state"), 

71 ), 

72 sa.ForeignKeyConstraint( 

73 ["flow_run_id"], 

74 ["flow_run.id"], 

75 name=op.f("fk_artifact__flow_run_id__flow_run"), 

76 ), 

77 sa.ForeignKeyConstraint( 

78 ["task_run_state_id"], 

79 ["task_run_state.id"], 

80 name=op.f("fk_artifact__task_run_state_id__task_run_state"), 

81 ), 

82 sa.ForeignKeyConstraint( 

83 ["task_run_id"], 

84 ["task_run.id"], 

85 name=op.f("fk_artifact__task_run_id__task_run"), 

86 ), 

87 sa.PrimaryKeyConstraint("id", name=op.f("pk_artifact")), 

88 ) 

89 with op.batch_alter_table("artifact", schema=None) as batch_op: 1a

90 batch_op.create_index( 1a

91 batch_op.f("ix_artifact__flow_run_id"), 

92 ["flow_run_id"], 

93 unique=False, 

94 ) 

95 batch_op.create_index( 1a

96 batch_op.f("ix_artifact__flow_run_state_id"), 

97 ["flow_run_state_id"], 

98 unique=False, 

99 ) 

100 batch_op.create_index(batch_op.f("ix_artifact__key"), ["key"], unique=True) 1a

101 batch_op.create_index( 1a

102 batch_op.f("ix_artifact__task_run_id"), 

103 ["task_run_id"], 

104 unique=False, 

105 ) 

106 batch_op.create_index( 1a

107 batch_op.f("ix_artifact__task_run_state_id"), 

108 ["task_run_state_id"], 

109 unique=False, 

110 ) 

111 batch_op.create_index( 1a

112 batch_op.f("ix_artifact__updated"), ["updated"], unique=False 

113 ) 

114 

115 with op.batch_alter_table("flow_run_state", schema=None) as batch_op: 1a

116 batch_op.add_column( 1a

117 sa.Column( 

118 "result_artifact_id", 

119 prefect.server.utilities.database.UUID(), 

120 nullable=True, 

121 ) 

122 ) 

123 batch_op.create_index( 1a

124 batch_op.f("ix_flow_run_state__result_artifact_id"), 

125 ["result_artifact_id"], 

126 unique=False, 

127 ) 

128 batch_op.create_foreign_key( 1a

129 batch_op.f("fk_flow_run_state__result_artifact_id__artifact"), 

130 "artifact", 

131 ["result_artifact_id"], 

132 ["id"], 

133 ondelete="SET NULL", 

134 use_alter=True, 

135 ) 

136 

137 with op.batch_alter_table("task_run_state", schema=None) as batch_op: 1a

138 batch_op.add_column( 1a

139 sa.Column( 

140 "result_artifact_id", 

141 prefect.server.utilities.database.UUID(), 

142 nullable=True, 

143 ) 

144 ) 

145 batch_op.create_index( 1a

146 batch_op.f("ix_task_run_state__result_artifact_id"), 

147 ["result_artifact_id"], 

148 unique=False, 

149 ) 

150 batch_op.create_foreign_key( 1a

151 batch_op.f("fk_task_run_state__result_artifact_id__artifact"), 

152 "artifact", 

153 ["result_artifact_id"], 

154 ["id"], 

155 ondelete="SET NULL", 

156 use_alter=True, 

157 ) 

158 

159 

160def downgrade(): 1a

161 op.execute("PRAGMA foreign_keys=OFF") 

162 

163 with op.batch_alter_table("task_run_state", schema=None) as batch_op: 

164 batch_op.drop_constraint( 

165 batch_op.f("fk_task_run_state__result_artifact_id__artifact"), 

166 type_="foreignkey", 

167 ) 

168 batch_op.drop_index(batch_op.f("ix_task_run_state__result_artifact_id")) 

169 batch_op.drop_column("result_artifact_id") 

170 

171 with op.batch_alter_table("flow_run_state", schema=None) as batch_op: 

172 batch_op.drop_constraint( 

173 batch_op.f("fk_flow_run_state__result_artifact_id__artifact"), 

174 type_="foreignkey", 

175 ) 

176 batch_op.drop_index(batch_op.f("ix_flow_run_state__result_artifact_id")) 

177 batch_op.drop_column("result_artifact_id") 

178 

179 with op.batch_alter_table("artifact", schema=None) as batch_op: 

180 batch_op.drop_index(batch_op.f("ix_artifact__updated")) 

181 batch_op.drop_index(batch_op.f("ix_artifact__task_run_id")) 

182 batch_op.drop_index(batch_op.f("ix_artifact__key")) 

183 batch_op.drop_index(batch_op.f("ix_artifact__flow_run_id")) 

184 

185 op.drop_table("artifact")