Skip to content

Commit 9c9e4b2

Browse files
fregataaclaude
fregataa
and
claude
committed
refactor(BA-6044,BA-6048): replace FK ref tuples with private dataclass
Replace the tuple-of-tuples FK reference tables in the domains and scaling_groups id-column migrations with a private frozen dataclass ``_FKRef``. Attribute access (``ref.constraint``, ``ref.column``, etc.) replaces positional unpacking and removes the need for the leading ``# (table, constraint, column, ...)`` comment. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent a3fb42b commit 9c9e4b2

2 files changed

Lines changed: 79 additions & 53 deletions

File tree

‎src/ai/backend/manager/models/alembic/versions/a1c3e5d7b9f2_add_id_column_to_domains.py‎

Lines changed: 44 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818

1919
# Part of: 26.5.0
2020

21+
from dataclasses import dataclass
22+
2123
import sqlalchemy as sa
2224
from alembic import op
2325

@@ -30,24 +32,38 @@
3032
depends_on = None
3133

3234

33-
# (referencing_table, fk_constraint_name, fk_column, on_update, on_delete)
34-
_DOMAIN_FK_REFS: tuple[tuple[str, str, str, str | None, str | None], ...] = (
35-
("endpoint_tokens", "fk_endpoint_tokens_domain_domains", "domain", None, "CASCADE"),
36-
("endpoints", "fk_endpoints_domain_domains", "domain", None, "RESTRICT"),
37-
("groups", "fk_groups_domain_name_domains", "domain_name", "CASCADE", "CASCADE"),
38-
("kernels", "fk_kernels_domain_name_domains", "domain_name", None, None),
39-
("model_cards", "fk_model_cards_domain_domains", "domain", None, "RESTRICT"),
40-
("routings", "fk_routings_domain_domains", "domain", None, "RESTRICT"),
41-
("session_templates", "fk_session_templates_domain_name_domains", "domain_name", None, None),
42-
("sessions", "fk_sessions_domain_name_domains", "domain_name", None, None),
43-
(
35+
@dataclass(frozen=True)
36+
class _FKRef:
37+
table: str
38+
constraint: str
39+
column: str
40+
on_update: str | None
41+
on_delete: str | None
42+
43+
44+
_DOMAIN_FK_REFS: tuple[_FKRef, ...] = (
45+
_FKRef("endpoint_tokens", "fk_endpoint_tokens_domain_domains", "domain", None, "CASCADE"),
46+
_FKRef("endpoints", "fk_endpoints_domain_domains", "domain", None, "RESTRICT"),
47+
_FKRef("groups", "fk_groups_domain_name_domains", "domain_name", "CASCADE", "CASCADE"),
48+
_FKRef("kernels", "fk_kernels_domain_name_domains", "domain_name", None, None),
49+
_FKRef("model_cards", "fk_model_cards_domain_domains", "domain", None, "RESTRICT"),
50+
_FKRef("routings", "fk_routings_domain_domains", "domain", None, "RESTRICT"),
51+
_FKRef(
52+
"session_templates",
53+
"fk_session_templates_domain_name_domains",
54+
"domain_name",
55+
None,
56+
None,
57+
),
58+
_FKRef("sessions", "fk_sessions_domain_name_domains", "domain_name", None, None),
59+
_FKRef(
4460
"sgroups_for_domains",
4561
"fk_sgroups_for_domains_domain_domains",
4662
"domain",
4763
"CASCADE",
4864
"CASCADE",
4965
),
50-
("users", "fk_users_domain_name_domains", "domain_name", None, None),
66+
_FKRef("users", "fk_users_domain_name_domains", "domain_name", None, None),
5167
)
5268

5369

@@ -62,41 +78,41 @@ def upgrade() -> None:
6278
),
6379
)
6480

65-
for table, fk_name, _column, _on_update, _on_delete in _DOMAIN_FK_REFS:
66-
op.drop_constraint(fk_name, table, type_="foreignkey")
81+
for ref in _DOMAIN_FK_REFS:
82+
op.drop_constraint(ref.constraint, ref.table, type_="foreignkey")
6783

6884
op.create_unique_constraint("uq_domains_name", "domains", ["name"])
6985
op.drop_constraint("pk_domains", "domains", type_="primary")
7086
op.create_primary_key("pk_domains", "domains", ["id"])
7187

72-
for table, fk_name, column, on_update, on_delete in _DOMAIN_FK_REFS:
88+
for ref in _DOMAIN_FK_REFS:
7389
op.create_foreign_key(
74-
fk_name,
75-
table,
90+
ref.constraint,
91+
ref.table,
7692
"domains",
77-
[column],
93+
[ref.column],
7894
["name"],
79-
onupdate=on_update,
80-
ondelete=on_delete,
95+
onupdate=ref.on_update,
96+
ondelete=ref.on_delete,
8197
)
8298

8399

84100
def downgrade() -> None:
85-
for table, fk_name, _column, _on_update, _on_delete in _DOMAIN_FK_REFS:
86-
op.drop_constraint(fk_name, table, type_="foreignkey")
101+
for ref in _DOMAIN_FK_REFS:
102+
op.drop_constraint(ref.constraint, ref.table, type_="foreignkey")
87103

88104
op.drop_constraint("pk_domains", "domains", type_="primary")
89105
op.drop_constraint("uq_domains_name", "domains", type_="unique")
90106
op.create_primary_key("pk_domains", "domains", ["name"])
91107
op.drop_column("domains", "id")
92108

93-
for table, fk_name, column, on_update, on_delete in _DOMAIN_FK_REFS:
109+
for ref in _DOMAIN_FK_REFS:
94110
op.create_foreign_key(
95-
fk_name,
96-
table,
111+
ref.constraint,
112+
ref.table,
97113
"domains",
98-
[column],
114+
[ref.column],
99115
["name"],
100-
onupdate=on_update,
101-
ondelete=on_delete,
116+
onupdate=ref.on_update,
117+
ondelete=ref.on_delete,
102118
)

‎src/ai/backend/manager/models/alembic/versions/b2d4f6e8c1a3_add_id_column_to_scaling_groups.py‎

Lines changed: 35 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717

1818
# Part of: 26.5.0
1919

20+
from dataclasses import dataclass
21+
2022
import sqlalchemy as sa
2123
from alembic import op
2224

@@ -29,39 +31,47 @@
2931
depends_on = None
3032

3133

32-
# (referencing_table, fk_constraint_name, fk_column, on_update, on_delete)
33-
_SCALING_GROUP_FK_REFS: tuple[tuple[str, str, str, str | None, str | None], ...] = (
34-
("agents", "fk_agents_scaling_group_scaling_groups", "scaling_group", None, None),
35-
(
34+
@dataclass(frozen=True)
35+
class _FKRef:
36+
table: str
37+
constraint: str
38+
column: str
39+
on_update: str | None
40+
on_delete: str | None
41+
42+
43+
_SCALING_GROUP_FK_REFS: tuple[_FKRef, ...] = (
44+
_FKRef("agents", "fk_agents_scaling_group_scaling_groups", "scaling_group", None, None),
45+
_FKRef(
3646
"endpoints",
3747
"fk_endpoints_resource_group_scaling_groups",
3848
"resource_group",
3949
None,
4050
"RESTRICT",
4151
),
42-
("kernels", "fk_kernels_scaling_group_scaling_groups", "scaling_group", None, None),
43-
(
52+
_FKRef("kernels", "fk_kernels_scaling_group_scaling_groups", "scaling_group", None, None),
53+
_FKRef(
4454
"sessions",
4555
"fk_sessions_scaling_group_name_scaling_groups",
4656
"scaling_group_name",
4757
None,
4858
None,
4959
),
50-
(
60+
_FKRef(
5161
"sgroups_for_domains",
5262
"fk_sgroups_for_domains_scaling_group_scaling_groups",
5363
"scaling_group",
5464
"CASCADE",
5565
"CASCADE",
5666
),
57-
(
67+
_FKRef(
5868
"sgroups_for_groups",
5969
"fk_sgroups_for_groups_scaling_group_scaling_groups",
6070
"scaling_group",
6171
"CASCADE",
6272
"CASCADE",
6373
),
64-
(
74+
_FKRef(
6575
"sgroups_for_keypairs",
6676
"fk_sgroups_for_keypairs_scaling_group_scaling_groups",
6777
"scaling_group",
@@ -82,41 +92,41 @@ def upgrade() -> None:
8292
),
8393
)
8494

85-
for table, fk_name, _column, _on_update, _on_delete in _SCALING_GROUP_FK_REFS:
86-
op.drop_constraint(fk_name, table, type_="foreignkey")
95+
for ref in _SCALING_GROUP_FK_REFS:
96+
op.drop_constraint(ref.constraint, ref.table, type_="foreignkey")
8797

8898
op.create_unique_constraint("uq_scaling_groups_name", "scaling_groups", ["name"])
8999
op.drop_constraint("pk_scaling_groups", "scaling_groups", type_="primary")
90100
op.create_primary_key("pk_scaling_groups", "scaling_groups", ["id"])
91101

92-
for table, fk_name, column, on_update, on_delete in _SCALING_GROUP_FK_REFS:
102+
for ref in _SCALING_GROUP_FK_REFS:
93103
op.create_foreign_key(
94-
fk_name,
95-
table,
104+
ref.constraint,
105+
ref.table,
96106
"scaling_groups",
97-
[column],
107+
[ref.column],
98108
["name"],
99-
onupdate=on_update,
100-
ondelete=on_delete,
109+
onupdate=ref.on_update,
110+
ondelete=ref.on_delete,
101111
)
102112

103113

104114
def downgrade() -> None:
105-
for table, fk_name, _column, _on_update, _on_delete in _SCALING_GROUP_FK_REFS:
106-
op.drop_constraint(fk_name, table, type_="foreignkey")
115+
for ref in _SCALING_GROUP_FK_REFS:
116+
op.drop_constraint(ref.constraint, ref.table, type_="foreignkey")
107117

108118
op.drop_constraint("pk_scaling_groups", "scaling_groups", type_="primary")
109119
op.drop_constraint("uq_scaling_groups_name", "scaling_groups", type_="unique")
110120
op.create_primary_key("pk_scaling_groups", "scaling_groups", ["name"])
111121
op.drop_column("scaling_groups", "id")
112122

113-
for table, fk_name, column, on_update, on_delete in _SCALING_GROUP_FK_REFS:
123+
for ref in _SCALING_GROUP_FK_REFS:
114124
op.create_foreign_key(
115-
fk_name,
116-
table,
125+
ref.constraint,
126+
ref.table,
117127
"scaling_groups",
118-
[column],
128+
[ref.column],
119129
["name"],
120-
onupdate=on_update,
121-
ondelete=on_delete,
130+
onupdate=ref.on_update,
131+
ondelete=ref.on_delete,
122132
)

0 commit comments

Comments
 (0)