diff --git a/src/sqlalchemy_cloudflare_d1/connection.py b/src/sqlalchemy_cloudflare_d1/connection.py index 005dce3..fd2791c 100644 --- a/src/sqlalchemy_cloudflare_d1/connection.py +++ b/src/sqlalchemy_cloudflare_d1/connection.py @@ -106,14 +106,22 @@ def _build_description( ) -> Optional[List[tuple]]: """Build cursor description from query result. + Column names sent by D1 are used whatever the statement starts with, so + a leading comment, EXPLAIN or VALUES still gets a description. + Args: operation: The SQL operation that was executed columns: Column names from the query result result_data: The result data rows Returns: - List of 7-tuples for SELECT-like statements, None otherwise + List of 7-tuples when D1 sent column names or the statement is + SELECT-like, None otherwise """ + # Build description from columns + if columns: + return [(name, None, None, None, None, None, None) for name in columns] + operation_upper = operation.strip().upper() is_select_like = ( operation_upper.startswith(("SELECT", "PRAGMA", "WITH")) @@ -123,10 +131,7 @@ def _build_description( if not is_select_like: return None - # Build description from columns - if columns: - return [(name, None, None, None, None, None, None) for name in columns] - elif result_data: + if result_data: # 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()] diff --git a/src/sqlalchemy_cloudflare_d1/dialect_async.py b/src/sqlalchemy_cloudflare_d1/dialect_async.py index 7ead83b..9d10a1e 100644 --- a/src/sqlalchemy_cloudflare_d1/dialect_async.py +++ b/src/sqlalchemy_cloudflare_d1/dialect_async.py @@ -84,11 +84,9 @@ def execute( self.await_(_cursor.execute(operation, parameters)) # Determine if this is a row-returning statement - operation_upper = operation.strip().upper() - is_select = ( - operation_upper.startswith(("SELECT", "PRAGMA", "WITH")) - or "RETURNING" in operation_upper - ) + # The cursor sets a description when D1 sent column names or the + # statement is SELECT-like, and None otherwise + is_select = _cursor.description is not None if is_select: # For SELECT statements, set description (may be empty list for no-column results) diff --git a/tests/integration/test_restapi_integration.py b/tests/integration/test_restapi_integration.py index d8ecf6f..91b5f44 100644 --- a/tests/integration/test_restapi_integration.py +++ b/tests/integration/test_restapi_integration.py @@ -3413,5 +3413,99 @@ class Task(Base): Base.metadata.drop_all(d1_engine) +# MARK: - Description Without SELECT Keyword Tests (Issue #32) + + +class TestDescriptionWithoutSelectKeyword: + """Test that statements returning rows get a description whatever they start with. + + Regression tests for bug where a leading comment, EXPLAIN or VALUES + returned rows with cursor.description set to None. + """ + + @pytest.mark.parametrize( + "sql", + [ + "-- note\nSELECT 1 AS x, 2 AS y", + "/* note */ SELECT 1 AS x, 2 AS y", + ], + ) + def test_comment_led_select_via_cursor(self, d1_connection, sql): + """Test SELECT with a leading comment has a description via DBAPI cursor.""" + cursor = d1_connection.cursor() + cursor.execute(sql) + + desc_names = [d[0] for d in cursor.description] if cursor.description else [] + assert desc_names == ["x", "y"] + assert cursor.fetchall() == [(1, 2)] + + def test_values_via_cursor(self, d1_connection): + """Test VALUES has a description via DBAPI cursor.""" + cursor = d1_connection.cursor() + cursor.execute("VALUES (1, 2)") + + desc_names = [d[0] for d in cursor.description] if cursor.description else [] + assert desc_names == ["column1", "column2"] + assert cursor.fetchall() == [(1, 2)] + + def test_explain_query_plan_via_cursor(self, d1_connection): + """Test EXPLAIN QUERY PLAN has a description via DBAPI cursor.""" + cursor = d1_connection.cursor() + cursor.execute("EXPLAIN QUERY PLAN SELECT 1") + + desc_names = [d[0] for d in cursor.description] if cursor.description else [] + assert "detail" in desc_names + assert len(cursor.fetchall()) >= 1 + + def test_insert_has_no_description(self, d1_connection, test_table_name): + """Test INSERT without RETURNING still has no description.""" + cursor = d1_connection.cursor() + + try: + cursor.execute( + f"CREATE TABLE IF NOT EXISTS {test_table_name} " + f"(id INTEGER PRIMARY KEY, name TEXT)" + ) + assert cursor.description is None + + cursor.execute(f"INSERT INTO {test_table_name} (name) VALUES (?)", ("a",)) + assert cursor.description is None + assert cursor.rowcount == 1 + finally: + cursor.execute(f"DROP TABLE IF EXISTS {test_table_name}") + + def test_comment_led_select_via_sqlalchemy(self, d1_engine): + """Test SELECT with a leading comment returns rows via SQLAlchemy engine.""" + from sqlalchemy import text + + with d1_engine.connect() as conn: + result = conn.execute(text("/* note */ SELECT 1 AS x, 2 AS y")) + columns = list(result.keys()) + rows = result.fetchall() + + assert columns == ["x", "y"] + assert rows == [(1, 2)] + + @pytest.mark.asyncio + async def test_comment_led_select_via_async_engine(self): + """Test SELECT with a leading comment returns rows via async engine.""" + from sqlalchemy import text + from sqlalchemy.ext.asyncio import create_async_engine + + 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("/* note */ SELECT 1 AS x, 2 AS y")) + columns = list(result.keys()) + rows = result.fetchall() + + assert columns == ["x", "y"] + assert rows == [(1, 2)] + finally: + await engine.dispose() + + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) diff --git a/tests/unit/test_dialect.py b/tests/unit/test_dialect.py index 1b55e1b..a561ea9 100644 --- a/tests/unit/test_dialect.py +++ b/tests/unit/test_dialect.py @@ -434,5 +434,132 @@ def test_version_matches_package_metadata(): assert sqlalchemy_cloudflare_d1.__version__ == installed +def _raw_transport(columns, rows): + """Build an httpx transport that answers every request like D1's /raw.""" + import httpx + + def handler(request): + return httpx.Response( + 200, + json={ + "success": True, + "result": [ + { + "results": {"columns": columns, "rows": rows}, + "meta": {}, + "success": True, + } + ], + }, + ) + + return httpx.MockTransport(handler) + + +@pytest.mark.parametrize( + "sql", + [ + "-- note\nSELECT 1 AS x, 2 AS y", + "/* note */ SELECT 1 AS x, 2 AS y", + "(SELECT 1 AS x, 2 AS y)", + "EXPLAIN QUERY PLAN SELECT 1 AS x, 2 AS y", + "VALUES (1, 2)", + ], +) +def test_description_uses_columns_for_any_statement(sql): + """Test that column names from D1 are used whatever the statement starts with.""" + import httpx + + from sqlalchemy_cloudflare_d1.connection import Connection + + conn = Connection(account_id="acct", database_id="db", api_token="token") + conn.client = httpx.Client(transport=_raw_transport(["x", "y"], [[1, 2]])) + + cursor = conn.cursor() + cursor.execute(sql) + + assert [d[0] for d in cursor.description] == ["x", "y"] + assert cursor.fetchall() == [(1, 2)] + + +def test_description_none_without_columns(): + """Test that a statement D1 sends no columns for still has no description.""" + import httpx + + from sqlalchemy_cloudflare_d1.connection import Connection + + conn = Connection(account_id="acct", database_id="db", api_token="token") + conn.client = httpx.Client(transport=_raw_transport([], [])) + + cursor = conn.cursor() + cursor.execute("INSERT INTO t (x) VALUES (1)") + + assert cursor.description is None + + +def test_description_empty_for_select_without_columns(): + """Test that a SELECT with no columns and no rows keeps an empty description.""" + import httpx + + from sqlalchemy_cloudflare_d1.connection import Connection + + conn = Connection(account_id="acct", database_id="db", api_token="token") + conn.client = httpx.Client(transport=_raw_transport([], [])) + + cursor = conn.cursor() + cursor.execute("SELECT * FROM t") + + assert cursor.description == [] + + +def test_engine_returns_rows_for_comment_led_select(): + """Test that a SELECT with a leading comment returns rows through SQLAlchemy.""" + import httpx + from sqlalchemy import create_engine, event, text + + engine = create_engine("cloudflare_d1://acct:token@db") + + @event.listens_for(engine, "connect") + def use_fake_raw(dbapi_connection, connection_record): + dbapi_connection.client = httpx.Client( + transport=_raw_transport(["x", "y"], [[1, 2]]) + ) + + with engine.connect() as conn: + result = conn.execute(text("/* note */ SELECT 1 AS x, 2 AS y")) + assert list(result.keys()) == ["x", "y"] + assert result.fetchall() == [(1, 2)] + + +def test_async_engine_returns_rows_for_comment_led_select(): + """Test that a SELECT with a leading comment returns rows through the async engine.""" + import asyncio + + import httpx + from sqlalchemy import event, text + from sqlalchemy.ext.asyncio import create_async_engine + + 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(["x", "y"], [[1, 2]]) + ) + + async def main(): + try: + async with engine.connect() as conn: + result = await conn.execute(text("/* note */ SELECT 1 AS x, 2 AS y")) + return list(result.keys()), result.fetchall() + finally: + await engine.dispose() + + keys, rows = asyncio.run(main()) + + assert keys == ["x", "y"] + assert rows == [(1, 2)] + + if __name__ == "__main__": pytest.main([__file__])