"""Search vector plugin for codegen_database dimensions.
:class:`SearchVectorPlugin` maintains a ``tsvector`` column on the
backing table, keeping it in sync with one or more source text columns
via a ``BEFORE INSERT OR UPDATE`` trigger.
PostgreSQL's full-text search infrastructure (``tsvector``,
``to_tsvector``, ``@@``, GIN indexes) is built into the database
engine — no extension is required for basic usage.
Usage::
products = CodegenDatabaseSimple(
"products",
"app",
metadata,
schema_items=[
Column("name", String, nullable=False),
Column("description", String),
Column("search_vector", TSVECTOR, nullable=True),
],
extra_plugins=[
SearchVectorPlugin(
source_columns=["name", "description"],
vector_column="search_vector",
),
],
)
The trigger function is registered under
``{schema}.{schema}_{tablename}_search_vector_update()`` and fires
``BEFORE INSERT OR UPDATE`` on the raw backing table.
You may optionally pass a GIN index via
:class:`~codegen_database.index.CodegenDatabaseIndex` on the vector column
(e.g. ``CodegenDatabaseIndex(columns=["search_vector"], using="gin")``),
which is the standard way to make full-text searches fast.
If the text search configuration (``ts_config``) references a
dictionary installed via a PostgreSQL extension (e.g.
``unaccent``, ``pg_trgm``), register that extension in metadata
via :func:`~codegen_database.pg_extension.register_pg_extension`.
"""
from __future__ import annotations
from pathlib import Path
from typing import TYPE_CHECKING
from sqlalchemy_declarative_extensions import (
register_function,
register_trigger,
)
from sqlalchemy_declarative_extensions.dialects.postgresql import (
Function,
FunctionSecurity,
Trigger,
)
from codegen_database.errors import CodegenDatabaseValidationError
from codegen_database.plugin import Dynamic, Plugin, produces, requires
from codegen_database.utils.naming import resolve_name
from codegen_database.utils.template import load_template
if TYPE_CHECKING:
from codegen_database.factory.context import FactoryContext
_TEMPLATES = Path(__file__).resolve().parent / "templates" / "search_vector"
_NAMING_DEFAULTS = {
"search_vector_function": (
"%(schema)s_%(table_name)s_search_vector_update"
),
"search_vector_trigger": ("%(table_name)s_search_vector_update"),
}
[docs]
@produces("search_vector_column")
@requires(Dynamic("table_key"))
class SearchVectorPlugin(Plugin):
r"""Maintain a ``tsvector`` column via a BEFORE trigger.
Generates and registers a ``BEFORE INSERT OR UPDATE`` trigger
function that recomputes::
NEW.vector_column := to_tsvector(
ts_config,
COALESCE(col1, '') || ' ' || COALESCE(col2, '') || ...
)
Source columns are concatenated with a space separator; ``NULL``
values are coalesced to an empty string so they do not suppress
the entire vector.
Args:
source_columns: One or more column names whose text content
feeds the search vector. At least one required.
vector_column: Name of the ``tsvector`` column to maintain
(default ``"search_vector"``). Must exist in
``schema_items`` with a compatible type.
ts_config: PostgreSQL text-search configuration name
(default ``"english"``). Use any configuration
installed in the database. If the configuration
depends on a PostgreSQL extension (e.g. ``"unaccent"``),
register that extension with
:func:`~codegen_database.pg_extension.register_pg_extension`.
table_key: Key in ``ctx`` for the backing table to attach
the trigger to (default ``"raw_table"``).
Raises:
CodegenDatabaseValidationError: If *source_columns* is empty, any
source column is absent from the table, or *vector_column*
is absent from the table.
"""
def __init__(
self,
source_columns: list[str],
vector_column: str = "search_vector",
ts_config: str = "english",
table_key: str = "raw_table",
) -> None:
"""Store configuration."""
if not source_columns:
msg = "source_columns must be a non-empty list"
raise CodegenDatabaseValidationError(msg)
self.source_columns = list(source_columns)
self.vector_column = vector_column
self.ts_config = ts_config
self.table_key = table_key
[docs]
def run(self, ctx: FactoryContext) -> None:
"""Validate columns, then register the trigger function."""
table = ctx[self.table_key]
col_names = {c.name for c in table.columns}
missing = [c for c in self.source_columns if c not in col_names]
if missing:
msg = (
f"SearchVectorPlugin: source column(s) {missing!r} not "
f"found in table {table.name!r}. Add them to schema_items."
)
raise CodegenDatabaseValidationError(msg)
if self.vector_column not in col_names:
msg = (
f"SearchVectorPlugin: vector column {self.vector_column!r} "
f"not found in table {table.name!r}. Add "
f"Column({self.vector_column!r}, TSVECTOR, nullable=True) "
f"to schema_items."
)
raise CodegenDatabaseValidationError(msg)
schema = ctx.schemaname
table_name = ctx.tablename
full_table = f"{schema}.{table.name}"
source_expr = " || ' ' || ".join(
f"COALESCE(NEW.{col}::text, '')" for col in self.source_columns
)
template = load_template(_TEMPLATES / "update.plpgsql.mako")
body = template.render(
vector_column=self.vector_column,
ts_config=self.ts_config,
source_expr=source_expr,
)
subs = {"table_name": table_name, "schema": schema}
fn_name = resolve_name(
ctx.metadata,
"search_vector_function",
subs,
_NAMING_DEFAULTS,
)
trigger_name = resolve_name(
ctx.metadata,
"search_vector_trigger",
subs,
_NAMING_DEFAULTS,
)
register_function(
ctx.metadata,
Function(
fn_name,
body,
returns="trigger",
language="plpgsql",
schema=schema,
security=FunctionSecurity.definer,
),
)
register_trigger(
ctx.metadata,
Trigger.before(
"insert",
"update",
on=full_table,
execute=f"{schema}.{fn_name}",
name=trigger_name,
).for_each_row(),
)
ctx["search_vector_column"] = self.vector_column