sqlspec 0.16.0__cp312-cp312-win_amd64.whl

sqlspec/builder/_select.py

@@ -0,0 +1,174 @@
+"""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 Any, Optional, Union
+
+from sqlglot import exp
+from typing_extensions import Self
+
+from sqlspec.builder._base import QueryBuilder, SafeQuery
+from sqlspec.builder.mixins import (
+    CommonTableExpressionMixin,
+    HavingClauseMixin,
+    JoinClauseMixin,
+    LimitOffsetClauseMixin,
+    OrderByClauseMixin,
+    PivotClauseMixin,
+    SelectClauseMixin,
+    SetOperationMixin,
+    UnpivotClauseMixin,
+    WhereClauseMixin,
+)
+from sqlspec.core.result import SQLResult
+
+__all__ = ("Select",)
+
+
+TABLE_HINT_PATTERN = r"\b{}\b(\s+AS\s+\w+)?"
+
+
+@dataclass
+class Select(
+    QueryBuilder,
+    WhereClauseMixin,
+    OrderByClauseMixin,
+    LimitOffsetClauseMixin,
+    SelectClauseMixin,
+    JoinClauseMixin,
+    HavingClauseMixin,
+    SetOperationMixin,
+    CommonTableExpressionMixin,
+    PivotClauseMixin,
+    UnpivotClauseMixin,
+):
+    """Type-safe builder for SELECT queries with schema/model integration.
+
+    This builder provides a fluent, safe interface for constructing SQL SELECT statements.
+
+    Example:
+        >>> class User(BaseModel):
+        ...     id: int
+        ...     name: str
+        >>> builder = Select("id", "name").from_("users")
+        >>> result = driver.execute(builder)
+    """
+
+    _with_parts: "dict[str, Union[exp.CTE, Select]]" = field(default_factory=dict, init=False)
+    _expression: Optional[exp.Expression] = field(default=None, init=False, repr=False, compare=False, hash=False)
+    _hints: "list[dict[str, object]]" = field(default_factory=list, init=False, repr=False)
+
+    def __init__(self, *columns: str, **kwargs: Any) -> None:
+        """Initialize SELECT with optional columns.
+
+        Args:
+            *columns: Column names to select (e.g., "id", "name", "u.email")
+            **kwargs: Additional QueryBuilder arguments (dialect, schema, etc.)
+
+        Examples:
+            Select("id", "name")  # Shorthand for Select().select("id", "name")
+            Select()              # Same as Select() - start empty
+        """
+        super().__init__(**kwargs)
+
+        self._with_parts = {}
+        self._expression = None
+        self._hints = []
+
+        self._create_base_expression()
+
+        if columns:
+            self.select(*columns)
+
+    @property
+    def _expected_result_type(self) -> "type[SQLResult]":
+        """Get the expected result type for SELECT operations.
+
+        Returns:
+            type: The SelectResult type.
+        """
+        return SQLResult
+
+    def _create_base_expression(self) -> "exp.Select":
+        if self._expression is None or not isinstance(self._expression, exp.Select):
+            self._expression = exp.Select()
+        return self._expression
+
+    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.
+        """
+        safe_query = super().build()
+
+        if not self._hints:
+            return safe_query
+
+        modified_expr = self._expression.copy() if self._expression else self._create_base_expression()
+
+        if isinstance(modified_expr, exp.Select):
+            statement_hints = [h["hint"] for h in self._hints if h.get("location") == "statement"]
+            if statement_hints:
+                hint_expressions = []
+
+                def parse_hint(hint: Any) -> exp.Expression:
+                    """Parse a single hint."""
+                    try:
+                        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:
+                            return hint_expr
+                        return exp.Anonymous(this=hint_str)
+                    except Exception:
+                        return exp.Anonymous(this=str(hint))
+
+                hint_expressions = [parse_hint(hint) for hint in statement_hints]
+
+                if hint_expressions:
+                    hint_node = exp.Hint(expressions=hint_expressions)
+                    modified_expr.set("hint", hint_node)
+
+        modified_sql = modified_expr.sql(dialect=self.dialect_name, pretty=True)
+
+        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"]
+                pattern = TABLE_HINT_PATTERN.format(re.escape(table))
+                compiled_pattern = re.compile(pattern, re.IGNORECASE)
+
+                def replacement_func(match: re.Match[str]) -> str:
+                    alias_part = match.group(1) or ""
+                    return f"/*+ {hint} */ {table}{alias_part}"  # noqa: B023
+
+                modified_sql = compiled_pattern.sub(replacement_func, modified_sql, count=1)
+
+        return SafeQuery(sql=modified_sql, parameters=safe_query.parameters, dialect=safe_query.dialect)

sqlspec/builder/_update.py

@@ -0,0 +1,186 @@
+"""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, Any, Optional, Union
+
+from sqlglot import exp
+from typing_extensions import Self
+
+from sqlspec.builder._base import QueryBuilder, SafeQuery
+from sqlspec.builder.mixins import (
+    ReturningClauseMixin,
+    UpdateFromClauseMixin,
+    UpdateSetClauseMixin,
+    UpdateTableClauseMixin,
+    WhereClauseMixin,
+)
+from sqlspec.core.result import SQLResult
+from sqlspec.exceptions import SQLBuilderError
+
+if TYPE_CHECKING:
+    from sqlspec.builder._select import Select
+
+__all__ = ("Update",)
+
+
+@dataclass(unsafe_hash=True)
+class Update(
+    QueryBuilder,
+    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
+        update_query = (
+            Update()
+            .table("users")
+            .set(name="John Doe")
+            .set(email="john@example.com")
+            .where("id = 1")
+        )
+
+        update_query = (
+            Update("users").set(name="John Doe").where("id = 1")
+        )
+
+        update_query = (
+            Update()
+            .table("users")
+            .set(status="active")
+            .where_eq("id", 123)
+        )
+
+        update_query = (
+            Update()
+            .table("users", "u")
+            .set(name="Updated Name")
+            .from_("profiles", "p")
+            .where("u.id = p.user_id AND p.is_verified = true")
+        )
+        ```
+    """
+
+    def __init__(self, table: Optional[str] = None, **kwargs: Any) -> None:
+        """Initialize UPDATE with optional table.
+
+        Args:
+            table: Target table name
+            **kwargs: Additional QueryBuilder arguments
+        """
+        super().__init__(**kwargs)
+
+        if table:
+            self.table(table)
+
+    @property
+    def _expected_result_type(self) -> "type[SQLResult]":
+        """Return the expected result type for this builder."""
+        return SQLResult
+
+    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, Select]",
+        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
+
+            subquery_parameters = table._parameters
+            if subquery_parameters:
+                for p_name, p_value in subquery_parameters.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)
+
+        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)
+
+        if getattr(self._expression, "this", None) is None:
+            msg = "No table specified for UPDATE statement."
+            raise SQLBuilderError(msg)
+
+        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()

sqlspec/builder/mixins/__init__.py

@@ -0,0 +1,55 @@
+"""SQL statement builder mixins."""
+
+from sqlspec.builder.mixins._cte_and_set_ops import CommonTableExpressionMixin, SetOperationMixin
+from sqlspec.builder.mixins._delete_operations import DeleteFromClauseMixin
+from sqlspec.builder.mixins._insert_operations import InsertFromSelectMixin, InsertIntoClauseMixin, InsertValuesMixin
+from sqlspec.builder.mixins._join_operations import JoinClauseMixin
+from sqlspec.builder.mixins._merge_operations import (
+    MergeIntoClauseMixin,
+    MergeMatchedClauseMixin,
+    MergeNotMatchedBySourceClauseMixin,
+    MergeNotMatchedClauseMixin,
+    MergeOnClauseMixin,
+    MergeUsingClauseMixin,
+)
+from sqlspec.builder.mixins._order_limit_operations import (
+    LimitOffsetClauseMixin,
+    OrderByClauseMixin,
+    ReturningClauseMixin,
+)
+from sqlspec.builder.mixins._pivot_operations import PivotClauseMixin, UnpivotClauseMixin
+from sqlspec.builder.mixins._select_operations import CaseBuilder, SelectClauseMixin
+from sqlspec.builder.mixins._update_operations import (
+    UpdateFromClauseMixin,
+    UpdateSetClauseMixin,
+    UpdateTableClauseMixin,
+)
+from sqlspec.builder.mixins._where_clause import HavingClauseMixin, WhereClauseMixin
+
+__all__ = (
+    "CaseBuilder",
+    "CommonTableExpressionMixin",
+    "DeleteFromClauseMixin",
+    "HavingClauseMixin",
+    "InsertFromSelectMixin",
+    "InsertIntoClauseMixin",
+    "InsertValuesMixin",
+    "JoinClauseMixin",
+    "LimitOffsetClauseMixin",
+    "MergeIntoClauseMixin",
+    "MergeMatchedClauseMixin",
+    "MergeNotMatchedBySourceClauseMixin",
+    "MergeNotMatchedClauseMixin",
+    "MergeOnClauseMixin",
+    "MergeUsingClauseMixin",
+    "OrderByClauseMixin",
+    "PivotClauseMixin",
+    "ReturningClauseMixin",
+    "SelectClauseMixin",
+    "SetOperationMixin",
+    "UnpivotClauseMixin",
+    "UpdateFromClauseMixin",
+    "UpdateSetClauseMixin",
+    "UpdateTableClauseMixin",
+    "WhereClauseMixin",
+)

sqlspec/builder/mixins/_cte_and_set_ops.py

@@ -0,0 +1,195 @@
+"""CTE (Common Table Expression) and Set Operations mixins for SQL builders."""
+
+from typing import Any, Optional, Union
+
+from sqlglot import exp
+from typing_extensions import Self
+
+from sqlspec.exceptions import SQLBuilderError
+
+__all__ = ("CommonTableExpressionMixin", "SetOperationMixin")
+
+
+class CommonTableExpressionMixin:
+    """Mixin providing WITH clause (Common Table Expressions) support for SQL builders."""
+
+    _expression: Optional[exp.Expression] = None
+
+    def with_(
+        self, name: str, query: Union[Any, str], recursive: bool = False, columns: Optional[list[str]] = None
+    ) -> Self:
+        """Add WITH clause (Common Table Expression).
+
+        Args:
+            name: The name of the CTE.
+            query: The query for the CTE (builder instance or SQL string).
+            recursive: Whether this is a recursive CTE.
+            columns: Optional column names for the CTE.
+
+        Raises:
+            SQLBuilderError: If the query type is unsupported.
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+        if self._expression is None:
+            msg = "Cannot add WITH clause: expression not initialized."
+            raise SQLBuilderError(msg)
+
+        if not isinstance(self._expression, (exp.Select, exp.Insert, exp.Update, exp.Delete)):
+            msg = f"Cannot add WITH clause to {type(self._expression).__name__} expression."
+            raise SQLBuilderError(msg)
+
+        cte_expr: Optional[exp.Expression] = None
+        if isinstance(query, str):
+            cte_expr = exp.maybe_parse(query, dialect=self.dialect)  # type: ignore[attr-defined]
+        elif isinstance(query, exp.Expression):
+            cte_expr = query
+        else:
+            built_query = query.to_statement()  # pyright: ignore
+            cte_sql = built_query.sql
+            cte_expr = exp.maybe_parse(cte_sql, dialect=self.dialect)  # type: ignore[attr-defined]
+
+            parameters = built_query.parameters
+            if parameters:
+                if isinstance(parameters, dict):
+                    for param_name, param_value in parameters.items():
+                        self.add_parameter(param_value, name=param_name)  # type: ignore[attr-defined]
+                elif isinstance(parameters, (list, tuple)):
+                    for param_value in parameters:
+                        self.add_parameter(param_value)  # type: ignore[attr-defined]
+
+        if not cte_expr:
+            msg = f"Could not parse CTE query: {query}"
+            raise SQLBuilderError(msg)
+
+        if columns:
+            cte_alias_expr = exp.alias_(cte_expr, name, table=[exp.to_identifier(col) for col in columns])
+        else:
+            cte_alias_expr = exp.alias_(cte_expr, name)
+
+        existing_with = self._expression.args.get("with")  # pyright: ignore
+        if existing_with:
+            existing_with.expressions.append(cte_alias_expr)
+            if recursive:
+                existing_with.set("recursive", recursive)
+        else:
+            self._expression = self._expression.with_(cte_alias_expr, as_=name, copy=False)  # type: ignore[union-attr]
+            if recursive:
+                with_clause = self._expression.find(exp.With)
+                if with_clause:
+                    with_clause.set("recursive", recursive)
+            self._with_ctes[name] = exp.CTE(this=cte_expr, alias=exp.to_table(name))  # type: ignore[attr-defined]
+
+        return self
+
+
+class SetOperationMixin:
+    """Mixin providing set operations (UNION, INTERSECT, EXCEPT) for SELECT builders."""
+
+    _expression: Any = None
+    _parameters: dict[str, Any]
+    dialect: Any = None
+
+    def union(self, other: Any, all_: bool = False) -> Self:
+        """Combine this query with another using UNION.
+
+        Args:
+            other: Another SelectBuilder or compatible builder to union with.
+            all_: If True, use UNION ALL instead of UNION.
+
+        Raises:
+            SQLBuilderError: If the current expression is not a SELECT statement.
+
+        Returns:
+            The new builder instance for the union query.
+        """
+        left_query = self.build()  # type: ignore[attr-defined]
+        right_query = other.build()
+        left_expr: Optional[exp.Expression] = exp.maybe_parse(left_query.sql, dialect=self.dialect)
+        right_expr: Optional[exp.Expression] = exp.maybe_parse(right_query.sql, dialect=self.dialect)
+        if not left_expr or not right_expr:
+            msg = "Could not parse queries for UNION operation"
+            raise SQLBuilderError(msg)
+        union_expr = exp.union(left_expr, right_expr, distinct=not all_)
+        new_builder = type(self)()
+        new_builder.dialect = self.dialect
+        new_builder._expression = union_expr
+        merged_parameters = dict(left_query.parameters)
+        for param_name, param_value in right_query.parameters.items():
+            if param_name in merged_parameters:
+                counter = 1
+                new_param_name = f"{param_name}_right_{counter}"
+                while new_param_name in merged_parameters:
+                    counter += 1
+                    new_param_name = f"{param_name}_right_{counter}"
+
+                def rename_parameter(node: exp.Expression) -> exp.Expression:
+                    if isinstance(node, exp.Placeholder) and node.name == param_name:  # noqa: B023
+                        return exp.Placeholder(this=new_param_name)  # noqa: B023
+                    return node
+
+                right_expr = right_expr.transform(rename_parameter)
+                union_expr = exp.union(left_expr, right_expr, distinct=not all_)
+                new_builder._expression = union_expr
+                merged_parameters[new_param_name] = param_value
+            else:
+                merged_parameters[param_name] = param_value
+        new_builder._parameters = merged_parameters
+        return new_builder
+
+    def intersect(self, other: Any) -> Self:
+        """Add INTERSECT clause.
+
+        Args:
+            other: Another SelectBuilder or compatible builder to intersect with.
+
+        Raises:
+            SQLBuilderError: If the current expression is not a SELECT statement.
+
+        Returns:
+            The new builder instance for the intersect query.
+        """
+        left_query = self.build()  # type: ignore[attr-defined]
+        right_query = other.build()
+        left_expr: Optional[exp.Expression] = exp.maybe_parse(left_query.sql, dialect=self.dialect)
+        right_expr: Optional[exp.Expression] = exp.maybe_parse(right_query.sql, dialect=self.dialect)
+        if not left_expr or not right_expr:
+            msg = "Could not parse queries for INTERSECT operation"
+            raise SQLBuilderError(msg)
+        intersect_expr = exp.intersect(left_expr, right_expr, distinct=True)
+        new_builder = type(self)()
+        new_builder.dialect = self.dialect
+        new_builder._expression = intersect_expr
+        merged_parameters = dict(left_query.parameters)
+        merged_parameters.update(right_query.parameters)
+        new_builder._parameters = merged_parameters
+        return new_builder
+
+    def except_(self, other: Any) -> Self:
+        """Combine this query with another using EXCEPT.
+
+        Args:
+            other: Another SelectBuilder or compatible builder to except with.
+
+        Raises:
+            SQLBuilderError: If the current expression is not a SELECT statement.
+
+        Returns:
+            The new builder instance for the except query.
+        """
+        left_query = self.build()  # type: ignore[attr-defined]
+        right_query = other.build()
+        left_expr: Optional[exp.Expression] = exp.maybe_parse(left_query.sql, dialect=self.dialect)
+        right_expr: Optional[exp.Expression] = exp.maybe_parse(right_query.sql, dialect=self.dialect)
+        if not left_expr or not right_expr:
+            msg = "Could not parse queries for EXCEPT operation"
+            raise SQLBuilderError(msg)
+        except_expr = exp.except_(left_expr, right_expr)
+        new_builder = type(self)()
+        new_builder.dialect = self.dialect
+        new_builder._expression = except_expr
+        merged_parameters = dict(left_query.parameters)
+        merged_parameters.update(right_query.parameters)
+        new_builder._parameters = merged_parameters
+        return new_builder

sqlspec/builder/mixins/_delete_operations.py

@@ -0,0 +1,36 @@
+"""Delete operation mixins for SQL builders."""
+
+from typing import Optional
+
+from sqlglot import exp
+from typing_extensions import Self
+
+from sqlspec.exceptions import SQLBuilderError
+
+__all__ = ("DeleteFromClauseMixin",)
+
+
+class DeleteFromClauseMixin:
+    """Mixin providing FROM clause for DELETE builders."""
+
+    _expression: Optional[exp.Expression] = None
+
+    def from_(self, table: str) -> Self:
+        """Set the target table for the DELETE statement.
+
+        Args:
+            table: The table name to delete from.
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+        if self._expression is None:
+            self._expression = exp.Delete()
+        if not isinstance(self._expression, exp.Delete):
+            current_expr_type = type(self._expression).__name__
+            msg = f"Base expression for Delete is {current_expr_type}, expected Delete."
+            raise SQLBuilderError(msg)
+
+        setattr(self, "_table", table)
+        self._expression.set("this", exp.to_table(table))
+        return self