add dml unit tests

Author: chelsea-lin

bigframes/core/compile/sqlglot/sql/base.py

@@ -1,4 +1,4 @@
-# Copyright 2025 Google LLC
+# Copyright 2026 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -58,6 +58,7 @@ def identifier(id: str) -> sge.Identifier:
 
 
 def literal(value: typing.Any, dtype: dtypes.Dtype) -> sge.Expression:
+    """Return a string representing column reference in a SQL."""
     sqlglot_type = sgt.from_bigframes_dtype(dtype) if dtype else None
     if sqlglot_type is None:
         if not pd.isna(value):
@@ -110,13 +111,15 @@ def literal(value: typing.Any, dtype: dtypes.Dtype) -> sge.Expression:
 
 
 def cast(arg: typing.Any, to: str, safe: bool = False) -> sge.Cast | sge.TryCast:
+    """Return a SQL expression that casts the given argument to the specified type."""
     if safe:
         return sge.TryCast(this=arg, to=to)
     else:
         return sge.Cast(this=arg, to=to)
 
 
 def table(table: bigquery.TableReference) -> sge.Table:
+    """Return a SQLGlot Table expression representing the given BigQuery table reference."""
     return sge.Table(
         this=sge.to_identifier(table.table_id, quoted=True),
         db=sge.to_identifier(table.dataset_id, quoted=True),

bigframes/core/compile/sqlglot/sql/dml.py

@@ -41,10 +41,12 @@ def replace(
         this=base.table(destination),
         using=_as_from_item(query_or_table),
         on=base.literal(False, dtypes.BOOL_DTYPE),
-        whens=[
-            sge.When(matched=False, source=True, then=sge.Delete()),
-            sge.When(matched=False, then=sge.Insert(this=sge.Var(this="ROW"))),
-        ],
+        whens=sge.Whens(
+            expressions=[
+                sge.When(matched=False, source=True, then=sge.Delete()),
+                sge.When(matched=False, then=sge.Insert(this=sge.Var(this="ROW"))),
+            ]
+        ),
     )
 
 

tests/unit/core/compile/sqlglot/snapshots/test_dml/test_insert_from_select/out.sql

@@ -0,0 +1,6 @@
+INSERT INTO `bigframes-dev`.`sqlglot_test`.`dest_table`
+(
+  SELECT
+    *
+  FROM `source_table`
+)

tests/unit/core/compile/sqlglot/snapshots/test_dml/test_insert_from_table/out.sql

@@ -0,0 +1,2 @@
+INSERT INTO `bigframes-dev`.`sqlglot_test`.`dest_table`
+`source_table`

tests/unit/core/compile/sqlglot/snapshots/test_dml/test_replace_from_select/out.sql

@@ -0,0 +1,9 @@
+MERGE INTO `bigframes-dev`.`sqlglot_test`.`dest_table`
+USING (
+  SELECT
+    *
+  FROM `source_table`
+)
+ON FALSE
+WHEN NOT MATCHED BY SOURCE THEN DELETE
+WHEN NOT MATCHED THEN INSERT ROW

tests/unit/core/compile/sqlglot/snapshots/test_dml/test_replace_from_table/out.sql

@@ -0,0 +1,5 @@
+MERGE INTO `bigframes-dev`.`sqlglot_test`.`dest_table`
+USING `source_table`
+ON FALSE
+WHEN NOT MATCHED BY SOURCE THEN DELETE
+WHEN NOT MATCHED THEN INSERT ROW

tests/unit/core/compile/sqlglot/test_dml.py

@@ -0,0 +1,73 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import bigframes_vendored.sqlglot.expressions as sge
+from google.cloud import bigquery
+import pytest
+
+from bigframes.core.compile.sqlglot.sql import base, dml
+
+pytest.importorskip("pytest_snapshot")
+
+
+def test_insert_from_select(snapshot):
+    query = sge.select("*").from_(
+        sge.Table(this=sge.Identifier(this="source_table", quoted=True))
+    )
+    destination = bigquery.TableReference.from_string(
+        "bigframes-dev.sqlglot_test.dest_table"
+    )
+
+    expr = dml.insert(query, destination)
+    sql = base.to_sql(expr)
+
+    snapshot.assert_match(sql, "out.sql")
+
+
+def test_insert_from_table(snapshot):
+    query = sge.Table(this=sge.Identifier(this="source_table", quoted=True))
+    destination = bigquery.TableReference.from_string(
+        "bigframes-dev.sqlglot_test.dest_table"
+    )
+
+    expr = dml.insert(query, destination)
+    sql = base.to_sql(expr)
+
+    snapshot.assert_match(sql, "out.sql")
+
+
+def test_replace_from_select(snapshot):
+    query = sge.select("*").from_(
+        sge.Table(this=sge.Identifier(this="source_table", quoted=True))
+    )
+    destination = bigquery.TableReference.from_string(
+        "bigframes-dev.sqlglot_test.dest_table"
+    )
+
+    expr = dml.replace(query, destination)
+    sql = base.to_sql(expr)
+
+    snapshot.assert_match(sql, "out.sql")
+
+
+def test_replace_from_table(snapshot):
+    query = sge.Table(this=sge.Identifier(this="source_table", quoted=True))
+    destination = bigquery.TableReference.from_string(
+        "bigframes-dev.sqlglot_test.dest_table"
+    )
+
+    expr = dml.replace(query, destination)
+    sql = base.to_sql(expr)
+
+    snapshot.assert_match(sql, "out.sql")