sqlspec 0.25.0__py3-none-any.whl → 0.27.0__py3-none-any.whl

sqlspec/builder/__init__.py

@@ -24,21 +24,64 @@ from sqlspec.builder._ddl import (
     Truncate,
 )
 from sqlspec.builder._delete import Delete
+from sqlspec.builder._dml import (
+    DeleteFromClauseMixin,
+    InsertFromSelectMixin,
+    InsertIntoClauseMixin,
+    InsertValuesMixin,
+    UpdateFromClauseMixin,
+    UpdateSetClauseMixin,
+    UpdateTableClauseMixin,
+)
+from sqlspec.builder._expression_wrappers import (
+    AggregateExpression,
+    ConversionExpression,
+    FunctionExpression,
+    MathExpression,
+    StringExpression,
+)
+from sqlspec.builder._factory import SQLFactory, sql
 from sqlspec.builder._insert import Insert
+from sqlspec.builder._join import JoinBuilder
 from sqlspec.builder._merge import Merge
-from sqlspec.builder._select import Select
+from sqlspec.builder._parsing_utils import (
+    extract_expression,
+    parse_column_expression,
+    parse_condition_expression,
+    parse_order_expression,
+    parse_table_expression,
+    to_expression,
+)
+from sqlspec.builder._select import (
+    Case,
+    CaseBuilder,
+    CommonTableExpressionMixin,
+    HavingClauseMixin,
+    LimitOffsetClauseMixin,
+    OrderByClauseMixin,
+    PivotClauseMixin,
+    ReturningClauseMixin,
+    Select,
+    SelectClauseMixin,
+    SetOperationMixin,
+    SubqueryBuilder,
+    UnpivotClauseMixin,
+    WhereClauseMixin,
+    WindowFunctionBuilder,
+)
 from sqlspec.builder._update import Update
-from sqlspec.builder.mixins import WhereClauseMixin
-from sqlspec.builder.mixins._join_operations import JoinBuilder
-from sqlspec.builder.mixins._select_operations import Case, SubqueryBuilder, WindowFunctionBuilder
 from sqlspec.exceptions import SQLBuilderError
 
 __all__ = (
+    "AggregateExpression",
     "AlterTable",
     "Case",
+    "CaseBuilder",
     "Column",
     "ColumnExpression",
     "CommentOn",
+    "CommonTableExpressionMixin",
+    "ConversionExpression",
     "CreateIndex",
     "CreateMaterializedView",
     "CreateSchema",
@@ -47,22 +90,48 @@ __all__ = (
     "CreateView",
     "DDLBuilder",
     "Delete",
+    "DeleteFromClauseMixin",
     "DropIndex",
     "DropSchema",
     "DropTable",
     "DropView",
     "FunctionColumn",
+    "FunctionExpression",
+    "HavingClauseMixin",
     "Insert",
+    "InsertFromSelectMixin",
+    "InsertIntoClauseMixin",
+    "InsertValuesMixin",
     "JoinBuilder",
+    "LimitOffsetClauseMixin",
+    "MathExpression",
     "Merge",
+    "OrderByClauseMixin",
+    "PivotClauseMixin",
     "QueryBuilder",
     "RenameTable",
+    "ReturningClauseMixin",
     "SQLBuilderError",
+    "SQLFactory",
     "SafeQuery",
     "Select",
+    "SelectClauseMixin",
+    "SetOperationMixin",
+    "StringExpression",
     "SubqueryBuilder",
     "Truncate",
+    "UnpivotClauseMixin",
     "Update",
+    "UpdateFromClauseMixin",
+    "UpdateSetClauseMixin",
+    "UpdateTableClauseMixin",
     "WhereClauseMixin",
     "WindowFunctionBuilder",
+    "extract_expression",
+    "parse_column_expression",
+    "parse_condition_expression",
+    "parse_order_expression",
+    "parse_table_expression",
+    "sql",
+    "to_expression",
 )

sqlspec/builder/_base.py

@@ -3,8 +3,10 @@
 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
@@ -19,13 +21,15 @@ from sqlspec.core.parameters import ParameterStyle, ParameterStyleConfig
 from sqlspec.core.statement import SQL, StatementConfig
 from sqlspec.exceptions import SQLBuilderError
 from sqlspec.utils.logging import get_logger
-from sqlspec.utils.type_guards import has_expression_and_parameters, has_sql_method, has_with_method
+from sqlspec.utils.type_guards import has_expression_and_parameters, has_sql_method, has_with_method, is_expression
 
 if TYPE_CHECKING:
     from sqlspec.core.result import SQLResult
 
 __all__ = ("QueryBuilder", "SafeQuery")
 
+MAX_PARAMETER_COLLISION_ATTEMPTS = 1000
+
 logger = get_logger(__name__)
 
 
@@ -34,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
@@ -64,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,
@@ -78,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] = {}
@@ -91,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:
@@ -104,13 +106,9 @@ class QueryBuilder(ABC):
 
         Args:
             expression: SQLGlot expression to set
-
-        Raises:
-            TypeError: If expression is not a SQLGlot Expression
         """
-        if not isinstance(expression, exp.Expression):
-            msg = f"Expected Expression, got {type(expression)}"
-            raise TypeError(msg)
+        if not is_expression(expression):
+            self._raise_invalid_expression_type(expression)
         self._expression = expression
 
     def has_expression(self) -> bool:
@@ -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:
@@ -151,7 +149,121 @@ class QueryBuilder(ABC):
         """
         raise SQLBuilderError(message) from cause
 
-    def _add_parameter(self, value: Any, context: Optional[str] = None) -> str:
+    @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 _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:
@@ -192,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
@@ -229,13 +341,11 @@ class QueryBuilder(ABC):
         if base_name not in self._parameters:
             return base_name
 
-        for i in range(1, 1000):
+        for i in range(1, MAX_PARAMETER_COLLISION_ATTEMPTS):
             name = f"{base_name}_{i}"
             if name not in self._parameters:
                 return name
 
-        import uuid
-
         return f"{base_name}_{uuid.uuid4().hex[:8]}"
 
     def _merge_cte_parameters(self, cte_name: str, parameters: dict[str, Any]) -> dict[str, str]:
@@ -275,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:
@@ -284,8 +394,6 @@ class QueryBuilder(ABC):
         Returns:
             A unique cache key representing the builder state and configuration
         """
-        import hashlib
-
         dialect_name: str = self.dialect_name or "default"
 
         if self._expression is None:
@@ -320,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:
@@ -334,41 +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_sql_builder_error("CTE query builder has no expression.")
-            if not isinstance(query_expr, exp.Select):
-                msg = f"CTE query builder expression must be a Select, got {type(query_expr).__name__}."
-                self._raise_sql_builder_error(msg)
-            cte_select_expression = query_expr
-            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):
-                msg = f"Updated CTE expression must be a Select, got {type(updated_expression).__name__}."
-                self._raise_sql_builder_error(msg)
-            cte_select_expression = updated_expression
-
-        elif isinstance(query, str):
-            try:
-                parsed_expression = sqlglot.parse_one(query, read=self.dialect_name)
-                if not isinstance(parsed_expression, exp.Select):
-                    msg = f"CTE query string must parse to a SELECT statement, got {type(parsed_expression).__name__}."
-                    self._raise_sql_builder_error(msg)
-                cte_select_expression = parsed_expression
-            except SQLGlotParseError as e:
-                self._raise_sql_builder_error(f"Failed to parse CTE query string: {e!s}", e)
-            except Exception as e:
-                msg = f"An unexpected error occurred while parsing CTE query string: {e!s}"
-                self._raise_sql_builder_error(msg, e)
-        elif isinstance(query, exp.Select):
-            cte_select_expression = query
-        else:
-            msg = f"Invalid query type for CTE: {type(query).__name__}"
-            self._raise_sql_builder_error(msg)
-
+        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
 
@@ -378,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)
@@ -438,15 +501,14 @@ class QueryBuilder(ABC):
             optimized = optimize(
                 expression, schema=self.schema, dialect=self.dialect_name, optimizer_settings=optimizer_settings
             )
-
             cache.put("optimized", cache_key, optimized)
-
         except Exception:
+            logger.debug("Expression optimization failed, using original expression")
             return expression
         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:
@@ -471,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:
@@ -482,18 +544,7 @@ class QueryBuilder(ABC):
         """
         safe_query = self.build()
 
-        if isinstance(safe_query.parameters, dict):
-            kwargs = safe_query.parameters
-            parameters: Optional[tuple[Any, ...]] = None
-        else:
-            kwargs = None
-            parameters = (
-                safe_query.parameters
-                if isinstance(safe_query.parameters, tuple)
-                else tuple(safe_query.parameters)
-                if safe_query.parameters
-                else None
-            )
+        kwargs, parameters = self._extract_statement_parameters(safe_query.parameters)
 
         if config is None:
             config = StatementConfig(
@@ -521,6 +572,28 @@ class QueryBuilder(ABC):
             return SQL(sql_string, *parameters, statement_config=config)
         return SQL(sql_string, statement_config=config)
 
+    def _extract_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 SafeQuery
+
+        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 __str__(self) -> str:
         """Return the SQL string representation of the query.
 
@@ -530,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