refactor(builder): share dialect and column helpers

Author: cofin

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,7 @@ 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 +989,7 @@ 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)
         )
 
 

sqlspec/builder/_factory.py

@@ -44,8 +44,10 @@
 from sqlspec.builder._join import JoinBuilder, create_join_builder
 from sqlspec.builder._merge import Merge
 from sqlspec.builder._parsing_utils import (
+    _coerce_column,
     _normalize_order_by,
     _normalize_partition_by,
+    _resolve_dialect,
     extract_expression,
     to_expression,
 )
@@ -283,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
@@ -300,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
 
@@ -317,7 +319,7 @@ 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):
@@ -337,7 +339,7 @@ def select(
         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):
@@ -355,7 +357,7 @@ def insert(self, table_or_sql: str | None = None, dialect: DialectType = None) -
         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):
@@ -372,7 +374,7 @@ def update(self, table_or_sql: str | None = None, dialect: DialectType = None) -
         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)
@@ -388,7 +390,7 @@ def delete(self, table_or_sql: str | None = None, dialect: DialectType = None) -
         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)
@@ -427,7 +429,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:
@@ -463,7 +465,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"}
@@ -483,7 +485,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.
@@ -494,7 +496,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.
@@ -506,7 +508,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.
@@ -518,7 +520,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.
@@ -530,7 +532,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.
@@ -542,7 +544,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.
@@ -554,7 +556,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.
@@ -566,7 +568,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.
@@ -578,7 +580,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.
@@ -590,7 +592,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.
@@ -602,7 +604,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.
@@ -614,7 +616,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.
@@ -626,7 +628,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.
@@ -637,7 +639,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,
@@ -650,7 +652,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(
@@ -664,7 +666,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(
@@ -1182,7 +1184,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
@@ -1195,7 +1197,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
@@ -1261,7 +1263,7 @@ 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
@@ -1408,7 +1410,7 @@ def coalesce(*expressions: str | exp.Expr) -> ConversionExpression:
         Returns:
             COALESCE expression.
         """
-        exprs = [exp.column(expr) if isinstance(expr, str) else expr for expr in expressions]
+        exprs = [_coerce_column(expr) for expr in expressions]
         return ConversionExpression(exp.Coalesce(expressions=exprs))
 
     @staticmethod

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):

sqlspec/builder/_parsing_utils.py

@@ -398,3 +398,11 @@ def _normalize_order_by(order_by: str | list[str] | exp.Expr | None) -> exp.Orde
     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

sqlspec/builder/_select.py

@@ -18,6 +18,7 @@
 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,
     parse_column_expression,
@@ -222,7 +223,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":
@@ -357,24 +358,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])
+            exp.Tuple(expressions=[_coerce_column(col) for col in column_set])
             for column_set in column_sets
         ]
         grouping_expr = exp.GroupingSets(expressions=grouping_sets)
@@ -1053,8 +1054,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)