diff --git a/.gitignore b/.gitignore index e276119..180fd4f 100644 --- a/.gitignore +++ b/.gitignore @@ -217,3 +217,5 @@ python_modules/ # pywrangler generated lock for Worker examples pylock.toml + +/.codex/ diff --git a/CHANGELOG.md b/CHANGELOG.md index 7cbe4c4..fa780fd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +## [0.4.1] + +### Fixed + +- Preserve distinct positional values when multiple result columns share a name in synchronous and asynchronous REST connections ([#31](https://github.com/CollierKing/sqlalchemy-cloudflare-d1/issues/31)). Thanks to [@danielalyoshin](https://github.com/danielalyoshin) for [PR #34](https://github.com/CollierKing/sqlalchemy-cloudflare-d1/pull/34). +- Recognize row-returning statements from D1's column metadata, including commented SQL, `VALUES`, and `EXPLAIN`, and make the async SQLAlchemy adapter follow the cursor description ([#32](https://github.com/CollierKing/sqlalchemy-cloudflare-d1/issues/32)). Thanks to [@danielalyoshin](https://github.com/danielalyoshin) for [PR #33](https://github.com/CollierKing/sqlalchemy-cloudflare-d1/pull/33). +- Extend positional read results to synchronous and asynchronous Worker bindings, preserving duplicate columns, NULLs, and empty-result column names. Use a native JavaScript options object for D1's `raw()` API; keep writes on `all()` to preserve affected-row counts and generated IDs, including `WITH`-prefixed inserts ([PR #37](https://github.com/CollierKing/sqlalchemy-cloudflare-d1/pull/37)). +- Correct Worker test-server readiness detection and process cleanup, and honor remote D1 bindings for the new integration tests. + +### Added + +- Unit and real-D1 regression coverage for Worker cursor results and duplicate columns through SQLAlchemy's async REST engine. + + ## [0.4.0] ### Added diff --git a/examples/workers/src/entry.py b/examples/workers/src/entry.py index c6b37f8..3a7b5d7 100644 --- a/examples/workers/src/entry.py +++ b/examples/workers/src/entry.py @@ -24,7 +24,9 @@ async def fetch(self, request, env): path = url.split("/")[-1].split("?")[0] if "/" in url else "" # Core test endpoints (matching REST API tests) - if path == "select": + if path == "cursor-results": + return await self.test_cursor_results() + elif path == "select": return await self.test_select() elif path == "sqlite-master": return await self.test_sqlite_master() @@ -175,6 +177,89 @@ async def fetch(self, request, env): else: return await self.index() + # MARK: - Cursor Result Regressions + + async def test_cursor_results(self): + """Exercise production cursors and SQLAlchemy against the bound database.""" + from sqlalchemy_cloudflare_d1.connection import SyncWorkerConnection + + queries = [ + "SELECT 1 AS id, 7 AS id", + "SELECT 1 AS id, NULL AS value, 7 AS id UNION ALL SELECT 2, 'x', 8", + "-- note\nSELECT 1 AS x WHERE 0", + "/* note */ SELECT 1 AS x WHERE 0", + 'WITH "select" AS (SELECT 1 AS x) SELECT x FROM "select" WHERE 0', + "VALUES (1, 2)", + "EXPLAIN SELECT 1", + ] + output = {} + for mode in ("async", "sync"): + connection = ( + WorkerConnection(self.env.DB) + if mode == "async" + else SyncWorkerConnection(self.env.DB) + ) + cursor = connection.cursor() + + async def execute(query, parameters=None): + if mode == "async": + await cursor.execute_async(query, parameters) + else: + cursor.execute(query, parameters) + + try: + results = [] + for query in queries: + await execute(query) + results.append( + { + "columns": [col[0] for col in cursor.description], + "rows": cursor.fetchall(), + } + ) + output[mode] = results + table = f"test_cursor_results_{uuid.uuid4().hex}" + await execute( + f"CREATE TABLE {table} (id INTEGER PRIMARY KEY, value TEXT)" + ) + try: + await execute( + f"/* note */ WITH select2 AS (SELECT ? AS value) " + f"INSERT INTO {table} (value) SELECT value FROM select2", + ("inserted once",), + ) + mutation = { + "rowcount": cursor.rowcount, + "lastrowid": cursor.lastrowid, + } + await execute(f"SELECT COUNT(*) FROM {table}") + mutation["count"] = cursor.fetchone()[0] + await execute(f"DELETE FROM {table}") + mutation["deleted"] = cursor.rowcount + output[f"{mode}_mutation"] = mutation + finally: + await execute(f"DROP TABLE {table}") + finally: + cursor.close() + connection.close() + + engine = create_engine_from_binding(self.env.DB) + try: + with engine.connect() as connection: + results = [] + for query in queries: + result = connection.exec_driver_sql(query) + results.append( + { + "columns": list(result.keys()), + "rows": [tuple(row) for row in result.fetchall()], + } + ) + output["engine"] = results + finally: + engine.dispose() + return Response.json(output) + def get_connection(self) -> WorkerConnection: """Get a WorkerConnection wrapping the D1 binding.""" return WorkerConnection(self.env.DB) diff --git a/src/sqlalchemy_cloudflare_d1/__init__.py b/src/sqlalchemy_cloudflare_d1/__init__.py index a1fd07f..0c46010 100644 --- a/src/sqlalchemy_cloudflare_d1/__init__.py +++ b/src/sqlalchemy_cloudflare_d1/__init__.py @@ -74,7 +74,7 @@ def __getattr__(name: str) -> Any: raise AttributeError(f"module {__name__!r} has no attribute {name!r}") -__version__ = "0.4.0" +__version__ = "0.4.1" __all__ = [ # Dialects "CloudflareD1Dialect", diff --git a/src/sqlalchemy_cloudflare_d1/connection.py b/src/sqlalchemy_cloudflare_d1/connection.py index d29877a..92e7ad3 100644 --- a/src/sqlalchemy_cloudflare_d1/connection.py +++ b/src/sqlalchemy_cloudflare_d1/connection.py @@ -7,6 +7,7 @@ """ import os +import re from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple, Union try: @@ -102,7 +103,7 @@ def _prepare_parameters(parameters: Optional[Sequence]) -> Optional[List]: def _build_description( - operation: str, columns: List[str], result_data: Optional[List[Dict[str, Any]]] + operation: str, columns: List[str], result_data: Optional[List[Any]] ) -> Optional[List[tuple]]: """Build cursor description from query result. @@ -131,7 +132,7 @@ def _build_description( if not is_select_like: return None - if result_data: + if result_data and isinstance(result_data[0], dict): # Fallback to first row keys if columns not available first_row = result_data[0] return [(name, None, None, None, None, None, None) for name in first_row.keys()] @@ -224,6 +225,68 @@ def _parse_all_result(all_result: Any) -> Dict[str, Any]: } +def _is_read_statement(query: str) -> bool: + """Recognize reads without mistaking a WITH-prefixed mutation for a read.""" + tokens = re.finditer( + r"--[^\n]*|/\*[\s\S]*?\*/|'(?:''|[^'])*'|" + r'"(?:""|[^"])*"|`(?:``|[^`])*`|\[[^\]]*\]|[(),]|[A-Za-z_][A-Za-z_0-9$]*', + query, + ) + first = None + depth = 0 + after_cte = False + for match in tokens: + token = match.group() + if token.startswith(("--", "/*", "'", '"', "`", "[")): + continue + if token == "(": + depth += 1 + continue + if token == ")": + depth -= 1 + after_cte = depth == 0 + continue + if depth: + continue + if token == ",": + after_cte = False + continue + keyword = token.upper() + if first is None: + first = keyword + if first != "WITH": + return first in ("SELECT", "VALUES", "EXPLAIN") + elif after_cte and keyword in ( + "SELECT", + "VALUES", + "INSERT", + "UPDATE", + "DELETE", + "REPLACE", + ): + return keyword in ("SELECT", "VALUES") + return False + + +async def _execute_worker_statement(stmt: Any, query: str) -> Dict[str, Any]: + """Keep read rows positional and execute mutations once with their metadata. + + A native JavaScript options object is required by D1; a Python dict does + not reliably expose columnNames to the binding in Pyodide. + """ + if not _is_read_statement(query): + return _parse_all_result(await stmt.all()) + + from js import JSON # type: ignore[import-not-found] + + raw_result = await stmt.raw(JSON.parse('{"columnNames":true}')) + if hasattr(raw_result, "to_py"): + raw_result = raw_result.to_py() + columns = list(raw_result[0]) if raw_result else [] + rows = [[_convert_js_null(value) for value in row] for row in raw_result[1:]] + return {"results": rows, "columns": columns, "meta": {}, "success": True} + + # MARK: - Row Class @@ -341,7 +404,7 @@ def fetchone(self) -> Optional[tuple]: row_data = self._result_data[self._position] self._position += 1 - # REST API rows are arrays in column order, Worker rows are dicts + # Reads preserve column order; mutation results may still contain dict rows if not isinstance(row_data, dict): return tuple(row_data) @@ -680,44 +743,7 @@ async def _execute_query_async( else: stmt = stmt.bind(parameters) - # MARK: - Execute using all() for reliable structured results - # all() returns {results: [{col: val, ...}, ...], meta: {...}} - # This avoids the raw() bug where single-row results lose the - # column names header row, causing data to end up in description. - all_result = await stmt.all() - - parsed = _parse_all_result(all_result) - - # MARK: - Fall back to raw() for column names on empty results - # all() doesn't return column info when results are empty. - # raw({columnNames: true}) works correctly for 0-row results. - # Only do this for SELECT queries to avoid re-executing mutations. - if ( - not parsed["columns"] - and not parsed["results"] - and query.strip().upper().startswith("SELECT") - ): - try: - stmt2 = self._d1.prepare(query) - if parameters: - if isinstance(parameters, (tuple, list)): - stmt2 = stmt2.bind(*parameters) - elif isinstance(parameters, dict): - stmt2 = stmt2.bind(*parameters.values()) - else: - stmt2 = stmt2.bind(parameters) - raw_result = await stmt2.raw({"columnNames": True}) - if hasattr(raw_result, "to_py"): - raw_result = raw_result.to_py() - if raw_result and len(raw_result) > 0: - first_row = raw_result[0] - if hasattr(first_row, "to_py"): - first_row = first_row.to_py() - parsed["columns"] = list(first_row) if first_row else [] - except Exception: - pass # Column names are best-effort for empty results - - return parsed + return await _execute_worker_statement(stmt, query) except Exception as e: raise OperationalError(f"D1 Worker query failed: {e}") @@ -1189,49 +1215,7 @@ def convert_param(val: Any) -> Any: else: stmt = stmt.bind(convert_param(parameters)) - # MARK: - Execute using all() for reliable structured results - # all() returns {results: [{col: val, ...}, ...], meta: {...}} - # This avoids the raw() bug where single-row results lose the - # column names header row, causing data to end up in description. - all_result = await stmt.all() - - # For empty results, fall back to raw() for column names - # all() doesn't return column info when results are empty. - # raw({columnNames: true}) works correctly for 0-row results. - # Only do this for SELECT queries to avoid re-executing mutations. - fallback_columns = None - parsed = _parse_all_result(all_result) - if ( - not parsed["columns"] - and not parsed["results"] - and query.strip().upper().startswith("SELECT") - ): - try: - stmt2 = self._d1.prepare(query) - if parameters: - if isinstance(parameters, (tuple, list)): - converted2 = [convert_param(p) for p in parameters] - stmt2 = stmt2.bind(*converted2) - elif isinstance(parameters, dict): - converted2 = [ - convert_param(v) for v in parameters.values() - ] - stmt2 = stmt2.bind(*converted2) - else: - stmt2 = stmt2.bind(convert_param(parameters)) - raw_result = await stmt2.raw({"columnNames": True}) - if hasattr(raw_result, "to_py"): - raw_result = raw_result.to_py() - if raw_result and len(raw_result) > 0: - first_row = raw_result[0] - if hasattr(first_row, "to_py"): - first_row = first_row.to_py() - fallback_columns = list(first_row) if first_row else [] - except Exception: - pass - if fallback_columns: - parsed["columns"] = fallback_columns - return parsed + return await _execute_worker_statement(stmt, query) return run_sync(_run()) # type: ignore[no-any-return] diff --git a/tests/conftest.py b/tests/conftest.py index deac01e..47a7c90 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,9 +2,10 @@ import os import shutil +import signal import socket import subprocess -import time +import threading import uuid from contextlib import closing from pathlib import Path @@ -72,7 +73,10 @@ def find_free_port() -> int: def pywrangler_dev_server( - project_dir: Path, timeout: int = 300, env: Optional[dict] = None + project_dir: Path, + timeout: int = 300, + env: Optional[dict] = None, + remote_bindings: bool = False, ) -> tuple[subprocess.Popen, int]: """Start a pywrangler dev server and return the process and port. @@ -80,48 +84,63 @@ def pywrangler_dev_server( project_dir: Path to the project directory containing wrangler.jsonc timeout: Maximum time to wait for server startup (default 300s for CI) env: Extra environment variables for the server process + remote_bindings: Honor configured remote bindings instead of forcing local D1 Returns: Tuple of (process, port) """ port = find_free_port() - # Start the dev server with --local flag - # Note: --remote has issues with Python Workers on Cloudflare edge + command = ["uv", "run", "pywrangler", "dev", "--port", str(port)] + if not remote_bindings: + command.append("--local") process = subprocess.Popen( - ["uv", "run", "pywrangler", "dev", "--local", "--port", str(port)], + command, cwd=project_dir, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, env={**os.environ, **(env or {})}, + start_new_session=True, ) - # Wait for server to be ready - start_time = time.time() - ready_message = "[wrangler:info] Ready on" - - while time.time() - start_time < timeout: - if process.poll() is not None: - # Process exited - output = process.stdout.read() if process.stdout else "" - raise RuntimeError(f"pywrangler dev exited unexpectedly: {output}") - - line = process.stdout.readline() if process.stdout else "" - if ready_message in line: + ready = threading.Event() + exited = threading.Event() + + def drain_output(): + # Keep draining after startup so a full pipe cannot stall the Worker. + if process.stdout: + for line in process.stdout: + # Local Explorer also announces localhost before bindings are ready. + if "[wrangler:info] Ready on" in line: + ready.set() + exited.set() + + threading.Thread(target=drain_output, daemon=True).start() + for _ in range(timeout * 10): + if ready.wait(0.1): return process, port + if exited.is_set(): + stop_dev_server(process) + raise RuntimeError("pywrangler dev exited before the Worker was ready") - # Also check for alternative ready messages - if f"localhost:{port}" in line.lower() or "ready" in line.lower(): - # Give it a moment to fully initialize - time.sleep(0.5) - return process, port - - # Timeout reached - process.terminate() + stop_dev_server(process) raise TimeoutError(f"pywrangler dev did not start within {timeout} seconds") +def stop_dev_server(process: subprocess.Popen) -> None: + """Stop only the process group created for this test server.""" + try: + os.killpg(process.pid, signal.SIGTERM) + except ProcessLookupError: + return + try: + process.wait(timeout=3) + except subprocess.TimeoutExpired: + os.killpg(process.pid, signal.SIGKILL) + process.wait() + + def get_worker_project_dir() -> Path: """Get the path to the examples/workers directory.""" return Path(__file__).parent.parent / "examples" / "workers" @@ -147,7 +166,7 @@ def initialized_worker(): @pytest.fixture(scope="session") -def dev_server(initialized_worker): +def dev_server(initialized_worker, request): """Session-scoped fixture that starts a single pywrangler dev server. The server is reused across all Worker tests and stopped after the session. @@ -161,29 +180,13 @@ def dev_server(initialized_worker): process = None try: - process, port = pywrangler_dev_server(project_dir) + process, port = pywrangler_dev_server( + project_dir, remote_bindings=getattr(request, "param", False) + ) yield port finally: if process is not None: - # Kill the entire process group to clean up workerd child processes - try: - # Try graceful termination first - process.terminate() - process.wait(timeout=3) - except subprocess.TimeoutExpired: - # Force kill if it doesn't terminate - process.kill() - process.wait() - - # Also kill any orphaned workerd processes on this port - try: - subprocess.run( - ["pkill", "-f", f"workerd.*{port}"], - capture_output=True, - timeout=5, - ) - except Exception: - pass + stop_dev_server(process) @pytest.fixture(scope="session") @@ -230,21 +233,7 @@ def hyperdrive_dev_server(): pass if process is not None: - try: - process.terminate() - process.wait(timeout=3) - except subprocess.TimeoutExpired: - process.kill() - process.wait() - - try: - subprocess.run( - ["pkill", "-f", f"workerd.*{port}"], - capture_output=True, - timeout=5, - ) - except Exception: - pass + stop_dev_server(process) @pytest.fixture diff --git a/tests/integration/test_restapi_integration.py b/tests/integration/test_restapi_integration.py index 7d41fcc..8013ab9 100644 --- a/tests/integration/test_restapi_integration.py +++ b/tests/integration/test_restapi_integration.py @@ -2520,6 +2520,30 @@ def test_join_via_sqlalchemy(self, d1_engine, joined_tables): assert rows == [(1, "x", 7)] + @pytest.mark.asyncio + async def test_join_via_async_sqlalchemy(self, joined_tables): + """Keep duplicate columns through the real async engine and adapter.""" + from sqlalchemy import text + from sqlalchemy.ext.asyncio import create_async_engine + + table_a, table_b = joined_tables + url = f"cloudflare_d1+async://{ACCOUNT_ID}:{API_TOKEN}@{DATABASE_ID}" + engine = create_async_engine(url) + + try: + async with engine.connect() as conn: + result = await conn.execute( + text( + f"/* note */ SELECT {table_a}.id, {table_a}.name, " + f"{table_b}.id FROM {table_a} JOIN {table_b} " + f"ON {table_b}.a_id = {table_a}.id" + ) + ) + assert list(result.keys()) == ["id", "name", "id"] + assert result.fetchall() == [(1, "x", 7)] + finally: + await engine.dispose() + # MARK: - Autoincrement Insert Tests (Issue #12) diff --git a/tests/integration/test_worker_integration.py b/tests/integration/test_worker_integration.py index b587b00..5bd9e3b 100644 --- a/tests/integration/test_worker_integration.py +++ b/tests/integration/test_worker_integration.py @@ -9,6 +9,7 @@ Requires: pywrangler dev running with --local flag for local D1 database. """ +import pytest import requests @@ -1274,3 +1275,50 @@ def test_hyperdrive_concurrent_queries_are_serialized(self, hyperdrive_dev_serve assert data["success"] is True, data assert data["results"] == [1, 2, 3, 4, 5], data + + +# MARK: - Real D1 Cursor Result Regressions + + +@pytest.mark.parametrize("dev_server", [True], indirect=True) +class TestRemoteWorkerCursorResults: + """Use the example's real D1 binding for sync, async and SQLAlchemy results.""" + + @pytest.fixture(scope="class") + def cursor_results(self, dev_server): + response = requests.get( + f"http://localhost:{dev_server}/cursor-results", timeout=120 + ) + assert response.status_code == 200, response.text + return response.json() + + @pytest.mark.parametrize("mode", ["async", "sync", "engine"]) + def test_duplicate_columns_and_nulls(self, cursor_results, mode): + results = cursor_results[mode] + assert results[0] == {"columns": ["id", "id"], "rows": [[1, 7]]} + assert results[1] == { + "columns": ["id", "value", "id"], + "rows": [[1, None, 7], [2, "x", 8]], + } + + @pytest.mark.parametrize("mode", ["async", "sync", "engine"]) + def test_empty_commented_and_cte_queries(self, cursor_results, mode): + for result in cursor_results[mode][2:5]: + assert result == {"columns": ["x"], "rows": []} + + @pytest.mark.parametrize("mode", ["async", "sync", "engine"]) + def test_values_and_explain(self, cursor_results, mode): + values, explain = cursor_results[mode][5:] + assert values == {"columns": ["column1", "column2"], "rows": [[1, 2]]} + assert explain["columns"][:3] == ["addr", "opcode", "p1"] + assert explain["rows"] + assert all(len(row) == len(explain["columns"]) for row in explain["rows"]) + + @pytest.mark.parametrize("mode", ["async", "sync"]) + def test_mutation_metadata_and_single_execution(self, cursor_results, mode): + assert cursor_results[f"{mode}_mutation"] == { + "rowcount": 1, + "lastrowid": 1, + "count": 1, + "deleted": 1, + } diff --git a/tests/unit/test_connection.py b/tests/unit/test_connection.py index 7c9d6ac..dcf7826 100644 --- a/tests/unit/test_connection.py +++ b/tests/unit/test_connection.py @@ -10,6 +10,7 @@ import httpx import pytest from sqlalchemy import create_engine, event, text +from sqlalchemy.ext.asyncio import create_async_engine from sqlalchemy_cloudflare_d1.connection import AsyncConnection, Connection @@ -92,5 +93,27 @@ def use_fake_raw(dbapi_connection, connection_record): assert rows == [(1, "x", 7)] +def test_async_engine_keeps_values_of_columns_with_same_name(): + """Keep positional values through the async SQLAlchemy adapter, too.""" + engine = create_async_engine("cloudflare_d1+async://acct:token@db") + + @event.listens_for(engine.sync_engine, "connect") + def use_fake_raw(dbapi_connection, connection_record): + dbapi_connection._connection.client = httpx.AsyncClient( + transport=raw_transport(["id", "name", "id"], [[1, "x", 7]]) + ) + + async def main(): + try: + async with engine.connect() as conn: + result = await conn.execute(text("/* note */ " + JOIN_SQL)) + assert list(result.keys()) == ["id", "name", "id"] + assert result.fetchall() == [(1, "x", 7)] + finally: + await engine.dispose() + + asyncio.run(main()) + + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/unit/test_worker_connection.py b/tests/unit/test_worker_connection.py new file mode 100644 index 0000000..189409d --- /dev/null +++ b/tests/unit/test_worker_connection.py @@ -0,0 +1,162 @@ +"""Worker cursor regressions using canned D1 binding responses.""" + +import asyncio +import json +import sys +from types import ModuleType, SimpleNamespace + +import pytest +from sqlalchemy import text + +from sqlalchemy_cloudflare_d1.connection import ( + SyncWorkerConnection, + WorkerConnection, + create_engine_from_binding, +) + + +# MARK: - Binding Fixtures + + +class Statement: + def __init__(self, raw_result, all_result): + self.raw_result = raw_result + self.all_result = all_result + self.calls = [] + self.parameters = None + + def bind(self, *parameters): + self.parameters = parameters + return self + + async def raw(self, options): + assert options["columnNames"] is True + self.calls.append("raw") + return self.raw_result + + async def all(self): + self.calls.append("all") + return self.all_result + + +class Binding: + def __init__(self, statement): + self.statement = statement + + def prepare(self, query): + return self.statement + + +@pytest.fixture(autouse=True) +def worker_runtime(monkeypatch): + """Supply only the Pyodide runtime boundary needed by these unit tests.""" + js = ModuleType("js") + js.JSON = SimpleNamespace(parse=json.loads) + pyodide = ModuleType("pyodide") + ffi = ModuleType("pyodide.ffi") + ffi.run_sync = asyncio.run + monkeypatch.setitem(sys.modules, "js", js) + monkeypatch.setitem(sys.modules, "pyodide", pyodide) + monkeypatch.setitem(sys.modules, "pyodide.ffi", ffi) + + +def execute(statement, sql, sync, parameters=None): + binding = Binding(statement) + connection = SyncWorkerConnection(binding) if sync else WorkerConnection(binding) + cursor = connection.cursor() + if sync: + cursor.execute(sql, parameters) + else: + asyncio.run(cursor.execute_async(sql, parameters)) + return cursor + + +# MARK: - Positional Results + + +@pytest.mark.parametrize("sync", [False, True]) +@pytest.mark.parametrize("rows", [[[1, "x", 7]], [[1, None, 7], [2, "y", 8]]]) +def test_duplicate_columns_keep_positional_values(sync, rows): + statement = Statement( + [["id", "name", "id"], *rows], + {"results": [{"id": 7, "name": "x"}], "meta": {}}, + ) + cursor = execute(statement, "SELECT a.id, a.name, b.id FROM a JOIN b", sync) + + assert [column[0] for column in cursor.description] == ["id", "name", "id"] + assert cursor.fetchone() == tuple(rows[0]) + assert cursor.fetchmany(2) == [tuple(row) for row in rows[1:]] + assert cursor.fetchone() is None + assert statement.calls == ["raw"] + + +def test_worker_engine_keeps_duplicate_columns(): + statement = Statement([["id", "id"], [1, 7]], {"results": [{"id": 7}], "meta": {}}) + engine = create_engine_from_binding(Binding(statement)) + try: + with engine.connect() as connection: + result = connection.execute(text("/* note */ SELECT 1 AS id, 7 AS id")) + assert list(result.keys()) == ["id", "id"] + assert result.fetchall() == [(1, 7)] + finally: + engine.dispose() + + +# MARK: - Empty Results + + +@pytest.mark.parametrize("sync", [False, True]) +@pytest.mark.parametrize( + "sql", + [ + "-- note\nSELECT 1 AS x WHERE ?", + "/* note */ SELECT 1 AS x WHERE ?", + "WITH c AS (SELECT 1 AS x) SELECT x FROM c WHERE ?", + 'WITH "select" AS (SELECT 1 AS x) SELECT x FROM "select" WHERE ?', + "WITH c(x) AS (SELECT '(') SELECT x FROM c WHERE ?", + "WITH a AS (SELECT 1), select2 AS (SELECT 1 AS x) SELECT x FROM select2 WHERE ?", + ], +) +def test_empty_select_preserves_columns_without_reexecution(sync, sql): + statement = Statement([["x"]], {"results": [], "meta": {}}) + cursor = execute(statement, sql, sync, parameters=(0,)) + + assert [column[0] for column in cursor.description] == ["x"] + assert cursor.fetchall() == [] + assert statement.parameters == (0,) + assert statement.calls == ["raw"] + + +# MARK: - Mutation Metadata + + +@pytest.mark.parametrize("sync", [False, True]) +@pytest.mark.parametrize( + "sql", + [ + "INSERT INTO t VALUES (?)", + "/* SELECT */ INSERT INTO t VALUES (?)", + "WITH c AS (SELECT ?) INSERT INTO t SELECT * FROM c", + "WITH select2 AS (SELECT ?) INSERT INTO t SELECT * FROM select2", + ], +) +def test_mutations_keep_metadata_and_execute_once(sync, sql): + statement = Statement( + [], {"results": [], "meta": {"changes": 1, "last_row_id": 41}} + ) + cursor = execute(statement, sql, sync, parameters=(1,)) + + assert cursor.rowcount == 1 + assert cursor.lastrowid == 41 + assert statement.parameters == (1,) + assert statement.calls == ["all"] + + +@pytest.mark.parametrize("sync", [False, True]) +@pytest.mark.parametrize("sql", ["VALUES (1, 2)", "EXPLAIN SELECT 1"]) +def test_other_read_statements_keep_column_header(sync, sql): + statement = Statement([["a", "b"], [1, 2]], {"results": [], "meta": {}}) + cursor = execute(statement, sql, sync) + assert [column[0] for column in cursor.description] == ["a", "b"] + assert cursor.fetchall() == [(1, 2)] + assert statement.calls == ["raw"]