Accept exp.Expr (incl. exp.Query) as sqlglot query type

README.md

@@ -543,7 +543,7 @@ for a `Response`, and the rows are *not* validated against the model.
   starts and released when it ends. A `con` must stay open until the response has been sent, i.e. come from a `Depends`
   with `yield` (FastAPI >= 0.118); `async with pool.connection() as conn: return PostgresJsonResponse(q, con=conn)`
   closes it too early. Prefer the pool.
-- **Query** — a literal string, a `psycopg.sql` composable, a sqlglot expression or a t-string, with values bound
+- **Query** — a literal string, a `psycopg.sql` composable, a sqlglot expression (any `exp.Expr`, so Select, Union, Subquery and other `exp.Query` values are fine) or a t-string, with values bound
   through `params` (`%(name)s`; a t-string carries its own values and takes no `params`). Passing only literals
   is enforced by the type checker, **not at runtime**: the query is wrapped as a subquery, never parsed, so never
   build it from user input. For a data-modifying CTE, or custom JSON, pass `query_produces_json=True` and return one

pgdevkit/db/fetch.py

@@ -28,7 +28,9 @@
 # The only things a query may be: a literal string (e.g. from `SqlLoader.load_sql()`), a sqlglot
 # expression, or psycopg's own safe-composition types (`psycopg.sql` / a t-string). A plain `str`
 # is rejected by the type checker on purpose, so user-controlled text can't be concatenated in.
-type SqlQueryNoTemplate = LiteralString | exp.Expression | Composable
+# The sqlglot type is `exp.Expr` (any query: Select, Union, Subquery...), not `exp.Expression`: in
+# sqlglot 30 `exp.Query` is an `exp.Expr` but not an `exp.Expression` (issue #56), so don't narrow it back.
+type SqlQueryNoTemplate = LiteralString | exp.Expr | Composable
 type SqlQuery = SqlQueryNoTemplate | Template
 
 

pyproject.toml

@@ -11,7 +11,7 @@ packages = ["pgdevkit"]
 
 [project]
 name = "pgdevkit"
-version = "0.15.1"
+version = "0.15.2"
 description = "A helper for developing with Postgres"
 readme = "README.md"
 requires-python = ">=3.14"

skills/pgdevkit/SKILL.md

@@ -234,7 +234,7 @@ deleted = await execute("DELETE FROM sessions WHERE user_id = %(id)s", {"id": 1}
 - Passing `con=` runs on *your* connection: `fetch_all` never commits, closes or releases it. Without it, a connection is borrowed from `pool=` (or the `set_default_pool()` pool) and released afterwards.
 - `statement_timeout=<seconds>` (optional) lets Postgres abort the query after that long, raising `QueryCanceled`; it is applied with `SET LOCAL` semantics for this call only (a borrowed `con` gets its previous value back). A plain value *replaces* the connection's own setting in both directions, so a per-call 900 s overrides a database-level safety net of 600 s; `statement_timeout=at_most(120)` (`from pgdevkit.db import at_most`) never loosens it: the effective timeout is `min(current, 120)`.
 - `cancel=` takes an `asyncio.Event`: set it (e.g. when the HTTP client disconnects) and the running query is cancelled on the server, raising `psycopg.errors.QueryCanceled`. If the awaiting task is cancelled instead, the server-side query is cancelled too and `CancelledError` propagates as usual (not `QueryCanceled`), so no query keeps running unattended. Don't use `cancel=` on a `con` that concurrent tasks share (it aborts whatever statement is running on it); if the cancel request itself can't be delivered, the connection is closed. After `QueryCanceled` a borrowed `con` is usable again once you `await con.rollback()` (autocommit connections need nothing). To stream a whole result as JSON from FastAPI (it wires the disconnect itself) see `PostgresJsonResponse` below.
-- `query` is typed `SqlQuery`: a literal string (`SqlLoader.load_sql()`), a sqlglot expression, a `psycopg.sql` composable or a t-string. A plain `str` fails type checking on purpose, so don't build SQL with f-strings/concatenation. A sqlglot expression is only as safe as the strings it was built from (its builders parse plain strings as SQL): pass user values as `exp.Placeholder` + `params`, never as literals/raw text. A literal `%` in a sqlglot expression (`LIKE 'a%'`, modulo) just works, with or without `params`; in a plain string or `psycopg.sql` composable it is psycopg's rule: write `%%` when you pass `params`. A t-string carries its own values, so don't also pass `params`.
+- `query` is typed `SqlQuery`: a literal string (`SqlLoader.load_sql()`), a sqlglot expression (any `exp.Expr`, so Select, Union, Subquery and other `exp.Query` values are fine), a `psycopg.sql` composable or a t-string. A plain `str` fails type checking on purpose, so don't build SQL with f-strings/concatenation. A sqlglot expression is only as safe as the strings it was built from (its builders parse plain strings as SQL): pass user values as `exp.Placeholder` + `params`, never as literals/raw text. A literal `%` in a sqlglot expression (`LIKE 'a%'`, modulo) just works, with or without `params`; in a plain string or `psycopg.sql` composable it is psycopg's rule: write `%%` when you pass `params`. A t-string carries its own values, so don't also pass `params`.
 
 ### Dynamic SQL
 

tests/db/test_fetch.py

@@ -120,6 +120,17 @@ async def test_accepts_sqlglot_expression_psycopg_sql_and_t_string(pool: PgPool)
     assert await fetch_all(t"SELECT name FROM widget WHERE id = {widget_id}", pool=pool) == [{"name": "sprocket"}]
 
 
+@requires_podman
+async def test_accepts_exp_query_typed_union_and_subquery(pool: PgPool):
+    # annotated exp.Query on purpose (issue #56)
+    union: exp.Query = select("name").from_("widget").where("id = 1").union(select("name").from_("widget").where("id = 2"))
+    assert sorted(r["name"] for r in await fetch_all(union, pool=pool)) == ["cog", "sprocket"]
+    sub: exp.Query = select("name").from_(select("name").from_("widget").subquery("w")).where("name = 'cog'")
+    assert await fetch_all(sub, pool=pool) == [{"name": "cog"}]
+    assert await fetch_one(union, pool=pool) is not None
+    assert await fetch_scalar(select("count(*)").from_(union.subquery("u")), pool=pool) == 2
+
+
 async def test_already_set_cancel_raises_before_touching_a_connection():
     cancel = asyncio.Event()
     cancel.set()

tests/db/test_query_expr_types.py

@@ -0,0 +1,69 @@
+"""sqlglot queries are typed as `exp.Expr`, so an `exp.Query` (base of Select/Union/Subquery) is accepted too.
+
+In sqlglot 30 `exp.Query` is an `exp.Expr` but not an `exp.Expression`; the annotations used to say
+`exp.Expression`, which a type checker such as pyright rejected for `exp.Query`-typed values (issue #56).
+"""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+
+import pytest
+from sqlglot import exp, select
+
+from pgdevkit.db.fetch import _render
+from pgdevkit.fastapi import PostgresJsonResponse
+
+
+def _union() -> exp.Query:
+    return select("1 AS a").union(select("2 AS a"))
+
+
+def _subquery() -> exp.Query:
+    return select("a").from_(select("1 AS a").subquery("t"))
+
+
+# --- runtime: no database needed for `_render` ---------------------------------------------------------------------
+
+
+def test_render_union_and_subquery():
+    assert _render(_union()) == "SELECT 1 AS a UNION SELECT 2 AS a"
+    assert _render(_subquery()) == "SELECT a FROM (SELECT 1 AS a) AS t"
+    assert _render(_union(), {}) == "SELECT 1 AS a UNION SELECT 2 AS a"
+
+
+def test_postgres_json_response_accepts_union_and_subquery():
+    for query in (_union(), _subquery()):
+        assert _render(query) in PostgresJsonResponse(query).query.as_string()
+    with pytest.raises(ValueError, match="Command"):
+        PostgresJsonResponse(select("1").union(select(exp.Command(this="SELECT", expression=" version()"))))
+
+
+def test_render_keeps_placeholders_and_escapes_percent_inside_a_union():
+    query = (
+        select("name").from_("w").where(exp.column("id").eq(exp.Placeholder(this="id"))).where("name LIKE '50%'")
+    ).union(select("'x%'"))
+    rendered = _render(query, {"id": 1})
+    assert "%(id)s" in rendered
+    assert "'50%%'" in rendered
+    assert "'x%%'" in rendered
+
+
+def test_render_rejects_command_and_bad_placeholder_nested_in_a_union():
+    command = exp.Command(this="SELECT", expression=" version()")
+    with pytest.raises(ValueError, match="Command"):
+        _render(select("1").union(select(command)))
+    with pytest.raises(ValueError, match="Command"):
+        _render(select("a").from_(select(command).subquery("t")))
+    bad = exp.Placeholder(this="x); DROP TABLE t;--")
+    with pytest.raises(ValueError, match="placeholder name"):
+        _render(select("1").union(select("2").where(bad)), {"x": 1})
+    with pytest.raises(ValueError, match="placeholder name"):
+        _render(select("a").from_(select("2 AS a").where(bad).subquery("t")), {"x": 1})
+
+
+def test_importing_pgdevkit_db_does_not_import_sqlglot():
+    code = "import sys, pgdevkit.db; sys.exit('sqlglot' in sys.modules)"
+    assert subprocess.run([sys.executable, "-c", code], check=False).returncode == 0
+