fix: handle ParseError gracefully in MERGE assignment validation and improve type annotations in parsing utilities

Author: cofin

sqlspec/builder/_merge.py

@@ -4,6 +4,7 @@
 parameter binding and validation.
 """
 
+import contextlib
 from collections.abc import Mapping, Sequence
 from datetime import datetime
 from decimal import Decimal
@@ -12,6 +13,7 @@
 
 from mypy_extensions import trait
 from sqlglot import exp
+from sqlglot.errors import ParseError
 from typing_extensions import Self
 
 from sqlspec.builder._base import QueryBuilder
@@ -45,32 +47,34 @@ def _is_column_reference(self, value: str) -> bool:
         if not isinstance(value, str):
             return False
 
-        parsed: exp.Expression | None = exp.maybe_parse(value.strip())
-        if parsed is None:
-            return False
-
-        if isinstance(parsed, exp.Column):
-            return parsed.table is not None and bool(parsed.table)
-
-        return isinstance(
-            parsed,
-            (
-                exp.Dot,
-                exp.Add,
-                exp.Sub,
-                exp.Mul,
-                exp.Div,
-                exp.Mod,
-                exp.Func,
-                exp.Anonymous,
-                exp.Null,
-                exp.CurrentTimestamp,
-                exp.CurrentDate,
-                exp.CurrentTime,
-                exp.Paren,
-                exp.Case,
-            ),
-        )
+        with contextlib.suppress(ParseError):
+            parsed: exp.Expression | None = exp.maybe_parse(value.strip())
+            if parsed is None:
+                return False
+
+            if isinstance(parsed, exp.Column):
+                return parsed.table is not None and bool(parsed.table)
+
+            return isinstance(
+                parsed,
+                (
+                    exp.Dot,
+                    exp.Add,
+                    exp.Sub,
+                    exp.Mul,
+                    exp.Div,
+                    exp.Mod,
+                    exp.Func,
+                    exp.Anonymous,
+                    exp.Null,
+                    exp.CurrentTimestamp,
+                    exp.CurrentDate,
+                    exp.CurrentTime,
+                    exp.Paren,
+                    exp.Case,
+                ),
+            )
+        return False
 
     def _process_assignment(self, target_column: str, value: Any) -> exp.Expression:
         column_identifier = exp.column(target_column) if isinstance(target_column, str) else target_column
@@ -571,7 +575,9 @@ def when_not_matched_then_insert(
             if values is None:
                 using_alias = None
                 using_expr = current_expr.args.get("using")
-                if using_expr is not None and (isinstance(using_expr, (exp.Subquery, exp.Table)) or hasattr(using_expr, "alias")):
+                if using_expr is not None and (
+                    isinstance(using_expr, (exp.Subquery, exp.Table)) or hasattr(using_expr, "alias")
+                ):
                     using_alias = using_expr.alias
                 column_values = [f"{using_alias}.{col}" for col in column_names] if using_alias else column_names
             else:

sqlspec/builder/_parsing_utils.py

@@ -28,7 +28,7 @@ def extract_column_name(column: str | exp.Column) -> str:
         Column name as string for use as parameter name
     """
     if isinstance(column, str):
-        col_expr = exp.maybe_parse(column)
+        col_expr: exp.Expression | None = exp.maybe_parse(column)
         if isinstance(col_expr, exp.Column):
             return col_expr.name
         return column.split(".")[-1] if "." in column else column
@@ -93,7 +93,7 @@ def parse_column_expression(column_input: str | exp.Expression | Any, builder: A
 def parse_table_expression(table_input: str, explicit_alias: str | None = None) -> exp.Expression:
     """Parses a table string that can be a name, a name with an alias, or a subquery string."""
     with contextlib.suppress(Exception):
-        parsed = exp.maybe_parse(f"SELECT * FROM {table_input}")
+        parsed: exp.Expression | None = exp.maybe_parse(f"SELECT * FROM {table_input}")
         if isinstance(parsed, exp.Select) and parsed.args.get("from"):
             from_clause = cast("exp.From", parsed.args.get("from"))
             table_expr = from_clause.this

tests/unit/test_sql_factory.py

@@ -1520,7 +1520,9 @@ def test_merge_complete_example() -> None:
     assert "WHEN NOT MATCHED THEN INSERT" in stmt.sql
     assert "WHEN NOT MATCHED BY SOURCE THEN UPDATE" in stmt.sql
     assert "NOW()" in stmt.sql
-    assert len(stmt.parameters) >= 6
+    assert "status" in stmt.parameters
+    assert stmt.parameters["status"] == "archived"
+    assert len(stmt.parameters) == 1
 
 
 def test_querybuilder_parameter_style_handling_regression() -> None: