sqlspec 0.13.1__py3-none-any.whl → 0.16.2__py3-none-any.whl

sqlspec/builder/mixins/_cte_and_set_ops.py

@@ -0,0 +1,222 @@
+"""CTE (Common Table Expression) and Set Operations mixins for SQL builders."""
+
+from typing import Any, Optional, Union
+
+from mypy_extensions import trait
+from sqlglot import exp
+from typing_extensions import Self
+
+from sqlspec.exceptions import SQLBuilderError
+
+__all__ = ("CommonTableExpressionMixin", "SetOperationMixin")
+
+
+@trait
+class CommonTableExpressionMixin:
+    """Mixin providing WITH clause (Common Table Expressions) support for SQL builders."""
+
+    __slots__ = ()
+    # Type annotation for PyRight - this will be provided by the base class
+    _expression: Optional[exp.Expression]
+
+    _with_ctes: Any  # Provided by QueryBuilder
+    dialect: Any  # Provided by QueryBuilder
+
+    def add_parameter(self, value: Any, name: Optional[str] = None) -> tuple[Any, str]:
+        """Add parameter - provided by QueryBuilder."""
+        msg = "Method must be provided by QueryBuilder subclass"
+        raise NotImplementedError(msg)
+
+    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)
+        elif isinstance(query, exp.Expression):
+            cte_expr = query
+        else:
+            built_query = query.to_statement()
+            cte_sql = built_query.sql
+            cte_expr = exp.maybe_parse(cte_sql, dialect=self.dialect)
+
+            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)
+                elif isinstance(parameters, (list, tuple)):
+                    for param_value in parameters:
+                        self.add_parameter(param_value)
+
+        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")
+        if existing_with:
+            existing_with.expressions.append(cte_alias_expr)
+            if recursive:
+                existing_with.set("recursive", recursive)
+        else:
+            # Only SELECT, INSERT, UPDATE support WITH clauses
+            if hasattr(self._expression, "with_") and isinstance(
+                self._expression, (exp.Select, exp.Insert, exp.Update)
+            ):
+                self._expression = self._expression.with_(cte_alias_expr, as_=name, copy=False)
+                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))
+
+        return self
+
+
+@trait
+class SetOperationMixin:
+    """Mixin providing set operations (UNION, INTERSECT, EXCEPT) for SELECT builders."""
+
+    __slots__ = ()
+    # Type annotation for PyRight - this will be provided by the base class
+    _expression: Optional[exp.Expression]
+
+    _parameters: dict[str, Any]
+    dialect: Any = None
+
+    def build(self) -> Any:
+        """Build the query - provided by QueryBuilder."""
+        msg = "Method must be provided by QueryBuilder subclass"
+        raise NotImplementedError(msg)
+
+    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()
+        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, old_name: str = param_name, new_name: str = new_param_name
+                ) -> exp.Expression:
+                    if isinstance(node, exp.Placeholder) and node.name == old_name:
+                        return exp.Placeholder(this=new_name)
+                    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()
+        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()
+        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/{statement/builder/mixins/_delete_from.py → builder/mixins/_delete_operations.py}

@@ -1,5 +1,8 @@
+"""Delete operation mixins for SQL builders."""
+
 from typing import Optional
 
+from mypy_extensions import trait
 from sqlglot import exp
 from typing_extensions import Self
 
@@ -8,10 +11,14 @@ from sqlspec.exceptions import SQLBuilderError
 __all__ = ("DeleteFromClauseMixin",)
 
 
+@trait
 class DeleteFromClauseMixin:
     """Mixin providing FROM clause for DELETE builders."""
 
-    _expression: Optional[exp.Expression] = None
+    __slots__ = ()
+
+    # Type annotation for PyRight - this will be provided by the base class
+    _expression: Optional[exp.Expression]
 
     def from_(self, table: str) -> Self:
         """Set the target table for the DELETE statement.

sqlspec/builder/mixins/_insert_operations.py

@@ -0,0 +1,244 @@
+"""Insert operation mixins for SQL builders."""
+
+from collections.abc import Sequence
+from typing import Any, Optional, TypeVar, Union
+
+from mypy_extensions import trait
+from sqlglot import exp
+from typing_extensions import Self
+
+from sqlspec.exceptions import SQLBuilderError
+from sqlspec.protocols import SQLBuilderProtocol
+
+BuilderT = TypeVar("BuilderT", bound=SQLBuilderProtocol)
+
+__all__ = ("InsertFromSelectMixin", "InsertIntoClauseMixin", "InsertValuesMixin")
+
+
+@trait
+class InsertIntoClauseMixin:
+    """Mixin providing INTO clause for INSERT builders."""
+
+    __slots__ = ()
+
+    # Type annotation for PyRight - this will be provided by the base class
+    _expression: Optional[exp.Expression]
+
+    def into(self, table: str) -> Self:
+        """Set the target table for the INSERT statement.
+
+        Args:
+            table: The name of the table to insert data into.
+
+        Raises:
+            SQLBuilderError: If the current expression is not an INSERT statement.
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+        if self._expression is None:
+            self._expression = exp.Insert()
+        if not isinstance(self._expression, exp.Insert):
+            msg = "Cannot set target table on a non-INSERT expression."
+            raise SQLBuilderError(msg)
+
+        setattr(self, "_table", table)
+        self._expression.set("this", exp.to_table(table))
+        return self
+
+
+@trait
+class InsertValuesMixin:
+    """Mixin providing VALUES and columns methods for INSERT builders."""
+
+    __slots__ = ()
+
+    # Type annotation for PyRight - this will be provided by the base class
+    _expression: Optional[exp.Expression]
+
+    _columns: Any  # Provided by QueryBuilder
+
+    def add_parameter(self, value: Any, name: Optional[str] = None) -> tuple[Any, str]:
+        """Add parameter - provided by QueryBuilder."""
+        msg = "Method must be provided by QueryBuilder subclass"
+        raise NotImplementedError(msg)
+
+    def _generate_unique_parameter_name(self, base_name: str) -> str:
+        """Generate unique parameter name - provided by QueryBuilder."""
+        msg = "Method must be provided by QueryBuilder subclass"
+        raise NotImplementedError(msg)
+
+    def columns(self, *columns: Union[str, exp.Expression]) -> Self:
+        """Set the columns for the INSERT statement and synchronize the _columns attribute on the builder."""
+        if self._expression is None:
+            self._expression = exp.Insert()
+        if not isinstance(self._expression, exp.Insert):
+            msg = "Cannot set columns on a non-INSERT expression."
+            raise SQLBuilderError(msg)
+        column_exprs = [exp.column(col) if isinstance(col, str) else col for col in columns]
+        self._expression.set("columns", column_exprs)
+        try:
+            cols = self._columns
+            if not columns:
+                cols.clear()
+            else:
+                cols[:] = [col.name if isinstance(col, exp.Column) else str(col) for col in columns]
+        except AttributeError:
+            pass
+        return self
+
+    def values(self, *values: Any, **kwargs: Any) -> Self:
+        """Add a row of values to the INSERT statement.
+
+        Supports:
+        - values(val1, val2, val3)
+        - values(col1=val1, col2=val2)
+        - values(mapping)
+
+        Args:
+            *values: Either positional values or a single mapping.
+            **kwargs: Column-value pairs.
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+        if self._expression is None:
+            self._expression = exp.Insert()
+        if not isinstance(self._expression, exp.Insert):
+            msg = "Cannot add values to a non-INSERT expression."
+            raise SQLBuilderError(msg)
+
+        if kwargs:
+            if values:
+                msg = "Cannot mix positional values with keyword values."
+                raise SQLBuilderError(msg)
+            try:
+                _columns = self._columns
+                if not _columns:
+                    self.columns(*kwargs.keys())
+            except AttributeError:
+                pass
+            row_exprs = []
+            for col, val in kwargs.items():
+                if isinstance(val, exp.Expression):
+                    row_exprs.append(val)
+                else:
+                    column_name = col if isinstance(col, str) else str(col)
+                    if "." in column_name:
+                        column_name = column_name.split(".")[-1]
+                    param_name = self._generate_unique_parameter_name(column_name)
+                    _, param_name = self.add_parameter(val, name=param_name)
+                    row_exprs.append(exp.var(param_name))
+        elif len(values) == 1 and hasattr(values[0], "items"):
+            mapping = values[0]
+            try:
+                _columns = self._columns
+                if not _columns:
+                    self.columns(*mapping.keys())
+            except AttributeError:
+                pass
+            row_exprs = []
+            for col, val in mapping.items():
+                if isinstance(val, exp.Expression):
+                    row_exprs.append(val)
+                else:
+                    column_name = col if isinstance(col, str) else str(col)
+                    if "." in column_name:
+                        column_name = column_name.split(".")[-1]
+                    param_name = self._generate_unique_parameter_name(column_name)
+                    _, param_name = self.add_parameter(val, name=param_name)
+                    row_exprs.append(exp.var(param_name))
+        else:
+            try:
+                _columns = self._columns
+                if _columns and len(values) != len(_columns):
+                    msg = f"Number of values ({len(values)}) does not match the number of specified columns ({len(_columns)})."
+                    raise SQLBuilderError(msg)
+            except AttributeError:
+                pass
+            row_exprs = []
+            for i, v in enumerate(values):
+                if isinstance(v, exp.Expression):
+                    row_exprs.append(v)
+                else:
+                    try:
+                        _columns = self._columns
+                        if _columns and i < len(_columns):
+                            column_name = (
+                                str(_columns[i]).split(".")[-1] if "." in str(_columns[i]) else str(_columns[i])
+                            )
+                            param_name = self._generate_unique_parameter_name(column_name)
+                        else:
+                            param_name = self._generate_unique_parameter_name(f"value_{i + 1}")
+                    except AttributeError:
+                        param_name = self._generate_unique_parameter_name(f"value_{i + 1}")
+                    _, param_name = self.add_parameter(v, name=param_name)
+                    row_exprs.append(exp.var(param_name))
+
+        values_expr = exp.Values(expressions=[row_exprs])
+        self._expression.set("expression", values_expr)
+        return self
+
+    def add_values(self, values: Sequence[Any]) -> Self:
+        """Add a row of values to the INSERT statement (alternative signature).
+
+        Args:
+            values: Sequence of values for the row.
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+        return self.values(*values)
+
+
+@trait
+class InsertFromSelectMixin:
+    """Mixin providing INSERT ... SELECT support for INSERT builders."""
+
+    __slots__ = ()
+
+    # Type annotation for PyRight - this will be provided by the base class
+    _expression: Optional[exp.Expression]
+
+    _table: Any  # Provided by QueryBuilder
+
+    def add_parameter(self, value: Any, name: Optional[str] = None) -> tuple[Any, str]:
+        """Add parameter - provided by QueryBuilder."""
+        msg = "Method must be provided by QueryBuilder subclass"
+        raise NotImplementedError(msg)
+
+    def from_select(self, select_builder: Any) -> Self:
+        """Sets the INSERT source to a SELECT statement.
+
+        Args:
+            select_builder: A SelectBuilder instance representing the SELECT query.
+
+        Returns:
+            The current builder instance for method chaining.
+
+        Raises:
+            SQLBuilderError: If the table is not set or the select_builder is invalid.
+        """
+        try:
+            if not self._table:
+                msg = "The target table must be set using .into() before adding values."
+                raise SQLBuilderError(msg)
+        except AttributeError:
+            msg = "The target table must be set using .into() before adding values."
+            raise SQLBuilderError(msg)
+        if self._expression is None:
+            self._expression = exp.Insert()
+        if not isinstance(self._expression, exp.Insert):
+            msg = "Cannot set INSERT source on a non-INSERT expression."
+            raise SQLBuilderError(msg)
+        subquery_parameters = select_builder._parameters
+        if subquery_parameters:
+            for p_name, p_value in subquery_parameters.items():
+                self.add_parameter(p_value, name=p_name)
+        select_expr = select_builder._expression
+        if select_expr and isinstance(select_expr, exp.Select):
+            self._expression.set("expression", select_expr.copy())
+        else:
+            msg = "SelectBuilder must have a valid SELECT expression."
+            raise SQLBuilderError(msg)
+        return self

sqlspec/{statement/builder/mixins/_join.py → builder/mixins/_join_operations.py}

@@ -1,25 +1,33 @@
 from typing import TYPE_CHECKING, Any, Optional, Union, cast
 
+from mypy_extensions import trait
 from sqlglot import exp
 from typing_extensions import Self
 
+from sqlspec.builder._parsing_utils import parse_table_expression
 from sqlspec.exceptions import SQLBuilderError
-from sqlspec.statement.builder._parsing_utils import parse_table_expression
 from sqlspec.utils.type_guards import has_query_builder_parameters
 
 if TYPE_CHECKING:
+    from sqlspec.core.statement import SQL
     from sqlspec.protocols import SQLBuilderProtocol
 
 __all__ = ("JoinClauseMixin",)
 
 
+@trait
 class JoinClauseMixin:
     """Mixin providing JOIN clause methods for SELECT builders."""
 
+    __slots__ = ()
+
+    # Type annotation for PyRight - this will be provided by the base class
+    _expression: Optional[exp.Expression]
+
     def join(
         self,
         table: Union[str, exp.Expression, Any],
-        on: Optional[Union[str, exp.Expression]] = None,
+        on: Optional[Union[str, exp.Expression, "SQL"]] = None,
         alias: Optional[str] = None,
         join_type: str = "INNER",
     ) -> Self:
@@ -33,25 +41,49 @@ class JoinClauseMixin:
         if isinstance(table, str):
             table_expr = parse_table_expression(table, alias)
         elif has_query_builder_parameters(table):
-            # Work directly with AST when possible to avoid string parsing
             if hasattr(table, "_expression") and getattr(table, "_expression", None) is not None:
                 table_expr_value = getattr(table, "_expression", None)
                 if table_expr_value is not None:
-                    subquery_exp = exp.paren(table_expr_value.copy())  # pyright: ignore
+                    subquery_exp = exp.paren(table_expr_value)
                 else:
                     subquery_exp = exp.paren(exp.Anonymous(this=""))
                 table_expr = exp.alias_(subquery_exp, alias) if alias else subquery_exp
             else:
-                subquery = table.build()  # pyright: ignore
+                subquery = table.build()
                 sql_str = subquery.sql if hasattr(subquery, "sql") and not callable(subquery.sql) else str(subquery)
                 subquery_exp = exp.paren(exp.maybe_parse(sql_str, dialect=getattr(builder, "dialect", None)))
                 table_expr = exp.alias_(subquery_exp, alias) if alias else subquery_exp
-            # Parameter merging logic can be added here if needed
         else:
             table_expr = table
         on_expr: Optional[exp.Expression] = None
         if on is not None:
-            on_expr = exp.condition(on) if isinstance(on, str) else on
+            if isinstance(on, str):
+                on_expr = exp.condition(on)
+            elif hasattr(on, "expression") and hasattr(on, "sql"):
+                # Handle SQL objects (from sql.raw with parameters)
+                expression = getattr(on, "expression", None)
+                if expression is not None and isinstance(expression, exp.Expression):
+                    # Merge parameters from SQL object into builder
+                    if hasattr(on, "parameters") and hasattr(builder, "add_parameter"):
+                        sql_parameters = getattr(on, "parameters", {})
+                        for param_name, param_value in sql_parameters.items():
+                            builder.add_parameter(param_value, name=param_name)
+                    on_expr = expression
+                else:
+                    # If expression is None, fall back to parsing the raw SQL
+                    sql_text = getattr(on, "sql", "")
+                    # Merge parameters even when parsing raw SQL
+                    if hasattr(on, "parameters") and hasattr(builder, "add_parameter"):
+                        sql_parameters = getattr(on, "parameters", {})
+                        for param_name, param_value in sql_parameters.items():
+                            builder.add_parameter(param_value, name=param_name)
+                    on_expr = exp.maybe_parse(sql_text) or exp.condition(str(sql_text))
+            # For other types (should be exp.Expression)
+            elif isinstance(on, exp.Expression):
+                on_expr = on
+            else:
+                # Last resort - convert to string and parse
+                on_expr = exp.condition(str(on))
         join_type_upper = join_type.upper()
         if join_type_upper == "INNER":
             join_expr = exp.Join(this=table_expr, on=on_expr)
@@ -68,22 +100,22 @@ class JoinClauseMixin:
         return cast("Self", builder)
 
     def inner_join(
-        self, table: Union[str, exp.Expression, Any], on: Union[str, exp.Expression], alias: Optional[str] = None
+        self, table: Union[str, exp.Expression, Any], on: Union[str, exp.Expression, "SQL"], alias: Optional[str] = None
     ) -> Self:
         return self.join(table, on, alias, "INNER")
 
     def left_join(
-        self, table: Union[str, exp.Expression, Any], on: Union[str, exp.Expression], alias: Optional[str] = None
+        self, table: Union[str, exp.Expression, Any], on: Union[str, exp.Expression, "SQL"], alias: Optional[str] = None
     ) -> Self:
         return self.join(table, on, alias, "LEFT")
 
     def right_join(
-        self, table: Union[str, exp.Expression, Any], on: Union[str, exp.Expression], alias: Optional[str] = None
+        self, table: Union[str, exp.Expression, Any], on: Union[str, exp.Expression, "SQL"], alias: Optional[str] = None
     ) -> Self:
         return self.join(table, on, alias, "RIGHT")
 
     def full_join(
-        self, table: Union[str, exp.Expression, Any], on: Union[str, exp.Expression], alias: Optional[str] = None
+        self, table: Union[str, exp.Expression, Any], on: Union[str, exp.Expression, "SQL"], alias: Optional[str] = None
     ) -> Self:
         return self.join(table, on, alias, "FULL")
 
@@ -101,12 +133,12 @@ class JoinClauseMixin:
             if hasattr(table, "_expression") and getattr(table, "_expression", None) is not None:
                 table_expr_value = getattr(table, "_expression", None)
                 if table_expr_value is not None:
-                    subquery_exp = exp.paren(table_expr_value.copy())  # pyright: ignore
+                    subquery_exp = exp.paren(table_expr_value)
                 else:
                     subquery_exp = exp.paren(exp.Anonymous(this=""))
                 table_expr = exp.alias_(subquery_exp, alias) if alias else subquery_exp
             else:
-                subquery = table.build()  # pyright: ignore
+                subquery = table.build()
                 sql_str = subquery.sql if hasattr(subquery, "sql") and not callable(subquery.sql) else str(subquery)
                 subquery_exp = exp.paren(exp.maybe_parse(sql_str, dialect=getattr(builder, "dialect", None)))
                 table_expr = exp.alias_(subquery_exp, alias) if alias else subquery_exp