test(sql): verify recursive CTE handling for #12041 (do not merge)

ibis/backends/sql/compilers/__init__.py

@@ -19,6 +19,7 @@
     "PySparkCompiler",
     "RisingWaveCompiler",
     "SQLiteCompiler",
+    "SingleStoreDBCompiler",
     "SnowflakeCompiler",
     "TrinoCompiler",
 ]
@@ -40,6 +41,7 @@
 from ibis.backends.sql.compilers.postgres import PostgresCompiler
 from ibis.backends.sql.compilers.pyspark import PySparkCompiler
 from ibis.backends.sql.compilers.risingwave import RisingWaveCompiler
+from ibis.backends.sql.compilers.singlestoredb import SingleStoreDBCompiler
 from ibis.backends.sql.compilers.snowflake import SnowflakeCompiler
 from ibis.backends.sql.compilers.sqlite import SQLiteCompiler
 from ibis.backends.sql.compilers.trino import TrinoCompiler

ibis/backends/sql/compilers/base.py

@@ -679,16 +679,21 @@ def fn(node, __unused__, **kwargs):
             if "alias" in this.args:
                 this = this.this
             modified_cte = sge.CTE(
-                alias=sg.to_identifier(aliases[cte], quoted=self.quoted), this=this
+                alias=sge.TableAlias(
+                    this=sg.to_identifier(aliases[cte], quoted=self.quoted)
+                ),
+                this=this,
             )
             merged_ctes.append(modified_cte)
         merged_ctes.extend(out.ctes)
-        out.args.pop(WITH_ARG, None)
+        with_ = out.args.pop(WITH_ARG, None)
+        recursive = with_ is not None and with_.recursive
 
         out = reduce(
             lambda parsed, cte: parsed.with_(
                 cte.args["alias"],
                 as_=cte.args["this"],
+                recursive=recursive,
                 dialect=self.dialect,
                 copy=False,
             ),
@@ -1612,18 +1617,24 @@ def add_query_to_expr(self, *, name: str, table: ir.Table, query: str) -> str:
         ctes = [
             *compiled_ibis_expr.ctes,
             sge.CTE(
-                alias=sg.to_identifier(name, quoted=self.quoted),
+                alias=sge.TableAlias(this=sg.to_identifier(name, quoted=self.quoted)),
                 this=compiled_ibis_expr,
             ),
             *compiled_query.ctes,
         ]
-        compiled_ibis_expr.args.pop(WITH_ARG, None)
-        compiled_query.args.pop(WITH_ARG, None)
+        # either side may contain a recursive CTE, and RECURSIVE applies to the
+        # whole WITH clause, so keep it if either one had it
+        recursive = False
+        for expr in (compiled_ibis_expr, compiled_query):
+            if (with_ := expr.args.pop(WITH_ARG, None)) is not None:
+                recursive |= with_.recursive
 
         # pull existing CTEs from the compiled Ibis expression and combine them
         # with the new query
         parsed = reduce(
-            lambda parsed, cte: parsed.with_(cte.args["alias"], as_=cte.args["this"]),
+            lambda parsed, cte: parsed.with_(
+                cte.args["alias"], as_=cte.args["this"], recursive=recursive
+            ),
             ctes,
             compiled_query,
         )

ibis/backends/sql/compilers/singlestoredb.py

@@ -825,18 +825,24 @@ def add_query_to_expr(self, *, name: str, table, query: str) -> str:
         ctes = [
             *compiled_ibis_expr.ctes,
             sge.CTE(
-                alias=sg.to_identifier(name, quoted=self.quoted),
+                alias=sge.TableAlias(this=sg.to_identifier(name, quoted=self.quoted)),
                 this=compiled_ibis_expr,
             ),
             *compiled_query.ctes,
         ]
-        compiled_ibis_expr.args.pop(WITH_ARG, None)
-        compiled_query.args.pop(WITH_ARG, None)
+        # either side may contain a recursive CTE, and RECURSIVE applies to the
+        # whole WITH clause, so keep it if either one had it
+        recursive = False
+        for expr in (compiled_ibis_expr, compiled_query):
+            if (with_ := expr.args.pop(WITH_ARG, None)) is not None:
+                recursive |= with_.recursive
 
         # pull existing CTEs from the compiled Ibis expression and combine them
         # with the new query
         parsed = reduce(
-            lambda parsed, cte: parsed.with_(cte.args["alias"], as_=cte.args["this"]),
+            lambda parsed, cte: parsed.with_(
+                cte.args["alias"], as_=cte.args["this"], recursive=recursive
+            ),
             ctes,
             compiled_query,
         )

ibis/backends/sql/tests/test_compiler.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
+import pytest
 import sqlglot as sg
 
 import ibis
 from ibis import _
+from ibis.backends import _get_backend_names
 from ibis.backends.sql.dialects import Trino
 
 
@@ -26,3 +28,11 @@ def test_transpile_join():
         "SELECT * FROM t1 JOIN t2 ON x = y", read="duckdb", write=Trino
     )
     assert "CROSS JOIN" not in result
+
+
+@pytest.mark.parametrize("dialect", sorted(_get_backend_names(exclude=("polars",))))
+def test_to_sql_knows_every_sql_backend(dialect):
+    # `ibis.to_sql` looks the dialect up on `ibis.backends.sql.compilers`, so
+    # every SQL backend's compiler module has to be imported there
+    expr = ibis.table({"a": "int64"}, name="t")
+    assert ibis.to_sql(expr, dialect=dialect)

ibis/backends/tests/snapshots/test_sql/test_cte_refs_in_topo_order/mssql/out.sql

@@ -1,6 +1,6 @@
 WITH [t1] AS (
   SELECT
-    [t0].[key]
+    [t0].[key] AS [key]
   FROM [leaf] AS [t0]
   WHERE
     (1 = 1)

ibis/backends/tests/test_dot_sql.py

@@ -5,20 +5,31 @@
 import pytest
 import sqlglot as sg
 import sqlglot.expressions as sge
+import toolz
 from pytest import param
 
 import ibis
+import ibis.backends.sql.compilers as sc
 import ibis.backends.sql.dialects  # to load dialects
 import ibis.common.exceptions as com
+import ibis.expr.operations as ops
 from ibis import _
 from ibis.backends import _get_backend_names
+from ibis.backends.conftest import _get_backends_to_test
 from ibis.backends.tests.base import PYTHON_SHORT_VERSION
 from ibis.backends.tests.errors import (
     ExaQueryError,
     GoogleBadRequest,
+    ImpalaHiveServer2Error,
     OracleDatabaseError,
+    PolarsSQLInterfaceError,
+    PsycoPg2InternalError,
+    PsycoPgSyntaxError,
+    Py4JJavaError,
     PyAthenaOperationalError,
+    PyDruidProgrammingError,
     PyODBCProgrammingError,
+    SingleStoreDBOperationalError,
 )
 
 pd = pytest.importorskip("pandas")
@@ -386,3 +397,204 @@ def test_scalar_dot_sql(con):
     sql = sg.select(sge.convert(1).as_("a")).sql(con.dialect)
     expr = con.sql(sql).as_scalar()
     assert expr.type().is_numeric()
+
+
+RECURSIVE_COUNTER = """
+WITH RECURSIVE "counter" AS (
+  SELECT 1 AS "n"
+  UNION ALL
+  SELECT "n" + 1 FROM "counter" WHERE "n" < 3
+)
+SELECT "n" FROM "counter"
+"""
+
+RECURSIVE_COUNTER_FROM_T = """
+WITH RECURSIVE "counter" AS (
+  SELECT "x" AS "n" FROM "t"
+  UNION ALL
+  SELECT "n" + 1 FROM "counter" WHERE "n" < 3
+)
+SELECT "n" FROM "counter"
+"""
+
+
+@pytest.mark.parametrize(
+    ("make_expr", "expected"),
+    [
+        param(
+            lambda con: con.sql(RECURSIVE_COUNTER, dialect="duckdb"),
+            [1, 2, 3],
+            id="con_sql",
+        ),
+        param(
+            lambda con: con.sql(RECURSIVE_COUNTER, dialect="duckdb").filter(_.n > 1),
+            [2, 3],
+            marks=pytest.mark.notyet(
+                ["bigquery"],
+                raises=GoogleBadRequest,
+                reason="BigQuery rejects WITH RECURSIVE nested in a subquery or CTE",
+            ),
+            id="con_sql_filtered",
+        ),
+        param(
+            lambda con: (
+                con.sql('SELECT 1 AS "x"', dialect="duckdb")
+                .alias("t")
+                .sql(RECURSIVE_COUNTER_FROM_T, dialect="duckdb")
+            ),
+            [1, 2, 3],
+            id="table_sql_recursive_query",
+        ),
+        param(
+            lambda con: (
+                con.sql(RECURSIVE_COUNTER, dialect="duckdb")
+                .alias("t")
+                .sql('SELECT "n" FROM "t"', dialect="duckdb")
+            ),
+            [1, 2, 3],
+            marks=pytest.mark.notyet(
+                ["bigquery"],
+                raises=GoogleBadRequest,
+                reason="BigQuery rejects WITH RECURSIVE nested in a subquery or CTE",
+            ),
+            id="table_sql_recursive_parent",
+        ),
+    ],
+)
+@pytest.mark.notyet(
+    ["polars"],
+    raises=PolarsSQLInterfaceError,
+    reason="Polars SQL doesn't support recursive CTEs",
+)
+@pytest.mark.notyet(
+    ["exasol"], raises=ExaQueryError, reason="Exasol doesn't support WITH RECURSIVE"
+)
+@pytest.mark.notyet(
+    ["risingwave"],
+    raises=PsycoPg2InternalError,
+    reason="RisingWave doesn't support recursive CTEs in batch queries",
+)
+@pytest.mark.notyet(
+    ["druid"],
+    raises=PyDruidProgrammingError,
+    reason="Druid doesn't support recursive CTEs",
+)
+@pytest.mark.notyet(
+    ["flink"],
+    raises=Py4JJavaError,
+    reason="Flink doesn't support recursive CTEs; sqlglot drops RECURSIVE",
+)
+@pytest.mark.notyet(
+    ["impala"],
+    raises=ImpalaHiveServer2Error,
+    reason="Impala doesn't support recursive CTEs; sqlglot drops RECURSIVE",
+)
+@pytest.mark.notyet(
+    ["materialize"],
+    raises=PsycoPgSyntaxError,
+    reason="Materialize only supports WITH MUTUALLY RECURSIVE",
+)
+@pytest.mark.notyet(
+    ["oracle"],
+    raises=OracleDatabaseError,
+    reason="sqlglot emits WITH RECURSIVE, which Oracle doesn't accept",
+)
+@pytest.mark.notyet(
+    ["singlestoredb"],
+    raises=SingleStoreDBOperationalError,
+    reason="unsupported recursive common table expression query shape",
+)
+@pytest.mark.notyet(
+    ["pyspark"],
+    reason="Spark < 4.1 doesn't support recursive CTEs (3.5 fails to parse, 4.0 to plan)",
+)
+def test_dot_sql_recursive_cte(con, make_expr, expected):
+    expr = make_expr(con)
+    result = expr.execute()
+    assert result["n"].astype("int64").sort_values().tolist() == expected
+
+
+def _sql_query_result(query, compiler):
+    """Build what `con.sql(query, dialect="duckdb")` builds, without a connection."""
+    query = sg.transpile(query, read="duckdb", write=compiler.dialect)[0]
+    return ops.SQLQueryResult(query, ibis.schema({"n": "int64"}), None).to_expr()
+
+
+no_recursive_keyword = toolz.compose(
+    pytest.mark.never(
+        ["mssql"], raises=AssertionError, reason="T-SQL has no RECURSIVE keyword"
+    ),
+    pytest.mark.notyet(
+        ["flink", "impala"],
+        raises=AssertionError,
+        reason="recursive CTEs aren't supported, so sqlglot drops RECURSIVE",
+    ),
+    pytest.mark.xfail_version(
+        pyspark=["sqlglot<30"],
+        databricks=["sqlglot<30"],
+        raises=AssertionError,
+        reason="sqlglot drops RECURSIVE for Spark before 30.0",
+    ),
+)
+
+
+@no_recursive_keyword
+@pytest.mark.parametrize("backend_name", _get_backends_to_test(discard=("polars",)))
+def test_recursive_cte_con_sql_compiles(backend_name):
+    compiler = getattr(sc, backend_name).compiler
+    expr = _sql_query_result(RECURSIVE_COUNTER, compiler)
+    assert "WITH RECURSIVE" in compiler.to_sqlglot(expr).sql(compiler.dialect)
+
+
+@no_recursive_keyword
+@pytest.mark.parametrize(
+    ("make_table", "query"),
+    [
+        param(
+            lambda _compiler: ibis.table({"x": "int64"}, name="src"),
+            RECURSIVE_COUNTER_FROM_T,
+            id="recursive_query",
+        ),
+        param(
+            lambda compiler: _sql_query_result(RECURSIVE_COUNTER, compiler),
+            'SELECT "n" FROM "t"',
+            id="recursive_parent",
+        ),
+    ],
+)
+@pytest.mark.parametrize("backend_name", _get_backends_to_test(discard=("polars",)))
+def test_recursive_cte_table_sql_compiles(backend_name, make_table, query):
+    compiler = getattr(sc, backend_name).compiler
+    query = sg.transpile(query, read="duckdb", write=compiler.dialect)[0]
+    sql = compiler.add_query_to_expr(name="t", table=make_table(compiler), query=query)
+    assert "WITH RECURSIVE" in sql
+
+
+_REC = (
+    "WITH RECURSIVE c AS (SELECT 1 AS n UNION ALL SELECT n + 1 FROM c WHERE n < 3) "
+    "SELECT n FROM c"
+)
+
+
+@pytest.mark.parametrize(
+    "sql",
+    [
+        param(_REC, id="top_level"),
+        param(f"SELECT * FROM ({_REC}) AS t0 WHERE n > 1", id="nested_in_subquery"),
+        param(f"WITH t AS ({_REC}) SELECT n FROM t", id="nested_in_cte"),
+        param(
+            "WITH RECURSIVE c AS (SELECT 1 AS n UNION ALL SELECT n + 1 FROM c "
+            "WHERE n < 3), t AS (SELECT n FROM c) SELECT n FROM t",
+            id="hoisted",
+        ),
+        param(
+            "SELECT * FROM (WITH c AS (SELECT 1 AS n) SELECT n FROM c) AS t0",
+            id="control_nonrecursive_nested",
+        ),
+    ],
+)
+def test_recursive_cte_bigquery_probe(con, sql):
+    """Temporary diagnostic: which placements of WITH RECURSIVE BigQuery accepts."""
+    if con.name != "bigquery":
+        pytest.skip("BigQuery-only diagnostic")
+    assert len(list(con.raw_sql(sql))) > 0

pyproject.toml

@@ -383,6 +383,9 @@ xfail_strict = true
 addopts = [
   "--strict-markers",
   "--strict-config",
+  "-v",
+  "-k",
+  "recursive_cte or test_dot_sql or test_sql",
   "--benchmark-disable",
   "--benchmark-group-by=name",
   "--benchmark-sort=name",