Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/database/_migrations/versions/sqlite/2024_05_21_123101_2ac65f1758c2_json_variables.py: 56%

30 statements  

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

1"""json_variables 

2 

3Revision ID: 2ac65f1758c2 

4Revises: 20fbd53b3cef 

5Create Date: 2024-05-21 12:29:43.948758 

6 

7""" 

8 

9import json 1a

10 

11import sqlalchemy as sa 1a

12from alembic import op 1a

13 

14from prefect.server.utilities.database import JSON 1a

15 

16# revision identifiers, used by Alembic. 

17revision = "2ac65f1758c2" 1a

18down_revision = "20fbd53b3cef" 1a

19branch_labels = None 1a

20depends_on = None 1a

21 

22 

23def upgrade(): 1a

24 op.add_column("variable", sa.Column("json_value", JSON, nullable=True)) 1a

25 

26 conn = op.get_bind() 1a

27 

28 result = conn.execute(sa.text("SELECT id, value FROM variable")) 1a

29 rows = result.fetchall() 1a

30 

31 for variable_id, value in rows: 31 ↛ 33line 31 didn't jump to line 33 because the loop on line 31 never started1a

32 # these values need to be json compatible strings 

33 json_value = json.dumps(value) 

34 conn.execute( 

35 sa.text("UPDATE variable SET json_value = :json_value WHERE id = :id"), 

36 {"json_value": json_value, "id": variable_id}, 

37 ) 

38 

39 with op.batch_alter_table("variable") as batch_op: 1a

40 batch_op.drop_column("value") 1a

41 batch_op.alter_column("json_value", new_column_name="value") 1a

42 

43 

44def downgrade(): 1a

45 op.add_column("variable", sa.Column("string_value", sa.String, nullable=True)) 

46 

47 conn = op.get_bind() 

48 

49 result = conn.execute(sa.text("SELECT id, value FROM variable")) 

50 rows = result.fetchall() 

51 

52 for variable_id, value in rows: 

53 string_value = json.loads(str(value)) 

54 conn.execute( 

55 sa.text("UPDATE variable SET string_value = :string_value WHERE id = :id"), 

56 {"string_value": string_value, "id": variable_id}, 

57 ) 

58 

59 with op.batch_alter_table("variable") as batch_op: 

60 batch_op.drop_column("value") 

61 batch_op.alter_column("string_value", new_column_name="value", nullable=False)