diff --git a/ibis/backends/sql/compilers/__init__.py b/ibis/backends/sql/compilers/__init__.py index 4a85c54594cc..57d96cd6bf08 100644 --- a/ibis/backends/sql/compilers/__init__.py +++ b/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 diff --git a/ibis/backends/sql/compilers/base.py b/ibis/backends/sql/compilers/base.py index 2e7333ce1bc4..c49d4b918b83 100644 --- a/ibis/backends/sql/compilers/base.py +++ b/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, ) diff --git a/ibis/backends/sql/compilers/singlestoredb.py b/ibis/backends/sql/compilers/singlestoredb.py index 8d429137deb2..c1baf4fafc3d 100644 --- a/ibis/backends/sql/compilers/singlestoredb.py +++ b/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, ) diff --git a/ibis/backends/sql/tests/test_compiler.py b/ibis/backends/sql/tests/test_compiler.py index 3596a7287b8c..f9248b664f6b 100644 --- a/ibis/backends/sql/tests/test_compiler.py +++ b/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) diff --git a/ibis/backends/tests/snapshots/test_sql/test_cte_refs_in_topo_order/mssql/out.sql b/ibis/backends/tests/snapshots/test_sql/test_cte_refs_in_topo_order/mssql/out.sql index 7a46232efae6..38448f048aef 100644 --- a/ibis/backends/tests/snapshots/test_sql/test_cte_refs_in_topo_order/mssql/out.sql +++ b/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) diff --git a/ibis/backends/tests/test_dot_sql.py b/ibis/backends/tests/test_dot_sql.py index e5bdfe120d90..f8f2e3c9a20a 100644 --- a/ibis/backends/tests/test_dot_sql.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 983305ee9931..71248f41a1b0 100644 --- a/pyproject.toml +++ b/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",