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

1"""State data migration cleanup 

2 

3Revision ID: 2882cd2df466 

4Revises: 2882cd2df465 

5Create Date: 2023-01-26 04:55:02.358638 

6 

7""" 

8 

9import sqlalchemy as sa 

10from alembic import op 

11 

12import prefect 

13 

14# revision identifiers, used by Alembic. 

15revision = "2882cd2df466" 

16down_revision = "2882cd2df465" 

17branch_labels = None 

18depends_on = None 

19 

20 

21def upgrade(): 

22 # drop state id columns after data migration 

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

24 batch_op.drop_index(batch_op.f("ix_artifact__task_run_state_id")) 

25 batch_op.drop_column("task_run_state_id") 

26 batch_op.drop_index(batch_op.f("ix_artifact__flow_run_state_id")) 

27 batch_op.drop_column("flow_run_state_id") 

28 

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

30 batch_op.drop_index(batch_op.f("ix_flow_run_state__has_data")) 

31 batch_op.drop_column("has_data") 

32 

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

34 batch_op.drop_index(batch_op.f("ix_task_run_state__has_data")) 

35 batch_op.drop_column("has_data") 

36 

37 

38def downgrade(): 

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

40 batch_op.add_column( 

41 sa.Column( 

42 "flow_run_state_id", 

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

44 nullable=True, 

45 ), 

46 ) 

47 batch_op.create_index( 

48 batch_op.f("ix_artifact__flow_run_state_id"), 

49 ["flow_run_state_id"], 

50 unique=False, 

51 ) 

52 batch_op.add_column( 

53 sa.Column( 

54 "task_run_state_id", 

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

56 nullable=True, 

57 ), 

58 ) 

59 batch_op.create_index( 

60 batch_op.f("ix_artifact__task_run_state_id"), 

61 ["task_run_state_id"], 

62 unique=False, 

63 ) 

64 

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

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

67 batch_op.create_index( 

68 batch_op.f("ix_flow_run_state__has_data"), 

69 ["has_data"], 

70 unique=False, 

71 ) 

72 

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

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

75 batch_op.create_index( 

76 batch_op.f("ix_task_run_state__has_data"), 

77 ["has_data"], 

78 unique=False, 

79 )