|
21 | 21 | from typing import Literal |
22 | 22 |
|
23 | 23 | import sqlalchemy as sa |
| 24 | +from alembic.operations.ops import CreateIndexOp, CreateTableOp, DropIndexOp, DropTableOp |
24 | 25 | from simple_module_core import ModuleBase |
25 | 26 | from simple_module_core.discovery import discover_modules, get_module_package_name |
26 | | -from sqlalchemy import MetaData |
| 27 | +from sqlalchemy import Column, Index, MetaData |
27 | 28 | from sqlalchemy.schema import SchemaItem |
28 | 29 |
|
29 | 30 | from simple_module_db.base import all_module_bases |
|
36 | 37 | "schema", "table", "column", "index", "unique_constraint", "foreign_key_constraint" |
37 | 38 | ] |
38 | 39 | IncludeObjectFn = Callable[[SchemaItem, str | None, _SchemaItemType, bool, SchemaItem | None], bool] |
| 40 | +ProcessRevisionDirectivesFn = Callable[[object, object, list], None] |
39 | 41 |
|
40 | 42 |
|
41 | 43 | def build_module_metadata(modules: Sequence[ModuleBase] | None = None) -> MetaData: |
@@ -110,6 +112,97 @@ def include_object( |
110 | 112 | return include_object |
111 | 113 |
|
112 | 114 |
|
| 115 | +def make_process_revision_directives( |
| 116 | + metadata: MetaData, |
| 117 | +) -> ProcessRevisionDirectivesFn: |
| 118 | + """Return an Alembic ``process_revision_directives`` hook that re-adds |
| 119 | + expression-based indexes silently dropped by autogenerate. |
| 120 | +
|
| 121 | + SQLAlchemy 2.0 can't reflect expression-based indexes (functional indexes |
| 122 | + like ``CREATE INDEX ... ON t (lower(email))``) under the SQLite dialect. |
| 123 | + Autogenerate guards against false-positive diffs there by *skipping* the |
| 124 | + index entirely — which is correct for an "existing table, can't tell if |
| 125 | + the index already exists" diff, but disastrous on initial CREATE TABLE: |
| 126 | + SQLite dev DBs end up without an index that production Postgres has. |
| 127 | +
|
| 128 | + This hook walks each generated ``MigrationScript`` and, for every |
| 129 | + ``CreateTableOp`` whose target table has expression-based indexes in the |
| 130 | + metadata, appends a matching ``CreateIndexOp``. The reverse ``DropIndexOp`` |
| 131 | + is inserted into ``downgrade_ops`` before the table drop for symmetry — |
| 132 | + not strictly required (dropping the table drops the index) but it keeps |
| 133 | + autogen output readable. |
| 134 | +
|
| 135 | + Call as:: |
| 136 | +
|
| 137 | + context.configure( |
| 138 | + ..., |
| 139 | + process_revision_directives=make_process_revision_directives(target_metadata), |
| 140 | + ) |
| 141 | + """ |
| 142 | + expression_indexes: dict[str, list[Index]] = {} |
| 143 | + for table in metadata.tables.values(): |
| 144 | + for index in table.indexes: |
| 145 | + if _index_is_expression_based(index): |
| 146 | + expression_indexes.setdefault(table.name, []).append(index) |
| 147 | + |
| 148 | + def process_revision_directives(context, revision, directives): |
| 149 | + if not expression_indexes: |
| 150 | + return |
| 151 | + for script in directives: |
| 152 | + upgrade_ops = getattr(script, "upgrade_ops", None) |
| 153 | + if upgrade_ops is not None: |
| 154 | + _inject_create_index_after_create_table(upgrade_ops, expression_indexes) |
| 155 | + downgrade_ops = getattr(script, "downgrade_ops", None) |
| 156 | + if downgrade_ops is not None: |
| 157 | + _inject_drop_index_before_drop_table(downgrade_ops, expression_indexes) |
| 158 | + |
| 159 | + return process_revision_directives |
| 160 | + |
| 161 | + |
| 162 | +def _index_is_expression_based(index: Index) -> bool: |
| 163 | + """An index is expression-based when any of its expressions is not a plain ``Column``.""" |
| 164 | + return any(not isinstance(expr, Column) for expr in index.expressions) |
| 165 | + |
| 166 | + |
| 167 | +def _inject_create_index_after_create_table(upgrade_ops, expression_indexes) -> None: |
| 168 | + existing_index_names = { |
| 169 | + getattr(op, "index_name", None) for op in upgrade_ops.ops if isinstance(op, CreateIndexOp) |
| 170 | + } |
| 171 | + new_ops: list = [] |
| 172 | + for op in upgrade_ops.ops: |
| 173 | + new_ops.append(op) |
| 174 | + if not isinstance(op, CreateTableOp): |
| 175 | + continue |
| 176 | + for index in expression_indexes.get(op.table_name, []): |
| 177 | + if index.name in existing_index_names: |
| 178 | + continue |
| 179 | + new_ops.append(CreateIndexOp.from_index(index)) |
| 180 | + existing_index_names.add(index.name) |
| 181 | + logger.info( |
| 182 | + "Re-emitting expression-based index %r on %r — autogenerate " |
| 183 | + "skipped it (dialect can't reflect functional indexes).", |
| 184 | + index.name, |
| 185 | + op.table_name, |
| 186 | + ) |
| 187 | + upgrade_ops.ops = new_ops |
| 188 | + |
| 189 | + |
| 190 | +def _inject_drop_index_before_drop_table(downgrade_ops, expression_indexes) -> None: |
| 191 | + existing_drop_names = { |
| 192 | + getattr(op, "index_name", None) for op in downgrade_ops.ops if isinstance(op, DropIndexOp) |
| 193 | + } |
| 194 | + new_ops: list = [] |
| 195 | + for op in downgrade_ops.ops: |
| 196 | + if isinstance(op, DropTableOp): |
| 197 | + for index in expression_indexes.get(op.table_name, []): |
| 198 | + if index.name in existing_drop_names: |
| 199 | + continue |
| 200 | + new_ops.append(DropIndexOp(index.name, table_name=op.table_name)) |
| 201 | + existing_drop_names.add(index.name) |
| 202 | + new_ops.append(op) |
| 203 | + downgrade_ops.ops = new_ops |
| 204 | + |
| 205 | + |
113 | 206 | def render_item(type_, obj, autogen_context): |
114 | 207 | """Alembic ``render_item`` callback for SQLModel + extension types. |
115 | 208 |
|
|
0 commit comments