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

17 statements  

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

1"""Add deployment to global concurrency limit FK 

2 

3Revision ID: 4ad4658cbefe 

4Revises: 7d6350aea855 

5Create Date: 2024-09-16 16:27:19.451150 

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 = "4ad4658cbefe" 1a

16down_revision = "7d6350aea855" 1a

17branch_labels = None 1a

18depends_on = None 1a

19 

20 

21def upgrade(): 1a

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

23 batch_op.add_column( 1a

24 sa.Column( 

25 "concurrency_limit_id", 

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

27 nullable=True, 

28 ) 

29 ) 

30 batch_op.create_foreign_key( 1a

31 batch_op.f("fk_deployment__concurrency_limit_id__concurrency_limit_v2"), 

32 "concurrency_limit_v2", 

33 ["concurrency_limit_id"], 

34 ["id"], 

35 ondelete="SET NULL", 

36 ) 

37 

38 # migrate existing data 

39 sql = sa.text( 1a

40 """ 

41 UPDATE deployment 

42 SET concurrency_limit_id = ( 

43 SELECT l.id 

44 FROM concurrency_limit_v2 l 

45 WHERE l.name = 'deployment:' || deployment.id 

46 ) 

47 WHERE EXISTS ( 

48 SELECT 1 

49 FROM concurrency_limit_v2 l 

50 WHERE l.name = 'deployment:' || deployment.id 

51 ); 

52 """ 

53 ) 

54 op.execute(sql) 1a

55 

56 

57def downgrade(): 1a

58 with op.batch_alter_table("deployment", schema=None) as batch_op: 

59 batch_op.drop_constraint( 

60 batch_op.f("fk_deployment__concurrency_limit_id__concurrency_limit_v2"), 

61 type_="foreignkey", 

62 ) 

63 batch_op.drop_column("concurrency_limit_id")