Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/database/_migrations/versions/sqlite/2022_05_19_165808_33439667aeea_add_block_schema_capabilities.py: 53%

39 statements  

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

1"""Add BlockSchema capabilities 

2 

3Revision ID: 33439667aeea 

4Revises: 888a0bb0df7b 

5Create Date: 2022-05-19 16:58:08.802305 

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 = "33439667aeea" 1a

16down_revision = "888a0bb0df7b" 1a

17branch_labels = None 1a

18depends_on = None 1a

19 

20 

21def upgrade(): 1a

22 # ### commands auto generated by Alembic - please adjust! ### 

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

24 batch_op.drop_constraint("fk_block__block_schema_id__block_schema") 1a

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

26 batch_op.add_column( 1a

27 sa.Column( 

28 "capabilities", 

29 prefect.server.utilities.database.JSON(astext_type=sa.Text()), 

30 server_default="[]", 

31 nullable=False, 

32 ), 

33 ) 

34 

35 connection = op.get_bind() 1a

36 meta_data = sa.MetaData() 1a

37 meta_data.reflect(connection) 1a

38 BLOCK_SCHEMA = meta_data.tables["block_schema"] 1a

39 

40 results = connection.execute(sa.select(BLOCK_SCHEMA.c.id, BLOCK_SCHEMA.c.type)) 1a

41 

42 for id, type in results: 42 ↛ 43line 42 didn't jump to line 43 because the loop on line 42 never started1a

43 if type == "STORAGE": 

44 connection.execute( 

45 sa.update(BLOCK_SCHEMA) 

46 .where(BLOCK_SCHEMA.c.id == id) 

47 .values(capabilities=["writeable", "readable", "storage"]) 

48 ) 

49 

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

51 batch_op.drop_index("ix_block_schema__type") 1a

52 batch_op.drop_column("type") 1a

53 

54 with op.batch_alter_table( 1a

55 "block_document", 

56 schema=None, 

57 ) as batch_op: 

58 batch_op.create_foreign_key( 1a

59 batch_op.f("fk_block__block_schema_id__block_schema"), 

60 "block_schema", 

61 ["block_schema_id"], 

62 ["id"], 

63 ondelete="cascade", 

64 ) 

65 # ### end Alembic commands ### 

66 

67 

68def downgrade(): 1a

69 # ### commands auto generated by Alembic - please adjust! ### 

70 with op.batch_alter_table("block_schema", schema=None) as batch_op: 

71 batch_op.add_column(sa.Column("type", sa.VARCHAR(), nullable=True)) 

72 batch_op.create_index("ix_block_schema__type", ["type"], unique=False) 

73 

74 connection = op.get_bind() 

75 meta_data = sa.MetaData() 

76 meta_data.reflect(connection) 

77 BLOCK_SCHEMA = meta_data.tables["block_schema"] 

78 

79 results = connection.execute( 

80 sa.select(BLOCK_SCHEMA.c.id, BLOCK_SCHEMA.c.capabilities) 

81 ) 

82 

83 for id, capabilities in results: 

84 if "storage" in capabilities: 

85 connection.execute( 

86 sa.update(BLOCK_SCHEMA) 

87 .where(BLOCK_SCHEMA.c.id == id) 

88 .values(type="STORAGE") 

89 ) 

90 

91 with op.batch_alter_table("block_schema", schema=None) as batch_op: 

92 batch_op.drop_column("capabilities") 

93 # ### end Alembic commands ###