"""Index support for codegen_database dimensions.
Provides :class:`CodegenDatabaseIndex`, a declarative index definition
that mirrors ``sqlalchemy.Index`` and uses ``{column_name}``
markers for column references.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Any
from codegen_database.errors import CodegenDatabaseValidationError
from codegen_database.validation import extract_column_names, resolve_markers
if TYPE_CHECKING:
from collections.abc import Callable
[docs]
class CodegenDatabaseIndex:
"""A declarative index definition with ``{col}`` markers.
Mirrors the ``sqlalchemy.Index`` constructor signature::
CodegenDatabaseIndex("idx_name", "{col1}", "{col2}",
unique=True, postgresql_using="btree")
Simple column references (``"{name}"``) and functional
expressions (``"lower({name})"``) are both supported.
Extra keyword arguments are passed through to the
underlying ``sqlalchemy.Index``.
Args:
name: Required index name.
*expressions: Index expressions using ``{column_name}``
markers.
unique: Whether to create a unique index.
**kw: Passed through to ``sqlalchemy.Index``
(e.g. ``postgresql_using``, ``postgresql_where``).
"""
__slots__ = ("expressions", "kw", "name", "unique")
name: str
expressions: list[str]
unique: bool
kw: dict[str, Any]
def __init__(
self,
name: str,
*expressions: str,
unique: bool = False,
**kw: Any, # noqa: ANN401
) -> None:
"""Create a new index definition."""
object.__setattr__(self, "name", name)
object.__setattr__(self, "expressions", list(expressions))
object.__setattr__(self, "unique", unique)
object.__setattr__(self, "kw", dict(kw))
def __setattr__(
self,
key: str,
value: object,
) -> None:
"""Prevent mutation after construction."""
msg = "CodegenDatabaseIndex instances are immutable"
raise AttributeError(msg)
def __repr__(self) -> str:
"""Return a constructor-style repr."""
parts = [repr(self.name)]
parts.extend(repr(e) for e in self.expressions)
if self.unique:
parts.append("unique=True")
for k, v in self.kw.items():
parts.append(f"{k}={v!r}")
return f"CodegenDatabaseIndex({', '.join(parts)})"
def __eq__(self, other: object) -> bool:
"""Compare by name, expressions, unique, and kw."""
if not isinstance(other, CodegenDatabaseIndex):
return NotImplemented
return (
self.name == other.name
and self.expressions == other.expressions
and self.unique == other.unique
and self.kw == other.kw
)
def __hash__(self) -> int:
"""Hash by name, expressions, and unique flag."""
return hash(
(
self.name,
tuple(self.expressions),
self.unique,
)
)
[docs]
def column_names(self) -> list[str]:
"""Extract ``{name}`` markers from all expressions.
Returns:
Column names in order of first appearance,
deduplicated across all expressions.
"""
seen: set[str] = set()
result: list[str] = []
for expr in self.expressions:
for name in extract_column_names(expr):
if name not in seen:
seen.add(name)
result.append(name)
return result
[docs]
def resolve(self, mapping: Callable[[str], str]) -> list[str]:
"""Replace ``{col}`` markers in each expression.
Args:
mapping: A callable that maps column names to
their resolved form.
Returns:
List of resolved expression strings.
"""
return [resolve_markers(expr, mapping) for expr in self.expressions]
[docs]
def collect_indices(
schema_items: list,
) -> list[CodegenDatabaseIndex]:
"""Filter :class:`CodegenDatabaseIndex` instances from schema items.
Args:
schema_items: Mixed list of schema items.
Returns:
Only the ``CodegenDatabaseIndex`` items, in original order.
"""
return [i for i in schema_items if isinstance(i, CodegenDatabaseIndex)]
[docs]
def trigram_indexes(
*columns: str,
table: str | None = None,
method: str = "gin",
) -> list[CodegenDatabaseIndex]:
"""Build one ``pg_trgm`` index per column for fuzzy text search.
A ``pg_trgm`` operator-class index is what makes the ``%``
similarity operator and ``ILIKE`` substring matches fast --
without one, a trigram search falls back to a sequential
scan. Pass the text columns a resource searches and splat the
result into a dimension's ``schema_items``; the
``TableIndexPlugin`` materializes each into a real index::
schema_items=[
Column("name", String),
Column("sku", String),
*trigram_indexes("name", "sku", table="product"),
]
Every index is ``USING <method> (<col> <method>_trgm_ops)`` --
the operator class is what binds the index to the trigram
operators, so a plain ``USING gin (col)`` would not serve
``%`` / ``ILIKE`` lookups.
The ``pg_trgm`` extension is registered by default --
``configure_metadata`` adds it via
:func:`codegen_database.pg_extension.register_default_pg_extensions`,
so the next autogenerated migration emits the
``CREATE EXTENSION`` when the database lacks it.
Args:
*columns: Text column names to index for trigram search.
table: Optional table name woven into each index name so
it stays unique schema-wide -- index names are
database-global. Omitted yields ``ix__<col>__trgm``;
given, ``ix__<table>__<col>__trgm``.
method: Index access method -- ``"gin"`` (default, the
usual choice) or ``"gist"``. The matching
``<method>_trgm_ops`` operator class is applied.
Returns:
One :class:`CodegenDatabaseIndex` per column, in input order.
Raises:
CodegenDatabaseValidationError: If *columns* is empty or contains
a duplicate, or *method* is neither ``"gin"`` nor
``"gist"``.
"""
if not columns:
msg = "trigram_indexes: pass at least one column name"
raise CodegenDatabaseValidationError(msg)
if len(set(columns)) != len(columns):
msg = f"trigram_indexes: duplicate column name in {list(columns)!r}"
raise CodegenDatabaseValidationError(msg)
if method not in ("gin", "gist"):
msg = f"trigram_indexes: method must be 'gin' or 'gist', got {method!r}"
raise CodegenDatabaseValidationError(msg)
ops_class = f"{method}_trgm_ops"
return [
CodegenDatabaseIndex(
f"ix__{table}__{col}__trgm" if table else f"ix__{col}__trgm",
f"{{{col}}}",
postgresql_using=method,
postgresql_ops={col: ops_class},
)
for col in columns
]