fix: Raise error instead of panicking for unsupported SQL operations (#…

…20789)
pola-rs · Feb 11, 2025 · c1e493a · c1e493a
1 parent e143e7e
commit c1e493a
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 4 deletions.
diff --git a/crates/polars-sql/src/context.rs b/crates/polars-sql/src/context.rs
@@ -346,10 +346,6 @@ impl SQLContext {
                 }
             },
             op => {
-                let op = match op {
-                    SetExpr::SetOperation { op, .. } => op,
-                    _ => unreachable!(),
-                };
                 polars_bail!(SQLInterface: "'{}' operation is currently unsupported", op)
             },
         }

diff --git a/py-polars/tests/unit/sql/test_set_ops.py b/py-polars/tests/unit/sql/test_set_ops.py
@@ -89,6 +89,54 @@ def test_except_intersect_all_unsupported(op: str, op_subtype: str) -> None:
         pl.sql(f"SELECT * FROM df1 {op} {op_subtype} SELECT * FROM df2")
 
 
+def test_update_statement_error() -> None:
+    df_large = pl.DataFrame(
+        {
+            "FQDN": ["c.ORG.na", "a.COM.na"],
+            "NS1": ["ns1.c.org.na", "ns1.d.net.na"],
+            "NS2": ["ns2.c.org.na", "ns2.d.net.na"],
+            "NS3": ["ns3.c.org.na", "ns3.d.net.na"],
+        }
+    )
+    df_small = pl.DataFrame(
+        {
+            "FQDN": ["c.org.na"],
+            "NS1": ["ns1.c.org.na|127.0.0.1"],
+            "NS2": ["ns2.c.org.na|127.0.0.1"],
+            "NS3": ["ns3.c.org.na|127.0.0.1"],
+        }
+    )
+
+    # Create a context and register the tables
+    ctx = pl.SQLContext()
+    ctx.register("large", df_large)
+    ctx.register("small", df_small)
+
+    with pytest.raises(
+        SQLInterfaceError,
+        match="'UPDATE large SET FQDN = u.FQDN, NS1 = u.NS1, NS2 = u.NS2, NS3 = u.NS3 FROM u WHERE large.FQDN = u.FQDN' operation is currently unsupported",
+    ):
+        ctx.execute("""
+            WITH u AS (
+                SELECT
+                    small.FQDN,
+                    small.NS1,
+                    small.NS2,
+                    small.NS3
+                FROM small
+                INNER JOIN large ON small.FQDN = large.FQDN
+            )
+            UPDATE large
+            SET
+                FQDN = u.FQDN,
+                NS1 = u.NS1,
+                NS2 = u.NS2,
+                NS3 = u.NS3
+            FROM u
+            WHERE large.FQDN = u.FQDN
+        """)
+
+
 @pytest.mark.parametrize("op", ["EXCEPT", "INTERSECT", "UNION"])
 def test_except_intersect_errors(op: str) -> None:
     df1 = pl.DataFrame({"x": [1, 9, 1, 1], "y": [2, 3, 4, 4], "z": [5, 5, 5, 5]})  # noqa: F841