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

1"""Add `trigger_id` to the unique index for `automation_bucket` 

2 

3Revision ID: 954db7517015 

4Revises: 916718e8330f 

5Create Date: 2024-04-09 13:20:36.838767 

6 

7""" 

8 

9from alembic import op 

10 

11import prefect.server.utilities.database 

12 

13# revision identifiers, used by Alembic. 

14revision = "954db7517015" 

15down_revision = "916718e8330f" 

16branch_labels = None 

17depends_on = None 

18 

19 

20def upgrade(): 

21 op.alter_column( 

22 "automation_bucket", 

23 "trigger_id", 

24 existing_type=prefect.server.utilities.database.UUID(), 

25 nullable=False, 

26 ) 

27 op.alter_column( 

28 "automation_bucket", 

29 "triggered_at", 

30 existing_type=prefect.server.utilities.database.Timestamp(timezone=True), 

31 nullable=True, 

32 ) 

33 op.drop_index( 

34 "uq_automation_bucket__automation_id__bucketing_key", 

35 table_name="automation_bucket", 

36 ) 

37 op.create_index( 

38 "uq_automation_bucket__automation_id__trigger_id__bucketing_key", 

39 "automation_bucket", 

40 ["automation_id", "trigger_id", "bucketing_key"], 

41 unique=True, 

42 ) 

43 

44 

45def downgrade(): 

46 op.drop_index( 

47 "uq_automation_bucket__automation_id__trigger_id__bucketing_key", 

48 table_name="automation_bucket", 

49 ) 

50 op.create_index( 

51 "uq_automation_bucket__automation_id__bucketing_key", 

52 "automation_bucket", 

53 ["automation_id", "bucketing_key"], 

54 unique=True, 

55 ) 

56 op.alter_column( 

57 "automation_bucket", 

58 "triggered_at", 

59 existing_type=prefect.server.utilities.database.Timestamp(timezone=True), 

60 nullable=False, 

61 ) 

62 op.alter_column( 

63 "automation_bucket", 

64 "trigger_id", 

65 existing_type=prefect.server.utilities.database.UUID(), 

66 nullable=True, 

67 )