sqlspec 0.47.0__cp314-cp314-win_amd64.whl

sqlspec/builder/_temporal.py

@@ -0,0 +1,167 @@
+"""Temporal query expressions.
+
+Overrides sqlglot's built-in version_sql generators to produce dialect-specific
+time-travel query syntax. Uses sqlglot's exp.Version and Table.version slot.
+
+Supported syntax by dialect:
+- BigQuery: table FOR SYSTEM_TIME AS OF timestamp (built-in)
+- Oracle: table AS OF TIMESTAMP timestamp / AS OF SCN scn
+- Snowflake: table AT (TIMESTAMP => timestamp) / BEFORE (...)
+- DuckDB: table AT (TIMESTAMP => timestamp)
+- CockroachDB (Postgres): table AS OF SYSTEM TIME timestamp
+"""
+
+from sqlglot import exp
+from sqlglot.dialects.bigquery import BigQuery
+from sqlglot.dialects.duckdb import DuckDB
+from sqlglot.dialects.oracle import Oracle
+from sqlglot.dialects.postgres import Postgres
+from sqlglot.dialects.snowflake import Snowflake
+from sqlglot.generator import (
+    _DISPATCH_CACHE,  # pyright: ignore[reportPrivateUsage]
+    Generator,
+)
+from sqlglot.generators.bigquery import BigQueryGenerator
+from sqlglot.generators.duckdb import DuckDBGenerator
+from sqlglot.generators.oracle import OracleGenerator
+from sqlglot.generators.postgres import PostgresGenerator
+from sqlglot.generators.snowflake import SnowflakeGenerator
+
+__all__ = ("create_temporal_table", "register_version_generators")
+
+
+def _oracle_version_sql(self: OracleGenerator, expression: exp.Version) -> str:
+    """Oracle: AS OF TIMESTAMP timestamp or AS OF SCN scn."""
+    expr = self.sql(expression, "expression")
+    this = expression.name or "TIMESTAMP"
+    return f"AS OF {this} {expr}"
+
+
+def _bigquery_version_sql(self: BigQueryGenerator, expression: exp.Version) -> str:
+    """BigQuery: FOR SYSTEM_TIME AS OF timestamp."""
+    expr = self.sql(expression, "expression")
+    return f"FOR SYSTEM_TIME AS OF {expr}"
+
+
+def _snowflake_version_sql(self: SnowflakeGenerator, expression: exp.Version) -> str:
+    """Snowflake: AT (TIMESTAMP => timestamp) or BEFORE (TIMESTAMP => ...).
+
+    AS OF is mapped to AT, and BEFORE is supported for point-before queries.
+    """
+    kind = expression.text("kind")
+    expr = self.sql(expression, "expression")
+    this = expression.name or "TIMESTAMP"
+    if kind and "BEFORE" in kind.upper():
+        return f"BEFORE ({this} => {expr})"
+    return f"AT ({this} => {expr})"
+
+
+def _duckdb_version_sql(self: DuckDBGenerator, expression: exp.Version) -> str:
+    """DuckDB: AT (TIMESTAMP => timestamp)."""
+    expr = self.sql(expression, "expression")
+    return f"AT (TIMESTAMP => {expr})"
+
+
+def _cockroachdb_version_sql(self: PostgresGenerator, expression: exp.Version) -> str:
+    """CockroachDB (via Postgres dialect): AS OF SYSTEM TIME timestamp."""
+    expr = self.sql(expression, "expression")
+    return f"AS OF SYSTEM TIME {expr}"
+
+
+def _default_version_sql(self: Generator, expression: exp.Version) -> str:
+    """Default: AS OF SYSTEM TIME timestamp (CockroachDB style).
+
+    When no dialect is specified, we default to CockroachDB/Postgres style
+    which is commonly expected for time-travel queries.
+    """
+    expr = self.sql(expression, "expression")
+    return f"AS OF SYSTEM TIME {expr}"
+
+
+def create_temporal_table(
+    table: "str | exp.Table | exp.Expr", as_of: "exp.Expr | str", kind: "str | None" = None
+) -> exp.Table:
+    """Create a table expression with temporal (time-travel) version clause.
+
+    Args:
+        table: Table name or expression.
+        as_of: Timestamp or SCN expression for the point-in-time.
+        kind: Optional version kind (e.g., "TIMESTAMP", "SCN", "SYSTEM TIME").
+              Defaults to dialect-appropriate value.
+
+    Returns:
+        Table expression with version clause that generates dialect-specific SQL.
+
+    Notes:
+        Inputs are normalized before building the ``exp.Version`` clause so both string table names and literal timestamps
+        work consistently.
+
+    Example:
+        >>> from sqlspec.builder import create_temporal_table
+        >>> from sqlglot import exp
+        >>> t = create_temporal_table(
+        ...     "orders", exp.Literal.string("2024-01-01")
+        ... )
+        >>> t.sql(dialect="oracle")
+        "orders AS OF TIMESTAMP '2024-01-01'"
+        >>> t.sql(dialect="bigquery")
+        "orders FOR SYSTEM_TIME AS OF '2024-01-01'"
+    """
+    if isinstance(table, str):
+        table_expr = exp.to_table(table)
+    elif isinstance(table, exp.Table):
+        table_expr = table.copy()
+    else:
+        table_expr = exp.to_table(str(table))
+
+    as_of_expr = exp.Literal.string(as_of) if isinstance(as_of, str) else as_of
+
+    version = exp.Version(this=kind or "TIMESTAMP", kind="AS OF", expression=as_of_expr)
+
+    table_expr.set("version", version)
+    return table_expr
+
+
+_VERSION_GENERATORS_REGISTERED = False
+
+
+def register_version_generators() -> None:
+    """Register dialect-specific version_sql generators.
+
+    This function is idempotent - calling it multiple times has no effect
+    after the first call. It is called automatically when the builder
+    module is imported.
+
+    Registers custom SQL generators for temporal (time-travel) queries:
+    - Default (no dialect): AS OF SYSTEM TIME (CockroachDB style)
+    - BigQuery: FOR SYSTEM_TIME AS OF
+    - Oracle: AS OF TIMESTAMP / AS OF SCN
+    - Snowflake: AT (TIMESTAMP => ...) / BEFORE (...)
+    - DuckDB: AT (TIMESTAMP => ...)
+    - Postgres/CockroachDB: AS OF SYSTEM TIME
+    """
+    global _VERSION_GENERATORS_REGISTERED
+    if _VERSION_GENERATORS_REGISTERED:
+        return
+
+    Generator.TRANSFORMS[exp.Version] = _default_version_sql
+
+    BigQuery.Generator.TRANSFORMS[exp.Version] = _bigquery_version_sql
+    Oracle.Generator.TRANSFORMS[exp.Version] = _oracle_version_sql
+    Snowflake.Generator.TRANSFORMS[exp.Version] = _snowflake_version_sql
+    DuckDB.Generator.TRANSFORMS[exp.Version] = _duckdb_version_sql
+    Postgres.Generator.TRANSFORMS[exp.Version] = _cockroachdb_version_sql
+
+    # Invalidate sqlglot's per-class dispatch cache so new TRANSFORMS entries
+    # are picked up by the next Generator instantiation.
+    for gen_cls in (
+        Generator,
+        BigQuery.Generator,
+        Oracle.Generator,
+        Snowflake.Generator,
+        DuckDB.Generator,
+        Postgres.Generator,
+    ):
+        _DISPATCH_CACHE.pop(gen_cls, None)
+
+    _VERSION_GENERATORS_REGISTERED = True

sqlspec/builder/_update.py

@@ -0,0 +1,173 @@
+"""UPDATE statement builder.
+
+Provides a fluent interface for building SQL UPDATE queries with
+parameter binding and validation.
+"""
+
+from typing import TYPE_CHECKING, Any, cast
+
+from sqlglot import exp
+from typing_extensions import Self
+
+from sqlspec.builder._base import BuiltQuery, QueryBuilder
+from sqlspec.builder._dml import UpdateFromClauseMixin, UpdateSetClauseMixin, UpdateTableClauseMixin
+from sqlspec.builder._explain import ExplainMixin
+from sqlspec.builder._join import build_join_clause
+from sqlspec.builder._select import ReturningClauseMixin, WhereClauseMixin
+from sqlspec.core import SQLResult
+from sqlspec.exceptions import SQLBuilderError
+
+if TYPE_CHECKING:
+    from sqlglot.dialects.dialect import DialectType
+
+    from sqlspec.builder._select import Select
+    from sqlspec.protocols import SQLBuilderProtocol
+
+__all__ = ("Update",)
+
+
+class Update(
+    QueryBuilder,
+    WhereClauseMixin,
+    ReturningClauseMixin,
+    UpdateSetClauseMixin,
+    UpdateFromClauseMixin,
+    UpdateTableClauseMixin,
+    ExplainMixin,
+):
+    """Builder for UPDATE statements.
+
+    Constructs SQL UPDATE statements with parameter binding and validation.
+
+    Example:
+        ```python
+        update_query = (
+            Update()
+            .table("users")
+            .set_(name="John Doe")
+            .set_(email="john@example.com")
+            .where("id = 1")
+        )
+
+        update_query = (
+            Update("users").set_(name="John Doe").where("id = 1")
+        )
+
+        update_query = (
+            Update()
+            .table("users")
+            .set_(status="active")
+            .where_eq("id", 123)
+        )
+
+        update_query = (
+            Update()
+            .table("users", "u")
+            .set_(name="Updated Name")
+            .from_("profiles", "p")
+            .where("u.id = p.user_id AND p.is_verified = true")
+        )
+        ```
+    """
+
+    __slots__ = ()
+    _expression: exp.Expr | None
+
+    def __init__(self, table: str | None = None, **kwargs: Any) -> None:
+        """Initialize UPDATE with optional table.
+
+        Args:
+            table: Target table name
+            **kwargs: Additional QueryBuilder arguments
+        """
+        (dialect, schema, enable_optimization, optimize_joins, optimize_predicates, simplify_expressions) = (
+            self._parse_query_builder_kwargs(kwargs)
+        )
+        super().__init__(
+            dialect=dialect,
+            schema=schema,
+            enable_optimization=enable_optimization,
+            optimize_joins=optimize_joins,
+            optimize_predicates=optimize_predicates,
+            simplify_expressions=simplify_expressions,
+        )
+        self._initialize_expression()
+
+        if table:
+            self.table(table)
+
+    @property
+    def _expected_result_type(self) -> "type[SQLResult]":
+        """Return the expected result type for this builder."""
+        return SQLResult
+
+    def _create_base_expression(self) -> exp.Update:
+        """Create a base UPDATE expression.
+
+        Returns:
+            A new sqlglot Update expression with empty clauses.
+        """
+        return exp.Update(this=None, expressions=[], joins=[])
+
+    def join(
+        self,
+        table: "str | exp.Expr | Select",
+        on: "str | exp.Expr",
+        alias: "str | None" = None,
+        join_type: str = "INNER",
+    ) -> "Self":
+        """Add JOIN clause to the UPDATE statement.
+
+        Args:
+            table: The table name, expression, or subquery to join.
+            on: The JOIN condition.
+            alias: Optional alias for the joined table.
+            join_type: Type of join (INNER, LEFT, RIGHT, FULL).
+
+        Returns:
+            The current builder instance for method chaining.
+
+        Raises:
+            SQLBuilderError: If the current expression is not an UPDATE statement.
+        """
+        if self._expression is None or not isinstance(self._expression, exp.Update):
+            msg = "Cannot add JOIN clause to non-UPDATE expression."
+            raise SQLBuilderError(msg)
+
+        join_expr = build_join_clause(cast("SQLBuilderProtocol", self), table, on, alias, join_type)
+
+        if not self._expression.args.get("joins"):
+            self._expression.set("joins", [])
+        self._expression.args["joins"].append(join_expr)
+
+        return self
+
+    def build(self, dialect: "DialectType" = None) -> "BuiltQuery":
+        """Build the UPDATE query with validation.
+
+        Args:
+            dialect: Optional dialect override for SQL generation.
+
+        Returns:
+            BuiltQuery: The built query with SQL and parameters.
+
+        Raises:
+            SQLBuilderError: If no table is set or expression is not an UPDATE.
+        """
+        if self._expression is None:
+            msg = "UPDATE expression not initialized."
+            raise SQLBuilderError(msg)
+
+        if not isinstance(self._expression, exp.Update):
+            msg = "No UPDATE expression to build or expression is of the wrong type."
+            raise SQLBuilderError(msg)
+
+        if self._expression.this is None:
+            msg = "No table specified for UPDATE statement."
+            raise SQLBuilderError(msg)
+
+        if not self._expression.args.get("expressions"):
+            msg = "At least one SET clause must be specified for UPDATE statement."
+            raise SQLBuilderError(msg)
+
+        return super().build(dialect=dialect)

sqlspec/builder/_vector_distance.py

@@ -0,0 +1,330 @@
+"""Vector distance helpers and SQL generator registration."""
+
+# ruff: noqa: N802
+# pyright: ignore[reportConstantRedefinition]
+from collections.abc import Callable, MutableMapping
+from typing import TYPE_CHECKING, Any, Final, TypeAlias, cast
+
+from sqlglot import exp
+
+if TYPE_CHECKING:
+    from sqlglot.dialects.bigquery import BigQuery
+    from sqlglot.dialects.duckdb import DuckDB
+    from sqlglot.dialects.mysql import MySQL
+    from sqlglot.dialects.oracle import Oracle
+    from sqlglot.dialects.postgres import Postgres
+    from sqlglot.generator import Generator
+
+    from sqlspec.dialects.spanner import Spangres, Spanner
+
+SupportedVectorDistanceDialect: TypeAlias = "BigQuery | DuckDB | MySQL | Oracle | Postgres | Spangres | Spanner"
+
+__all__ = (
+    "VectorDistance",
+    "has_vector_distance_ancestor",
+    "is_vector_distance_expression",
+    "render_vector_distance_bigquery",
+    "render_vector_distance_duckdb",
+    "render_vector_distance_generic",
+    "render_vector_distance_mysql",
+    "render_vector_distance_oracle",
+    "render_vector_distance_postgres",
+    "vector_distance_metric",
+)
+
+_VECTOR_DISTANCE_META_KEY: Final[str] = "sqlspec_vector_distance_metric"
+_OperatorTransform = Callable[[Any, exp.Operator], str]
+_SQLGLOT_VECTOR_DISTANCE_REGISTERED = False
+_BASE_OPERATOR_TRANSFORM: _OperatorTransform | None = None
+_POSTGRES_OPERATOR_TRANSFORM: _OperatorTransform | None = None
+_MYSQL_OPERATOR_TRANSFORM: _OperatorTransform | None = None
+_ORACLE_OPERATOR_TRANSFORM: _OperatorTransform | None = None
+_BIGQUERY_OPERATOR_TRANSFORM: _OperatorTransform | None = None
+_DUCKDB_OPERATOR_TRANSFORM: _OperatorTransform | None = None
+_SPANNER_OPERATOR_TRANSFORM: _OperatorTransform | None = None
+_SPANGRES_OPERATOR_TRANSFORM: _OperatorTransform | None = None
+
+
+def _normalize_metric(metric: Any) -> str:
+    """Normalize vector metrics to a lowercase string."""
+    if isinstance(metric, exp.Literal):
+        return str(metric.this).lower()
+    if isinstance(metric, exp.Identifier):
+        identifier = metric.this
+        return identifier.lower() if isinstance(identifier, str) else "euclidean"
+    if isinstance(metric, str):
+        return metric.lower()
+    return "euclidean"
+
+
+def is_vector_distance_expression(expression: object) -> bool:
+    """Return True when an Operator node is a SQLSpec vector-distance expression."""
+    return isinstance(expression, exp.Operator) and _VECTOR_DISTANCE_META_KEY in expression.meta
+
+
+def has_vector_distance_ancestor(expression: exp.Expr) -> bool:
+    """Return True when any ancestor is a SQLSpec vector-distance expression."""
+    parent = expression.parent
+    while parent is not None:
+        if is_vector_distance_expression(parent):
+            return True
+        parent = parent.parent
+    return False
+
+
+def vector_distance_metric(expression: object) -> str:
+    """Get the normalized vector-distance metric from an Operator node."""
+    if not isinstance(expression, exp.Operator):
+        msg = f"Expected sqlglot Operator, got {type(expression)}"
+        raise TypeError(msg)
+    metric = expression.meta.get(_VECTOR_DISTANCE_META_KEY)
+    if isinstance(metric, str):
+        return metric
+    operator = expression.args.get("operator")
+    return str(operator).lower() if operator is not None else "euclidean"
+
+
+def _build_vector_distance(this: exp.Expr, expression: exp.Expr, metric: Any = "euclidean") -> exp.Operator:
+    normalized_metric = _normalize_metric(metric)
+    node = exp.Operator(this=this, expression=expression, operator=normalized_metric)
+    node.meta[_VECTOR_DISTANCE_META_KEY] = normalized_metric
+    return node
+
+
+def VectorDistance(*, this: exp.Expr, expression: exp.Expr, metric: Any = "euclidean") -> exp.Operator:
+    """Build a SQLSpec vector-distance expression."""
+    _register_with_sqlglot()
+    return _build_vector_distance(this=this, expression=expression, metric=metric)
+
+
+def render_vector_distance_postgres(left: str, right: str, metric: str) -> str:
+    """Render PostgreSQL pgvector operator syntax."""
+    operator_map = {"euclidean": "<->", "cosine": "<=>", "inner_product": "<#>"}
+
+    operator = operator_map.get(metric)
+    if operator:
+        return f"{left} {operator} {right}"
+
+    return render_vector_distance_generic(left, right, metric)
+
+
+def render_vector_distance_mysql(left: str, right: str, metric: str) -> str:
+    """Render MySQL DISTANCE function syntax."""
+    metric_map = {"euclidean": "EUCLIDEAN", "cosine": "COSINE", "inner_product": "DOT"}
+
+    mysql_metric = metric_map.get(metric, "EUCLIDEAN")
+
+    if ("ARRAY" in right or "[" in right) and "STRING_TO_VECTOR" not in right:
+        right = f"STRING_TO_VECTOR({right})"
+
+    return f"DISTANCE({left}, {right}, '{mysql_metric}')"
+
+
+def render_vector_distance_oracle(left: str, right: str, metric: str) -> str:
+    """Render Oracle VECTOR_DISTANCE function syntax."""
+    metric_map = {
+        "euclidean": "EUCLIDEAN",
+        "cosine": "COSINE",
+        "inner_product": "DOT",
+        "euclidean_squared": "EUCLIDEAN_SQUARED",
+    }
+
+    oracle_metric = metric_map.get(metric, "EUCLIDEAN")
+
+    if ("[" in right or "ARRAY" in right) and "TO_VECTOR" not in right:
+        right = f"TO_VECTOR({right})"
+
+    return f"VECTOR_DISTANCE({left}, {right}, {oracle_metric})"
+
+
+def render_vector_distance_bigquery(left: str, right: str, metric: str) -> str:
+    """Render BigQuery vector distance function syntax."""
+    function_map = {"euclidean": "EUCLIDEAN_DISTANCE", "cosine": "COSINE_DISTANCE", "inner_product": "DOT_PRODUCT"}
+
+    function_name = function_map.get(metric)
+    if function_name:
+        return f"{function_name}({left}, {right})"
+
+    return render_vector_distance_generic(left, right, metric)
+
+
+def render_vector_distance_duckdb(left: str, right: str, metric: str) -> str:
+    """Render DuckDB VSS extension function syntax."""
+    function_map = {
+        "euclidean": "array_distance",
+        "cosine": "array_cosine_distance",
+        "inner_product": "array_negative_inner_product",
+    }
+    function_name = function_map.get(metric)
+    if function_name:
+        return f"{function_name}({left}, CAST({right} AS DOUBLE[]))"
+
+    return render_vector_distance_generic(left, right, metric)
+
+
+def render_vector_distance_generic(left: str, right: str, metric: str) -> str:
+    """Render generic VECTOR_DISTANCE function syntax."""
+    return f"VECTOR_DISTANCE({left}, {right}, '{metric.upper()}')"
+
+
+def _render_with_metric(generator: "Generator", expression: exp.Operator, dialect: str) -> str:
+    left_sql = generator.sql(expression, "this")
+    right_sql = generator.sql(expression, "expression")
+    metric = vector_distance_metric(expression)
+
+    if dialect == "postgres":
+        return render_vector_distance_postgres(left_sql, right_sql, metric)
+    if dialect == "mysql":
+        return render_vector_distance_mysql(left_sql, right_sql, metric)
+    if dialect == "oracle":
+        if isinstance(expression.expression, exp.Array):
+            values = [
+                str(item.this) if isinstance(item, exp.Literal) else generator.sql(item)
+                for item in expression.expression.expressions
+            ]
+            right_sql = f"TO_VECTOR('[{', '.join(values)}]')"
+        return render_vector_distance_oracle(left_sql, right_sql, metric)
+    if dialect == "bigquery":
+        return render_vector_distance_bigquery(left_sql, right_sql, metric)
+    if dialect == "duckdb":
+        if isinstance(expression.expression, exp.Array) and expression.expression.expressions:
+            target_type = f"DOUBLE[{len(expression.expression.expressions)}]"
+        else:
+            target_type = "DOUBLE[]"
+        function_map = {
+            "euclidean": "array_distance",
+            "cosine": "array_cosine_distance",
+            "inner_product": "array_negative_inner_product",
+        }
+        function_name = function_map.get(metric)
+        if function_name:
+            return f"{function_name}({left_sql}, CAST({right_sql} AS {target_type}))"
+        return render_vector_distance_generic(left_sql, right_sql, metric)
+
+    return render_vector_distance_generic(left_sql, right_sql, metric)
+
+
+def _operator_sql_base(generator: "Generator", expression: exp.Operator) -> str:
+    if is_vector_distance_expression(expression):
+        return _render_with_metric(generator, expression, "generic")
+    return _require_operator_transform(_BASE_OPERATOR_TRANSFORM)(generator, expression)
+
+
+def _operator_sql_postgres(generator: "Generator", expression: exp.Operator) -> str:
+    if is_vector_distance_expression(expression):
+        return _render_with_metric(generator, expression, "postgres")
+    return _require_operator_transform(_POSTGRES_OPERATOR_TRANSFORM)(generator, expression)
+
+
+def _operator_sql_mysql(generator: "Generator", expression: exp.Operator) -> str:
+    if is_vector_distance_expression(expression):
+        return _render_with_metric(generator, expression, "mysql")
+    return _require_operator_transform(_MYSQL_OPERATOR_TRANSFORM)(generator, expression)
+
+
+def _operator_sql_oracle(generator: "Generator", expression: exp.Operator) -> str:
+    if is_vector_distance_expression(expression):
+        return _render_with_metric(generator, expression, "oracle")
+    return _require_operator_transform(_ORACLE_OPERATOR_TRANSFORM)(generator, expression)
+
+
+def _operator_sql_bigquery(generator: "Generator", expression: exp.Operator) -> str:
+    if is_vector_distance_expression(expression):
+        return _render_with_metric(generator, expression, "bigquery")
+    return _require_operator_transform(_BIGQUERY_OPERATOR_TRANSFORM)(generator, expression)
+
+
+def _operator_sql_duckdb(generator: "Generator", expression: exp.Operator) -> str:
+    if is_vector_distance_expression(expression):
+        return _render_with_metric(generator, expression, "duckdb")
+    return _require_operator_transform(_DUCKDB_OPERATOR_TRANSFORM)(generator, expression)
+
+
+def _operator_sql_spanner(generator: "Generator", expression: exp.Operator) -> str:
+    if is_vector_distance_expression(expression):
+        return _render_with_metric(generator, expression, "bigquery")
+    return _require_operator_transform(_SPANNER_OPERATOR_TRANSFORM)(generator, expression)
+
+
+def _operator_sql_spangres(generator: "Generator", expression: exp.Operator) -> str:
+    if is_vector_distance_expression(expression):
+        return _render_with_metric(generator, expression, "postgres")
+    return _require_operator_transform(_SPANGRES_OPERATOR_TRANSFORM)(generator, expression)
+
+
+def _require_operator_transform(transform: _OperatorTransform | None) -> _OperatorTransform:
+    """Return a registered fallback transform or fail loudly."""
+    if transform is None:
+        msg = "Vector-distance SQLGlot transforms have not been registered"
+        raise RuntimeError(msg)
+    return transform
+
+
+def _register_operator_transform(
+    transforms: MutableMapping[type[exp.Expr], "_OperatorTransform"], wrapper: "_OperatorTransform"
+) -> None:
+    """Install an operator transform wrapper."""
+    transforms[exp.Operator] = wrapper
+
+
+def _register_with_sqlglot() -> None:
+    """Register vector-distance Operator rendering with SQLGlot generators."""
+    global \
+        _SQLGLOT_VECTOR_DISTANCE_REGISTERED, \
+        _BASE_OPERATOR_TRANSFORM, \
+        _POSTGRES_OPERATOR_TRANSFORM, \
+        _MYSQL_OPERATOR_TRANSFORM, \
+        _ORACLE_OPERATOR_TRANSFORM, \
+        _BIGQUERY_OPERATOR_TRANSFORM, \
+        _DUCKDB_OPERATOR_TRANSFORM, \
+        _SPANNER_OPERATOR_TRANSFORM, \
+        _SPANGRES_OPERATOR_TRANSFORM
+
+    if _SQLGLOT_VECTOR_DISTANCE_REGISTERED:
+        return
+
+    from sqlglot.dialects.bigquery import BigQuery
+    from sqlglot.dialects.duckdb import DuckDB
+    from sqlglot.dialects.mysql import MySQL
+    from sqlglot.dialects.oracle import Oracle
+    from sqlglot.dialects.postgres import Postgres
+    from sqlglot.generator import Generator
+
+    from sqlspec.dialects.spanner import Spangres, Spanner
+
+    _BASE_OPERATOR_TRANSFORM = cast("_OperatorTransform", Generator.TRANSFORMS[exp.Operator])
+    _POSTGRES_OPERATOR_TRANSFORM = cast("_OperatorTransform", Postgres.Generator.TRANSFORMS[exp.Operator])
+    _MYSQL_OPERATOR_TRANSFORM = cast("_OperatorTransform", MySQL.Generator.TRANSFORMS[exp.Operator])
+    _ORACLE_OPERATOR_TRANSFORM = cast("_OperatorTransform", Oracle.Generator.TRANSFORMS[exp.Operator])
+    _BIGQUERY_OPERATOR_TRANSFORM = cast("_OperatorTransform", BigQuery.Generator.TRANSFORMS[exp.Operator])
+    _DUCKDB_OPERATOR_TRANSFORM = cast("_OperatorTransform", DuckDB.Generator.TRANSFORMS[exp.Operator])
+    _SPANNER_OPERATOR_TRANSFORM = cast("_OperatorTransform", Spanner.Generator.TRANSFORMS[exp.Operator])
+    _SPANGRES_OPERATOR_TRANSFORM = cast("_OperatorTransform", Spangres.Generator.TRANSFORMS[exp.Operator])
+
+    _register_operator_transform(Generator.TRANSFORMS, _operator_sql_base)
+    _register_operator_transform(Postgres.Generator.TRANSFORMS, _operator_sql_postgres)
+    _register_operator_transform(MySQL.Generator.TRANSFORMS, _operator_sql_mysql)
+    _register_operator_transform(Oracle.Generator.TRANSFORMS, _operator_sql_oracle)
+    _register_operator_transform(BigQuery.Generator.TRANSFORMS, _operator_sql_bigquery)
+    _register_operator_transform(DuckDB.Generator.TRANSFORMS, _operator_sql_duckdb)
+    _register_operator_transform(Spanner.Generator.TRANSFORMS, _operator_sql_spanner)
+    _register_operator_transform(Spangres.Generator.TRANSFORMS, _operator_sql_spangres)
+
+    # sqlglot caches the dispatch table (built from TRANSFORMS) per Generator class
+    # in _DISPATCH_CACHE. We must invalidate stale entries so the next instantiation
+    # picks up our new Operator transforms.
+    from sqlglot.generator import _DISPATCH_CACHE  # pyright: ignore[reportPrivateUsage]
+
+    for gen_cls in (
+        Generator,
+        Postgres.Generator,
+        MySQL.Generator,
+        Oracle.Generator,
+        BigQuery.Generator,
+        DuckDB.Generator,
+        Spanner.Generator,
+        Spangres.Generator,
+    ):
+        _DISPATCH_CACHE.pop(gen_cls, None)
+
+    _SQLGLOT_VECTOR_DISTANCE_REGISTERED = True