sqlspec 0.11.1__py3-none-any.whl → 0.12.0__py3-none-any.whl

sqlspec/statement/builder/_ddl_utils.py

@@ -0,0 +1,119 @@
+"""DDL builder utilities."""
+
+from typing import TYPE_CHECKING, Optional
+
+from sqlglot import exp
+
+if TYPE_CHECKING:
+    from sqlspec.statement.builder.ddl import ColumnDefinition, ConstraintDefinition
+
+__all__ = ("build_column_expression", "build_constraint_expression")
+
+
+def build_column_expression(col: "ColumnDefinition") -> "exp.Expression":
+    """Build SQLGlot expression for a column definition."""
+    # Start with column name and type
+    col_def = exp.ColumnDef(this=exp.to_identifier(col.name), kind=exp.DataType.build(col.dtype))
+
+    # Add constraints
+    constraints: list[exp.ColumnConstraint] = []
+
+    if col.not_null:
+        constraints.append(exp.ColumnConstraint(kind=exp.NotNullColumnConstraint()))
+
+    if col.primary_key:
+        constraints.append(exp.ColumnConstraint(kind=exp.PrimaryKeyColumnConstraint()))
+
+    if col.unique:
+        constraints.append(exp.ColumnConstraint(kind=exp.UniqueColumnConstraint()))
+
+    if col.default is not None:
+        # Handle different default value types
+        default_expr: Optional[exp.Expression] = None
+        if isinstance(col.default, str):
+            # Check if it's a function/expression or a literal string
+            if col.default.upper() in {"CURRENT_TIMESTAMP", "CURRENT_DATE", "CURRENT_TIME"} or "(" in col.default:
+                default_expr = exp.maybe_parse(col.default)
+            else:
+                default_expr = exp.Literal.string(col.default)
+        elif isinstance(col.default, (int, float)):
+            default_expr = exp.Literal.number(col.default)
+        elif col.default is True:
+            default_expr = exp.true()
+        elif col.default is False:
+            default_expr = exp.false()
+        else:
+            default_expr = exp.Literal.string(str(col.default))
+
+        constraints.append(exp.ColumnConstraint(kind=default_expr))
+
+    if col.check:
+        check_expr = exp.Check(this=exp.maybe_parse(col.check))
+        constraints.append(exp.ColumnConstraint(kind=check_expr))
+
+    if col.comment:
+        constraints.append(exp.ColumnConstraint(kind=exp.CommentColumnConstraint(this=exp.Literal.string(col.comment))))
+
+    if col.generated:
+        # Handle generated columns (computed columns)
+        generated_expr = exp.GeneratedAsIdentityColumnConstraint(this=exp.maybe_parse(col.generated))
+        constraints.append(exp.ColumnConstraint(kind=generated_expr))
+
+    if col.collate:
+        constraints.append(exp.ColumnConstraint(kind=exp.CollateColumnConstraint(this=exp.to_identifier(col.collate))))
+
+    # Set constraints on column definition
+    if constraints:
+        col_def.set("constraints", constraints)
+
+    return col_def
+
+
+def build_constraint_expression(constraint: "ConstraintDefinition") -> "Optional[exp.Expression]":
+    """Build SQLGlot expression for a table constraint."""
+    if constraint.constraint_type == "PRIMARY KEY":
+        # Build primary key constraint
+        pk_cols = [exp.to_identifier(col) for col in constraint.columns]
+        pk_constraint = exp.PrimaryKey(expressions=pk_cols)
+
+        if constraint.name:
+            return exp.Constraint(this=exp.to_identifier(constraint.name), expression=pk_constraint)
+        return pk_constraint
+
+    if constraint.constraint_type == "FOREIGN KEY":
+        # Build foreign key constraint
+        fk_cols = [exp.to_identifier(col) for col in constraint.columns]
+        ref_cols = [exp.to_identifier(col) for col in constraint.references_columns]
+
+        fk_constraint = exp.ForeignKey(
+            expressions=fk_cols,
+            reference=exp.Reference(
+                this=exp.to_table(constraint.references_table) if constraint.references_table else None,
+                expressions=ref_cols,
+                on_delete=constraint.on_delete,
+                on_update=constraint.on_update,
+            ),
+        )
+
+        if constraint.name:
+            return exp.Constraint(this=exp.to_identifier(constraint.name), expression=fk_constraint)
+        return fk_constraint
+
+    if constraint.constraint_type == "UNIQUE":
+        # Build unique constraint
+        unique_cols = [exp.to_identifier(col) for col in constraint.columns]
+        unique_constraint = exp.UniqueKeyProperty(expressions=unique_cols)
+
+        if constraint.name:
+            return exp.Constraint(this=exp.to_identifier(constraint.name), expression=unique_constraint)
+        return unique_constraint
+
+    if constraint.constraint_type == "CHECK":
+        # Build check constraint
+        check_expr = exp.Check(this=exp.maybe_parse(constraint.condition) if constraint.condition else None)
+
+        if constraint.name:
+            return exp.Constraint(this=exp.to_identifier(constraint.name), expression=check_expr)
+        return check_expr
+
+    return None

sqlspec/statement/builder/_parsing_utils.py

@@ -0,0 +1,135 @@
+"""Centralized parsing utilities for SQLSpec builders.
+
+This module provides common parsing functions to handle complex SQL expressions
+that users might pass as strings to various builder methods.
+"""
+
+import contextlib
+from typing import Any, Optional, Union, cast
+
+from sqlglot import exp, maybe_parse, parse_one
+
+
+def parse_column_expression(column_input: Union[str, exp.Expression]) -> exp.Expression:
+    """Parse a column input that might be a complex expression.
+
+    Handles cases like:
+    - Simple column names: "name" -> Column(this=name)
+    - Qualified names: "users.name" -> Column(table=users, this=name)
+    - Aliased columns: "name AS user_name" -> Alias(this=Column(name), alias=user_name)
+    - Function calls: "MAX(price)" -> Max(this=Column(price))
+    - Complex expressions: "CASE WHEN ... END" -> Case(...)
+
+    Args:
+        column_input: String or SQLGlot expression representing a column/expression
+
+    Returns:
+        exp.Expression: Parsed SQLGlot expression
+    """
+    if isinstance(column_input, exp.Expression):
+        return column_input
+    return exp.maybe_parse(column_input) or exp.column(str(column_input))
+
+
+def parse_table_expression(table_input: str, explicit_alias: Optional[str] = None) -> exp.Expression:
+    """Parses a table string that can be a name, a name with an alias, or a subquery string."""
+    with contextlib.suppress(Exception):
+        # Wrapping in a SELECT statement is a robust way to parse various table-like syntaxes
+        parsed = parse_one(f"SELECT * FROM {table_input}")
+        if isinstance(parsed, exp.Select) and parsed.args.get("from"):
+            from_clause = cast("exp.From", parsed.args.get("from"))
+            table_expr = from_clause.this
+
+            if explicit_alias:
+                return exp.alias_(table_expr, explicit_alias)  # type:ignore[no-any-return]
+            return table_expr  # type:ignore[no-any-return]
+
+    return exp.to_table(table_input, alias=explicit_alias)
+
+
+def parse_order_expression(order_input: Union[str, exp.Expression]) -> exp.Expression:
+    """Parse an ORDER BY expression that might include direction.
+
+    Handles cases like:
+    - Simple column: "name" -> Column(this=name)
+    - With direction: "name DESC" -> Ordered(this=Column(name), desc=True)
+    - Qualified: "users.name ASC" -> Ordered(this=Column(table=users, this=name), desc=False)
+    - Function: "COUNT(*) DESC" -> Ordered(this=Count(this=Star), desc=True)
+
+    Args:
+        order_input: String or SQLGlot expression for ORDER BY
+
+    Returns:
+        exp.Expression: Parsed SQLGlot expression (usually Ordered or Column)
+    """
+    if isinstance(order_input, exp.Expression):
+        return order_input
+
+    with contextlib.suppress(Exception):
+        parsed = maybe_parse(str(order_input), into=exp.Ordered)
+        if parsed:
+            return parsed
+
+    return parse_column_expression(order_input)
+
+
+def parse_condition_expression(
+    condition_input: Union[str, exp.Expression, tuple[str, Any]], builder: "Any" = None
+) -> exp.Expression:
+    """Parse a condition that might be complex SQL.
+
+    Handles cases like:
+    - Simple conditions: "name = 'John'" -> EQ(Column(name), Literal('John'))
+    - Tuple format: ("name", "John") -> EQ(Column(name), Literal('John'))
+    - Complex conditions: "age > 18 AND status = 'active'" -> And(GT(...), EQ(...))
+    - Function conditions: "LENGTH(name) > 5" -> GT(Length(Column(name)), Literal(5))
+
+    Args:
+        condition_input: String, tuple, or SQLGlot expression for condition
+        builder: Optional builder instance for parameter binding
+
+    Returns:
+        exp.Expression: Parsed SQLGlot expression (usually a comparison or logical op)
+    """
+    if isinstance(condition_input, exp.Expression):
+        return condition_input
+
+    tuple_condition_parts = 2
+    if isinstance(condition_input, tuple) and len(condition_input) == tuple_condition_parts:
+        # Handle (column, value) tuple format with proper parameter binding
+        column, value = condition_input
+        column_expr = parse_column_expression(column)
+        if value is None:
+            return exp.Is(this=column_expr, expression=exp.null())
+        # Use builder's parameter system if available
+        if builder and hasattr(builder, "add_parameter"):
+            _, param_name = builder.add_parameter(value)
+            return exp.EQ(this=column_expr, expression=exp.Placeholder(this=param_name))
+        # Fallback to literal value
+        if isinstance(value, str):
+            return exp.EQ(this=column_expr, expression=exp.Literal.string(value))
+        if isinstance(value, (int, float)):
+            return exp.EQ(this=column_expr, expression=exp.Literal.number(str(value)))
+        return exp.EQ(this=column_expr, expression=exp.Literal.string(str(value)))
+
+    if not isinstance(condition_input, str):
+        condition_input = str(condition_input)
+
+    try:
+        # Parse as condition using SQLGlot's condition parser
+        return exp.condition(condition_input)
+    except Exception:
+        # If that fails, try parsing as a general expression
+        try:
+            parsed = exp.maybe_parse(condition_input)  # type: ignore[var-annotated]
+            if parsed:
+                return parsed  # type:ignore[no-any-return]
+        except Exception:  # noqa: S110
+            # SQLGlot condition parsing failed, will use raw condition
+            pass
+
+    # Ultimate fallback: treat as raw condition string
+    return exp.condition(condition_input)
+
+
+__all__ = ("parse_column_expression", "parse_condition_expression", "parse_order_expression", "parse_table_expression")

sqlspec/statement/builder/base.py

@@ -0,0 +1,328 @@
+"""Safe SQL query builder with validation and parameter binding.
+
+This module provides a fluent interface for building SQL queries safely,
+with automatic parameter binding and validation. Enhanced with SQLGlot's
+advanced builder patterns and optimization capabilities.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Generic, NoReturn, Optional, Union
+
+import sqlglot
+from sqlglot import Dialect, exp
+from sqlglot.dialects.dialect import DialectType
+from sqlglot.errors import ParseError as SQLGlotParseError
+from sqlglot.optimizer import optimize
+from typing_extensions import Self
+
+from sqlspec.exceptions import SQLBuilderError
+from sqlspec.statement.sql import SQL, SQLConfig
+from sqlspec.typing import RowT
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from sqlspec.statement.result import SQLResult
+
+__all__ = ("QueryBuilder", "SafeQuery")
+
+logger = get_logger(__name__)
+
+
+@dataclass(frozen=True)
+class SafeQuery:
+    """A safely constructed SQL query with bound parameters."""
+
+    sql: str
+    parameters: dict[str, Any] = field(default_factory=dict)
+    dialect: DialectType = field(default=None)
+
+
+@dataclass
+class QueryBuilder(ABC, Generic[RowT]):
+    """Abstract base class for SQL query builders with SQLGlot optimization.
+
+    Provides common functionality for dialect handling, parameter management,
+    query construction, and automatic query optimization using SQLGlot's
+    advanced capabilities.
+
+    New features:
+    - Automatic query optimization (join reordering, predicate pushdown)
+    - Query complexity analysis
+    - Smart parameter naming based on context
+    - Expression caching for performance
+    """
+
+    dialect: DialectType = field(default=None)
+    schema: Optional[dict[str, dict[str, str]]] = field(default=None)
+    _expression: Optional[exp.Expression] = field(default=None, init=False, repr=False, compare=False, hash=False)
+    _parameters: dict[str, Any] = field(default_factory=dict, init=False, repr=False, compare=False, hash=False)
+    _parameter_counter: int = field(default=0, init=False, repr=False, compare=False, hash=False)
+    _with_ctes: dict[str, exp.CTE] = field(default_factory=dict, init=False, repr=False, compare=False, hash=False)
+    enable_optimization: bool = field(default=True, init=True)
+    optimize_joins: bool = field(default=True, init=True)
+    optimize_predicates: bool = field(default=True, init=True)
+    simplify_expressions: bool = field(default=True, init=True)
+
+    def __post_init__(self) -> None:
+        self._expression = self._create_base_expression()
+        if not self._expression:
+            # This path should be unreachable if _raise_sql_builder_error has NoReturn
+            self._raise_sql_builder_error(
+                "QueryBuilder._create_base_expression must return a valid sqlglot expression."
+            )
+
+    @abstractmethod
+    def _create_base_expression(self) -> exp.Expression:
+        """Create the base sqlglot expression for the specific query type.
+
+        Examples:
+            For a SELECT query, this would return `exp.Select()`.
+            For an INSERT query, this would return `exp.Insert()`.
+
+        Returns:
+            exp.Expression: A new sqlglot expression.
+        """
+
+    @property
+    @abstractmethod
+    def _expected_result_type(self) -> "type[SQLResult[RowT]]":
+        """The expected result type for the query being built.
+
+        Returns:
+            type[ResultT]: The type of the result.
+        """
+
+    @staticmethod
+    def _raise_sql_builder_error(message: str, cause: Optional[BaseException] = 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
+
+    def _add_parameter(self, value: Any, context: Optional[str] = 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 (e.g., "where", "join")
+
+        Returns:
+            str: The placeholder name for the parameter (e.g., :param_1 or :where_param_1).
+        """
+        self._parameter_counter += 1
+
+        # Use context-aware naming if provided
+        param_name = f"{context}_param_{self._parameter_counter}" if context else f"param_{self._parameter_counter}"
+
+        self._parameters[param_name] = value
+        return param_name
+
+    def add_parameter(self: Self, value: Any, name: Optional[str] = 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_sql_builder_error(f"Parameter name '{name}' already exists.")
+            param_name_to_use = name
+        else:
+            self._parameter_counter += 1
+            param_name_to_use = f"param_{self._parameter_counter}"
+
+        self._parameters[param_name_to_use] = value
+        return self, param_name_to_use
+
+    def _generate_unique_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
+        """
+        if base_name not in self._parameters:
+            return base_name
+
+        i = 1
+        while True:
+            name = f"{base_name}_{i}"
+            if name not in self._parameters:
+                return name
+            i += 1
+
+    def with_cte(self: Self, alias: str, query: "Union[QueryBuilder[Any], exp.Select, str]") -> 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 expression.
+
+        Returns:
+            Self: The current builder instance for method chaining.
+        """
+        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):
+            if query._expression is None:
+                self._raise_sql_builder_error("CTE query builder has no expression.")
+            if not isinstance(query._expression, exp.Select):
+                msg = f"CTE query builder expression must be a Select, got {type(query._expression).__name__}."
+                self._raise_sql_builder_error(msg)
+            cte_select_expression = query._expression.copy()
+            for p_name, p_value in query._parameters.items():
+                self.add_parameter(p_value, f"cte_{alias}_{p_name}")
+
+        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)
+                # parsed_expression is now known to be exp.Select
+                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.copy()
+        else:
+            msg = f"Invalid query type for CTE: {type(query).__name__}"
+            self._raise_sql_builder_error(msg)
+            return self  # This line won't be reached but satisfies type checkers
+
+        self._with_ctes[alias] = exp.CTE(this=cte_select_expression, alias=exp.to_table(alias))
+        return self
+
+    def build(self) -> "SafeQuery":
+        """Builds the SQL query string and parameters.
+
+        Returns:
+            SafeQuery: A dataclass containing the SQL string and parameters.
+        """
+        if self._expression is None:
+            self._raise_sql_builder_error("QueryBuilder expression not initialized.")
+
+        final_expression = self._expression.copy()
+
+        if self._with_ctes:
+            if hasattr(final_expression, "with_") and callable(getattr(final_expression, "with_", None)):
+                for alias, cte_node in self._with_ctes.items():
+                    final_expression = final_expression.with_(  # pyright: ignore
+                        cte_node.args["this"], as_=alias, copy=False
+                    )
+            elif (
+                isinstance(final_expression, (exp.Select, exp.Insert, exp.Update, exp.Delete, exp.Union))
+                and self._with_ctes
+            ):
+                final_expression = exp.With(expressions=list(self._with_ctes.values()), this=final_expression)
+
+        # Apply SQLGlot optimizations if enabled
+        if self.enable_optimization:
+            final_expression = self._optimize_expression(final_expression)
+
+        try:
+            sql_string = final_expression.sql(dialect=self.dialect_name, pretty=True)
+        except Exception as e:
+            err_msg = f"Error generating SQL from expression: {e!s}"
+            logger.exception("SQL generation failed")
+            self._raise_sql_builder_error(err_msg, e)
+
+        return SafeQuery(sql=sql_string, parameters=self._parameters.copy(), dialect=self.dialect)
+
+    def _optimize_expression(self, expression: exp.Expression) -> exp.Expression:
+        """Apply SQLGlot optimizations to the expression.
+
+        Args:
+            expression: The expression to optimize
+
+        Returns:
+            The optimized expression
+        """
+        if not self.enable_optimization:
+            return expression
+
+        try:
+            # Use SQLGlot's comprehensive optimizer
+            return optimize(
+                expression.copy(),
+                schema=self.schema,
+                dialect=self.dialect_name,
+                optimizer_settings={
+                    "optimize_joins": self.optimize_joins,
+                    "pushdown_predicates": self.optimize_predicates,
+                    "simplify_expressions": self.simplify_expressions,
+                },
+            )
+        except Exception:
+            # Continue with unoptimized query on failure
+            return expression
+
+    def to_statement(self, config: "Optional[SQLConfig]" = None) -> "SQL":
+        """Converts the built query into a SQL statement object.
+
+        Args:
+            config: Optional SQL configuration.
+
+        Returns:
+            SQL: A SQL statement object.
+        """
+        safe_query = self.build()
+
+        return SQL(
+            statement=safe_query.sql,
+            parameters=safe_query.parameters,
+            _dialect=safe_query.dialect,
+            _config=config,
+            _builder_result_type=self._expected_result_type,
+        )
+
+    def __str__(self) -> str:
+        """Return the SQL string representation of the query.
+
+        Returns:
+            str: The SQL string for this query.
+        """
+        try:
+            return self.build().sql
+        except Exception:
+            # Fallback to default representation if build fails
+            return super().__str__()
+
+    @property
+    def dialect_name(self) -> "Optional[str]":
+        """Returns the name of the dialect, if set."""
+        if isinstance(self.dialect, str):
+            return self.dialect
+        if self.dialect is not None:
+            if isinstance(self.dialect, type) and issubclass(self.dialect, Dialect):
+                return self.dialect.__name__.lower()
+            if isinstance(self.dialect, Dialect):
+                return type(self.dialect).__name__.lower()
+            # Handle case where dialect might have a __name__ attribute
+            if hasattr(self.dialect, "__name__"):
+                return self.dialect.__name__.lower()
+        return None