Coverage for slidge/db/alembic/versions/e365e04e9ac3_fix_missing_existing_constraints_in_.py: 72%

1"""Fix missing existing constraints in SQLite 

2 

3Revision ID: e365e04e9ac3 

4Revises: abea9b63fa34 

5Create Date: 2026-02-11 19:57:06.448033 

6 

7""" 

8 

9import logging 

10from collections.abc import Sequence 

11 

12import sqlalchemy as sa 

13from alembic import op 

14 

15# revision identifiers, used by Alembic. 

16revision: str = "e365e04e9ac3" 

17down_revision: str | None = "abea9b63fa34" 

18branch_labels: str | Sequence[str] | None = None 

19depends_on: str | Sequence[str] | None = None 

20 

21 

22def upgrade() -> None: 

23 fix_constraints("contact", "user_account_id", "jid", "legacy_id") 

24 fix_constraints("contact_sent", "contact_id", "msg_id") 

25 fix_constraints("mam", "room_id", "stanza_id") 

26 fix_constraints("attachment", "user_account_id", "legacy_file_id") 

27 fix_constraints("participant", "room_id", "resource", "contact_id", "occupant_id") 

28 

29 

30def fix_constraints(table: str, col1: str, *cols: str) -> None: 

31 bogus_name = f"uq_{table}_{col1}" 

32 

33 inspector = sa.inspect(op.get_bind()) 

34 constraints = inspector.get_unique_constraints(table) 

35 constraint_exists = any( 

36 constraint["name"] == bogus_name for constraint in constraints 

37 ) 

38 if not constraint_exists: 

39 # If this runs on a newly created DB, this is unnecessary because we actually 

40 # edited past migration to allow running with postgresql. 

41 logging.info(f"No need to fix constraints for {bogus_name}") # noqa: LOG015 

42 return 

43 

44 with op.batch_alter_table(table, schema=None) as batch_op: 

45 logging.info(f"Dropping constraint {bogus_name}") # noqa: LOG015 

46 batch_op.drop_constraint(bogus_name, type_="unique") 

47 for col2 in cols: 

48 op.execute(f""" 

49 DELETE FROM '{table}' 

50 WHERE id NOT IN ( 

51 SELECT MIN(id) 

52 FROM '{table}' 

53 GROUP BY '{col1}', '{col2}' 

54 ); 

55 """) 

56 proper_name = f"uq_{table}_{col1}_{col2}" 

57 logging.info(f"Creating constraint {proper_name}") # noqa: LOG015 

58 batch_op.create_unique_constraint(proper_name, [col1, col2]) 

59 

60 

61def downgrade() -> None: 

62 raise RuntimeError("Downgrade not supported!")