fix(server): rowcount

tekumara · Dec 15, 2024 · 58e15c5 · 58e15c5
1 parent 7a6f9a3
commit 58e15c5
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 2 deletions.
diff --git a/fakesnow/server.py b/fakesnow/server.py
@@ -87,7 +87,7 @@ async def query_request(request: Request) -> JSONResponse:
                 "data": {
                     "rowtype": rowtype,
                     "rowsetBase64": rowset_b64,
-                    "total": 1,
+                    "total": cur._rowcount,  # noqa: SLF001
                     "queryId": cur.sfqid,
                     "queryResultFormat": "arrow",
                 },

diff --git a/tests/test_fakes.py b/tests/test_fakes.py
@@ -1210,7 +1210,7 @@ def test_variables(conn: snowflake.connector.SnowflakeConnection):
 
 def test_values(conn: snowflake.connector.SnowflakeConnection):
     with conn.cursor(snowflake.connector.cursor.DictCursor) as cur:
-        cur.execute("select * from VALUES ('Amsterdam', 1), ('London', 2)")
+        cur.execute("select * from values ('Amsterdam', 1), ('London', 2)")
 
         assert cur.fetchall() == [
             {"COLUMN1": "Amsterdam", "COLUMN2": 1},

diff --git a/tests/test_server.py b/tests/test_server.py
@@ -12,6 +12,13 @@
 from tests.utils import indent
 
 
+def test_server_rowcount(scur: snowflake.connector.cursor.SnowflakeCursor):
+    cur = scur
+
+    cur.execute("select * from values ('Salted'), ('Caramel')")
+    assert cur.rowcount == 2
+
+
 def test_server_sfid(scur: snowflake.connector.cursor.SnowflakeCursor) -> None:
     cur = scur
     assert not cur.sfqid