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

sqlspec/statement/builder/select.py

@@ -0,0 +1,206 @@
+"""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.
+"""
+
+import re
+from dataclasses import dataclass, field
+from typing import Optional, Union, cast
+
+from sqlglot import exp
+from typing_extensions import Self
+
+from sqlspec.statement.builder.base import QueryBuilder, SafeQuery
+from sqlspec.statement.builder.mixins import (
+    AggregateFunctionsMixin,
+    CaseBuilderMixin,
+    CommonTableExpressionMixin,
+    FromClauseMixin,
+    GroupByClauseMixin,
+    HavingClauseMixin,
+    JoinClauseMixin,
+    LimitOffsetClauseMixin,
+    OrderByClauseMixin,
+    PivotClauseMixin,
+    SelectColumnsMixin,
+    SetOperationMixin,
+    UnpivotClauseMixin,
+    WhereClauseMixin,
+    WindowFunctionsMixin,
+)
+from sqlspec.statement.result import SQLResult
+from sqlspec.typing import RowT
+
+__all__ = ("SelectBuilder",)
+
+
+@dataclass
+class SelectBuilder(
+    QueryBuilder[RowT],
+    WhereClauseMixin,
+    OrderByClauseMixin,
+    LimitOffsetClauseMixin,
+    SelectColumnsMixin,
+    JoinClauseMixin,
+    FromClauseMixin,
+    GroupByClauseMixin,
+    HavingClauseMixin,
+    SetOperationMixin,
+    CommonTableExpressionMixin,
+    AggregateFunctionsMixin,
+    WindowFunctionsMixin,
+    CaseBuilderMixin,
+    PivotClauseMixin,
+    UnpivotClauseMixin,
+):
+    """Type-safe builder for SELECT queries with schema/model integration.
+
+    This builder provides a fluent, safe interface for constructing SQL SELECT statements.
+    It supports type-safe result mapping via the `as_schema()` method, allowing users to
+    associate a schema/model (such as a Pydantic model, dataclass, or msgspec.Struct) with
+    the query for static type checking and IDE support.
+
+    Example:
+        >>> class User(BaseModel):
+        ...     id: int
+        ...     name: str
+        >>> builder = (
+        ...     SelectBuilder()
+        ...     .select("id", "name")
+        ...     .from_("users")
+        ...     .as_schema(User)
+        ... )
+        >>> result: list[User] = driver.execute(builder)
+
+    Attributes:
+        _schema: The schema/model class for row typing, if set via as_schema().
+    """
+
+    _with_parts: "dict[str, Union[exp.CTE, SelectBuilder]]" = field(default_factory=dict, init=False)
+    _expression: Optional[exp.Expression] = field(default=None, init=False, repr=False, compare=False, hash=False)
+    _schema: Optional[type[RowT]] = None
+    _hints: "list[dict[str, object]]" = field(default_factory=list, init=False, repr=False)
+
+    def __post_init__(self) -> "None":
+        super().__post_init__()
+        if self._expression is None:
+            self._create_base_expression()
+
+    @property
+    def _expected_result_type(self) -> "type[SQLResult[RowT]]":
+        """Get the expected result type for SELECT operations.
+
+        Returns:
+            type: The SelectResult type.
+        """
+        return SQLResult[RowT]
+
+    def _create_base_expression(self) -> "exp.Select":
+        if self._expression is None or not isinstance(self._expression, exp.Select):
+            self._expression = exp.Select()
+        # At this point, self._expression is exp.Select
+        return self._expression
+
+    def as_schema(self, schema: "type[RowT]") -> "SelectBuilder[RowT]":
+        """Return a new SelectBuilder instance parameterized with the given schema/model type.
+
+        This enables type-safe result mapping: the returned builder will carry the schema type
+        for static analysis and IDE autocompletion. The schema should be a class such as a Pydantic
+        model, dataclass, or msgspec.Struct that describes the expected row shape.
+
+        Args:
+            schema: The schema/model class to use for row typing (e.g., a Pydantic model, dataclass, or msgspec.Struct).
+
+        Returns:
+            SelectBuilder[RowT]: A new SelectBuilder instance with RowT set to the provided schema/model type.
+        """
+        new_builder = SelectBuilder()
+        new_builder._expression = self._expression.copy() if self._expression is not None else None
+        new_builder._parameters = self._parameters.copy()
+        new_builder._parameter_counter = self._parameter_counter
+        new_builder.dialect = self.dialect
+        new_builder._schema = schema  # type: ignore[assignment]
+        return cast("SelectBuilder[RowT]", new_builder)
+
+    def with_hint(
+        self,
+        hint: "str",
+        *,
+        location: "str" = "statement",
+        table: "Optional[str]" = None,
+        dialect: "Optional[str]" = None,
+    ) -> "Self":
+        """Attach an optimizer or dialect-specific hint to the query.
+
+        Args:
+            hint: The raw hint string (e.g., 'INDEX(users idx_users_name)').
+            location: Where to apply the hint ('statement', 'table').
+            table: Table name if the hint is for a specific table.
+            dialect: Restrict the hint to a specific dialect (optional).
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+        self._hints.append({"hint": hint, "location": location, "table": table, "dialect": dialect})
+        return self
+
+    def build(self) -> "SafeQuery":
+        """Builds the SQL query string and parameters with hint injection.
+
+        Returns:
+            SafeQuery: A dataclass containing the SQL string and parameters.
+        """
+        # Call parent build method which handles CTEs and optimization
+        safe_query = super().build()
+
+        # Apply hints using SQLGlot's proper hint support (more robust than regex)
+        if hasattr(self, "_hints") and self._hints:
+            modified_expr = self._expression.copy() if self._expression else None
+
+            if modified_expr and isinstance(modified_expr, exp.Select):
+                # Apply statement-level hints using SQLGlot's Hint expression
+                statement_hints = [h["hint"] for h in self._hints if h.get("location") == "statement"]
+                if statement_hints:
+                    # Parse each hint and create proper hint expressions
+                    hint_expressions = []
+                    for hint in statement_hints:
+                        try:
+                            # Try to parse hint as an expression (e.g., "INDEX(users idx_name)")
+                            hint_str = str(hint)  # Ensure hint is a string
+                            hint_expr: Optional[exp.Expression] = exp.maybe_parse(hint_str, dialect=self.dialect_name)
+                            if hint_expr:
+                                hint_expressions.append(hint_expr)
+                            else:
+                                # Create a raw identifier for unparsable hints
+                                hint_expressions.append(exp.Anonymous(this=hint_str))
+                        except Exception:  # noqa: PERF203
+                            hint_expressions.append(exp.Anonymous(this=str(hint)))
+
+                    # Create a Hint node and attach to SELECT
+                    if hint_expressions:
+                        hint_node = exp.Hint(expressions=hint_expressions)
+                        modified_expr.set("hint", hint_node)
+
+                # For table-level hints, we'll fall back to comment injection in SQL
+                # since SQLGlot doesn't have a standard way to attach hints to individual tables
+                modified_sql = modified_expr.sql(dialect=self.dialect_name, pretty=True)
+
+                # Apply table-level hints via string manipulation (as fallback)
+                table_hints = [h for h in self._hints if h.get("location") == "table" and h.get("table")]
+                if table_hints:
+                    for th in table_hints:
+                        table = str(th["table"])
+                        hint = th["hint"]
+                        # More precise regex that captures the table and optional alias
+                        pattern = rf"\b{re.escape(table)}\b(\s+AS\s+\w+)?"
+
+                        def replacement_func(match: re.Match[str]) -> str:
+                            alias_part = match.group(1) or ""
+                            return f"/*+ {hint} */ {table}{alias_part}"  # noqa: B023
+
+                        modified_sql = re.sub(pattern, replacement_func, modified_sql, flags=re.IGNORECASE, count=1)
+
+                return SafeQuery(sql=modified_sql, parameters=safe_query.parameters, dialect=safe_query.dialect)
+
+        return safe_query

sqlspec/statement/builder/update.py

@@ -0,0 +1,178 @@
+"""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.
+"""
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional, Union
+
+from sqlglot import exp
+from typing_extensions import Self
+
+from sqlspec.exceptions import SQLBuilderError
+from sqlspec.statement.builder.base import QueryBuilder, SafeQuery
+from sqlspec.statement.builder.mixins import (
+    ReturningClauseMixin,
+    UpdateFromClauseMixin,
+    UpdateSetClauseMixin,
+    UpdateTableClauseMixin,
+    WhereClauseMixin,
+)
+from sqlspec.statement.result import SQLResult
+from sqlspec.typing import RowT
+
+if TYPE_CHECKING:
+    from sqlspec.statement.builder.select import SelectBuilder
+
+__all__ = ("UpdateBuilder",)
+
+
+@dataclass(unsafe_hash=True)
+class UpdateBuilder(
+    QueryBuilder[RowT],
+    WhereClauseMixin,
+    ReturningClauseMixin,
+    UpdateSetClauseMixin,
+    UpdateFromClauseMixin,
+    UpdateTableClauseMixin,
+):
+    """Builder for UPDATE statements.
+
+    This builder provides a fluent interface for constructing SQL UPDATE statements
+    with automatic parameter binding and validation.
+
+    Example:
+        ```python
+        # Basic UPDATE
+        update_query = (
+            UpdateBuilder()
+            .table("users")
+            .set(name="John Doe")
+            .set(email="john@example.com")
+            .where("id = 1")
+        )
+
+        # UPDATE with parameterized conditions
+        update_query = (
+            UpdateBuilder()
+            .table("users")
+            .set(status="active")
+            .where_eq("id", 123)
+        )
+
+        # UPDATE with FROM clause (PostgreSQL style)
+        update_query = (
+            UpdateBuilder()
+            .table("users", "u")
+            .set(name="Updated Name")
+            .from_("profiles", "p")
+            .where("u.id = p.user_id AND p.is_verified = true")
+        )
+        ```
+    """
+
+    @property
+    def _expected_result_type(self) -> "type[SQLResult[RowT]]":
+        """Return the expected result type for this builder."""
+        return SQLResult[RowT]
+
+    def _create_base_expression(self) -> exp.Update:
+        """Create a base UPDATE expression.
+
+        Returns:
+            A new sqlglot Update expression with empty clauses.
+        """
+        return exp.Update(this=None, expressions=[], joins=[])
+
+    def join(
+        self,
+        table: "Union[str, exp.Expression, SelectBuilder[RowT]]",
+        on: "Union[str, exp.Expression]",
+        alias: "Optional[str]" = None,
+        join_type: str = "INNER",
+    ) -> "Self":
+        """Add JOIN clause to the UPDATE statement.
+
+        Args:
+            table: The table name, expression, or subquery to join.
+            on: The JOIN condition.
+            alias: Optional alias for the joined table.
+            join_type: Type of join (INNER, LEFT, RIGHT, FULL).
+
+        Returns:
+            The current builder instance for method chaining.
+
+        Raises:
+            SQLBuilderError: If the current expression is not an UPDATE statement.
+        """
+        if self._expression is None or not isinstance(self._expression, exp.Update):
+            msg = "Cannot add JOIN clause to non-UPDATE expression."
+            raise SQLBuilderError(msg)
+
+        table_expr: exp.Expression
+        if isinstance(table, str):
+            table_expr = exp.table_(table, alias=alias)
+        elif isinstance(table, QueryBuilder):
+            subquery = table.build()
+            subquery_exp = exp.paren(exp.maybe_parse(subquery.sql, dialect=self.dialect))
+            table_expr = exp.alias_(subquery_exp, alias) if alias else subquery_exp
+
+            # Merge parameters
+            subquery_params = table._parameters
+            if subquery_params:
+                for p_name, p_value in subquery_params.items():
+                    self.add_parameter(p_value, name=p_name)
+        else:
+            table_expr = table
+
+        on_expr: exp.Expression = exp.condition(on) if isinstance(on, str) else on
+
+        join_type_upper = join_type.upper()
+        if join_type_upper == "INNER":
+            join_expr = exp.Join(this=table_expr, on=on_expr)
+        elif join_type_upper == "LEFT":
+            join_expr = exp.Join(this=table_expr, on=on_expr, side="LEFT")
+        elif join_type_upper == "RIGHT":
+            join_expr = exp.Join(this=table_expr, on=on_expr, side="RIGHT")
+        elif join_type_upper == "FULL":
+            join_expr = exp.Join(this=table_expr, on=on_expr, side="FULL", kind="OUTER")
+        else:
+            msg = f"Unsupported join type: {join_type}"
+            raise SQLBuilderError(msg)
+
+        # Add join to the UPDATE expression
+        if not self._expression.args.get("joins"):
+            self._expression.set("joins", [])
+        self._expression.args["joins"].append(join_expr)
+
+        return self
+
+    def build(self) -> "SafeQuery":
+        """Build the UPDATE query with validation.
+
+        Returns:
+            SafeQuery: The built query with SQL and parameters.
+
+        Raises:
+            SQLBuilderError: If no table is set or expression is not an UPDATE.
+        """
+        if self._expression is None:
+            msg = "UPDATE expression not initialized."
+            raise SQLBuilderError(msg)
+
+        if not isinstance(self._expression, exp.Update):
+            msg = "No UPDATE expression to build or expression is of the wrong type."
+            raise SQLBuilderError(msg)
+
+        # Check that the table is set
+        if getattr(self._expression, "this", None) is None:
+            msg = "No table specified for UPDATE statement."
+            raise SQLBuilderError(msg)
+
+        # Check that at least one SET expression exists
+        if not self._expression.args.get("expressions"):
+            msg = "At least one SET clause must be specified for UPDATE statement."
+            raise SQLBuilderError(msg)
+
+        return super().build()