sqlspec 0.26.0__py3-none-any.whl → 0.28.0__py3-none-any.whl

sqlspec/builder/_base.py

@@ -6,7 +6,7 @@ Provides abstract base classes and core functionality for SQL query builders.
 import hashlib
 import uuid
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, NoReturn, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, NoReturn, cast
 
 import sqlglot
 from sqlglot import Dialect, exp
@@ -38,9 +38,7 @@ class SafeQuery:
 
     __slots__ = ("dialect", "parameters", "sql")
 
-    def __init__(
-        self, sql: str, parameters: Optional[dict[str, Any]] = None, dialect: Optional[DialectType] = None
-    ) -> None:
+    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
@@ -68,8 +66,8 @@ class QueryBuilder(ABC):
 
     def __init__(
         self,
-        dialect: Optional[DialectType] = None,
-        schema: Optional[dict[str, dict[str, str]]] = None,
+        dialect: DialectType | None = None,
+        schema: dict[str, dict[str, str]] | None = None,
         enable_optimization: bool = True,
         optimize_joins: bool = True,
         optimize_predicates: bool = True,
@@ -82,7 +80,7 @@ class QueryBuilder(ABC):
         self.optimize_predicates = optimize_predicates
         self.simplify_expressions = simplify_expressions
 
-        self._expression: Optional[exp.Expression] = None
+        self._expression: exp.Expression | None = None
         self._parameters: dict[str, Any] = {}
         self._parameter_counter: int = 0
         self._with_ctes: dict[str, exp.CTE] = {}
@@ -95,7 +93,7 @@ class QueryBuilder(ABC):
                 "QueryBuilder._create_base_expression must return a valid sqlglot expression."
             )
 
-    def get_expression(self) -> Optional[exp.Expression]:
+    def get_expression(self) -> exp.Expression | None:
         """Get expression reference (no copy).
 
         Returns:
@@ -139,7 +137,7 @@ class QueryBuilder(ABC):
         """
 
     @staticmethod
-    def _raise_sql_builder_error(message: str, cause: Optional[BaseException] = None) -> NoReturn:
+    def _raise_sql_builder_error(message: str, cause: BaseException | None = None) -> NoReturn:
         """Helper to raise SQLBuilderError, potentially with a cause.
 
         Args:
@@ -191,7 +189,81 @@ class QueryBuilder(ABC):
         msg = f"Failed to parse CTE query: {cause!s}"
         raise SQLBuilderError(msg) from cause
 
-    def _add_parameter(self, value: Any, context: Optional[str] = None) -> str:
+    def _build_final_expression(self, *, copy: bool = False) -> exp.Expression:
+        """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_sql_builder_error("QueryBuilder expression not initialized.")
+
+        base_expression = self._expression.copy() if copy else self._expression
+
+        if not self._with_ctes:
+            return base_expression
+
+        final_expression: exp.Expression = base_expression
+        if has_with_method(final_expression):
+            for alias, cte_node in self._with_ctes.items():
+                final_expression = cast("Any", final_expression).with_(cte_node.args["this"], as_=alias, copy=False)
+            return cast("exp.Expression", final_expression)
+
+        if isinstance(final_expression, (exp.Select, exp.Insert, exp.Update, exp.Delete, exp.Union)):
+            return exp.With(expressions=list(self._with_ctes.values()), this=final_expression)
+
+        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 _resolve_cte_query(self, alias: str, query: "QueryBuilder | exp.Select | str") -> exp.Select:
+        """Resolve a CTE query into a Select expression with merged parameters."""
+        if isinstance(query, QueryBuilder):
+            query_expr = query.get_expression()
+            if query_expr is None:
+                self._raise_cte_query_error(alias, "query builder has no expression")
+            if not isinstance(query_expr, exp.Select):
+                self._raise_cte_query_error(alias, f"expression must be a Select, got {type(query_expr).__name__}")
+            cte_select_expression = query_expr.copy()
+            param_mapping = self._merge_cte_parameters(alias, query.parameters)
+            updated_expression = self._update_placeholders_in_expression(cte_select_expression, param_mapping)
+            if not isinstance(updated_expression, exp.Select):  # pragma: no cover - defensive
+                msg = "CTE placeholder update produced non-select expression"
+                raise SQLBuilderError(msg)
+            return updated_expression
+
+        if isinstance(query, str):
+            try:
+                parsed_expression = sqlglot.parse_one(query, read=self.dialect_name)
+            except SQLGlotParseError as e:  # pragma: no cover - defensive
+                self._raise_cte_parse_error(e)
+            if not isinstance(parsed_expression, exp.Select):
+                self._raise_cte_query_error(
+                    alias, f"query string must parse to SELECT, got {type(parsed_expression).__name__}"
+                )
+            return parsed_expression
+
+        if isinstance(query, exp.Select):
+            return query
+
+        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:
@@ -232,7 +304,7 @@ class QueryBuilder(ABC):
 
         return expression.transform(replacer, copy=False)
 
-    def add_parameter(self: Self, value: Any, name: Optional[str] = None) -> tuple[Self, str]:
+    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
@@ -313,7 +385,7 @@ class QueryBuilder(ABC):
 
         return expression.transform(placeholder_replacer, copy=False)
 
-    def _generate_builder_cache_key(self, config: "Optional[StatementConfig]" = None) -> str:
+    def _generate_builder_cache_key(self, config: "StatementConfig | None" = None) -> str:
         """Generate cache key based on builder state and configuration.
 
         Args:
@@ -356,7 +428,7 @@ class QueryBuilder(ABC):
         state_string = "|".join(state_parts)
         return f"builder:{hashlib.sha256(state_string.encode()).hexdigest()[:16]}"
 
-    def with_cte(self: Self, alias: str, query: "Union[QueryBuilder, exp.Select, str]") -> Self:
+    def with_cte(self: Self, alias: str, query: "QueryBuilder | exp.Select | str") -> Self:
         """Adds a Common Table Expression (CTE) to the query.
 
         Args:
@@ -370,35 +442,7 @@ class QueryBuilder(ABC):
         if alias in self._with_ctes:
             self._raise_sql_builder_error(f"CTE with alias '{alias}' already exists.")
 
-        cte_select_expression: exp.Select
-
-        if isinstance(query, QueryBuilder):
-            query_expr = query.get_expression()
-            if query_expr is None:
-                self._raise_cte_query_error(alias, "query builder has no expression")
-            if not isinstance(query_expr, exp.Select):
-                self._raise_cte_query_error(alias, f"expression must be a Select, got {type(query_expr).__name__}")
-            cte_select_expression = query_expr
-            param_mapping = self._merge_cte_parameters(alias, query.parameters)
-            cte_select_expression = cast(
-                "exp.Select", self._update_placeholders_in_expression(cte_select_expression, param_mapping)
-            )
-
-        elif isinstance(query, str):
-            try:
-                parsed_expression = sqlglot.parse_one(query, read=self.dialect_name)
-                if not isinstance(parsed_expression, exp.Select):
-                    self._raise_cte_query_error(
-                        alias, f"query string must parse to SELECT, got {type(parsed_expression).__name__}"
-                    )
-                cte_select_expression = parsed_expression
-            except SQLGlotParseError as e:
-                self._raise_cte_parse_error(e)
-        elif isinstance(query, exp.Select):
-            cte_select_expression = query
-        else:
-            self._raise_cte_query_error(alias, f"invalid query type: {type(query).__name__}")
-
+        cte_select_expression = self._resolve_cte_query(alias, query)
         self._with_ctes[alias] = exp.CTE(this=cte_select_expression, alias=exp.to_table(alias))
         return self
 
@@ -408,18 +452,7 @@ class QueryBuilder(ABC):
         Returns:
             SafeQuery: A dataclass containing the SQL string and parameters.
         """
-        if self._expression is None:
-            self._raise_sql_builder_error("QueryBuilder expression not initialized.")
-
-        if self._with_ctes:
-            final_expression = self._expression
-            if has_with_method(final_expression):
-                for alias, cte_node in self._with_ctes.items():
-                    final_expression = cast("Any", final_expression).with_(cte_node.args["this"], as_=alias, copy=False)
-            elif isinstance(final_expression, (exp.Select, exp.Insert, exp.Update, exp.Delete, exp.Union)):
-                final_expression = exp.With(expressions=list(self._with_ctes.values()), this=final_expression)
-        else:
-            final_expression = self._expression
+        final_expression = self._build_final_expression()
 
         if self.enable_optimization and isinstance(final_expression, exp.Expression):
             final_expression = self._optimize_expression(final_expression)
@@ -475,7 +508,7 @@ class QueryBuilder(ABC):
         else:
             return optimized
 
-    def to_statement(self, config: "Optional[StatementConfig]" = None) -> "SQL":
+    def to_statement(self, config: "StatementConfig | None" = None) -> "SQL":
         """Converts the built query into a SQL statement object.
 
         Args:
@@ -500,7 +533,7 @@ class QueryBuilder(ABC):
 
         return sql_statement
 
-    def _to_statement(self, config: "Optional[StatementConfig]" = None) -> "SQL":
+    def _to_statement(self, config: "StatementConfig | None" = None) -> "SQL":
         """Internal method to create SQL statement.
 
         Args:
@@ -541,7 +574,7 @@ class QueryBuilder(ABC):
 
     def _extract_statement_parameters(
         self, raw_parameters: Any
-    ) -> "tuple[Optional[dict[str, Any]], Optional[tuple[Any, ...]]]":
+    ) -> "tuple[dict[str, Any] | None, tuple[Any, ...] | None]":
         """Extract parameters for SQL statement creation.
 
         Args:
@@ -570,7 +603,7 @@ class QueryBuilder(ABC):
         return self.build().sql
 
     @property
-    def dialect_name(self) -> "Optional[str]":
+    def dialect_name(self) -> "str | None":
         """Returns the name of the dialect, if set."""
         if isinstance(self.dialect, str):
             return self.dialect

sqlspec/builder/_column.py

@@ -6,7 +6,7 @@ SQL conditions with parameter binding.
 
 from collections.abc import Iterable
 from datetime import date, datetime
-from typing import Any, Optional, cast
+from typing import Any, cast
 
 from sqlglot import exp
 
@@ -77,7 +77,7 @@ class Column:
 
     __slots__ = ("_expression", "name", "table")
 
-    def __init__(self, name: str, table: Optional[str] = None) -> None:
+    def __init__(self, name: str, table: str | None = None) -> None:
         self.name = name
         self.table = table
 
@@ -122,7 +122,7 @@ class Column:
         """Apply NOT operator (~)."""
         return ColumnExpression(exp.Not(this=self._expression))
 
-    def like(self, pattern: str, escape: Optional[str] = None) -> ColumnExpression:
+    def like(self, pattern: str, escape: str | None = None) -> ColumnExpression:
         """SQL LIKE pattern matching."""
         if escape:
             like_expr = exp.Like(
@@ -159,7 +159,7 @@ class Column:
         """SQL IS NOT NULL."""
         return ColumnExpression(exp.Not(this=exp.Is(this=self._expression, expression=exp.Null())))
 
-    def not_like(self, pattern: str, escape: Optional[str] = None) -> ColumnExpression:
+    def not_like(self, pattern: str, escape: str | None = None) -> ColumnExpression:
         """SQL NOT LIKE pattern matching."""
         return ~self.like(pattern, escape)
 
@@ -211,7 +211,7 @@ class Column:
         """SQL CEIL() function."""
         return FunctionColumn(exp.Ceil(this=self._expression))
 
-    def substring(self, start: int, length: Optional[int] = None) -> "FunctionColumn":
+    def substring(self, start: int, length: int | None = None) -> "FunctionColumn":
         """SQL SUBSTRING() function."""
         args = [self._convert_value(start)]
         if length is not None: