Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 17 additions & 13 deletions src/pycel2sql/_converter.py
Original file line number Diff line number Diff line change
Expand Up @@ -1674,9 +1674,10 @@ def _visit_comprehension(self, source: Branch, macro_name: str, args: list[Branc
self._comprehension_depth -= 1

def _write_unnest_source(self, source: Branch, iter_var: str) -> None:
"""Write the UNNEST(source) AS var clause."""
self._dialect.write_unnest(self._w, lambda: self._visit_child(source))
self._w.write(f" AS {iter_var}")
"""Write the comprehension source clause, binding iter_var."""
self._dialect.write_comprehension_source(
self._w, lambda: self._visit_child(source), iter_var
)

def _visit_comp_all(self, source: Branch, args: list[Branch]) -> None:
"""all(x, pred) -> NOT EXISTS (SELECT 1 FROM UNNEST(src) AS x WHERE NOT (pred))"""
Expand All @@ -1686,11 +1687,13 @@ def _visit_comp_all(self, source: Branch, args: list[Branch]) -> None:
pred = args[1]
self._comprehension_vars.add(iter_var)
try:
self._w.write("NOT EXISTS (SELECT 1 FROM ")
self._write_unnest_source(source, iter_var)
self._w.write(" WHERE NOT (")
self._visit_child(pred)
self._w.write("))")
def _body() -> None:
self._write_unnest_source(source, iter_var)
self._w.write(" WHERE NOT (")
self._visit_child(pred)
self._w.write(")")

self._dialect.write_comprehension_not_exists(self._w, _body)
finally:
self._comprehension_vars.discard(iter_var)

Expand All @@ -1702,11 +1705,12 @@ def _visit_comp_exists(self, source: Branch, args: list[Branch]) -> None:
pred = args[1]
self._comprehension_vars.add(iter_var)
try:
self._w.write("EXISTS (SELECT 1 FROM ")
self._write_unnest_source(source, iter_var)
self._w.write(" WHERE ")
self._visit_child(pred)
self._w.write(")")
def _body() -> None:
self._write_unnest_source(source, iter_var)
self._w.write(" WHERE ")
self._visit_child(pred)

self._dialect.write_comprehension_exists(self._w, _body)
finally:
self._comprehension_vars.discard(iter_var)

Expand Down
26 changes: 26 additions & 0 deletions src/pycel2sql/dialect/_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -223,6 +223,32 @@ def write_format(
@abstractmethod
def write_unnest(self, w: StringIO, write_source: WriteFunc) -> None: ...

def write_comprehension_source(
self, w: StringIO, write_source: WriteFunc, iter_var: str
) -> None:
"""Write the FROM source of a comprehension subquery, binding iter_var.

The default binds the unnest alias directly, which is only correct
where the source is row-valued (PostgreSQL/BigQuery/Spark UNNEST).
Table-valued sources (SQLite json_each, MySQL JSON_TABLE) and
DuckDB's struct-producing UNNEST override this to rename the value
column to iter_var, so bare references to it resolve.
"""
self.write_unnest(w, write_source)
w.write(f" AS {iter_var}")

def write_comprehension_exists(self, w: StringIO, write_body: WriteFunc) -> None:
"""Wrap a comprehension's existential subquery: EXISTS (SELECT 1 FROM <body>)."""
w.write("EXISTS (SELECT 1 FROM ")
write_body()
w.write(")")

def write_comprehension_not_exists(self, w: StringIO, write_body: WriteFunc) -> None:
"""Negation of write_comprehension_exists: NOT EXISTS (SELECT 1 FROM <body>)."""
w.write("NOT EXISTS (SELECT 1 FROM ")
write_body()
w.write(")")

@abstractmethod
def write_array_subquery_open(self, w: StringIO) -> None: ...

Expand Down
8 changes: 8 additions & 0 deletions src/pycel2sql/dialect/duckdb.py
Original file line number Diff line number Diff line change
Expand Up @@ -248,6 +248,14 @@ def write_unnest(self, w: StringIO, write_source: WriteFunc) -> None:
write_source()
w.write(")")

def write_comprehension_source(
self, w: StringIO, write_source: WriteFunc, iter_var: str
) -> None:
# DuckDB's FROM UNNEST(arr) AS a binds a to a STRUCT(unnest ...) row,
# so a = 'x' is a cast error; the column-alias form binds the value.
self.write_unnest(w, write_source)
w.write(f" AS _t({iter_var})")

def write_array_subquery_open(self, w: StringIO) -> None:
w.write("ARRAY(SELECT ")

Expand Down
22 changes: 22 additions & 0 deletions src/pycel2sql/dialect/mysql.py
Original file line number Diff line number Diff line change
Expand Up @@ -283,6 +283,28 @@ def write_unnest(self, w: StringIO, write_source: WriteFunc) -> None:
write_source()
w.write(", '$[*]' COLUMNS(value TEXT PATH '$'))")

def write_comprehension_source(
self, w: StringIO, write_source: WriteFunc, iter_var: str
) -> None:
# JSON_TABLE is table-valued, so the value column is renamed to
# iter_var through a derived table (JSON_TABLE itself needs an alias).
w.write(f"(SELECT value AS {iter_var} FROM ")
self.write_unnest(w, write_source)
w.write(" AS jt) AS _t")

def write_comprehension_exists(self, w: StringIO, write_body: WriteFunc) -> None:
# Not EXISTS: the MySQL 8.x optimizer turns a correlated EXISTS into a
# semijoin and loses the correlation to JSON_TABLE, silently matching
# nothing (works from 9.x). COUNT comparisons are never transformed.
w.write("(SELECT COUNT(*) FROM ")
write_body()
w.write(") > 0")

def write_comprehension_not_exists(self, w: StringIO, write_body: WriteFunc) -> None:
w.write("(SELECT COUNT(*) FROM ")
write_body()
w.write(") = 0")

def write_array_subquery_open(self, w: StringIO) -> None:
w.write("(SELECT JSON_ARRAYAGG(")

Expand Down
9 changes: 9 additions & 0 deletions src/pycel2sql/dialect/sqlite.py
Original file line number Diff line number Diff line change
Expand Up @@ -313,6 +313,15 @@ def write_unnest(self, w: StringIO, write_source: WriteFunc) -> None:
write_source()
w.write(")")

def write_comprehension_source(
self, w: StringIO, write_source: WriteFunc, iter_var: str
) -> None:
# json_each is table-valued (rows of key, value, type, ...), so a bare
# reference to the alias is "no such column"; rename value to iter_var.
w.write(f"(SELECT value AS {iter_var} FROM ")
self.write_unnest(w, write_source)
w.write(") AS _t")

def write_array_subquery_open(self, w: StringIO) -> None:
w.write("(SELECT json_group_array(")

Expand Down
74 changes: 74 additions & 0 deletions tests/integration/test_comprehensions.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,74 @@
"""Integration tests for comprehensions — executed against real databases.

These exist because comprehension SQL can parse and still be wrong: SQLite's
json_each and MySQL's JSON_TABLE are table-valued, so a bare reference to the
iteration variable is "no such column"; DuckDB's UNNEST binds the alias to a
STRUCT row; and MySQL 8.x turns a correlated EXISTS into a semijoin that loses
the JSON_TABLE correlation and silently matches nothing.
"""

from __future__ import annotations

import pytest

from pycel2sql.dialect.duckdb import DuckDBDialect
from pycel2sql.dialect.mysql import MySQLDialect
from pycel2sql.dialect.postgres import PostgresDialect
from pycel2sql.dialect.sqlite import SQLiteDialect
from pycel2sql.schema import FieldSchema, Schema

from tests.integration.conftest import execute_cel, get_names


pytestmark = pytest.mark.integration

COMPREHENSION_DBS = [
pytest.param(("pg", PostgresDialect()), id="pg", marks=pytest.mark.postgres),
pytest.param(("duckdb", DuckDBDialect()), id="duckdb", marks=pytest.mark.duckdb),
pytest.param(("sqlite", SQLiteDialect()), id="sqlite", marks=pytest.mark.sqlite),
pytest.param(("mysql", MySQLDialect()), id="mysql", marks=pytest.mark.mysql),
]

SCHEMAS = {"t": Schema([FieldSchema("tags", repeated=True)])}


@pytest.fixture(params=COMPREHENSION_DBS)
def comp_db(request):
db_name, dialect = request.param
conn = request.getfixturevalue(f"{db_name}_db")
return conn, dialect, db_name


class TestComprehensions:
def test_exists(self, comp_db):
conn, dialect, name = comp_db
rows = execute_cel(
conn, 't.tags.exists(x, x == "python")', dialect, name,
schemas=SCHEMAS, table_alias="t",
)
assert get_names(rows) == {"Alice", "Charlie"}

def test_exists_no_match(self, comp_db):
conn, dialect, name = comp_db
rows = execute_cel(
conn, 't.tags.exists(x, x == "cobol")', dialect, name,
schemas=SCHEMAS, table_alias="t",
)
assert get_names(rows) == set()

def test_all(self, comp_db):
conn, dialect, name = comp_db
rows = execute_cel(
conn, 't.tags.all(x, x != "rust")', dialect, name,
schemas=SCHEMAS, table_alias="t",
)
# Eve has no tags: all() over an empty array is vacuously true.
assert get_names(rows) == {"Alice", "Charlie", "Eve"}

def test_exists_one(self, comp_db):
conn, dialect, name = comp_db
rows = execute_cel(
conn, 't.tags.exists_one(x, x == "go")', dialect, name,
schemas=SCHEMAS, table_alias="t",
)
assert get_names(rows) == {"Bob", "Charlie"}
6 changes: 5 additions & 1 deletion tests/test_mysql.py
Original file line number Diff line number Diff line change
Expand Up @@ -153,8 +153,12 @@ def test_filter(self, d):
def test_exists(self, d):
schemas = {"t": Schema([FieldSchema("arr", repeated=True)])}
result = convert("t.arr.exists(x, x > 5)", dialect=d, schemas=schemas)
assert "EXISTS" in result
# COUNT comparison, not EXISTS: MySQL 8.x turns a correlated EXISTS
# into a semijoin and loses the JSON_TABLE correlation.
assert result.startswith("(SELECT COUNT(*) FROM ")
assert result.endswith(") > 0")
assert "JSON_TABLE(" in result
assert "value AS x" in result


class TestMySQLStruct:
Expand Down
Loading