diff --git a/README.md b/README.md index d803cbf..e14c89e 100644 --- a/README.md +++ b/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 diff --git a/pgdevkit/db/fetch.py b/pgdevkit/db/fetch.py index b3b94e8..3a234dd 100644 --- a/pgdevkit/db/fetch.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index f38552f..a505512 100644 --- a/pyproject.toml +++ b/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" diff --git a/skills/pgdevkit/SKILL.md b/skills/pgdevkit/SKILL.md index e1d7674..61c37e0 100644 --- a/skills/pgdevkit/SKILL.md +++ b/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=` (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 diff --git a/tests/db/test_fetch.py b/tests/db/test_fetch.py index 63e5428..b4972e8 100644 --- a/tests/db/test_fetch.py +++ b/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() diff --git a/tests/db/test_query_expr_types.py b/tests/db/test_query_expr_types.py new file mode 100644 index 0000000..166869f --- /dev/null +++ b/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 + diff --git a/uv.lock b/uv.lock index 21dd842..9889250 100644 --- a/uv.lock +++ b/uv.lock @@ -451,7 +451,7 @@ wheels = [ [[package]] name = "pgdevkit" -version = "0.15.1" +version = "0.15.2" source = { editable = "." } dependencies = [ { name = "docker" },