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

1"""json_variables 

2 

3Revision ID: 94622c1663e8 

4Revises: b23c83a12cb4 

5Create Date: 2024-05-21 10:14:57.246286 

6 

7""" 

8 

9import json 

10 

11import sqlalchemy as sa 

12from alembic import op 

13 

14from prefect.server.utilities.database import JSON 

15 

16# revision identifiers, used by Alembic. 

17revision = "94622c1663e8" 

18down_revision = "b23c83a12cb4" 

19branch_labels = None 

20depends_on = None 

21 

22 

23def upgrade(): 

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

25 

26 conn = op.get_bind() 

27 

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

29 rows = result.fetchall() 

30 

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

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 op.drop_column("variable", "value") 

40 op.alter_column("variable", "json_value", new_column_name="value") 

41 

42 

43def downgrade(): 

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

45 

46 conn = op.get_bind() 

47 

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

49 rows = result.fetchall() 

50 

51 for variable_id, value in rows: 

52 string_value = str(value) 

53 conn.execute( 

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

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

56 ) 

57 

58 op.drop_column("variable", "value") 

59 op.alter_column("variable", "string_value", new_column_name="value", nullable=False)