sqlspec 0.36.0__cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl

sqlspec/builder/_temporal.py

@@ -0,0 +1,139 @@
+"""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.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 Generator
+
+__all__ = ("create_temporal_table", "register_version_generators")
+
+
+def _oracle_version_sql(self: "Oracle.Generator", 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 _snowflake_version_sql(self: "Snowflake.Generator", 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: "DuckDB.Generator", expression: exp.Version) -> str:
+    """DuckDB: AT (TIMESTAMP => timestamp)."""
+    expr = self.sql(expression, "expression")
+    return f"AT (TIMESTAMP => {expr})"
+
+
+def _cockroachdb_version_sql(self: "Postgres.Generator", 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.Expression", as_of: "exp.Expression | 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)
+    - Oracle: AS OF TIMESTAMP / AS OF SCN
+    - Snowflake: AT (TIMESTAMP => ...) / BEFORE (...)
+    - DuckDB: AT (TIMESTAMP => ...)
+    - Postgres/CockroachDB: AS OF SYSTEM TIME
+    - BigQuery: Uses built-in FOR SYSTEM_TIME AS OF
+    """
+    global _VERSION_GENERATORS_REGISTERED
+    if _VERSION_GENERATORS_REGISTERED:
+        return
+
+    Generator.TRANSFORMS[exp.Version] = _default_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
+
+    _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.Expression | 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.Expression | Select",
+        on: "str | exp.Expression",
+        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_expressions.py

@@ -0,0 +1,267 @@
+"""Custom SQLGlot expressions for vector distance operations.
+
+Provides dialect-specific SQL generation for vector similarity search
+across PostgreSQL (pgvector), MySQL 9+, Oracle 23ai+, BigQuery, and Spanner.
+"""
+
+from contextlib import suppress
+from typing import Any
+
+from sqlglot import exp
+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
+
+__all__ = ("VectorDistance",)
+
+
+class VectorDistance(exp.Expression):
+    """Vector distance expression with dialect-specific generation.
+
+    Generates database-specific SQL for vector distance calculations:
+    - PostgreSQL (pgvector): Operators <->, <=>, <#>
+    - MySQL 9+: DISTANCE(col, vec, 'METRIC') function
+    - Oracle 23ai+: VECTOR_DISTANCE(col, vec, METRIC) function
+    - Generic: VECTOR_DISTANCE(col, vec, 'METRIC') function
+
+    The metric is stored as a raw string attribute (not parametrized) and drives
+    dialect-specific generation at SQL build time.
+    """
+
+    arg_types = {"this": True, "expression": True, "metric": False}
+
+    def __init__(self, **args: Any) -> None:
+        """Initialize VectorDistance with metric stored in args."""
+        metric_value = args.get("metric", "euclidean")
+        if isinstance(metric_value, exp.Literal):
+            metric_value = str(metric_value.this).lower()
+        elif isinstance(metric_value, exp.Identifier):
+            metric_value = metric_value.this.lower()
+        elif isinstance(metric_value, str):
+            metric_value = metric_value.lower()
+        else:
+            metric_value = "euclidean"
+
+        args["metric"] = exp.Identifier(this=metric_value)
+        super().__init__(**args)
+
+    @property
+    def left(self) -> "exp.Expression":
+        """Get the left operand (column)."""
+        result: exp.Expression = self.this
+        return result
+
+    @property
+    def right(self) -> "exp.Expression":
+        """Get the right operand (vector value)."""
+        result: exp.Expression = self.expression
+        return result
+
+    @property
+    def metric(self) -> str:
+        """Get the distance metric as raw string (not parametrized)."""
+        metric_expr = self.args.get("metric")
+        if isinstance(metric_expr, exp.Identifier):
+            metric_name: str = metric_expr.this
+            return metric_name.lower()
+        return "euclidean"
+
+    def sql(self, dialect: "Any | None" = None, **opts: Any) -> str:
+        """Generate dialect-specific SQL.
+
+        This overrides the default sql() method to provide custom
+        dialect-specific generation for vector distance operations.
+
+        Args:
+            dialect: Target SQL dialect (postgres, mysql, oracle, bigquery, duckdb, etc.)
+            **opts: Additional SQL generation options
+
+        Returns:
+            Dialect-specific SQL string
+        """
+        dialect_name = str(dialect).lower() if dialect else "generic"
+
+        left_sql = self.left.sql(dialect=dialect, **opts)
+        right_sql = self.right.sql(dialect=dialect, **opts)
+        metric = self.metric
+
+        if dialect_name in {"postgres", "postgresql"}:
+            return self._sql_postgres(left_sql, right_sql, metric)
+
+        if dialect_name == "mysql":
+            return self._sql_mysql(left_sql, right_sql, metric)
+
+        if dialect_name == "oracle":
+            return self._sql_oracle(left_sql, right_sql, metric)
+
+        if dialect_name == "bigquery":
+            return self._sql_bigquery(left_sql, right_sql, metric)
+
+        if dialect_name == "duckdb":
+            return self._sql_duckdb(left_sql, right_sql, metric)
+
+        return self._sql_generic(left_sql, right_sql, metric)
+
+    def _sql_postgres(self, left: str, right: str, metric: str) -> str:
+        """Generate PostgreSQL pgvector operator syntax."""
+        operator_map = {"euclidean": "<->", "cosine": "<=>", "inner_product": "<#>"}
+
+        operator = operator_map.get(metric)
+        if operator:
+            return f"{left} {operator} {right}"
+
+        return self._sql_generic(left, right, metric)
+
+    def _sql_mysql(self, left: str, right: str, metric: str) -> str:
+        """Generate 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 _sql_oracle(self, left: str, right: str, metric: str) -> str:
+        """Generate 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 isinstance(self.expression, exp.Array):
+            values = []
+            for expr in self.expression.expressions:
+                if isinstance(expr, exp.Literal):
+                    values.append(str(expr.this))
+                else:  # pragma: no cover - defensive
+                    values.append(expr.sql(dialect="oracle"))
+            right = f"TO_VECTOR('[{', '.join(values)}]')"
+        elif ("ARRAY" in right or "[" in right) and "TO_VECTOR" not in right:
+            right = f"TO_VECTOR({right})"
+
+        return f"VECTOR_DISTANCE({left}, {right}, {oracle_metric})"
+
+    def _sql_bigquery(self, left: str, right: str, metric: str) -> str:
+        """Generate 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 self._sql_generic(left, right, metric)
+
+    def _sql_duckdb(self, left: str, right: str, metric: str) -> str:
+        """Generate DuckDB VSS extension function syntax.
+
+        DuckDB's VSS extension provides:
+        - array_distance(): L2 squared distance (euclidean)
+        - array_cosine_distance(): Cosine distance (1 - cosine_similarity)
+        - array_negative_inner_product(): Negative inner product
+
+        Note: Array literals must be cast to DOUBLE[] since DuckDB infers
+        decimal literals as DECIMAL type, but VSS functions require DOUBLE[].
+        """
+        function_map = {
+            "euclidean": "array_distance",
+            "cosine": "array_cosine_distance",
+            "inner_product": "array_negative_inner_product",
+        }
+        target_type = "DOUBLE[]"
+        if isinstance(self.expression, exp.Array) and self.expression.expressions:
+            target_type = f"DOUBLE[{len(self.expression.expressions)}]"
+
+        function_name = function_map.get(metric)
+        if function_name:
+            right_cast = f"CAST({right} AS {target_type})"
+            return f"{function_name}({left}, {right_cast})"
+
+        return self._sql_generic(left, right, metric)
+
+    def _sql_generic(self, left: str, right: str, metric: str) -> str:
+        """Generate generic VECTOR_DISTANCE function syntax."""
+        return f"VECTOR_DISTANCE({left}, {right}, '{metric.upper()}')"
+
+
+def _vector_distance_sql_base(generator: "Generator", expression: "VectorDistance") -> str:
+    """Base generator for VectorDistance expressions."""
+    return expression._sql_generic(  # pyright: ignore[reportPrivateUsage]
+        generator.sql(expression.left), generator.sql(expression.right), expression.metric
+    )
+
+
+def _vector_distance_sql_postgres(generator: "Generator", expression: "VectorDistance") -> str:
+    """PostgreSQL generator for VectorDistance expressions."""
+    return expression._sql_postgres(  # pyright: ignore[reportPrivateUsage]
+        generator.sql(expression.left), generator.sql(expression.right), expression.metric
+    )
+
+
+def _vector_distance_sql_mysql(generator: "Generator", expression: "VectorDistance") -> str:
+    """MySQL generator for VectorDistance expressions."""
+    return expression._sql_mysql(  # pyright: ignore[reportPrivateUsage]
+        generator.sql(expression.left), generator.sql(expression.right), expression.metric
+    )
+
+
+def _vector_distance_sql_oracle(generator: "Generator", expression: "VectorDistance") -> str:
+    """Oracle generator for VectorDistance expressions."""
+    return expression._sql_oracle(  # pyright: ignore[reportPrivateUsage]
+        generator.sql(expression.left), generator.sql(expression.right), expression.metric
+    )
+
+
+def _vector_distance_sql_bigquery(generator: "Generator", expression: "VectorDistance") -> str:
+    """BigQuery generator for VectorDistance expressions."""
+    return expression._sql_bigquery(  # pyright: ignore[reportPrivateUsage]
+        generator.sql(expression.left), generator.sql(expression.right), expression.metric
+    )
+
+
+def _vector_distance_sql_spanner(generator: "Generator", expression: "VectorDistance") -> str:
+    """Spanner generator for VectorDistance expressions (same as BigQuery)."""
+    return expression._sql_bigquery(  # pyright: ignore[reportPrivateUsage]
+        generator.sql(expression.left), generator.sql(expression.right), expression.metric
+    )
+
+
+def _vector_distance_sql_duckdb(generator: "Generator", expression: "VectorDistance") -> str:
+    """DuckDB generator for VectorDistance expressions."""
+    return expression._sql_duckdb(  # pyright: ignore[reportPrivateUsage]
+        generator.sql(expression.left), generator.sql(expression.right), expression.metric
+    )
+
+
+def _register_with_sqlglot() -> None:
+    """Register VectorDistance with SQLGlot's generator dispatch system."""
+    spanner_dialect: type | None = None
+    spangres_dialect: type | None = None
+    with suppress(ImportError):
+        from sqlspec.adapters.spanner.dialect import Spangres, Spanner
+
+        spanner_dialect = Spanner
+        spangres_dialect = Spangres
+
+    Generator.TRANSFORMS[VectorDistance] = _vector_distance_sql_base
+
+    Postgres.Generator.TRANSFORMS[VectorDistance] = _vector_distance_sql_postgres
+    MySQL.Generator.TRANSFORMS[VectorDistance] = _vector_distance_sql_mysql
+    Oracle.Generator.TRANSFORMS[VectorDistance] = _vector_distance_sql_oracle
+    BigQuery.Generator.TRANSFORMS[VectorDistance] = _vector_distance_sql_bigquery
+    DuckDB.Generator.TRANSFORMS[VectorDistance] = _vector_distance_sql_duckdb
+    if spanner_dialect is not None:
+        spanner_dialect.Generator.TRANSFORMS[VectorDistance] = _vector_distance_sql_spanner  # type: ignore[attr-defined]
+    if spangres_dialect is not None:
+        spangres_dialect.Generator.TRANSFORMS[VectorDistance] = _vector_distance_sql_postgres  # type: ignore[attr-defined]
+
+
+_register_with_sqlglot()