diff --git a/sqlspec/builder/_base.py b/sqlspec/builder/_base.py index 6d5e9157f..df9696ae1 100644 --- a/sqlspec/builder/_base.py +++ b/sqlspec/builder/_base.py @@ -20,6 +20,7 @@ from typing_extensions import Self from sqlspec.builder._locking import register_lock_generator +from sqlspec.builder._parsing_utils import _resolve_dialect from sqlspec.builder._vector_distance import has_vector_distance_ancestor from sqlspec.core import ( SQL, @@ -647,7 +648,9 @@ def build(self, dialect: DialectType = None) -> "BuiltQuery": err_msg = f"Error generating SQL from expression: {e!s}" self._raise_builder_error(err_msg, e) - return BuiltQuery(sql=sql_string, parameters=self._parameters.copy(), dialect=dialect or self.dialect) + return BuiltQuery( + sql=sql_string, parameters=self._parameters.copy(), dialect=_resolve_dialect(dialect, self.dialect) + ) def to_sql(self, show_parameters: bool = False, dialect: DialectType = None) -> str: """Return SQL string with optional parameter substitution. @@ -988,7 +991,9 @@ def build_static_expression( identify = self._should_identify(target_dialect) sql_string = expr.sql(dialect=target_dialect, pretty=True, identify=identify) return BuiltQuery( - sql=sql_string, parameters=parameters.copy() if parameters else {}, dialect=dialect or self.dialect + sql=sql_string, + parameters=parameters.copy() if parameters else {}, + dialect=_resolve_dialect(dialect, self.dialect), ) diff --git a/sqlspec/builder/_column.py b/sqlspec/builder/_column.py index 859452482..1d87e739b 100644 --- a/sqlspec/builder/_column.py +++ b/sqlspec/builder/_column.py @@ -143,11 +143,11 @@ def alias(self, alias_name: str) -> exp.Expr: def asc(self) -> exp.Ordered: """Create an ASC ordering expression.""" - return exp.Ordered(this=self._expression, desc=False) + return exp.Ordered(this=self._expression, desc=False, nulls_first=False) def desc(self) -> exp.Ordered: """Create a DESC ordering expression.""" - return exp.Ordered(this=self._expression, desc=True) + return exp.Ordered(this=self._expression, desc=True, nulls_first=False) def as_(self, alias: str) -> exp.Alias: """Create an aliased expression.""" @@ -248,7 +248,11 @@ def substring(self, start: int, length: int | None = None) -> "FunctionColumn": def coalesce(self, *values: Any) -> "FunctionColumn": """SQL COALESCE() function.""" expressions = [self._expression] + [self._convert_value(v) for v in values] - return FunctionColumn(exp.Coalesce(expressions=expressions)) + return FunctionColumn( + exp.Coalesce(this=expressions[0], expressions=expressions[1:]) + if expressions + else exp.Coalesce(this=exp.Var(this="")) + ) def count(self) -> "FunctionColumn": """SQL COUNT() function.""" diff --git a/sqlspec/builder/_ddl.py b/sqlspec/builder/_ddl.py index 4b3593c00..48ab2546b 100644 --- a/sqlspec/builder/_ddl.py +++ b/sqlspec/builder/_ddl.py @@ -19,7 +19,6 @@ from sqlglot.dialects.dialect import DialectType from sqlspec.builder._column import ColumnExpression - from sqlspec.core import StatementConfig __all__ = ( "AlterOperation", @@ -90,7 +89,7 @@ def build_column_expression(col: "ColumnDefinition") -> "exp.Expr": constraints.append(exp.ColumnConstraint(kind=exp.UniqueColumnConstraint())) if col.auto_increment: - constraints.append(exp.ColumnConstraint(kind=exp.AutoIncrementColumnConstraint())) + constraints.append(exp.ColumnConstraint(kind=exp.AutoIncrementColumnConstraint(this=False))) if col.default is not None: default_expr: exp.Expr | None = None @@ -204,6 +203,10 @@ def _create_base_expression(self) -> exp.Expr: msg = "Subclasses must implement _create_base_expression." raise NotImplementedError(msg) + def _require(self, value: object, message: str) -> None: + if not value: + self._raise_builder_error(message) + def _resolve_select_query(self, query: object, context: str, *, require_select_type: bool = True) -> exp.Expr: select_parameters: dict[str, Any] | None = None @@ -238,9 +241,6 @@ def build(self, dialect: "DialectType" = None) -> "BuiltQuery": self._expression = self._create_base_expression() return super().build(dialect=dialect) - def to_statement(self, config: "StatementConfig | None" = None) -> "SQL": - return super().to_statement(config=config) - @trait class _IfExistsDDLMixin: @@ -589,8 +589,7 @@ def with_option(self, key: str, value: "Any") -> "Self": def _create_base_expression(self) -> "exp.Expr": """Create the SQLGlot expression for CREATE TABLE.""" - if not self._columns and not self._like_table: - self._raise_builder_error("Table must have at least one column or use LIKE clause") + self._require(self._columns or self._like_table, "Table must have at least one column or use LIKE clause") column_defs: list[exp.Expr] = [] for col in self._columns: @@ -632,7 +631,7 @@ def _create_base_expression(self) -> "exp.Expr": props.append(exp.TemporaryProperty()) if self._like_table: props.append(exp.LikeProperty(this=exp.to_table(self._like_table))) - properties_node = exp.Properties(expressions=props) if props else None + properties_node = _wrap_properties(props) create_target: exp.Expr = table_identifier if self._like_table and not column_defs else schema_expr return exp.Create(kind="TABLE", this=create_target, exists=self._if_not_exists, properties=properties_node) @@ -686,8 +685,7 @@ def _drop_expression_args(self) -> "dict[str, Any]": return {} def _create_base_expression(self) -> exp.Expr: - if not self._name: - self._raise_builder_error(f"{self._object_label} name must be set for DROP {self._drop_kind}.") + self._require(self._name, f"{self._object_label} name must be set for DROP {self._drop_kind}.") return exp.Drop( kind=self._drop_kind, this=self._build_drop_this(), @@ -874,8 +872,7 @@ def _create_base_expression(self) -> exp.Expr: Columns are turned into raw expressions (not ``Ordered``) to preserve natural NULL ordering, string ``where`` clauses become expressions, and the final ``exp.Index`` is wrapped in an ``exp.Create`` with the configured flags. """ - if not self._index_name or not self._table_name: - self._raise_builder_error("Index name and table name must be set for CREATE INDEX.") + self._require(self._index_name and self._table_name, "Index name and table name must be set for CREATE INDEX.") cols: list[exp.Expr] = [] for col in self._columns: @@ -901,7 +898,9 @@ def _create_base_expression(self) -> exp.Expr: index_params.set("where", exp.Where(this=where_expr)) index_expr = exp.Index( - this=exp.to_identifier(self._index_name), table=exp.to_table(self._table_name), params=index_params + this=exp.to_identifier(self._index_name), + table=exp.to_table(cast("str", self._table_name)), + params=index_params, ) return exp.Create(kind="INDEX", this=index_expr, unique=self._unique, exists=self._if_not_exists) @@ -937,8 +936,7 @@ def continue_identity(self) -> Self: return self def _create_base_expression(self) -> exp.Expr: - if not self._table_name: - self._raise_builder_error("Table name must be set for TRUNCATE TABLE.") + self._require(self._table_name, "Table name must be set for TRUNCATE TABLE.") identity_expr = exp.Var(this=self._identity) if self._identity else None option_expr = exp.Var(this="CASCADE") if self._cascade else None return exp.TruncateTable( @@ -1013,14 +1011,13 @@ def authorization(self, user_name: str) -> Self: return self def _create_base_expression(self) -> exp.Expr: - if not self._schema_name: - self._raise_builder_error("Schema name must be set for CREATE SCHEMA.") + self._require(self._schema_name, "Schema name must be set for CREATE SCHEMA.") props: list[exp.Property] = [] if self._authorization: props.append( exp.Property(this=exp.to_identifier("AUTHORIZATION"), value=exp.to_identifier(self._authorization)) ) - properties_node = exp.Properties(expressions=props) if props else None + properties_node = _wrap_properties(props) return exp.Create( kind="SCHEMA", this=exp.to_identifier(self._schema_name), @@ -1061,8 +1058,7 @@ def as_select(self, select_query: "str | exp.Expr") -> Self: return self def _create_base_expression(self) -> exp.Expr: - if not self._table_name: - self._raise_builder_error("Table name must be set for CREATE TABLE AS SELECT.") + self._require(self._table_name, "Table name must be set for CREATE TABLE AS SELECT.") if self._select_query is None: self._raise_builder_error("SELECT query must be set for CREATE TABLE AS SELECT.") @@ -1074,7 +1070,7 @@ def _create_base_expression(self) -> exp.Expr: if has_with_method(select_expr): select_expr = select_expr.with_(cte.this, as_=alias, copy=False) - create_target: exp.Expr = exp.to_table(self._table_name) + create_target: exp.Expr = exp.to_table(cast("str", self._table_name)) if self._columns: create_target = exp.Schema(this=create_target, expressions=[exp.to_identifier(c) for c in self._columns]) @@ -1157,16 +1153,13 @@ def with_hint(self, hint: str) -> Self: return self def _create_base_expression(self) -> exp.Expr: - if not self._view_name: - self._raise_builder_error("View name must be set for CREATE MATERIALIZED VIEW.") + self._require(self._view_name, "View name must be set for CREATE MATERIALIZED VIEW.") if self._select_query is None: self._raise_builder_error("SELECT query must be set for CREATE MATERIALIZED VIEW.") select_expr = self._resolve_select_query(self._select_query, "materialized view") - schema_expr = None - if self._columns: - schema_expr = exp.Schema(expressions=[exp.column(c) for c in self._columns]) + schema_expr = _build_optional_column_schema(self._columns) props: list[exp.Property] = [] if self._refresh_mode: @@ -1180,9 +1173,9 @@ def _create_base_expression(self) -> exp.Expr: for k, v in self._storage_parameters.items(): props.append(exp.Property(this=exp.to_identifier(k), value=exp.convert(str(v)))) if self._with_data is not None: - props.append(exp.Property(this=exp.to_identifier("WITH_DATA" if self._with_data else "NO_DATA"))) + props.append(exp.WithDataProperty(no=not self._with_data)) props.extend(exp.Property(this=exp.to_identifier("HINT"), value=exp.convert(hint)) for hint in self._hints) - properties_node = exp.Properties(expressions=props) if props else None + properties_node = _wrap_properties(props) create_target: exp.Expr = exp.to_table(self._view_name) if schema_expr is not None: @@ -1233,21 +1226,18 @@ def with_hint(self, hint: str) -> Self: return self def _create_base_expression(self) -> exp.Expr: - if not self._view_name: - self._raise_builder_error("View name must be set for CREATE VIEW.") + self._require(self._view_name, "View name must be set for CREATE VIEW.") if self._select_query is None: self._raise_builder_error("SELECT query must be set for CREATE VIEW.") select_expr = self._resolve_select_query(self._select_query, "view") - schema_expr = None - if self._columns: - schema_expr = exp.Schema(expressions=[exp.column(c) for c in self._columns]) + schema_expr = _build_optional_column_schema(self._columns) props: list[exp.Property] = [ exp.Property(this=exp.to_identifier("HINT"), value=exp.convert(h)) for h in self._hints ] - properties_node = exp.Properties(expressions=props) if props else None + properties_node = _wrap_properties(props) create_target: exp.Expr = exp.to_table(self._view_name) if schema_expr is not None: @@ -1453,8 +1443,7 @@ def drop_column_default(self, column: str) -> "Self": def _create_base_expression(self) -> "exp.Expr": """Create the SQLGlot expression for ALTER TABLE.""" - if not self._operations: - self._raise_builder_error("At least one operation must be specified for ALTER TABLE") + self._require(self._operations, "At least one operation must be specified for ALTER TABLE") if self._schema: table = exp.Table(this=exp.to_identifier(self._table_name), db=exp.to_identifier(self._schema)) @@ -1616,10 +1605,19 @@ def to(self, new_name: str) -> Self: return self def _create_base_expression(self) -> exp.Expr: - if not self._old_name or not self._new_name: - self._raise_builder_error("Both old and new table names must be set for RENAME TABLE.") + self._require(self._old_name and self._new_name, "Both old and new table names must be set for RENAME TABLE.") return exp.Alter( this=exp.to_table(self._old_name), kind="TABLE", actions=[exp.AlterRename(this=exp.to_identifier(self._new_name))], ) + + +def _build_optional_column_schema(columns: list[str]) -> exp.Schema | None: + if not columns: + return None + return exp.Schema(expressions=[exp.column(column) for column in columns]) + + +def _wrap_properties(properties: list[exp.Property]) -> exp.Properties | None: + return exp.Properties(expressions=properties) if properties else None diff --git a/sqlspec/builder/_factory.py b/sqlspec/builder/_factory.py index 4be1c4af7..c08fb77aa 100644 --- a/sqlspec/builder/_factory.py +++ b/sqlspec/builder/_factory.py @@ -5,7 +5,7 @@ import hashlib import logging -from typing import TYPE_CHECKING, Any, TypeVar, Union, cast, final +from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar, Union, cast, final import sqlglot from sqlglot import exp @@ -43,7 +43,14 @@ from sqlspec.builder._insert import Insert from sqlspec.builder._join import JoinBuilder, create_join_builder from sqlspec.builder._merge import Merge -from sqlspec.builder._parsing_utils import extract_expression, to_expression +from sqlspec.builder._parsing_utils import ( + _coerce_column, + _normalize_order_by, + _normalize_partition_by, + _resolve_dialect, + extract_expression, + to_expression, +) from sqlspec.builder._select import Case, Select, SubqueryBuilder, WindowFunctionBuilder from sqlspec.builder._update import Update from sqlspec.core import SQL @@ -93,6 +100,7 @@ logger = get_logger("sqlspec.builder.factory") BuilderT = TypeVar("BuilderT", bound=QueryBuilder) +ColumnLike: TypeAlias = Union[str, exp.Expr, "ExpressionWrapper", "Case", "Column"] MIN_SQL_LIKE_STRING_LENGTH = 6 MIN_DECODE_ARGS = 2 @@ -277,7 +285,7 @@ def __call__(self, statement: str, dialect: DialectType = None) -> "Any": """ try: - parsed_expr = sqlglot.parse_one(statement, read=dialect or self.dialect) + parsed_expr = sqlglot.parse_one(statement, read=_resolve_dialect(dialect, self.dialect)) except Exception as e: msg = f"Failed to parse SQL: {e}" raise SQLBuilderError(msg) from e @@ -294,7 +302,7 @@ def __call__(self, statement: str, dialect: DialectType = None) -> "Any": if actual_type_str == "SELECT" or ( actual_type_str == "WITH" and parsed_expr.this and isinstance(parsed_expr.this, exp.Select) ): - builder = Select(dialect=dialect or self.dialect) + builder = Select(dialect=_resolve_dialect(dialect, self.dialect)) builder.set_expression(parsed_expr) return builder @@ -311,89 +319,83 @@ def __call__(self, statement: str, dialect: DialectType = None) -> "Any": def select( self, *columns_or_sql: Union[str, exp.Expr, Column, "SQL", "Case"], dialect: DialectType = None ) -> "Select": - builder_dialect = dialect or self.dialect + builder_dialect = _resolve_dialect(dialect, self.dialect) if len(columns_or_sql) == 1 and isinstance(columns_or_sql[0], str): sql_candidate = columns_or_sql[0].strip() if self._looks_like_sql(sql_candidate): - parsed_expr = self._parse_sql_expression(sql_candidate, builder_dialect) - detected = "COMMAND" if parsed_expr is None else self._detect_type_from_expression(parsed_expr) - if detected not in {"SELECT", "WITH"}: - msg = ( - f"sql.select() expects a SELECT or WITH statement, got {detected}. " - f"Use sql.{detected.lower()}() if a dedicated builder exists, or ensure the SQL is SELECT/WITH." - ) - raise SQLBuilderError(msg) + parsed_expr = self._detect_or_raise( + sql_candidate, + builder_dialect, + {"SELECT", "WITH"}, + "sql.select() expects a SELECT or WITH statement, got {detected}. " + "Use sql.{detected_lower}() if a dedicated builder exists, or ensure the SQL is SELECT/WITH.", + ) select_builder = Select(dialect=builder_dialect) - return self._populate_select_from_sql(select_builder, sql_candidate, parsed_expr) + return self._populate_builder_from_sql(select_builder, sql_candidate, exp.Select, parsed_expr) select_builder = Select(dialect=builder_dialect) if columns_or_sql: select_builder.select(*columns_or_sql) return select_builder def insert(self, table_or_sql: str | None = None, dialect: DialectType = None) -> "Insert": - builder_dialect = dialect or self.dialect + builder_dialect = _resolve_dialect(dialect, self.dialect) builder = Insert(dialect=builder_dialect) if table_or_sql: if self._looks_like_sql(table_or_sql): - parsed_expr = self._parse_sql_expression(table_or_sql, builder_dialect) - detected = "COMMAND" if parsed_expr is None else self._detect_type_from_expression(parsed_expr) - if detected not in {"INSERT", "SELECT"}: - msg = ( - f"sql.insert() expects INSERT or SELECT (for insert-from-select), got {detected}. " - f"Use sql.{detected.lower()}() if a dedicated builder exists, " - f"or ensure the SQL is INSERT/SELECT." - ) - raise SQLBuilderError(msg) - return self._populate_insert_from_sql(builder, table_or_sql, parsed_expr) + parsed_expr = self._detect_or_raise( + table_or_sql, + builder_dialect, + {"INSERT", "SELECT"}, + "sql.insert() expects INSERT or SELECT (for insert-from-select), got {detected}. " + "Use sql.{detected_lower}() if a dedicated builder exists, or ensure the SQL is INSERT/SELECT.", + ) + return self._populate_builder_from_sql(builder, table_or_sql, exp.Insert, parsed_expr) return builder.into(table_or_sql) return builder def update(self, table_or_sql: str | None = None, dialect: DialectType = None) -> "Update": - builder_dialect = dialect or self.dialect + builder_dialect = _resolve_dialect(dialect, self.dialect) builder = Update(dialect=builder_dialect) if table_or_sql: if self._looks_like_sql(table_or_sql): - parsed_expr = self._parse_sql_expression(table_or_sql, builder_dialect) - detected = "COMMAND" if parsed_expr is None else self._detect_type_from_expression(parsed_expr) - if detected != "UPDATE": - msg = ( - f"sql.update() expects UPDATE statement, got {detected}. " - f"Use sql.{detected.lower()}() if a dedicated builder exists." - ) - raise SQLBuilderError(msg) - return self._populate_update_from_sql(builder, table_or_sql, parsed_expr) + parsed_expr = self._detect_or_raise( + table_or_sql, + builder_dialect, + {"UPDATE"}, + "sql.update() expects UPDATE statement, got {detected}. " + "Use sql.{detected_lower}() if a dedicated builder exists.", + ) + return self._populate_builder_from_sql(builder, table_or_sql, exp.Update, parsed_expr) return builder.table(table_or_sql) return builder def delete(self, table_or_sql: str | None = None, dialect: DialectType = None) -> "Delete": - builder_dialect = dialect or self.dialect + builder_dialect = _resolve_dialect(dialect, self.dialect) if table_or_sql and self._looks_like_sql(table_or_sql): builder = Delete(dialect=builder_dialect) - parsed_expr = self._parse_sql_expression(table_or_sql, builder_dialect) - detected = "COMMAND" if parsed_expr is None else self._detect_type_from_expression(parsed_expr) - if detected != "DELETE": - msg = ( - f"sql.delete() expects DELETE statement, got {detected}. " - f"Use sql.{detected.lower()}() if a dedicated builder exists." - ) - raise SQLBuilderError(msg) - return self._populate_delete_from_sql(builder, table_or_sql, parsed_expr) + parsed_expr = self._detect_or_raise( + table_or_sql, + builder_dialect, + {"DELETE"}, + "sql.delete() expects DELETE statement, got {detected}. " + "Use sql.{detected_lower}() if a dedicated builder exists.", + ) + return self._populate_builder_from_sql(builder, table_or_sql, exp.Delete, parsed_expr) return Delete(table_or_sql, dialect=builder_dialect) if table_or_sql else Delete(dialect=builder_dialect) def merge(self, table_or_sql: str | None = None, dialect: DialectType = None) -> "Merge": - builder_dialect = dialect or self.dialect + builder_dialect = _resolve_dialect(dialect, self.dialect) if table_or_sql and self._looks_like_sql(table_or_sql): builder = Merge(dialect=builder_dialect) - parsed_expr = self._parse_sql_expression(table_or_sql, builder_dialect) - detected = "COMMAND" if parsed_expr is None else self._detect_type_from_expression(parsed_expr) - if detected != "MERGE": - msg = ( - f"sql.merge() expects MERGE statement, got {detected}. " - f"Use sql.{detected.lower()}() if a dedicated builder exists." - ) - raise SQLBuilderError(msg) - return self._populate_merge_from_sql(builder, table_or_sql, parsed_expr) + parsed_expr = self._detect_or_raise( + table_or_sql, + builder_dialect, + {"MERGE"}, + "sql.merge() expects MERGE statement, got {detected}. " + "Use sql.{detected_lower}() if a dedicated builder exists.", + ) + return self._populate_builder_from_sql(builder, table_or_sql, exp.Merge, parsed_expr) return Merge(table_or_sql, dialect=builder_dialect) if table_or_sql else Merge(dialect=builder_dialect) @@ -421,7 +423,7 @@ def explain( Returns: Explain builder for further configuration """ - builder_dialect = dialect or self.dialect + builder_dialect = _resolve_dialect(dialect, self.dialect) fmt = None if format is not None: @@ -457,7 +459,7 @@ def upsert(self, table: str, dialect: DialectType = None) -> "Merge | Insert": Returns: MERGE builder for supported databases, INSERT builder otherwise """ - builder_dialect = dialect or self.dialect + builder_dialect = _resolve_dialect(dialect, self.dialect) dialect_str = str(builder_dialect).lower() if builder_dialect else None merge_supported = {"postgres", "postgresql", "oracle", "bigquery"} @@ -477,7 +479,7 @@ def create_table(self, table_name: str, dialect: DialectType = None) -> "CreateT Returns: CreateTable builder instance """ - return CreateTable(table_name, dialect=dialect or self.dialect) + return CreateTable(table_name, dialect=_resolve_dialect(dialect, self.dialect)) def create_table_as_select(self, dialect: DialectType = None) -> "CreateTableAsSelect": """Create a CREATE TABLE AS SELECT builder. @@ -488,7 +490,7 @@ def create_table_as_select(self, dialect: DialectType = None) -> "CreateTableAsS Returns: CreateTableAsSelect builder instance """ - return CreateTableAsSelect(dialect=dialect or self.dialect) + return CreateTableAsSelect(dialect=_resolve_dialect(dialect, self.dialect)) def create_view(self, view_name: str, dialect: DialectType = None) -> "CreateView": """Create a CREATE VIEW builder. @@ -500,7 +502,7 @@ def create_view(self, view_name: str, dialect: DialectType = None) -> "CreateVie Returns: CreateView builder instance """ - return CreateView(view_name, dialect=dialect or self.dialect) + return CreateView(view_name, dialect=_resolve_dialect(dialect, self.dialect)) def create_materialized_view(self, view_name: str, dialect: DialectType = None) -> "CreateMaterializedView": """Create a CREATE MATERIALIZED VIEW builder. @@ -512,7 +514,7 @@ def create_materialized_view(self, view_name: str, dialect: DialectType = None) Returns: CreateMaterializedView builder instance """ - return CreateMaterializedView(view_name, dialect=dialect or self.dialect) + return CreateMaterializedView(view_name, dialect=_resolve_dialect(dialect, self.dialect)) def create_index(self, index_name: str, dialect: DialectType = None) -> "CreateIndex": """Create a CREATE INDEX builder. @@ -524,7 +526,7 @@ def create_index(self, index_name: str, dialect: DialectType = None) -> "CreateI Returns: CreateIndex builder instance """ - return CreateIndex(index_name, dialect=dialect or self.dialect) + return CreateIndex(index_name, dialect=_resolve_dialect(dialect, self.dialect)) def create_schema(self, schema_name: str, dialect: DialectType = None) -> "CreateSchema": """Create a CREATE SCHEMA builder. @@ -536,7 +538,7 @@ def create_schema(self, schema_name: str, dialect: DialectType = None) -> "Creat Returns: CreateSchema builder instance """ - return CreateSchema(schema_name, dialect=dialect or self.dialect) + return CreateSchema(schema_name, dialect=_resolve_dialect(dialect, self.dialect)) def drop_table(self, table_name: str, dialect: DialectType = None) -> "DropTable": """Create a DROP TABLE builder. @@ -548,7 +550,7 @@ def drop_table(self, table_name: str, dialect: DialectType = None) -> "DropTable Returns: DropTable builder instance """ - return DropTable(table_name, dialect=dialect or self.dialect) + return DropTable(table_name, dialect=_resolve_dialect(dialect, self.dialect)) def drop_view(self, view_name: str, dialect: DialectType = None) -> "DropView": """Create a DROP VIEW builder. @@ -560,7 +562,7 @@ def drop_view(self, view_name: str, dialect: DialectType = None) -> "DropView": Returns: DropView builder instance """ - return DropView(view_name, dialect=dialect or self.dialect) + return DropView(view_name, dialect=_resolve_dialect(dialect, self.dialect)) def drop_materialized_view(self, view_name: str, dialect: DialectType = None) -> "DropMaterializedView": """Create a DROP MATERIALIZED VIEW builder. @@ -572,7 +574,7 @@ def drop_materialized_view(self, view_name: str, dialect: DialectType = None) -> Returns: DropMaterializedView builder instance """ - return DropMaterializedView(view_name, dialect=dialect or self.dialect) + return DropMaterializedView(view_name, dialect=_resolve_dialect(dialect, self.dialect)) def drop_index(self, index_name: str, dialect: DialectType = None) -> "DropIndex": """Create a DROP INDEX builder. @@ -584,7 +586,7 @@ def drop_index(self, index_name: str, dialect: DialectType = None) -> "DropIndex Returns: DropIndex builder instance """ - return DropIndex(index_name, dialect=dialect or self.dialect) + return DropIndex(index_name, dialect=_resolve_dialect(dialect, self.dialect)) def drop_schema(self, schema_name: str, dialect: DialectType = None) -> "DropSchema": """Create a DROP SCHEMA builder. @@ -596,7 +598,7 @@ def drop_schema(self, schema_name: str, dialect: DialectType = None) -> "DropSch Returns: DropSchema builder instance """ - return DropSchema(schema_name, dialect=dialect or self.dialect) + return DropSchema(schema_name, dialect=_resolve_dialect(dialect, self.dialect)) def alter_table(self, table_name: str, dialect: DialectType = None) -> "AlterTable": """Create an ALTER TABLE builder. @@ -608,7 +610,7 @@ def alter_table(self, table_name: str, dialect: DialectType = None) -> "AlterTab Returns: AlterTable builder instance """ - return AlterTable(table_name, dialect=dialect or self.dialect) + return AlterTable(table_name, dialect=_resolve_dialect(dialect, self.dialect)) def rename_table(self, old_name: str, dialect: DialectType = None) -> "RenameTable": """Create a RENAME TABLE builder. @@ -620,7 +622,7 @@ def rename_table(self, old_name: str, dialect: DialectType = None) -> "RenameTab Returns: RenameTable builder instance """ - return RenameTable(old_name, dialect=dialect or self.dialect) + return RenameTable(old_name, dialect=_resolve_dialect(dialect, self.dialect)) def comment_on(self, dialect: DialectType = None) -> "CommentOn": """Create a COMMENT ON builder. @@ -631,7 +633,7 @@ def comment_on(self, dialect: DialectType = None) -> "CommentOn": Returns: CommentOn builder instance """ - return CommentOn(dialect=dialect or self.dialect) + return CommentOn(dialect=_resolve_dialect(dialect, self.dialect)) def copy_from( self, @@ -644,7 +646,7 @@ def copy_from( ) -> SQL: """Build a COPY ... FROM statement.""" - effective_dialect = dialect or self.dialect + effective_dialect = _resolve_dialect(dialect, self.dialect) return build_copy_from_statement(table, source, columns=columns, options=options, dialect=effective_dialect) def copy_to( @@ -658,7 +660,7 @@ def copy_to( ) -> SQL: """Build a COPY ... TO statement.""" - effective_dialect = dialect or self.dialect + effective_dialect = _resolve_dialect(dialect, self.dialect) return build_copy_to_statement(table, target, columns=columns, options=options, dialect=effective_dialect) def copy( @@ -749,35 +751,14 @@ def _populate_builder_from_sql( ) return builder - def _populate_insert_from_sql( - self, builder: "Insert", sql_string: str, parsed_expr: "exp.Expr | None" = None - ) -> "Insert": - """Parse SQL string and populate INSERT builder using SQLGlot directly.""" - return self._populate_builder_from_sql(builder, sql_string, exp.Insert, parsed_expr) - - def _populate_select_from_sql( - self, builder: "Select", sql_string: str, parsed_expr: "exp.Expr | None" = None - ) -> "Select": - """Parse SQL string and populate SELECT builder using SQLGlot directly.""" - return self._populate_builder_from_sql(builder, sql_string, exp.Select, parsed_expr) - - def _populate_update_from_sql( - self, builder: "Update", sql_string: str, parsed_expr: "exp.Expr | None" = None - ) -> "Update": - """Parse SQL string and populate UPDATE builder using SQLGlot directly.""" - return self._populate_builder_from_sql(builder, sql_string, exp.Update, parsed_expr) - - def _populate_delete_from_sql( - self, builder: "Delete", sql_string: str, parsed_expr: "exp.Expr | None" = None - ) -> "Delete": - """Parse SQL string and populate DELETE builder using SQLGlot directly.""" - return self._populate_builder_from_sql(builder, sql_string, exp.Delete, parsed_expr) - - def _populate_merge_from_sql( - self, builder: "Merge", sql_string: str, parsed_expr: "exp.Expr | None" = None - ) -> "Merge": - """Parse SQL string and populate MERGE builder using SQLGlot directly.""" - return self._populate_builder_from_sql(builder, sql_string, exp.Merge, parsed_expr) + def _detect_or_raise( + self, sql_string: str, dialect: DialectType, expected_types: set[str], error_message: str + ) -> exp.Expr | None: + parsed_expr = self._parse_sql_expression(sql_string, dialect) + detected = "COMMAND" if parsed_expr is None else self._detect_type_from_expression(parsed_expr) + if detected not in expected_types: + raise SQLBuilderError(error_message.format(detected=detected, detected_lower=detected.lower())) + return parsed_expr def column(self, name: str, table: str | None = None) -> Column: """Create a column reference. @@ -967,9 +948,7 @@ def raw(sql_fragment: str, **parameters: Any) -> "exp.Expr | SQL": return SQL(sql_fragment, parameters) - def count( - self, column: Union[str, exp.Expr, "ExpressionWrapper", "Case", "Column"] = "*", distinct: bool = False - ) -> AggregateExpression: + def count(self, column: ColumnLike = "*", distinct: bool = False) -> AggregateExpression: """Create a COUNT expression. Args: @@ -989,7 +968,7 @@ def count( expr = exp.Count(this=exp.Distinct(expressions=[col_expr])) if distinct else exp.Count(this=col_expr) return AggregateExpression(expr) - def count_distinct(self, column: Union[str, exp.Expr, "ExpressionWrapper", "Case"]) -> AggregateExpression: + def count_distinct(self, column: ColumnLike) -> AggregateExpression: """Create a COUNT(DISTINCT column) expression. Args: @@ -1002,8 +981,9 @@ def count_distinct(self, column: Union[str, exp.Expr, "ExpressionWrapper", "Case def count_over( self, - column: Union[str, exp.Expr, "ExpressionWrapper", "Case", "Column"] = "*", + column: ColumnLike = "*", partition_by: str | list[str] | exp.Expr | None = None, + order_by: str | list[str] | exp.Expr | None = None, ) -> FunctionExpression: """Create a COUNT() OVER() window function for inline total counts. @@ -1013,6 +993,7 @@ def count_over( Args: column: Column to count (default "*" for COUNT(*)). partition_by: Optional columns to partition by. + order_by: Optional columns to order by. Returns: COUNT() OVER() window function expression. @@ -1024,19 +1005,18 @@ def count_over( count_expr = exp.Count(this=col_expr) over_args: dict[str, Any] = {} - if partition_by: - if isinstance(partition_by, str): - over_args["partition_by"] = [exp.column(partition_by)] - elif isinstance(partition_by, list): - over_args["partition_by"] = [exp.column(col) for col in partition_by] - elif isinstance(partition_by, exp.Expr): - over_args["partition_by"] = [partition_by] + normalized_partition = _normalize_partition_by(partition_by) + if normalized_partition is not None: + over_args["partition_by"] = normalized_partition + normalized_order = _normalize_order_by(order_by) + if normalized_order is not None: + over_args["order"] = normalized_order return FunctionExpression(exp.Window(this=count_expr, **over_args)) def sum_over( self, - column: Union[str, exp.Expr, "ExpressionWrapper", "Case"], + column: ColumnLike, partition_by: str | list[str] | exp.Expr | None = None, order_by: str | list[str] | exp.Expr | None = None, ) -> FunctionExpression: @@ -1055,7 +1035,7 @@ def sum_over( def avg_over( self, - column: Union[str, exp.Expr, "ExpressionWrapper", "Case"], + column: ColumnLike, partition_by: str | list[str] | exp.Expr | None = None, order_by: str | list[str] | exp.Expr | None = None, ) -> FunctionExpression: @@ -1074,7 +1054,7 @@ def avg_over( def max_over( self, - column: Union[str, exp.Expr, "ExpressionWrapper", "Case"], + column: ColumnLike, partition_by: str | list[str] | exp.Expr | None = None, order_by: str | list[str] | exp.Expr | None = None, ) -> FunctionExpression: @@ -1093,7 +1073,7 @@ def max_over( def min_over( self, - column: Union[str, exp.Expr, "ExpressionWrapper", "Case"], + column: ColumnLike, partition_by: str | list[str] | exp.Expr | None = None, order_by: str | list[str] | exp.Expr | None = None, ) -> FunctionExpression: @@ -1111,7 +1091,7 @@ def min_over( return self._create_window_function("MIN", [col_expr], partition_by, order_by) @staticmethod - def sum(column: Union[str, exp.Expr, "ExpressionWrapper", "Case"], distinct: bool = False) -> AggregateExpression: + def sum(column: ColumnLike, distinct: bool = False) -> AggregateExpression: """Create a SUM expression. Args: @@ -1127,7 +1107,7 @@ def sum(column: Union[str, exp.Expr, "ExpressionWrapper", "Case"], distinct: boo return AggregateExpression(exp.Sum(this=col_expr)) @staticmethod - def avg(column: Union[str, exp.Expr, "ExpressionWrapper", "Case"]) -> AggregateExpression: + def avg(column: ColumnLike) -> AggregateExpression: """Create an AVG expression. Args: @@ -1140,7 +1120,7 @@ def avg(column: Union[str, exp.Expr, "ExpressionWrapper", "Case"]) -> AggregateE return AggregateExpression(exp.Avg(this=col_expr)) @staticmethod - def max(column: Union[str, exp.Expr, "ExpressionWrapper", "Case"]) -> AggregateExpression: + def max(column: ColumnLike) -> AggregateExpression: """Create a MAX expression. Args: @@ -1153,7 +1133,7 @@ def max(column: Union[str, exp.Expr, "ExpressionWrapper", "Case"]) -> AggregateE return AggregateExpression(exp.Max(this=col_expr)) @staticmethod - def min(column: Union[str, exp.Expr, "ExpressionWrapper", "Case"]) -> AggregateExpression: + def min(column: ColumnLike) -> AggregateExpression: """Create a MIN expression. Args: @@ -1175,7 +1155,7 @@ def rollup(*columns: str | exp.Expr) -> FunctionExpression: Returns: ROLLUP expression. """ - column_exprs = [exp.column(col) if isinstance(col, str) else col for col in columns] + column_exprs = [_coerce_column(col) for col in columns] return FunctionExpression(exp.Rollup(expressions=column_exprs)) @staticmethod @@ -1188,7 +1168,7 @@ def cube(*columns: str | exp.Expr) -> FunctionExpression: Returns: CUBE expression. """ - column_exprs = [exp.column(col) if isinstance(col, str) else col for col in columns] + column_exprs = [_coerce_column(col) for col in columns] return FunctionExpression(exp.Cube(expressions=column_exprs)) @staticmethod @@ -1254,11 +1234,11 @@ def concat(*expressions: str | exp.Expr) -> StringExpression: Returns: CONCAT expression. """ - exprs = [exp.column(expr) if isinstance(expr, str) else expr for expr in expressions] + exprs = [_coerce_column(expr) for expr in expressions] return StringExpression(exp.Concat(expressions=exprs)) @staticmethod - def upper(column: str | exp.Expr) -> StringExpression: + def upper(column: ColumnLike) -> StringExpression: """Create an UPPER expression. Args: @@ -1267,11 +1247,11 @@ def upper(column: str | exp.Expr) -> StringExpression: Returns: UPPER expression. """ - col_expr = exp.column(column) if isinstance(column, str) else column + col_expr = extract_expression(column) return StringExpression(exp.Upper(this=col_expr)) @staticmethod - def lower(column: str | exp.Expr) -> StringExpression: + def lower(column: ColumnLike) -> StringExpression: """Create a LOWER expression. Args: @@ -1280,11 +1260,11 @@ def lower(column: str | exp.Expr) -> StringExpression: Returns: LOWER expression. """ - col_expr = exp.column(column) if isinstance(column, str) else column + col_expr = extract_expression(column) return StringExpression(exp.Lower(this=col_expr)) @staticmethod - def length(column: str | exp.Expr) -> StringExpression: + def length(column: ColumnLike) -> StringExpression: """Create a LENGTH expression. Args: @@ -1293,11 +1273,11 @@ def length(column: str | exp.Expr) -> StringExpression: Returns: LENGTH expression. """ - col_expr = exp.column(column) if isinstance(column, str) else column + col_expr = extract_expression(column) return StringExpression(exp.Length(this=col_expr)) @staticmethod - def round(column: str | exp.Expr, decimals: int = 0) -> MathExpression: + def round(column: ColumnLike, decimals: int = 0) -> MathExpression: """Create a ROUND expression. Args: @@ -1307,7 +1287,7 @@ def round(column: str | exp.Expr, decimals: int = 0) -> MathExpression: Returns: ROUND expression. """ - col_expr = exp.column(column) if isinstance(column, str) else column + col_expr = extract_expression(column) if decimals == 0: return MathExpression(exp.Round(this=col_expr)) return MathExpression(exp.Round(this=col_expr, decimals=exp.Literal.number(decimals))) @@ -1335,7 +1315,7 @@ def to_literal(value: Any) -> FunctionExpression: return FunctionExpression(exp.convert(value)) @staticmethod - def decode(column: str | exp.Expr, *args: str | exp.Expr | Any) -> FunctionExpression: + def decode(column: ColumnLike, *args: str | exp.Expr | Any) -> FunctionExpression: """Create a DECODE expression (Oracle-style conditional logic). DECODE compares column to each search value and returns the corresponding result. @@ -1352,7 +1332,7 @@ def decode(column: str | exp.Expr, *args: str | exp.Expr | Any) -> FunctionExpre Returns: CASE expression equivalent to DECODE. """ - col_expr = exp.column(column) if isinstance(column, str) else column + col_expr = extract_expression(column) if len(args) < MIN_DECODE_ARGS: msg = "DECODE requires at least one search/result pair" @@ -1378,7 +1358,7 @@ def decode(column: str | exp.Expr, *args: str | exp.Expr | Any) -> FunctionExpre return FunctionExpression(exp.Case(ifs=conditions, default=default)) @staticmethod - def cast(column: str | exp.Expr, data_type: str) -> ConversionExpression: + def cast(column: ColumnLike, data_type: str) -> ConversionExpression: """Create a CAST expression for type conversion. Args: @@ -1388,7 +1368,7 @@ def cast(column: str | exp.Expr, data_type: str) -> ConversionExpression: Returns: CAST expression. """ - col_expr = exp.column(column) if isinstance(column, str) else column + col_expr = extract_expression(column) return ConversionExpression(exp.Cast(this=col_expr, to=exp.DataType.build(data_type))) @staticmethod @@ -1401,11 +1381,13 @@ def coalesce(*expressions: str | exp.Expr) -> ConversionExpression: Returns: COALESCE expression. """ - exprs = [exp.column(expr) if isinstance(expr, str) else expr for expr in expressions] - return ConversionExpression(exp.Coalesce(expressions=exprs)) + exprs = [_coerce_column(expr) for expr in expressions] + return ConversionExpression( + exp.Coalesce(this=exprs[0], expressions=exprs[1:]) if exprs else exp.Coalesce(this=exp.Var(this="")) + ) @staticmethod - def nvl(column: str | exp.Expr, substitute_value: str | exp.Expr | Any) -> ConversionExpression: + def nvl(column: ColumnLike, substitute_value: str | exp.Expr | Any) -> ConversionExpression: """Create an NVL (Oracle-style) expression using COALESCE. Args: @@ -1415,13 +1397,13 @@ def nvl(column: str | exp.Expr, substitute_value: str | exp.Expr | Any) -> Conve Returns: COALESCE expression equivalent to NVL. """ - col_expr = exp.column(column) if isinstance(column, str) else column + col_expr = extract_expression(column) sub_expr = to_expression(substitute_value) - return ConversionExpression(exp.Coalesce(expressions=[col_expr, sub_expr])) + return ConversionExpression(exp.Coalesce(this=col_expr, expressions=[sub_expr])) @staticmethod def nvl2( - column: str | exp.Expr, value_if_not_null: str | exp.Expr | Any, value_if_null: str | exp.Expr | Any + column: ColumnLike, value_if_not_null: str | exp.Expr | Any, value_if_null: str | exp.Expr | Any ) -> ConversionExpression: """Create an NVL2 (Oracle-style) expression using CASE. @@ -1436,7 +1418,7 @@ def nvl2( Returns: CASE expression equivalent to NVL2. """ - col_expr = exp.column(column) if isinstance(column, str) else column + col_expr = extract_expression(column) not_null_expr = to_expression(value_if_not_null) null_expr = to_expression(value_if_null) @@ -1536,7 +1518,7 @@ def dense_rank( def lag( self, - column: str | exp.Expr, + column: ColumnLike, offset: int = 1, default: Any = None, partition_by: str | list[str] | exp.Expr | None = None, @@ -1556,7 +1538,7 @@ def lag( Returns: LAG window function expression. """ - col_expr = exp.column(column) if isinstance(column, str) else column + col_expr = extract_expression(column) func_args: list[exp.Expr] = [col_expr, exp.Literal.number(offset)] if default is not None: func_args.append(exp.convert(default)) @@ -1564,7 +1546,7 @@ def lag( def lead( self, - column: str | exp.Expr, + column: ColumnLike, offset: int = 1, default: Any = None, partition_by: str | list[str] | exp.Expr | None = None, @@ -1584,7 +1566,7 @@ def lead( Returns: LEAD window function expression. """ - col_expr = exp.column(column) if isinstance(column, str) else column + col_expr = extract_expression(column) func_args: list[exp.Expr] = [col_expr, exp.Literal.number(offset)] if default is not None: func_args.append(exp.convert(default)) @@ -1612,21 +1594,12 @@ def _create_window_function( over_args: dict[str, Any] = {} - if partition_by: - if isinstance(partition_by, str): - over_args["partition_by"] = [exp.column(partition_by)] - elif isinstance(partition_by, list): - over_args["partition_by"] = [exp.column(col) for col in partition_by] - elif isinstance(partition_by, exp.Expr): - over_args["partition_by"] = [partition_by] - - if order_by: - if isinstance(order_by, str): - over_args["order"] = exp.Order(expressions=[exp.column(order_by).asc()]) - elif isinstance(order_by, list): - over_args["order"] = exp.Order(expressions=[exp.column(col).asc() for col in order_by]) - elif isinstance(order_by, exp.Expr): - over_args["order"] = exp.Order(expressions=[order_by]) + normalized_partition = _normalize_partition_by(partition_by) + if normalized_partition is not None: + over_args["partition_by"] = normalized_partition + normalized_order = _normalize_order_by(order_by) + if normalized_order is not None: + over_args["order"] = normalized_order return FunctionExpression(exp.Window(this=func_expr, **over_args)) diff --git a/sqlspec/builder/_join.py b/sqlspec/builder/_join.py index 82eda97ce..c5b378abc 100644 --- a/sqlspec/builder/_join.py +++ b/sqlspec/builder/_join.py @@ -11,7 +11,7 @@ from typing_extensions import Self from sqlspec.builder._base import BuiltQuery, QueryBuilder -from sqlspec.builder._parsing_utils import parse_table_expression +from sqlspec.builder._parsing_utils import extract_sql_object_expression, parse_table_expression from sqlspec.builder._temporal import register_version_generators from sqlspec.exceptions import SQLBuilderError from sqlspec.utils.type_guards import has_expression_and_parameters, has_expression_and_sql, has_parameter_builder @@ -23,28 +23,17 @@ __all__ = ("JoinBuilder", "JoinClauseMixin", "create_join_builder") -def _condition_from_sql_object(on: Any, builder: "SQLBuilderProtocol") -> exp.Expr: - if has_expression_and_parameters(on) and on.expression is not None: - for param_name, param_value in on.parameters.items(): - builder.add_parameter(param_value, name=param_name) - return cast("exp.Expr", on.expression) - if has_expression_and_parameters(on): - for param_name, param_value in on.parameters.items(): - builder.add_parameter(param_value, name=param_name) - raw_sql = getattr(on, "raw_sql", None) - if raw_sql is None: - raw_sql = on.sql # pyright: ignore[reportAttributeAccessIssue] - parsed_expr = exp.maybe_parse(raw_sql, dialect=builder.dialect) - return parsed_expr if parsed_expr is not None else exp.condition(raw_sql) - - def _parse_join_condition(builder: "SQLBuilderProtocol", on: Union[str, exp.Expr, "SQL"] | None) -> exp.Expr | None: if on is None: return None if isinstance(on, str): return exp.condition(on) - if has_expression_and_sql(on): - return _condition_from_sql_object(on, builder) + if has_expression_and_sql(on) or has_expression_and_parameters(on): + + def parse_join_condition(sql_text: str) -> exp.Expr: + return exp.maybe_parse(sql_text, dialect=builder.dialect) or exp.condition(sql_text) + + return extract_sql_object_expression(on, builder=builder, parse_sql=parse_join_condition) if isinstance(on, exp.Expr): return on return exp.condition(str(on)) diff --git a/sqlspec/builder/_merge.py b/sqlspec/builder/_merge.py index 8c07ba7e6..a427b8b4c 100644 --- a/sqlspec/builder/_merge.py +++ b/sqlspec/builder/_merge.py @@ -18,7 +18,7 @@ from sqlspec.builder._base import QueryBuilder from sqlspec.builder._explain import ExplainMixin -from sqlspec.builder._parsing_utils import extract_sql_object_expression +from sqlspec.builder._parsing_utils import _coerce_column, _resolve_dialect, extract_sql_object_expression from sqlspec.builder._select import is_explicitly_quoted from sqlspec.core import SQLResult from sqlspec.exceptions import DialectNotSupportedError, SQLBuilderError @@ -137,7 +137,7 @@ def _is_column_reference(self, value: str) -> bool: return False def _process_assignment(self, target_column: str, value: Any) -> exp.Expr: - column_identifier = exp.column(target_column) if isinstance(target_column, str) else target_column + column_identifier = _coerce_column(target_column) if has_expression_and_sql(value): value_expr = extract_sql_object_expression(value, builder=self) @@ -724,7 +724,7 @@ def build(self, dialect: "DialectType" = None) -> "Any": Built statement object. """ self._validate_dialect_support() - target_dialect = dialect or self.dialect + target_dialect = _resolve_dialect(dialect, self.dialect) if isinstance(target_dialect, str): dialect_name = target_dialect elif isinstance(target_dialect, type): diff --git a/sqlspec/builder/_parsing_utils.py b/sqlspec/builder/_parsing_utils.py index 43a7326ca..29c6be247 100644 --- a/sqlspec/builder/_parsing_utils.py +++ b/sqlspec/builder/_parsing_utils.py @@ -6,6 +6,7 @@ import contextlib import re +from collections.abc import Callable from typing import TYPE_CHECKING, Any, Final from sqlglot import exp, maybe_parse @@ -301,7 +302,9 @@ def parse_condition_expression(condition_input: str | exp.Expr | tuple[str, Any] return exp.condition(condition_input) -def extract_sql_object_expression(value: Any, builder: Any | None = None) -> exp.Expr: +def extract_sql_object_expression( + value: Any, builder: Any | None = None, parse_sql: Callable[[str], exp.Expr] | None = None +) -> exp.Expr: """Extract SQLGlot expression from SQL object value with parameter merging. Handles the common pattern of: @@ -317,6 +320,7 @@ def extract_sql_object_expression(value: Any, builder: Any | None = None) -> exp Args: value: The SQL object value to process builder: Optional builder instance for parameter merging (must have add_parameter method) + parse_sql: Optional parser for raw SQL fallback text. Returns: SQLGlot Expression extracted from the SQL object @@ -324,19 +328,27 @@ def extract_sql_object_expression(value: Any, builder: Any | None = None) -> exp Raises: ValueError: If the value doesn't appear to be a SQL object """ - if not has_expression_and_sql(value): - msg = f"Value does not have both expression and sql attributes: {type(value)}" + has_sql = has_expression_and_sql(value) + has_parameters = has_expression_and_parameters(value) + if not has_sql and not has_parameters: + msg = f"Value does not expose an expression with SQL or parameters: {type(value)}" raise ValueError(msg) if value.expression is not None and isinstance(value.expression, exp.Expr): _merge_sql_parameters(value, builder) return value.expression + if not has_sql: + msg = f"Value has no SQL text fallback: {type(value)}" + raise ValueError(msg) + _merge_sql_parameters(value, builder) sql_text = getattr(value, "raw_sql", None) if sql_text is None: sql_text = value.sql if not callable(value.sql) else str(value) + if parse_sql is not None: + return parse_sql(sql_text) return exp.maybe_parse(sql_text) or exp.convert(str(sql_text)) @@ -376,3 +388,31 @@ def to_expression(value: Any) -> exp.Expr: if isinstance(value, exp.Expr): return value return exp.convert(value) + + +def _normalize_partition_by(partition_by: str | list[str] | exp.Expr | None) -> list[exp.Expr] | None: + if isinstance(partition_by, str): + return [exp.column(partition_by)] + if isinstance(partition_by, list): + return [exp.column(column) for column in partition_by] + if isinstance(partition_by, exp.Expr): + return [partition_by] + return None + + +def _normalize_order_by(order_by: str | list[str] | exp.Expr | None) -> exp.Order | None: + if isinstance(order_by, str): + return exp.Order(expressions=[exp.column(order_by).asc()]) + if isinstance(order_by, list): + return exp.Order(expressions=[exp.column(column).asc() for column in order_by]) + if isinstance(order_by, exp.Expr): + return exp.Order(expressions=[order_by]) + return None + + +def _coerce_column(value: str | exp.Expr) -> exp.Expr: + return exp.column(value) if isinstance(value, str) else value + + +def _resolve_dialect(dialect: "DialectType | None", default: "DialectType | None") -> "DialectType | None": + return dialect or default diff --git a/sqlspec/builder/_select.py b/sqlspec/builder/_select.py index 5dd85e7a6..8190260bb 100644 --- a/sqlspec/builder/_select.py +++ b/sqlspec/builder/_select.py @@ -18,8 +18,10 @@ from sqlspec.builder._join import JoinClauseMixin, _attach_as_of_version from sqlspec.builder._parsing_utils import ( _PARAMETER_VALIDATOR, + _coerce_column, extract_column_name, extract_expression, + extract_sql_object_expression, parse_column_expression, parse_condition_expression, parse_order_expression, @@ -29,6 +31,7 @@ from sqlspec.core import SQL, ParameterStyle, SQLResult from sqlspec.exceptions import SQLBuilderError from sqlspec.utils.type_guards import ( + has_expression_and_parameters, has_expression_and_sql, has_parameter_builder, has_sqlglot_expression, @@ -197,7 +200,7 @@ def __call__(self, subquery: Any) -> exp.Expr: if self._operation == "exists": return exp.Exists(this=subquery_expr) if self._operation == "in": - return exp.In(expressions=[subquery_expr]) + return exp.In(this=exp.Var(this=""), expressions=[subquery_expr]) if self._operation == "any": return exp.Any(this=subquery_expr) if self._operation == "all": @@ -222,7 +225,7 @@ def __call__(self, *function_args: Any) -> "WindowFunctionBuilder": return self def partition_by(self, *columns: str | exp.Expr) -> "WindowFunctionBuilder": - self._partition_by = [exp.column(column) if isinstance(column, str) else column for column in columns] + self._partition_by = [_coerce_column(column) for column in columns] return self def order_by(self, *columns: str | exp.Expr) -> "WindowFunctionBuilder": @@ -233,7 +236,7 @@ def order_by(self, *columns: str | exp.Expr) -> "WindowFunctionBuilder": elif isinstance(column, exp.Ordered): ordered_columns.append(column) else: - ordered_columns.append(exp.Ordered(this=column, desc=False)) + ordered_columns.append(exp.Ordered(this=column, desc=False, nulls_first=False)) self._order_by = ordered_columns return self @@ -357,25 +360,24 @@ def group_by(self, *columns: str | exp.Expr) -> Self: return cast("Self", builder) for column in columns: - column_expr = exp.column(column) if isinstance(column, str) else column + column_expr = _coerce_column(column) select_expr = select_expr.group_by(column_expr, copy=False) builder.set_expression(select_expr) return cast("Self", builder) def group_by_rollup(self, *columns: str | exp.Expr) -> Self: - column_exprs = [exp.column(column) if isinstance(column, str) else column for column in columns] + column_exprs = [_coerce_column(column) for column in columns] rollup_expr = exp.Rollup(expressions=column_exprs) return self.group_by(rollup_expr) def group_by_cube(self, *columns: str | exp.Expr) -> Self: - column_exprs = [exp.column(column) if isinstance(column, str) else column for column in columns] + column_exprs = [_coerce_column(column) for column in columns] cube_expr = exp.Cube(expressions=column_exprs) return self.group_by(cube_expr) def group_by_grouping_sets(self, *column_sets: tuple[str, ...] | list[str]) -> Self: grouping_sets = [ - exp.Tuple(expressions=[exp.column(col) if isinstance(col, str) else col for col in column_set]) - for column_set in column_sets + exp.Tuple(expressions=[_coerce_column(col) for col in column_set]) for column_set in column_sets ] grouping_expr = exp.GroupingSets(expressions=grouping_sets) return self.group_by(grouping_expr) @@ -402,7 +404,9 @@ def order_by(self, *items: Union[str, exp.Ordered, "Column"], desc: bool = False if isinstance(extracted_item, exp.Alias): alias_name = (extracted_item.alias or "").lower() if alias_name in {"asc", "desc"}: - extracted_item = exp.Ordered(this=extracted_item.this, desc=alias_name == "desc") + extracted_item = exp.Ordered( + this=extracted_item.this, desc=alias_name == "desc", nulls_first=False + ) order_item = ( extracted_item.desc() if desc and not isinstance(extracted_item, exp.Ordered) else extracted_item ) @@ -646,17 +650,16 @@ def _normalize_subquery_expression(self, subquery: Any, builder: "SQLBuilderProt builder.add_parameter(parameters) return subquery_expr - if has_expression_and_sql(subquery): - self._merge_parameters(subquery) - expression_attr = subquery.expression - if isinstance(expression_attr, exp.Expr): - return expression_attr - sql_text = subquery.sql - parsed_from_sql: exp.Expr | None = exp.maybe_parse(sql_text, dialect=builder.dialect) - if parsed_from_sql is None: - msg = f"Could not parse subquery SQL: {sql_text}" - raise SQLBuilderError(msg) - return parsed_from_sql + if has_expression_and_sql(subquery) or has_expression_and_parameters(subquery): + + def parse_subquery(sql_text: str) -> exp.Expr: + parsed_from_sql: exp.Expr | None = exp.maybe_parse(sql_text, dialect=builder.dialect) + if parsed_from_sql is None: + msg = f"Could not parse subquery SQL: {sql_text}" + raise SQLBuilderError(msg) + return parsed_from_sql + + return extract_sql_object_expression(subquery, builder=self, parse_sql=parse_subquery) if isinstance(subquery, exp.Expr): return subquery @@ -777,16 +780,10 @@ def _normalize_condition_expression(expression: exp.Expr) -> exp.Expr: if isinstance(raw_expr, exp.Expr): return builder._parameterize_expression(_normalize_condition_expression(raw_expr)) return parse_condition_expression(str(condition)) - if has_expression_and_sql(condition): - expression_attr = condition.expression - if isinstance(expression_attr, exp.Expr): - self._merge_parameters(condition) - return _normalize_condition_expression(expression_attr) - sql_text = getattr(condition, "raw_sql", None) - if sql_text is None: - sql_text = condition.sql - self._merge_parameters(condition) - return parse_condition_expression(sql_text) + if has_expression_and_sql(condition) or has_expression_and_parameters(condition): + return _normalize_condition_expression( + extract_sql_object_expression(condition, builder=self, parse_sql=parse_condition_expression) + ) msg = f"Unsupported condition type: {type(condition).__name__}" raise SQLBuilderError(msg) @@ -1053,8 +1050,8 @@ def pivot( raise TypeError(msg) agg_name = aggregate_function if isinstance(aggregate_function, str) else aggregate_function.name - agg_column = exp.column(aggregate_column) if isinstance(aggregate_column, str) else aggregate_column - pivot_col_expr = exp.column(pivot_column) if isinstance(pivot_column, str) else pivot_column + agg_column = _coerce_column(aggregate_column) + pivot_col_expr = _coerce_column(pivot_column) pivot_agg_expr = exp.func(agg_name, agg_column) diff --git a/tests/unit/builder/test_expression_arg_contracts.py b/tests/unit/builder/test_expression_arg_contracts.py index f198b8687..a16a2799c 100644 --- a/tests/unit/builder/test_expression_arg_contracts.py +++ b/tests/unit/builder/test_expression_arg_contracts.py @@ -6,10 +6,14 @@ expression and require the full clause text. """ +from typing import Any, cast + import pytest +from sqlglot import exp from sqlspec import sql from sqlspec.builder import Column +from sqlspec.builder._ddl import ColumnDefinition, build_column_expression from sqlspec.exceptions import SQLBuilderError @@ -71,3 +75,45 @@ def test_column_any_renders_operand_array() -> None: def test_column_not_any_renders_operand_array() -> None: rendered = Column("status").not_any_(["a", "b"]).sqlglot_expression.sql(dialect="postgres") assert rendered == "status <> ANY(ARRAY['a', 'b'])" + + +def test_builder_expressions_satisfy_required_sqlglot_arguments() -> None: + ordered_select = cast("Any", sql.select("a").from_("t")).order_by(exp.alias_(exp.column("a"), "asc")) + ordered_expression = ordered_select.get_expression() + assert ordered_expression is not None + expressions = [ + sql.coalesce("a", "b").expression, + sql.nvl("a", "b").expression, + Column("a").coalesce("b").sqlglot_expression, + Column("a").asc(), + Column("a").desc(), + next(sql.row_number_.order_by(exp.column("a")).build().find_all(exp.Ordered)), + next(ordered_expression.find_all(exp.Ordered)), + next( + build_column_expression(ColumnDefinition("id", "INT", auto_increment=True)).find_all( + exp.AutoIncrementColumnConstraint + ) + ), + next( + sql + .create_materialized_view("mv") + .as_select("SELECT 1") + .with_data() + ._create_base_expression() + .find_all(exp.Property) + ), + sql.in_(sql.select("id").from_("items")), + ] + assert [(type(expression).__name__, expression.error_messages()) for expression in expressions] == [ + (type(expression).__name__, []) for expression in expressions + ] + + +def test_materialized_view_with_data_renders_postgres_clause() -> None: + expression = sql.create_materialized_view("mv").as_select("SELECT 1").with_data()._create_base_expression() + assert expression.sql(dialect="postgres") == "CREATE MATERIALIZED_VIEW mv AS SELECT 1 WITH DATA" + + +def test_materialized_view_no_data_renders_postgres_clause() -> None: + expression = sql.create_materialized_view("mv").as_select("SELECT 1").no_data()._create_base_expression() + assert expression.sql(dialect="postgres") == "CREATE MATERIALIZED_VIEW mv AS SELECT 1 WITH NO DATA" diff --git a/tests/unit/builder/test_sql_factory.py b/tests/unit/builder/test_sql_factory.py index 8d12073d6..60c30eec3 100644 --- a/tests/unit/builder/test_sql_factory.py +++ b/tests/unit/builder/test_sql_factory.py @@ -3,6 +3,8 @@ import importlib.util import math +from collections.abc import Callable +from typing import Any, cast import pytest from sqlglot import exp, parse_one @@ -24,6 +26,7 @@ build_copy_from_statement, build_copy_to_statement, ) +from sqlspec.builder import _select as select_module from sqlspec.builder._expression_wrappers import ( ConversionExpression, ExpressionWrapper, @@ -32,10 +35,24 @@ StringExpression, ) from sqlspec.builder._join import create_join_builder +from sqlspec.builder._parsing_utils import extract_sql_object_expression from sqlspec.core import SQL from sqlspec.exceptions import SQLBuilderError +class _SQLExpressionObject: + def __init__(self, sql_text: str, expression: exp.Expr, parameters: dict[str, object]) -> None: + self.expression = expression + self.parameters = parameters + self.sql = sql_text + + +class _ParametersExpressionObject: + def __init__(self, _sql_text: str, expression: exp.Expr, parameters: dict[str, object]) -> None: + self.expression = expression + self.parameters = parameters + + def test_sql_factory_instance() -> None: """Test that sql is an instance of SQLFactory.""" assert isinstance(sql, SQLFactory) @@ -55,6 +72,50 @@ def test_decode_renders_trailing_default_else_clause() -> None: assert "ELSE 0" in rendered.upper() +@pytest.mark.parametrize( + "function", + [ + lambda value: sql.upper(value), + lambda value: sql.lower(value), + lambda value: sql.length(value), + lambda value: sql.round(value), + lambda value: sql.decode(value, "active", 1), + lambda value: sql.cast(value, "TEXT"), + lambda value: sql.nvl(value, 0), + lambda value: sql.nvl2(value, 1, 0), + lambda value: sql.lag(value), + lambda value: sql.lead(value), + ], +) +@pytest.mark.parametrize( + "value", [Column("amount"), ExpressionWrapper(exp.column("amount")), sql.case().when("amount > 0", 1).else_(0)] +) +def test_scalar_functions_accept_builder_expressions(function: object, value: object) -> None: + result = function(value) # type: ignore[operator] + assert isinstance(result.expression, exp.Expr) # type: ignore[attr-defined] + assert result.expression.sql() # type: ignore[attr-defined] + + +@pytest.mark.parametrize( + "function", + [ + lambda value: sql.count_distinct(value), + lambda value: sql.sum_over(value), + lambda value: sql.avg_over(value), + lambda value: sql.max_over(value), + lambda value: sql.min_over(value), + lambda value: sql.sum(value), + lambda value: sql.avg(value), + lambda value: sql.max(value), + lambda value: sql.min(value), + ], +) +def test_aggregate_functions_accept_columns(function: object) -> None: + result = function(Column("amount")) # type: ignore[operator] + assert isinstance(result.expression, exp.Expr) # type: ignore[attr-defined] + assert result.expression.sql() # type: ignore[attr-defined] + + def test_where_eq_uses_placeholder_not_var() -> None: """Test that where_eq uses Placeholder instead of var for parameters.""" query = sql.select("*").from_("users").where_eq("name", "John") @@ -933,6 +994,58 @@ def test_multiple_sql_raw_objects_parameter_merging() -> None: assert ":min_date" in stmt.sql +def test_select_conditions_use_shared_expression_extraction(monkeypatch: pytest.MonkeyPatch) -> None: + calls = 0 + + def extract( + value: object, builder: object | None = None, parse_sql: Callable[[str], exp.Expr] | None = None + ) -> exp.Expr: + nonlocal calls + calls += 1 + return extract_sql_object_expression(value, builder, parse_sql) + + monkeypatch.setattr(select_module, "extract_sql_object_expression", extract) + sql.select("*").from_("users").where_exists( + sql.raw("SELECT 1 FROM posts WHERE posts.user_id = :user_id", user_id=1) + ).where(sql.raw("id = :id", id=1)) + assert calls == 2 + + +@pytest.mark.parametrize("object_type", [_SQLExpressionObject, _ParametersExpressionObject]) +def test_sql_object_protocol_shapes_merge_where_parameters( + object_type: type[_SQLExpressionObject] | type[_ParametersExpressionObject], +) -> None: + condition = object_type("users.id = :user_id", exp.condition("users.id = :user_id"), {"user_id": 41}) + statement = sql.select("*").from_("users").where(cast("Any", condition)).build() + assert ":user_id" in statement.sql + assert statement.parameters["user_id"] == 41 + + +@pytest.mark.parametrize("object_type", [_SQLExpressionObject, _ParametersExpressionObject]) +def test_sql_object_protocol_shapes_merge_subquery_parameters( + object_type: type[_SQLExpressionObject] | type[_ParametersExpressionObject], +) -> None: + sql_text = "SELECT 1 FROM posts WHERE posts.user_id = :user_id" + subquery = object_type(sql_text, parse_one(sql_text), {"user_id": 42}) + statement = sql.select("*").from_("users").where_exists(subquery).build() + assert ":user_id" in statement.sql + assert statement.parameters["user_id"] == 42 + + +@pytest.mark.parametrize("object_type", [_SQLExpressionObject, _ParametersExpressionObject]) +def test_sql_object_protocol_shapes_merge_join_parameters( + object_type: type[_SQLExpressionObject] | type[_ParametersExpressionObject], +) -> None: + condition = object_type( + "users.id = posts.user_id AND posts.status = :status", + exp.condition("users.id = posts.user_id AND posts.status = :status"), + {"status": "published"}, + ) + statement = sql.select("*").from_("users").join("posts", cast("Any", condition)).build() + assert ":status" in statement.sql + assert statement.parameters["status"] == "published" + + def test_sql_raw_without_parameters_still_works() -> None: """Test that SQL objects without parameters still work correctly.""" raw_expr = sql.raw("NOW()") @@ -1712,6 +1825,14 @@ def test_count_over_method_with_multiple_partition_columns() -> None: assert "status" in built.lower() +def test_count_over_method_with_ordering() -> None: + """Test count_over() accepts the same ordering inputs as other window functions.""" + count_expr = sql.count_over(partition_by="department", order_by=["created_at", "id"]) + built = str(count_expr) + assert "PARTITION BY department" in built + assert "ORDER BY created_at, id" in built + + def test_count_over_property_basic() -> None: """Test count_over_ property returns WindowFunctionBuilder.""" builder = sql.count_over_