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

33 statements  

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

1"""State data migration cleanup 

2 

3Revision ID: f92143d30c27 

4Revises: f92143d30c26 

5Create Date: 2023-01-12 00:00:44.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 = "f92143d30c27" 1a

16down_revision = "f92143d30c26" 1a

17branch_labels = None 1a

18depends_on = None 1a

19 

20 

21def upgrade(): 1a

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

23 

24 # drop state id columns after data migration 

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

26 batch_op.drop_index(batch_op.f("ix_artifact__task_run_state_id")) 1a

27 batch_op.drop_column("task_run_state_id") 1a

28 batch_op.drop_index(batch_op.f("ix_artifact__flow_run_state_id")) 1a

29 batch_op.drop_column("flow_run_state_id") 1a

30 

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

32 batch_op.drop_index(batch_op.f("ix_flow_run_state__has_data")) 1a

33 batch_op.drop_column("has_data") 1a

34 

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

36 batch_op.drop_index(batch_op.f("ix_task_run_state__has_data")) 1a

37 batch_op.drop_column("has_data") 1a

38 

39 

40def downgrade(): 1a

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

42 

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

44 batch_op.add_column( 

45 sa.Column( 

46 "flow_run_state_id", 

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

48 nullable=True, 

49 ), 

50 ) 

51 batch_op.create_index( 

52 batch_op.f("ix_artifact__flow_run_state_id"), 

53 ["flow_run_state_id"], 

54 unique=False, 

55 ) 

56 batch_op.add_column( 

57 sa.Column( 

58 "task_run_state_id", 

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

60 nullable=True, 

61 ), 

62 ) 

63 batch_op.create_index( 

64 batch_op.f("ix_artifact__task_run_state_id"), 

65 ["task_run_state_id"], 

66 unique=False, 

67 ) 

68 

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

70 batch_op.add_column(sa.Column("has_data", sa.Boolean)) 

71 batch_op.create_index( 

72 batch_op.f("ix_flow_run_state__has_data"), 

73 ["has_data"], 

74 unique=False, 

75 ) 

76 

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

78 batch_op.add_column(sa.Column("has_data", sa.Boolean)) 

79 batch_op.create_index( 

80 batch_op.f("ix_task_run_state__has_data"), 

81 ["has_data"], 

82 unique=False, 

83 )