"""Base query builder with validation and parameter binding.
Provides abstract base classes and core functionality for SQL query builders.
"""
import re
from abc import abstractmethod
from collections.abc import Callable, Iterable, Mapping
from typing import Any, Final, NoReturn, Self, cast
import sqlglot
from sqlglot import Dialect, exp
from sqlglot import optimizer as sqlglot_optimizer
from sqlglot.dialects.dialect import DialectType
from sqlglot.errors import OptimizeError as SQLGlotOptimizeError
from sqlglot.errors import ParseError as SQLGlotParseError
from sqlglot.optimizer import RULES
from sqlglot.optimizer.eliminate_ctes import eliminate_ctes as _eliminate_ctes_rule
from sqlglot.optimizer.merge_subqueries import merge_subqueries as _merge_subqueries_rule
from sqlglot.optimizer.normalize_identifiers import normalize_identifiers as _normalize_identifiers_rule
from sqlglot.optimizer.optimize_joins import optimize_joins as _optimize_joins_rule
from sqlglot.optimizer.pushdown_predicates import pushdown_predicates as _pushdown_predicates_rule
from sqlglot.optimizer.pushdown_projections import pushdown_projections as _pushdown_projections_rule
from sqlglot.optimizer.qualify_columns import quote_identifiers as _quote_identifiers_rule
from sqlglot.optimizer.scope import build_scope
from sqlglot.optimizer.simplify import simplify as _simplify_rule
from sqlglot.schema import MappingSchema
from sqlspec.builder._locking import register_lock_generator
from sqlspec.builder._parsing_utils import _normalize_dialect, _resolve_dialect
from sqlspec.builder._vector_distance import has_vector_distance_ancestor
from sqlspec.core import (
SQL,
ParameterStyle,
ParameterStyleConfig,
StatementConfig,
get_cache,
get_cache_config,
hash_optimized_expression,
)
from sqlspec.core.filters import StatementFilter
from sqlspec.core.hashing import _expression_cache_fingerprint
from sqlspec.data_dictionary import DialectConfig, get_dialect_config
from sqlspec.exceptions import SQLBuilderError
from sqlspec.utils.logging import get_logger
from sqlspec.utils.type_guards import has_expression_and_parameters, is_expression
from sqlspec.utils.uuids import uuid4
__all__ = ("BuiltQuery", "ExpressionBuilder", "QueryBuilder")
MAX_PARAMETER_COLLISION_ATTEMPTS = 1000
PARAMETER_INDEX_PATTERN = re.compile(r"^param_(?P<index>\d+)$")
_PARAMETER_SUFFIX_PATTERN = re.compile(r"^(?P<base>.+)_(?P<index>\d+)$")
_UPPER_FOLDING_DIALECTS: Final[frozenset[str]] = frozenset({"oracle", "db2"})
_UNQUOTED_IDENTIFIER_PATTERN = re.compile(r"^[A-Za-z_][\w$#]*$")
logger = get_logger(__name__)
[docs]
class BuiltQuery:
"""SQL query with bound parameters."""
__slots__ = ("dialect", "parameters", "sql")
[docs]
def __init__(self, sql: str, parameters: dict[str, Any] | None = None, dialect: DialectType | None = None) -> None:
self.sql = sql
self.parameters = parameters if parameters is not None else {}
self.dialect = dialect
def __repr__(self) -> str:
parameter_keys = sorted(self.parameters.keys())
return f"BuiltQuery(sql={self.sql!r}, parameters={parameter_keys!r}, dialect={self.dialect!r})"
def __eq__(self, other: object) -> bool:
if not isinstance(other, BuiltQuery):
return NotImplemented
return self.sql == other.sql and self.parameters == other.parameters and self.dialect == other.dialect
def __hash__(self) -> int:
return hash((self.sql, frozenset(self.parameters.items()), self.dialect))
[docs]
class QueryBuilder:
"""Abstract base class for SQL query builders.
Provides common functionality for dialect handling, parameter management,
and query construction using SQLGlot.
"""
__slots__ = (
"_expression",
"_merge_target_quoted",
"_parameter_counter",
"_parameter_name_counters",
"_parameters",
"_with_ctes",
"dialect",
"enable_optimization",
"optimize_joins",
"optimize_predicates",
"schema",
"simplify_expressions",
)
[docs]
def __init__(
self,
dialect: DialectType | None = None,
schema: dict[str, dict[str, str]] | None = None,
enable_optimization: bool = True,
optimize_joins: bool = True,
optimize_predicates: bool = True,
simplify_expressions: bool = True,
) -> None:
self.dialect = dialect
self.schema = schema
self.enable_optimization = enable_optimization
self.optimize_joins = optimize_joins
self.optimize_predicates = optimize_predicates
self.simplify_expressions = simplify_expressions
self._expression: exp.Expr | None = None
self._parameter_name_counters: dict[str, int] = {}
self._parameters: dict[str, Any] = {}
self._parameter_counter: int = 0
self._with_ctes: dict[str, exp.CTE] = {}
self._merge_target_quoted = False
[docs]
def __str__(self) -> str:
"""Return the SQL string representation of the query.
Returns:
str: The SQL string for this query.
"""
return self.build().sql
@property
def dialect_name(self) -> "str | None":
"""Returns the normalized name of the dialect, if set."""
return _normalize_dialect(self.dialect)
@property
def parameters(self) -> dict[str, Any]:
"""Public access to query parameters."""
return self._parameters
@property
def with_ctes(self) -> "dict[str, exp.CTE]":
"""Get WITH clause CTEs (public API)."""
return dict(self._with_ctes)
[docs]
def get_expression(self) -> exp.Expr | None:
"""Get expression reference (no copy).
Returns:
The current SQLGlot expression or None if not set
"""
return self._expression
[docs]
def set_expression(self, expression: exp.Expr) -> None:
"""Set expression with validation.
Args:
expression: SQLGlot expression to set
"""
if not is_expression(expression):
self._raise_invalid_expression_type(expression)
self._expression = expression
[docs]
def has_expression(self) -> bool:
"""Check if expression exists.
Returns:
True if expression is set, False otherwise
"""
return self._expression is not None
[docs]
def add_parameter_for_expression(self, value: Any, context: str | None = None) -> str:
"""Add a parameter for expression parameterization.
Args:
value: The value of the parameter.
context: Optional context hint for parameter naming.
Returns:
Parameter placeholder name.
"""
return self._add_parameter(value, context=context)
[docs]
def add_parameter(self: Self, value: Any, name: str | None = None) -> tuple[Self, str]:
"""Explicitly adds a parameter to the query.
This is useful for parameters that are not directly tied to a
builder method like `where` or `values`.
Args:
value: The value of the parameter.
name: Optional explicit name for the parameter. If None, a name
will be generated.
Returns:
tuple[Self, str]: The builder instance and the parameter name.
"""
if name:
if name in self._parameters:
self._raise_builder_error(f"Parameter name '{name}' already exists.")
self._parameters[name] = value
self._sync_parameter_name_counter(name)
return self, name
self._parameter_counter += 1
param_name = f"param_{self._parameter_counter}"
self._parameters[param_name] = value
self._sync_parameter_name_counter(param_name)
return self, param_name
[docs]
def load_parameters(self, parameters: "Mapping[str, Any]") -> None:
"""Load a parameter mapping into the builder.
Args:
parameters: Mapping of parameter names to values.
"""
if not parameters:
return
for name, value in parameters.items():
if name in self._parameters:
self._raise_builder_error(f"Parameter name '{name}' already exists.")
self._parameters[name] = value
self._sync_parameter_name_counter(name)
[docs]
def load_ctes(self, ctes: "Iterable[exp.CTE]") -> None:
"""Load SQLGlot CTE nodes into the builder.
Args:
ctes: Iterable of CTE expressions to register.
"""
for cte in ctes:
alias = self._resolve_cte_alias(cte)
if alias in self._with_ctes:
self._raise_builder_error(f"CTE '{alias}' already exists.")
self._with_ctes[alias] = cte
[docs]
def create_placeholder(self, value: Any, base_name: str) -> tuple[exp.Placeholder, str]:
"""Create placeholder expression with a unique parameter name.
Args:
value: Parameter value to bind.
base_name: Seed for parameter naming.
Returns:
Tuple of placeholder expression and the final parameter name.
"""
param_name = self._next_parameter_name(base_name)
_, param_name = self.add_parameter(value, name=param_name)
return exp.Placeholder(this=param_name), param_name
[docs]
def with_cte(
self: Self,
alias: str,
query: "QueryBuilder | exp.Select | exp.SetOperation | exp.Values | str | Any",
recursive: bool = False,
columns: "list[str] | None" = None,
) -> Self:
"""Adds a Common Table Expression (CTE) to the query.
Args:
alias: The alias for the CTE.
query: The CTE query, which can be another QueryBuilder instance,
a raw SQL string, or a sqlglot Select or Values expression.
recursive: Whether the CTE is recursive.
columns: Optional list of column aliases for the CTE.
Returns:
Self: The current builder instance for method chaining.
"""
if alias in self._with_ctes:
self._raise_builder_error(f"CTE with alias '{alias}' already exists.")
cte_select_expression = self._resolve_cte_query(alias, query)
cte_columns = columns
if cte_columns is None and isinstance(query, exp.Values):
values_alias = query.args.get("alias")
if isinstance(values_alias, exp.TableAlias):
cte_columns = [column.name for column in values_alias.columns]
if cte_columns is None:
query_cols = getattr(query, "columns", None)
query_private_cols = getattr(query, "_columns", None)
if isinstance(query_cols, (list, tuple)):
cte_columns = [str(c) for c in query_cols]
elif isinstance(query_private_cols, (list, tuple)):
cte_columns = [str(c) for c in query_private_cols]
if cte_columns:
alias_node: exp.Expr = exp.TableAlias(
this=exp.to_identifier(alias), columns=[exp.to_identifier(c) for c in cte_columns]
)
else:
alias_node = exp.to_table(alias)
self._with_ctes[alias] = exp.CTE(this=cte_select_expression, alias=alias_node)
self._with_ctes[alias].meta["recursive"] = recursive
return self
[docs]
def with_(
self: Self,
name: str,
query: "QueryBuilder | exp.Select | exp.SetOperation | exp.Values | str | Any",
recursive: bool = False,
columns: "list[str] | None" = None,
) -> Self:
"""Alias for with_cte for parity across builders.
Args:
name: The alias/name for the CTE.
query: The CTE query expression or builder.
recursive: Whether the CTE is recursive.
columns: Optional list of column aliases for the CTE.
Returns:
Self: The current builder instance for method chaining.
"""
return self.with_cte(name, query, recursive=recursive, columns=columns)
[docs]
def build(self, dialect: DialectType = None) -> "BuiltQuery":
"""Builds the SQL query string and parameters.
Args:
dialect: Optional dialect override. If provided, generates SQL for this dialect
instead of the builder's default dialect.
Returns:
BuiltQuery: A dataclass containing the SQL string and parameters.
"""
final_expression = self._prepare_expression_for_build(dialect=dialect, copy=True)
target_dialect = self._build_dialect(dialect)
self._validate_update_from(final_expression, _resolve_dialect(dialect, self.dialect))
if self.enable_optimization and isinstance(final_expression, exp.Expr):
final_expression = self._optimize_expression(final_expression, dialect=target_dialect)
final_expression = self._prepare_dialect_expression(final_expression, target_dialect, dialect)
try:
if isinstance(final_expression, exp.Expr):
normalized_expression = (
self._unquote_identifiers(final_expression)
if self._folds_unquoted_to_upper(target_dialect)
else final_expression
)
identify = self._should_identify(target_dialect)
if normalized_expression.find(exp.Lock) and target_dialect != "db2":
register_lock_generator(target_dialect)
sql_string = normalized_expression.sql(
dialect=target_dialect, pretty=True, identify=identify, copy=False
)
sql_string = self._postprocess_built_sql(sql_string, dialect=dialect)
else:
sql_string = str(final_expression)
except Exception as e:
err_msg = f"Error generating SQL from expression: {e!s}"
self._raise_builder_error(err_msg, e)
return BuiltQuery(sql=sql_string, parameters=self._parameters.copy(), dialect=target_dialect)
[docs]
def to_sql(self, show_parameters: bool = False, dialect: DialectType = None) -> str:
"""Return SQL string with optional parameter substitution.
Args:
show_parameters: If True, replace parameter placeholders with actual values (for debugging).
If False (default), return SQL with parameter placeholders.
dialect: Optional dialect override. If provided, generates SQL for this dialect
instead of the builder's default dialect.
Returns:
SQL string with or without parameter values filled in
Warning:
SQL with show_parameters=True is for debugging ONLY.
Never execute SQL with interpolated parameters directly - use parameterized queries.
"""
safe_query = self.build(dialect=dialect)
if not show_parameters:
return safe_query.sql
sql = safe_query.sql
parameters = safe_query.parameters
if not parameters:
return sql
sorted_names = sorted(parameters, key=len, reverse=True)
pattern = re.compile(r":(" + "|".join(re.escape(name) for name in sorted_names) + r")\b")
def _replace_match(match: "re.Match[str]") -> str:
param_name = match.group(1)
param_value = parameters[param_name]
if isinstance(param_value, str):
escaped = param_value.replace("'", "''")
return f"'{escaped}'"
if param_value is None:
return "NULL"
if isinstance(param_value, bool):
return "TRUE" if param_value else "FALSE"
return str(param_value)
return pattern.sub(_replace_match, sql)
[docs]
def to_statement(self, config: "StatementConfig | None" = None) -> "SQL":
"""Converts the built query into a SQL statement object.
Args:
config: Optional SQL configuration.
Returns:
SQL: A SQL statement object.
"""
cache_config = get_cache_config()
if (
not cache_config.compiled_cache_enabled
or not cache_config.sql_cache_enabled
or (config is not None and not config.enable_caching)
):
return self._to_statement(config)
cache_key_str = self._cache_key(config)
cache = get_cache()
cached_entry = cache.get_builder(cache_key_str)
if cached_entry is None:
cache_entry = self._create_builder_cache_entry(config)
cache.put_builder(cache_key_str, cache_entry)
else:
cache_entry = cast("_BuilderCacheEntry", cached_entry)
return self._statement_from_cache_entry(cache_entry, config)
[docs]
def apply_filters(self, *filters: StatementFilter, config: "StatementConfig | None" = None) -> "SQL":
"""Apply statement filters to this query builder.
Converts the builder to a SQL statement, then applies each filter's
``append_to_statement`` method to produce the final SQL with all
WHERE clauses, ORDER BY, LIMIT/OFFSET, etc.
Args:
*filters: Statement filters.
config: Optional SQL configuration.
Returns:
SQL statement with all filters applied.
"""
sql_statement = self.to_statement(config)
for filter_obj in filters:
sql_statement = filter_obj.append_to_statement(sql_statement)
return sql_statement
[docs]
def set_parameters(self, parameters: dict[str, Any]) -> None:
"""Set query parameters (public API)."""
self._parameters = parameters.copy()
self._parameter_counter = 0
self._parameter_name_counters.clear()
for name in self._parameters:
self._sync_parameter_name_counter(name)
[docs]
def generate_unique_parameter_name(self, base_name: str) -> str:
"""Generate unique parameter name (public API)."""
return self._next_parameter_name(base_name)
[docs]
def build_static_expression(
self,
expression: exp.Expr | None = None,
parameters: dict[str, Any] | None = None,
*,
cache_key: str | None = None,
expression_factory: Callable[[], exp.Expr] | None = None,
copy: bool = True,
optimize_expression: bool | None = None,
dialect: DialectType | None = None,
) -> "BuiltQuery":
"""Compile a pre-built expression with optional caching and parameters.
Designed for hot paths that construct an AST once and reuse it with
different parameters, avoiding repeated parse/optimize cycles.
Args:
expression: Pre-built sqlglot expression to render (required when cache_key is not provided).
parameters: Optional parameter mapping to include in the result.
cache_key: When provided, the expression will be cached under this key.
expression_factory: Factory used to build the expression on cache miss.
copy: Copy the expression before rendering to avoid caller mutation.
optimize_expression: Override builder optimization toggle for this call.
dialect: Optional dialect override for SQL generation.
Returns:
BuiltQuery containing SQL and parameters.
"""
expr: exp.Expr | None = None
target_dialect = self._build_dialect(dialect)
if cache_key is not None:
cache = get_cache()
cached_expr = cache.get_expression(cache_key)
if cached_expr is None:
if expression_factory is None:
msg = "expression_factory is required when cache_key is provided"
self._raise_builder_error(msg)
expr_candidate = expression_factory()
if not is_expression(expr_candidate):
self._raise_invalid_expression_type(expr_candidate)
expr_to_store = expr_candidate.copy()
should_optimize = self.enable_optimization if optimize_expression is None else optimize_expression
if should_optimize:
expr_to_store = self._optimize_expression(
expr_to_store, force=optimize_expression is True, dialect=target_dialect
)
cache.put_expression(cache_key, expr_to_store)
cached_expr = expr_to_store
expr = cached_expr.copy()
else:
if expression is None:
msg = "expression must be provided when cache_key is not set"
self._raise_builder_error(msg)
expr = expression.copy() if copy else expression
should_optimize = self.enable_optimization if optimize_expression is None else optimize_expression
if should_optimize:
expr = self._optimize_expression(expr, force=optimize_expression is True, dialect=target_dialect)
if expr is None:
self._raise_builder_error("Static expression could not be resolved.")
self._validate_update_from(expr, _resolve_dialect(dialect, self.dialect))
expr = self._prepare_dialect_expression(expr, target_dialect, dialect)
if self._folds_unquoted_to_upper(target_dialect):
expr = self._unquote_identifiers(expr)
identify = self._should_identify(target_dialect)
if expr.find(exp.Lock) and target_dialect != "db2":
register_lock_generator(target_dialect)
sql_string = expr.sql(dialect=target_dialect, pretty=True, identify=identify, copy=not copy)
sql_string = self._postprocess_built_sql(sql_string, dialect=dialect)
return BuiltQuery(
sql=sql_string,
parameters=parameters.copy() if parameters else {},
dialect=_resolve_dialect(dialect, self.dialect),
)
@classmethod
def _parse_init_options(
cls, kwargs: "dict[str, Any]"
) -> "tuple[DialectType | None, dict[str, dict[str, str]] | None, bool, bool, bool, bool]":
dialect = kwargs.pop("dialect", None)
schema = kwargs.pop("schema", None)
enable_optimization = kwargs.pop("enable_optimization", True)
optimize_joins = kwargs.pop("optimize_joins", True)
optimize_predicates = kwargs.pop("optimize_predicates", True)
simplify_expressions = kwargs.pop("simplify_expressions", True)
if kwargs:
unknown = ", ".join(sorted(kwargs.keys()))
cls._raise_builder_error(f"Unexpected QueryBuilder arguments: {unknown}")
return (dialect, schema, enable_optimization, optimize_joins, optimize_predicates, simplify_expressions)
def _init_query_builder(self, kwargs: "dict[str, Any]") -> None:
"""Consume shared builder options from kwargs and initialize QueryBuilder state."""
(dialect, schema, enable_optimization, optimize_joins, optimize_predicates, simplify_expressions) = (
self._parse_init_options(kwargs)
)
QueryBuilder.__init__(
self,
dialect=dialect,
schema=schema,
enable_optimization=enable_optimization,
optimize_joins=optimize_joins,
optimize_predicates=optimize_predicates,
simplify_expressions=simplify_expressions,
)
def _initialize_expression(self) -> None:
"""Initialize the base expression. Called after __init__."""
self._expression = self._create_base_expression()
if not self._expression:
self._raise_builder_error("QueryBuilder._create_base_expression must return a valid sqlglot expression.")
@abstractmethod
def _create_base_expression(self) -> exp.Expr:
"""Create the base sqlglot expression for the specific query type.
Returns:
A new sqlglot expression appropriate for the query type.
"""
@staticmethod
def _raise_builder_error(message: str, cause: BaseException | None = None) -> NoReturn:
"""Helper to raise SQLBuilderError, potentially with a cause.
Args:
message: The error message.
cause: The optional original exception to chain.
Raises:
SQLBuilderError: Always raises this exception.
"""
raise SQLBuilderError(message) from cause
@staticmethod
def _raise_invalid_expression_type(expression: Any) -> NoReturn:
"""Raise error for invalid expression type.
Args:
expression: The invalid expression object
Raises:
TypeError: Always raised for type mismatch
"""
msg = f"Expected Expression, got {type(expression)}"
raise TypeError(msg)
@staticmethod
def _raise_cte_query_error(alias: str, message: str) -> NoReturn:
"""Raise error for CTE query issues.
Args:
alias: CTE alias name
message: Specific error message
Raises:
SQLBuilderError: Always raised for CTE errors
"""
msg = f"CTE '{alias}': {message}"
raise SQLBuilderError(msg)
@staticmethod
def _raise_cte_parse_error(cause: BaseException) -> NoReturn:
"""Raise error for CTE parsing failures.
Args:
cause: The original parsing exception
Raises:
SQLBuilderError: Always raised with chained cause
"""
msg = f"Failed to parse CTE query: {cause!s}"
raise SQLBuilderError(msg) from cause
def _prepare_expression_for_build(self, dialect: DialectType = None, *, copy: bool = True) -> exp.Expr:
"""Validate builder state and return the expression ready for optimization and rendering."""
return self._build_final_expression(copy=copy)
def _postprocess_built_sql(self, sql_string: str, dialect: DialectType = None) -> str:
"""Apply builder-specific post-processing to a rendered SQL string."""
return self._strip_merge_target_quotes(sql_string)
def _build_final_expression(self, *, copy: bool = False) -> exp.Expr:
"""Construct the current expression with attached CTEs.
Args:
copy: Whether to copy the underlying expression tree before
applying transformations.
Returns:
Expression representing the current builder state with CTEs applied.
"""
if self._expression is None:
self._raise_builder_error("QueryBuilder expression not initialized.")
base_expression = self._expression.copy() if copy or self._with_ctes else self._expression
if not self._with_ctes:
return base_expression
final_expression: exp.Expr = base_expression
existing_with = final_expression.args.get("with_")
if existing_with is None:
final_expression.set("with_", exp.With(expressions=[cte.copy() for cte in self._with_ctes.values()]))
else:
for cte_node in self._with_ctes.values():
if cte_node not in existing_with.expressions:
existing_with.append("expressions", cte_node.copy())
if any(cte.meta.get("recursive") for cte in self._with_ctes.values()):
final_expression.args["with_"].set("recursive", True)
return final_expression
def _spawn_like_self(self: Self) -> Self:
"""Create a new builder instance with matching configuration."""
return type(self)(
dialect=self.dialect,
schema=self.schema,
enable_optimization=self.enable_optimization,
optimize_joins=self.optimize_joins,
optimize_predicates=self.optimize_predicates,
simplify_expressions=self.simplify_expressions,
)
def _embed_subquery(self, subquery: Any, alias_hint: str = "subquery") -> exp.Expr:
"""Extract an AST expression from a subquery builder and merge its parameters without collision."""
subquery_expression: exp.Expr
parameters: dict[str, Any] = dict(getattr(subquery, "parameters", None) or {})
if isinstance(subquery, QueryBuilder):
subquery_expression = subquery._prepare_expression_for_build(copy=True)
elif hasattr(subquery, "_build_final_expression"):
subquery_expression = subquery._build_final_expression(copy=True)
elif hasattr(subquery, "get_expression") and subquery.get_expression() is not None:
subquery_expression = subquery.get_expression().copy()
else:
built = subquery.build() if hasattr(subquery, "build") else subquery
sql_text = built.sql if isinstance(built, BuiltQuery) else str(built)
if isinstance(built, BuiltQuery) and not parameters:
parameters = dict(built.parameters)
subquery_expression = exp.maybe_parse(sql_text, dialect=self.dialect_name) or exp.convert(sql_text)
if parameters:
param_mapping: dict[str, str] = {}
for old_name, value in parameters.items():
target_name = self._next_parameter_name(old_name) if old_name in self._parameters else old_name
_, actual_name = self.add_parameter(value, name=target_name)
if actual_name != old_name:
param_mapping[old_name] = actual_name
if param_mapping:
subquery_expression = self._update_placeholders(subquery_expression, param_mapping)
return subquery_expression
def _resolve_cte_query(
self, alias: str, query: "QueryBuilder | exp.Select | exp.SetOperation | exp.Values | str | Any"
) -> exp.Expr:
"""Resolve a CTE query into a Select or Values expression with merged parameters."""
if isinstance(query, QueryBuilder):
query_expr = query._build_final_expression(copy=True)
if query_expr is None:
self._raise_cte_query_error(alias, "query builder has no expression")
if not isinstance(query_expr, (exp.Select, exp.SetOperation, exp.Values)):
self._raise_cte_query_error(
alias, f"expression must be a Select or Values, got {type(query_expr).__name__}"
)
cte_select_expression: exp.Expr = query_expr
if isinstance(cte_select_expression, exp.Values) and cte_select_expression.args.get("alias"):
cte_select_expression.set("alias", None)
param_mapping = self._merge_cte_parameters(alias, query.parameters)
if param_mapping:
cte_select_expression = self._update_placeholders(cte_select_expression, param_mapping)
return cte_select_expression
if hasattr(query, "get_expression"):
raw_query = cast("Any", query)
raw_query_expr = (
raw_query._build_final_expression(copy=True)
if hasattr(raw_query, "_build_final_expression")
else raw_query.get_expression()
)
if raw_query_expr is None:
self._raise_cte_query_error(alias, "query builder has no expression")
if not isinstance(raw_query_expr, (exp.Select, exp.SetOperation, exp.Values)):
self._raise_cte_query_error(
alias, f"expression must be a Select or Values, got {type(raw_query_expr).__name__}"
)
cte_duck_expression: exp.Expr = raw_query_expr.copy()
if isinstance(cte_duck_expression, exp.Values) and cte_duck_expression.args.get("alias"):
cte_duck_expression.set("alias", None)
if hasattr(raw_query, "parameters"):
param_mapping = self._merge_cte_parameters(alias, raw_query.parameters)
if param_mapping:
cte_duck_expression = self._update_placeholders(cte_duck_expression, param_mapping)
return cte_duck_expression
if isinstance(query, str):
try:
parsed_expression = sqlglot.parse_one(query, read=self.dialect_name)
except SQLGlotParseError as e: # pragma: no cover
self._raise_cte_parse_error(e)
if not isinstance(parsed_expression, (exp.Select, exp.SetOperation, exp.Values)):
self._raise_cte_query_error(
alias, f"query string must parse to SELECT or VALUES, got {type(parsed_expression).__name__}"
)
return parsed_expression
if isinstance(query, (exp.Select, exp.SetOperation, exp.Values)):
cte_expression = query.copy()
if isinstance(cte_expression, exp.Values):
cte_expression.set("alias", None)
return cte_expression
self._raise_cte_query_error(alias, f"invalid query type: {type(query).__name__}")
msg = "Unreachable"
raise AssertionError(msg)
def _add_parameter(self, value: Any, context: str | None = None) -> str:
"""Adds a parameter to the query and returns its placeholder name.
Args:
value: The value of the parameter.
context: Optional context hint for parameter naming
Returns:
str: The placeholder name for the parameter.
"""
self._parameter_counter += 1
param_name = f"{context}_param_{self._parameter_counter}" if context else f"param_{self._parameter_counter}"
self._parameters[param_name] = value
self._sync_parameter_name_counter(param_name)
return param_name
def _parameterize_expression(self, expression: exp.Expr) -> exp.Expr:
"""Replace literal values in an expression with bound parameters.
This method traverses a SQLGlot expression tree and replaces literal
values with parameter placeholders, adding the values to the builder's
parameter collection.
Args:
expression: The SQLGlot expression to parameterize
Returns:
A new expression with literals replaced by parameter placeholders
"""
return expression.transform(_ExpressionParameterizer(self), copy=True)
def _resolve_cte_alias(self, cte: exp.CTE) -> str:
alias_name = cte.alias_or_name
if not alias_name:
self._raise_builder_error("CTE alias is required.")
return str(alias_name)
def _update_parameter_counter(self, name: str) -> None:
match = PARAMETER_INDEX_PATTERN.match(name)
if not match:
return
index = int(match.group("index"))
self._parameter_counter = max(self._parameter_counter, index)
def _sync_parameter_name_counter(self, name: str) -> None:
self._update_parameter_counter(name)
suffix_match = _PARAMETER_SUFFIX_PATTERN.match(name)
if suffix_match:
base_name = suffix_match.group("base")
index = int(suffix_match.group("index"))
self._parameter_name_counters[base_name] = max(self._parameter_name_counters.get(base_name, 0), index)
else:
self._parameter_name_counters.setdefault(name, 0)
def _next_parameter_name(self, base_name: str) -> str:
"""Generate unique parameter name when collision occurs.
Args:
base_name: The desired base name for the parameter
Returns:
A unique parameter name that doesn't exist in current parameters
"""
current_index = self._parameter_name_counters.get(base_name, 0)
if base_name not in self._parameters:
self._parameter_name_counters[base_name] = current_index
return base_name
next_index = current_index + 1
candidate = f"{base_name}_{next_index}"
while candidate in self._parameters:
next_index += 1
if next_index > MAX_PARAMETER_COLLISION_ATTEMPTS:
return f"{base_name}_{uuid4().hex[:8]}"
candidate = f"{base_name}_{next_index}"
self._parameter_name_counters[base_name] = next_index
return candidate
def _merge_cte_parameters(self, cte_name: str, parameters: dict[str, Any]) -> dict[str, str]:
"""Merge CTE parameters with unique naming to prevent collisions.
Args:
cte_name: The name of the CTE for parameter prefixing
parameters: The CTE's parameter dictionary
Returns:
Mapping of old parameter names to new unique names
"""
param_mapping = {}
for old_name, value in parameters.items():
new_name = self._next_parameter_name(f"{cte_name}_{old_name}")
param_mapping[old_name] = new_name
self.add_parameter(value, name=new_name)
return param_mapping
def _update_placeholders(self, expression: exp.Expr, param_mapping: dict[str, str]) -> exp.Expr:
"""Update parameter placeholders in expression to use new names.
Args:
expression: The SQLGlot expression to update
param_mapping: Mapping of old parameter names to new names
Returns:
Updated expression with new placeholder names
"""
return expression.transform(_PlaceholderReplacer(param_mapping), copy=False)
def _cache_key(self, config: "StatementConfig | None" = None) -> str:
"""Generate cache key based on builder state and configuration.
Args:
config: Optional SQL configuration that affects the generated SQL
Returns:
A unique cache key representing the builder state and configuration
"""
dialect_override = config.dialect if config is not None and config.dialect is not None else None
final_expression = self._prepare_expression_for_build(dialect=dialect_override, copy=True)
dialect = self._build_dialect(dialect_override)
hints = getattr(self, "_hints", None)
settings: dict[str, Any] = {
"enable_optimization": self.enable_optimization,
"optimize_joins": self.optimize_joins,
"optimize_predicates": self.optimize_predicates,
"simplify_expressions": self.simplify_expressions,
}
if hints:
settings["hints"] = tuple(
(h.get("hint"), h.get("location"), h.get("table"), h.get("dialect")) for h in hints
)
fingerprint = _expression_cache_fingerprint(
final_expression,
parameter_signature=tuple(self._parameters),
dialect=dialect,
schema=self.schema,
settings=settings,
)
return f"builder:{fingerprint}"
def _has_sql_postprocessing(self) -> bool:
return False
def _build_dialect(self, dialect: DialectType = None) -> str | None:
return _normalize_dialect(dialect or self.dialect)
def _validate_dialect_lock(self, lock: exp.Lock, dialect: str, config: DialectConfig) -> None:
"""Reject lock options the target dialect cannot express."""
if dialect in {"spanner", "spangres"} and (
not lock.args.get("update") or lock.args.get("wait") is not None or lock.expressions or lock.args.get("key")
):
self._raise_builder_error(f"Dialect '{dialect}' supports only plain FOR UPDATE without lock modifiers.")
if lock.args.get("key") and dialect != "postgres":
self._raise_builder_error(f"Dialect '{dialect}' does not support PostgreSQL key lock modes.")
if dialect == "oracle" and not lock.args.get("update"):
self._raise_builder_error("Dialect 'oracle' does not support FOR SHARE.")
if dialect not in {"spanner", "spangres"} and config.get_feature_flag("supports_for_update") is False:
self._raise_builder_error(f"Dialect '{dialect}' does not support FOR UPDATE / row locking.")
if lock.args.get("wait") is False and config.get_feature_flag("supports_skip_locked") is False:
self._raise_builder_error(f"Dialect '{dialect}' does not support SKIP LOCKED.")
if dialect == "db2":
wait = lock.args.get("wait")
if wait is True or isinstance(wait, exp.Literal):
self._raise_builder_error(
"Dialect 'db2' does not support NOWAIT or WAIT; set CURRENT LOCK TIMEOUT on the session."
)
if lock.expressions:
self._raise_builder_error("Dialect 'db2' does not support FOR UPDATE OF table targets.")
def _prepare_dialect_expression(
self, expression: exp.Expr, dialect: str | None, source_dialect: DialectType = None
) -> exp.Expr:
"""Validate and translate a copy so builds never mutate the builder AST."""
if dialect is None:
return expression
try:
config = get_dialect_config("spanner" if dialect == "spangres" else dialect)
except ValueError:
return expression
if str(source_dialect or self.dialect).lower() == "mariadb" and expression.find(exp.Lock):
expression = expression.copy()
for lock in expression.find_all(exp.Lock):
if lock.expressions:
self._raise_builder_error("MariaDB locking clauses do not support OF targets.")
if not lock.args.get("update"):
lock.set("sqlspec_share_mode", True)
for lock in expression.find_all(exp.Lock):
self._validate_dialect_lock(lock, dialect, config)
if dialect == "spangres":
self._validate_spangres_conflicts(expression)
if config.get_feature_flag("supports_on_conflict") is not False or not expression.find(exp.OnConflict):
return expression
if dialect != "mysql":
self._raise_builder_error(f"Dialect '{dialect}' does not support ON CONFLICT; use sql.merge() instead.")
expression = expression.copy()
for conflict in expression.find_all(exp.OnConflict):
if conflict.args.get("duplicate"):
continue
if any(conflict.args.get(key) for key in ("where", "index_predicate", "constraint")):
self._raise_builder_error("MySQL cannot preserve ON CONFLICT predicates or named constraints.")
assignments = conflict.args.get("expressions")
if str(conflict.args.get("action", "")).upper() == "DO NOTHING":
keys = conflict.args.get("conflict_keys")
insert = conflict.find_ancestor(exp.Insert)
schema = insert.this if insert is not None else None
columns = keys or (schema.expressions if isinstance(schema, exp.Schema) else None)
if not columns:
self._raise_builder_error("MySQL DO NOTHING requires a conflict column or explicit insert columns.")
column = exp.column(columns[0].name)
assignments = [exp.EQ(this=column, expression=column.copy())]
elif not assignments:
self._raise_builder_error("ON CONFLICT DO UPDATE requires at least one assignment.")
for assignment in assignments:
for column in list(assignment.find_all(exp.Column)):
if column.table.lower() == "excluded":
column.replace(exp.Anonymous(this="VALUES", expressions=[exp.column(column.name)]))
conflict.replace(exp.OnConflict(duplicate=True, action=exp.var("UPDATE"), expressions=assignments))
return expression
def _validate_spangres_conflicts(self, expression: exp.Expr) -> None:
"""Check PostgreSQL-mode Spanner upsert restrictions visible in the AST."""
for conflict in expression.find_all(exp.OnConflict):
if any(conflict.args.get(key) for key in ("where", "index_predicate", "constraint", "duplicate")):
self._raise_builder_error("Spanner PostgreSQL does not support these ON CONFLICT modifiers.")
assignments = conflict.args.get("expressions") or []
if assignments:
insert = conflict.find_ancestor(exp.Insert)
schema = insert.this if insert is not None else None
assigned_columns = set()
for assignment in assignments:
value = assignment.expression
if (
not isinstance(value, exp.Column)
or value.table.lower() != "excluded"
or value.name != assignment.this.name
):
self._raise_builder_error("Spanner PostgreSQL conflict updates require excluded column values.")
assigned_columns.add(assignment.this.name)
if isinstance(schema, exp.Schema) and assigned_columns != {
column.name for column in schema.expressions
}:
self._raise_builder_error("Spanner PostgreSQL conflict updates must assign every inserted column.")
def _optimize_expression(
self, expression: exp.Expr, *, force: bool = False, dialect: DialectType = None
) -> exp.Expr:
"""Apply SQLGlot optimizations to the expression.
The ``ON CONFLICT`` / ``ON DUPLICATE KEY UPDATE`` clause of an INSERT only receives
identifier normalization and quoting, so each ``SET`` assignment keeps its target
column on the left.
Args:
expression: The expression to optimize
force: Optimize even when the builder-level toggle is disabled.
dialect: Optional dialect override for optimization rules.
Returns:
The optimized expression
"""
if not force and not self.enable_optimization:
return expression
if not self.optimize_joins and not self.optimize_predicates and not self.simplify_expressions:
return expression
target_dialect = _normalize_dialect(dialect) if dialect is not None else self.dialect_name
optimizer_dialect = None if target_dialect in {"oracle", "db2", "snowflake"} else target_dialect
conflict = expression.args.get("conflict") if isinstance(expression, exp.Insert) else None
if isinstance(conflict, exp.OnConflict):
return self._optimize_insert_with_conflict(expression, conflict, force=force, dialect=optimizer_dialect)
optimizer_settings = {
"optimize_joins": self.optimize_joins,
"pushdown_predicates": self.optimize_predicates,
"simplify_expressions": self.simplify_expressions,
}
dialect_key = optimizer_dialect or "default"
cache_key = hash_optimized_expression(
expression, dialect=dialect_key, schema=self.schema, optimizer_settings=optimizer_settings
)
cache = get_cache()
cached_optimized = cache.get_optimized(cache_key)
if cached_optimized is not None:
return cast("exp.Expr", cached_optimized).copy()
if any(isinstance(cte.this, exp.Values) and cte.alias_column_names for cte in expression.find_all(exp.CTE)):
return expression
excluded_rules = set()
if not self.optimize_joins:
excluded_rules.add(_optimize_joins_rule)
if not self.optimize_predicates:
excluded_rules.add(_pushdown_predicates_rule)
if not self.simplify_expressions:
excluded_rules.add(_simplify_rule)
if expression.args.get("with_") is not None or self._with_ctes:
excluded_rules.add(_eliminate_ctes_rule)
excluded_rules.add(_merge_subqueries_rule)
excluded_rules.add(_pushdown_projections_rule)
elif self.schema is None and (
(isinstance(expression, exp.Select) and expression.is_star) or expression.find(exp.Star) is not None
):
excluded_rules.add(_pushdown_projections_rule)
rules = RULES if not excluded_rules else tuple(rule for rule in RULES if rule not in excluded_rules)
rule_sets = [rules]
if _pushdown_projections_rule not in excluded_rules:
rule_sets.append(tuple(rule for rule in rules if rule is not _pushdown_projections_rule))
schema: Any = (
cast("dict[str, object]", self.schema)
if self.schema is not None
else _infer_single_table_schema(expression, dialect=optimizer_dialect)
)
for rule_set in rule_sets:
try:
optimized = sqlglot_optimizer.optimize(
expression, schema=schema, dialect=optimizer_dialect, rules=rule_set
)
cache.put_optimized(cache_key, optimized.copy())
except SQLGlotOptimizeError:
continue
except Exception:
break
return optimized
logger.debug("Expression optimization failed, using original expression")
return expression
def _optimize_insert_with_conflict(
self, expression: exp.Expr, conflict: exp.OnConflict, *, force: bool, dialect: str | None = None
) -> exp.Expr:
"""Optimize an INSERT without its conflict clause, then attach the clause with quoted identifiers.
Args:
expression: The INSERT expression owning ``conflict``.
conflict: The INSERT's ``ON CONFLICT`` / ``ON DUPLICATE KEY UPDATE`` clause.
force: Optimize even when the builder-level toggle is disabled.
dialect: Optional normalized dialect name for identifier normalization.
Returns:
The optimized INSERT with the normalized and quoted conflict clause, or the
original INSERT with its clause unchanged when optimization fails.
"""
expression.set("conflict", None)
try:
optimized = self._optimize_expression(expression, force=force, dialect=dialect)
finally:
expression.set("conflict", conflict)
if optimized is expression:
return expression
quoted_conflict = _quote_identifiers_rule(
_normalize_identifiers_rule(conflict.copy(), dialect=dialect), dialect=dialect
)
optimized.set("conflict", quoted_conflict)
return optimized
def _to_statement(self, config: "StatementConfig | None" = None) -> "SQL":
"""Internal method to create SQL statement.
Args:
config: Optional SQL configuration.
Returns:
SQL: A SQL statement object.
"""
cache_entry = self._create_builder_cache_entry(config)
return self._statement_from_cache_entry(cache_entry, config)
def _validate_update_from(self, expression: exp.Expr, dialect: DialectType) -> None:
if not dialect or not any(node.args.get("from_") is not None for node in expression.find_all(exp.Update)):
return
dialect_name = (
dialect.lower() if isinstance(dialect, str) else type(Dialect.get_or_raise(dialect)).__name__.lower()
)
try:
config = get_dialect_config("spanner" if dialect_name == "spangres" else dialect_name)
except ValueError:
return
if not config.feature_flags.get("supports_update_from", True):
msg = (
f"Dialect '{dialect_name}' does not support UPDATE ... FROM clauses. "
"Consider using MERGE or a JOIN-based UPDATE instead."
)
raise SQLBuilderError(msg)
def _create_builder_cache_entry(self, config: "StatementConfig | None") -> "_BuilderCacheEntry":
dialect_override = config.dialect if config is not None else None
resolved_dialect = self._build_dialect(dialect_override)
statement_expression = self._prepare_expression_for_build(dialect=dialect_override, copy=True)
self._validate_update_from(statement_expression, _resolve_dialect(dialect_override, self.dialect))
if self.enable_optimization and isinstance(statement_expression, exp.Expr):
statement_expression = self._optimize_expression(statement_expression, dialect=resolved_dialect)
statement_expression = self._prepare_dialect_expression(
statement_expression, resolved_dialect, dialect_override
)
if statement_expression.find(exp.Lock) and resolved_dialect != "db2":
register_lock_generator(resolved_dialect)
if self._folds_unquoted_to_upper(resolved_dialect):
statement_expression = self._unquote_identifiers(statement_expression)
return _BuilderCacheEntry(statement_expression, resolved_dialect)
def _statement_from_cache_entry(self, cache_entry: "_BuilderCacheEntry", config: "StatementConfig | None") -> "SQL":
statement_target: exp.Expr | str
if self._has_sql_postprocessing():
statement_target = self.build(dialect=cache_entry.dialect).sql
else:
statement_target = cache_entry.expression.copy()
kwargs, parameters = self._statement_parameters(self._parameters.copy())
statement_config = config
if statement_config is not None and statement_config.dialect != cache_entry.dialect:
statement_config = statement_config.replace(dialect=cache_entry.dialect)
if statement_config is None:
statement_config = StatementConfig(
parameter_config=ParameterStyleConfig(
default_parameter_style=ParameterStyle.QMARK, supported_parameter_styles={ParameterStyle.QMARK}
),
dialect=cache_entry.dialect,
)
if kwargs:
return SQL(statement_target, statement_config=statement_config, **kwargs)
if parameters:
return SQL(statement_target, *parameters, statement_config=statement_config)
return SQL(statement_target, statement_config=statement_config)
def _statement_parameters(self, raw_parameters: Any) -> "tuple[dict[str, Any] | None, tuple[Any, ...] | None]":
"""Extract parameters for SQL statement creation.
Args:
raw_parameters: Raw parameter data from BuiltQuery
Returns:
Tuple of (kwargs, parameters) for SQL statement construction
"""
if isinstance(raw_parameters, dict):
return raw_parameters, None
if isinstance(raw_parameters, tuple):
return None, raw_parameters
if raw_parameters:
return None, tuple(raw_parameters)
return None, None
def _merge_parameters(self, sql_obj: Any) -> dict[str, str]:
"""Merge parameters from a SQL object into the builder.
Args:
sql_obj: Object with parameters attribute containing parameter mappings
Returns:
Mapping of original parameter names to renamed parameter names.
"""
param_mapping: dict[str, str] = {}
if not has_expression_and_parameters(sql_obj):
return param_mapping
sql_parameters = sql_obj.parameters
for param_name, param_value in sql_parameters.items():
target_name = self._next_parameter_name(param_name) if param_name in self._parameters else param_name
_, actual_name = self.add_parameter(param_value, name=target_name)
if actual_name != param_name:
param_mapping[param_name] = actual_name
return param_mapping
def _folds_unquoted_to_upper(self, dialect: "DialectType | str | None") -> bool:
"""Check if the target dialect folds unquoted identifiers to uppercase."""
if dialect is None:
return False
return str(dialect).lower() in _UPPER_FOLDING_DIALECTS
def _unquote_identifiers(self, expression: exp.Expr) -> exp.Expr:
"""Remove generated quoting while preserving explicit table identifiers.
Upper-folding dialects resolve quoted lowercase names case-sensitively, so quoting is
removed to keep lookups aligned with how unquoted DDL created the objects.
"""
return expression.copy().transform(_unquote_identifier, copy=False)
def _strip_merge_target_quotes(self, sql_string: str) -> str:
if sql_string.startswith('MERGE INTO "') and not self._merge_target_quoted:
end_quote = sql_string.find('"', len('MERGE INTO "'))
if end_quote > 0:
table_name = sql_string[len('MERGE INTO "') : end_quote]
remainder = sql_string[end_quote + 1 :]
return f"MERGE INTO {table_name}{remainder}"
return sql_string
def _should_identify(self, dialect: "DialectType | str | None") -> bool:
"""Determine whether to quote identifiers for the given dialect."""
if dialect is None:
return True
return str(dialect).lower() not in _UPPER_FOLDING_DIALECTS
[docs]
class ExpressionBuilder(QueryBuilder):
"""Builder wrapper for a pre-parsed SQLGlot expression."""
__slots__ = ()
[docs]
def __init__(self, expression: exp.Expr, **kwargs: Any) -> None:
self._init_query_builder(kwargs)
if not is_expression(expression):
self._raise_invalid_expression_type(expression)
self._expression = expression
def _create_base_expression(self) -> exp.Expr:
if self._expression is None:
msg = "ExpressionBuilder requires an expression at construction."
self._raise_builder_error(msg)
return self._expression
class _BuilderCacheEntry:
__slots__ = ("dialect", "expression")
def __init__(self, expression: exp.Expr, dialect: "DialectType | None") -> None:
self.expression = expression
self.dialect = dialect
class _ExpressionParameterizer:
__slots__ = ("_builder",)
def __init__(self, builder: "QueryBuilder") -> None:
self._builder = builder
def __call__(self, node: exp.Expr) -> exp.Expr:
if isinstance(node, exp.Literal):
if node.this in {True, False, None}:
return node
parent = node.parent
if isinstance(parent, exp.Array) and has_vector_distance_ancestor(node):
return node
value = node.this
if node.is_number and isinstance(node.this, str):
try:
value = float(node.this) if "." in node.this or "e" in node.this.lower() else int(node.this)
except ValueError:
value = node.this
param_name = self._builder.add_parameter_for_expression(value, context="where")
return exp.Placeholder(this=param_name)
return node
class _PlaceholderReplacer:
__slots__ = ("_param_mapping",)
def __init__(self, param_mapping: dict[str, str]) -> None:
self._param_mapping = param_mapping
def __call__(self, node: exp.Expr) -> exp.Expr:
if isinstance(node, exp.Placeholder) and str(node.this) in self._param_mapping:
return exp.Placeholder(this=self._param_mapping[str(node.this)])
return node
def _infer_single_table_schema(expression: exp.Expr, dialect: DialectType = None) -> MappingSchema | None:
"""Synthesize a lightweight MappingSchema for single-table scopes when no schema is provided."""
try:
root = build_scope(expression)
if root is None:
return None
inferred: dict[str, Any] = {}
visible: dict[str, Any] = {}
for scope in root.traverse():
if len(scope.selected_sources) != 1 or scope.lateral_sources:
continue
alias_name, (_, source) = next(iter(scope.selected_sources.items()))
if not isinstance(source, exp.Table) or not source.name:
continue
select_aliases = {
s.alias
for s in (scope.expression.selects if isinstance(scope.expression, exp.Select) else ())
if isinstance(s, exp.Alias)
and s.alias
and not (isinstance(s.this, exp.Column) and s.this.name == s.alias)
}
cols = {
col.name: "UNKNOWN"
for col in scope.columns
if col.name
and col.name != "*"
and (col.table in {alias_name, source.name} or (not col.table and col.name not in select_aliases))
}
if not cols:
continue
parts = [p.name for p in source.parts]
target = inferred
vis_target = visible
for part in parts[:-1]:
target = target.setdefault(part, {})
vis_target = vis_target.setdefault(part, {})
leaf = parts[-1]
existing = target.get(leaf)
if isinstance(existing, dict):
existing.update(cols)
else:
target[leaf] = cols
vis_target[leaf] = []
if not inferred:
return None
return MappingSchema(inferred, visible=visible, dialect=dialect, normalize=False)
except Exception:
return None
def _unquote_identifier(node: exp.Expr) -> exp.Expr:
if (
isinstance(node, exp.Identifier)
and not node.meta.get("sqlspec_explicit_table_quote")
and _UNQUOTED_IDENTIFIER_PATTERN.fullmatch(node.name)
):
node.set("quoted", False)
return node
