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

45 statements  

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

1"""Adds block type slug 

2 

3Revision ID: f335f9633eec 

4Revises: 2fe8ef6a6514 

5Create Date: 2022-07-25 14:25:15.809720 

6 

7""" 

8 

9import sqlalchemy as sa 1a

10from alembic import op 1a

11 

12from prefect.blocks.core import Block 1a

13from prefect.utilities.slugify import slugify 1a

14 

15# revision identifiers, used by Alembic. 

16revision = "f335f9633eec" 1a

17down_revision = "2fe8ef6a6514" 1a

18branch_labels = None 1a

19depends_on = None 1a

20 

21 

22def replace_name_with_slug(fields): 1a

23 if isinstance(fields, list): 

24 return [replace_name_with_slug(x) for x in fields] 

25 if isinstance(fields, dict): 

26 new_dict = {} 

27 for k, v in fields.items(): 

28 if k == "block_type_name": 

29 new_dict["block_type_slug"] = slugify(v) 

30 else: 

31 new_dict[k] = v 

32 return new_dict 

33 return fields 

34 

35 

36def upgrade(): 1a

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

38 

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

40 batch_op.add_column(sa.Column("slug", sa.String(), nullable=True)) 1a

41 batch_op.drop_index("uq_block_type__name") 1a

42 batch_op.create_index("uq_block_type__slug", ["slug"], unique=True) 1a

43 

44 # Add slugs to existing block types 

45 connection = op.get_bind() 1a

46 meta_data = sa.MetaData() 1a

47 meta_data.reflect(connection) 1a

48 BLOCK_TYPE = meta_data.tables["block_type"] 1a

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

50 

51 block_types_result = connection.execute( 1a

52 sa.select(BLOCK_TYPE.c.id, BLOCK_TYPE.c.name) 

53 ).all() 

54 for block_type_id, block_type_name in block_types_result: 54 ↛ 55line 54 didn't jump to line 55 because the loop on line 54 never started1a

55 connection.execute( 

56 sa.update(BLOCK_TYPE) 

57 .where(BLOCK_TYPE.c.id == block_type_id) 

58 .values(slug=slugify(block_type_name)) 

59 ) 

60 

61 block_schemas_result = connection.execute( 

62 sa.select(BLOCK_SCHEMA.c.id, BLOCK_SCHEMA.c.fields).where( 

63 BLOCK_SCHEMA.c.block_type_id == block_type_id 

64 ) 

65 ).all() 

66 for block_schema_id, block_schema_fields in block_schemas_result: 

67 new_fields = replace_name_with_slug(block_schema_fields) 

68 connection.execute( 

69 sa.update(BLOCK_SCHEMA) 

70 .where(BLOCK_SCHEMA.c.id == block_schema_id) 

71 .values( 

72 fields=replace_name_with_slug(block_schema_fields), 

73 checksum=Block._calculate_schema_checksum(new_fields), 

74 ) 

75 ) 

76 

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

78 batch_op.alter_column("slug", existing_type=sa.VARCHAR(), nullable=False) 1a

79 

80 op.execute("PRAGMA foreign_keys=ON") 1a

81 

82 

83def downgrade(): 1a

84 with op.batch_alter_table("block_type", schema=None) as batch_op: 

85 batch_op.drop_index("uq_block_type__slug") 

86 batch_op.create_index("uq_block_type__name", ["name"], unique=False) 

87 batch_op.drop_column("slug")