fix(dot-sql): ensure that CTEs can be used in .sql

Author: cpcloud

ibis/backends/polars/__init__.py

@@ -449,17 +449,12 @@ def compile(
 
         return translate(node, ctx=self._context)
 
-    def _get_sql_string_view_schema(self, name, table, query) -> sch.Schema:
-        import sqlglot as sg
-
-        cte = sg.parse_one(str(ibis.to_sql(table, dialect="postgres")), read="postgres")
-        parsed = sg.parse_one(query, read=self.dialect)
-        parsed.args["with"] = cte.args.pop("with", [])
-        parsed = parsed.with_(
-            sg.to_identifier(name, quoted=True), as_=cte, dialect=self.dialect
-        )
+    def _get_sql_string_view_schema(
+        self, *, name: str, table: ir.Table, query: str
+    ) -> sch.Schema:
+        from ibis.backends.sql.compilers.postgres import compiler
 
-        sql = parsed.sql(self.dialect)
+        sql = compiler.add_query_to_expr(name=name, table=table, query=query)
         return self._get_schema_using_query(sql)
 
     def _get_schema_using_query(self, query: str) -> sch.Schema:

ibis/backends/sql/__init__.py

@@ -191,18 +191,10 @@ def _get_schema_using_query(self, query: str) -> sch.Schema:
             The schema inferred from `query`
         """
 
-    def _get_sql_string_view_schema(self, name, table, query) -> sch.Schema:
-        compiler = self.compiler
-        dialect = compiler.dialect
-
-        cte = compiler.to_sqlglot(table)
-        parsed = sg.parse_one(query, read=dialect)
-        parsed.args["with"] = cte.args.pop("with", [])
-        parsed = parsed.with_(
-            sg.to_identifier(name, quoted=compiler.quoted), as_=cte, dialect=dialect
-        )
-
-        sql = parsed.sql(dialect)
+    def _get_sql_string_view_schema(
+        self, *, name: str, table: ir.Table, query: str
+    ) -> sch.Schema:
+        sql = self.compiler.add_query_to_expr(name=name, table=table, query=query)
         return self._get_schema_using_query(sql)
 
     def _register_udfs(self, expr: ir.Expr) -> None:

ibis/backends/sql/compilers/base.py

@@ -1613,6 +1613,31 @@ def visit_DropColumns(self, op, *, parent, columns_to_drop):
         )
         return sg.select(*columns_to_keep).from_(parent)
 
+    def add_query_to_expr(self, *, name: str, table: ir.Table, query: str) -> str:
+        dialect = self.dialect
+
+        compiled_ibis_expr = self.to_sqlglot(table)
+
+        # 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"]),
+            compiled_ibis_expr.ctes,
+            sg.parse_one(query, read=dialect),
+        )
+
+        # remove all ctes from the compiled expression, since they're now in
+        # our larger expression
+        compiled_ibis_expr.args.pop("with", None)
+
+        # add the new str query as a CTE
+        parsed = parsed.with_(
+            sg.to_identifier(name, quoted=self.quoted), as_=compiled_ibis_expr
+        )
+
+        # generate the SQL string
+        return parsed.sql(dialect)
+
 
 # `__init_subclass__` is uncalled for subclasses - we manually call it here to
 # autogenerate the base class implementations as well.

ibis/backends/tests/test_dot_sql.py

@@ -5,6 +5,7 @@
 
 import pytest
 import sqlglot as sg
+import sqlglot.expressions as sge
 from pytest import param
 
 import ibis
@@ -24,11 +25,21 @@
 
 _NAMES = {
     "bigquery": f"ibis_gbq_testing_{getpass.getuser()}_{PYTHON_SHORT_VERSION}.functional_alltypes",
-    "exasol": '"functional_alltypes"',
 }
 
 
-@pytest.mark.notyet(["oracle"], reason="table quoting behavior")
+@pytest.fixture(scope="module")
+def ftname_raw(con):
+    return _NAMES.get(con.name, "functional_alltypes")
+
+
+@pytest.fixture(scope="module")
+def ftname(con, ftname_raw):
+    table = sg.parse_one(ftname_raw, into=sge.Table)
+    table = sg.table(table.name, db=table.db, catalog=table.catalog, quoted=True)
+    return table.sql(con.dialect)
+
+
 @dot_sql_never
 @pytest.mark.parametrize(
     "schema",
@@ -37,10 +48,9 @@
         param({"s": "string", "new_col": "double"}, id="explicit_schema"),
     ],
 )
-def test_con_dot_sql(backend, con, schema):
+def test_con_dot_sql(backend, con, schema, ftname):
     alltypes = backend.functional_alltypes
     # pull out the quoted name
-    name = _NAMES.get(con.name, "functional_alltypes")
     quoted = True
     cols = [
         sg.column("string_col", quoted=quoted).as_("s", quoted=quoted).sql(con.dialect),
@@ -50,7 +60,7 @@ def test_con_dot_sql(backend, con, schema):
     ]
     t = (
         con.sql(
-            f"SELECT {', '.join(cols)} FROM {name}",
+            f"SELECT {', '.join(cols)} FROM {ftname}",
             schema=schema,
         )
         .group_by("s")  # group by a column from SQL
@@ -325,9 +335,16 @@ def test_cte(alltypes, df):
 
 
 @dot_sql_never
-def test_bare_minimum(con, alltypes, df):
+def test_bare_minimum(alltypes, df, ftname_raw):
     """Test that a backend that supports dot sql can do the most basic thing."""
 
-    name = _NAMES.get(con.name, "functional_alltypes").replace('"', "")
-    expr = alltypes.sql(f'SELECT COUNT(*) AS "n" FROM "{name}"', dialect="duckdb")
+    expr = alltypes.sql(f'SELECT COUNT(*) AS "n" FROM "{ftname_raw}"', dialect="duckdb")
     assert expr.to_pandas().iat[0, 0] == len(df)
+
+
+@dot_sql_never
+def test_embedded_cte(alltypes, ftname_raw):
+    sql = f'WITH "x" AS (SELECT * FROM "{ftname_raw}") SELECT * FROM "x"'
+    expr = alltypes.sql(sql, dialect="duckdb")
+    result = expr.head(1).execute()
+    assert len(result) == 1

ibis/expr/types/relations.py

@@ -3553,7 +3553,7 @@ def sql(self, query: str, dialect: str | None = None) -> ir.Table:
             name = util.gen_name("sql_query")
             expr = self
 
-        schema = backend._get_sql_string_view_schema(name, expr, query)
+        schema = backend._get_sql_string_view_schema(name=name, table=expr, query=query)
         node = ops.SQLStringView(child=self.op(), query=query, schema=schema)
         return node.to_expr()