Skip to content

Commit bb69e9d

Browse files
committed
make cursor client-side
1 parent c70c252 commit bb69e9d

File tree

6 files changed

+268
-184
lines changed

6 files changed

+268
-184
lines changed

.pre-commit-config.yaml

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
files: ydb_dbapi
12
repos:
23
- repo: https://github.com/pre-commit/pre-commit-hooks
34
rev: v4.6.0
@@ -23,5 +24,3 @@ repos:
2324
hooks:
2425
- id: mypy
2526
name: mypy
26-
27-
exclude: '.github/'

poetry.lock

Lines changed: 182 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ types-protobuf = "^5.28.0.20240924"
1818
testcontainers = "^4.8.2"
1919
pytest = "^8.3.3"
2020
pytest-asyncio = "^0.24.0"
21+
sqlalchemy = "^2.0.36"
2122

2223
[build-system]
2324
requires = ["poetry-core"]

tests/test_connection.py

Lines changed: 18 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ def _test_isolation_level_read_only(
3434
if read_only:
3535
with pytest.raises(dbapi.DatabaseError):
3636
cursor.execute(query)
37-
cursor.finish_query()
37+
cursor._scroll_stream()
3838

3939
else:
4040
cursor.execute(query)
@@ -53,14 +53,14 @@ def _test_connection(self, connection: dbapi.Connection) -> None:
5353
cur = connection.cursor()
5454
with suppress(dbapi.DatabaseError):
5555
cur.execute("DROP TABLE foo")
56-
cur.finish_query()
56+
cur._scroll_stream()
5757

5858
assert not connection.check_exists("/local/foo")
5959
with pytest.raises(dbapi.ProgrammingError):
6060
connection.describe("/local/foo")
6161

6262
cur.execute("CREATE TABLE foo(id Int64 NOT NULL, PRIMARY KEY (id))")
63-
cur.finish_query()
63+
cur._scroll_stream()
6464

6565
assert connection.check_exists("/local/foo")
6666

@@ -77,12 +77,12 @@ def _test_cursor_raw_query(self, connection: dbapi.Connection) -> None:
7777

7878
with suppress(dbapi.DatabaseError):
7979
cur.execute("DROP TABLE test")
80-
cur.finish_query()
80+
cur._scroll_stream()
8181

8282
cur.execute(
8383
"CREATE TABLE test(id Int64 NOT NULL, text Utf8, PRIMARY KEY (id))"
8484
)
85-
cur.finish_query()
85+
cur._scroll_stream()
8686

8787
cur.execute(
8888
"""
@@ -104,7 +104,7 @@ def _test_cursor_raw_query(self, connection: dbapi.Connection) -> None:
104104
)
105105
},
106106
)
107-
cur.finish_query()
107+
cur._scroll_stream()
108108

109109
cur.execute("DROP TABLE test")
110110

@@ -121,7 +121,7 @@ def _test_errors(self, connection: dbapi.Connection) -> None:
121121

122122
with suppress(dbapi.DatabaseError):
123123
cur.execute("DROP TABLE test")
124-
cur.finish_query()
124+
cur._scroll_stream()
125125

126126
with pytest.raises(dbapi.DataError):
127127
cur.execute("SELECT 18446744073709551616")
@@ -136,10 +136,10 @@ def _test_errors(self, connection: dbapi.Connection) -> None:
136136
cur.execute("SELECT * FROM test")
137137

138138
cur.execute("CREATE TABLE test(id Int64, PRIMARY KEY (id))")
139-
cur.finish_query()
139+
cur._scroll_stream()
140140

141141
cur.execute("INSERT INTO test(id) VALUES(1)")
142-
cur.finish_query()
142+
cur._scroll_stream()
143143

144144
with pytest.raises(dbapi.IntegrityError):
145145
cur.execute("INSERT INTO test(id) VALUES(1)")
@@ -214,7 +214,7 @@ async def _test_isolation_level_read_only(
214214
if read_only:
215215
with pytest.raises(dbapi.DatabaseError):
216216
await cursor.execute(query)
217-
await cursor.finish_query()
217+
await cursor._scroll_stream()
218218

219219
else:
220220
await cursor.execute(query)
@@ -235,7 +235,7 @@ async def _test_connection(
235235
cur = connection.cursor()
236236
with suppress(dbapi.DatabaseError):
237237
await cur.execute("DROP TABLE foo")
238-
await cur.finish_query()
238+
await cur._scroll_stream()
239239

240240
assert not await connection.check_exists("/local/foo")
241241
with pytest.raises(dbapi.ProgrammingError):
@@ -244,7 +244,7 @@ async def _test_connection(
244244
await cur.execute(
245245
"CREATE TABLE foo(id Int64 NOT NULL, PRIMARY KEY (id))"
246246
)
247-
await cur.finish_query()
247+
await cur._scroll_stream()
248248

249249
assert await connection.check_exists("/local/foo")
250250

@@ -263,12 +263,12 @@ async def _test_cursor_raw_query(
263263

264264
with suppress(dbapi.DatabaseError):
265265
await cur.execute("DROP TABLE test")
266-
await cur.finish_query()
266+
await cur._scroll_stream()
267267

268268
await cur.execute(
269269
"CREATE TABLE test(id Int64 NOT NULL, text Utf8, PRIMARY KEY (id))"
270270
)
271-
await cur.finish_query()
271+
await cur._scroll_stream()
272272

273273
await cur.execute(
274274
"""
@@ -290,7 +290,7 @@ async def _test_cursor_raw_query(
290290
)
291291
},
292292
)
293-
await cur.finish_query()
293+
await cur._scroll_stream()
294294

295295
await cur.execute("DROP TABLE test")
296296

@@ -307,7 +307,7 @@ async def _test_errors(self, connection: dbapi.AsyncConnection) -> None:
307307

308308
with suppress(dbapi.DatabaseError):
309309
await cur.execute("DROP TABLE test")
310-
await cur.finish_query()
310+
await cur._scroll_stream()
311311

312312
with pytest.raises(dbapi.DataError):
313313
await cur.execute("SELECT 18446744073709551616")
@@ -322,10 +322,10 @@ async def _test_errors(self, connection: dbapi.AsyncConnection) -> None:
322322
await cur.execute("SELECT * FROM test")
323323

324324
await cur.execute("CREATE TABLE test(id Int64, PRIMARY KEY (id))")
325-
await cur.finish_query()
325+
await cur._scroll_stream()
326326

327327
await cur.execute("INSERT INTO test(id) VALUES(1)")
328-
await cur.finish_query()
328+
await cur._scroll_stream()
329329

330330
with pytest.raises(dbapi.IntegrityError):
331331
await cur.execute("INSERT INTO test(id) VALUES(1)")

tests/test_cursor.py

Lines changed: 23 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -119,48 +119,19 @@ async def test_cursor_fetch_all(
119119
assert await cursor.fetchall() == []
120120

121121
@pytest.mark.asyncio
122-
async def test_cursor_next_set(
122+
async def test_cursor_fetch_one_multiple_result_sets(
123123
self, session: ydb.aio.QuerySession
124124
) -> None:
125125
async with ydb_dbapi.AsyncCursor(session=session) as cursor:
126-
yql_text = """SELECT 1 as val; SELECT 2 as val;"""
127-
await cursor.execute(query=yql_text)
128-
129-
res = await cursor.fetchall()
130-
assert res is not None
131-
assert len(res) == 1
132-
assert res[0][0] == 1
133-
134-
nextset = await cursor.nextset()
135-
assert nextset
136-
137-
res = await cursor.fetchall()
138-
assert res is not None
139-
assert len(res) == 1
140-
assert res[0][0] == 2
141-
142-
nextset = await cursor.nextset()
143-
assert nextset
144-
145-
assert await cursor.fetchall() == []
146-
147-
nextset = await cursor.nextset()
148-
assert not nextset
149-
150-
@pytest.mark.asyncio
151-
async def test_cursor_fetch_one_autoscroll(
152-
self, session: ydb.aio.QuerySession
153-
) -> None:
154-
async with ydb_dbapi.AsyncCursor(
155-
session=session, auto_scroll_result_sets=True
156-
) as cursor:
157126
yql_text = """
158127
SELECT id, val FROM table;
159128
SELECT id, val FROM table1;
160129
SELECT id, val FROM table2;
161130
"""
162131
await cursor.execute(query=yql_text)
163132

133+
assert cursor.rowcount == 12
134+
164135
for i in range(RESULT_SET_LENGTH * RESULT_SET_COUNT):
165136
res = await cursor.fetchone()
166137
assert res is not None
@@ -170,19 +141,19 @@ async def test_cursor_fetch_one_autoscroll(
170141
assert not await cursor.nextset()
171142

172143
@pytest.mark.asyncio
173-
async def test_cursor_fetch_many_autoscroll(
144+
async def test_cursor_fetch_many_multiple_result_sets(
174145
self, session: ydb.aio.QuerySession
175146
) -> None:
176-
async with ydb_dbapi.AsyncCursor(
177-
session=session, auto_scroll_result_sets=True
178-
) as cursor:
147+
async with ydb_dbapi.AsyncCursor(session=session) as cursor:
179148
yql_text = """
180149
SELECT id, val FROM table;
181150
SELECT id, val FROM table1;
182151
SELECT id, val FROM table2;
183152
"""
184153
await cursor.execute(query=yql_text)
185154

155+
assert cursor.rowcount == 12
156+
186157
halfsize = (RESULT_SET_LENGTH * RESULT_SET_COUNT) // 2
187158
for _ in range(2):
188159
res = await cursor.fetchmany(size=halfsize)
@@ -193,19 +164,19 @@ async def test_cursor_fetch_many_autoscroll(
193164
assert not await cursor.nextset()
194165

195166
@pytest.mark.asyncio
196-
async def test_cursor_fetch_all_autoscroll(
167+
async def test_cursor_fetch_all_multiple_result_sets(
197168
self, session: ydb.aio.QuerySession
198169
) -> None:
199-
async with ydb_dbapi.AsyncCursor(
200-
session=session, auto_scroll_result_sets=True
201-
) as cursor:
170+
async with ydb_dbapi.AsyncCursor(session=session) as cursor:
202171
yql_text = """
203172
SELECT id, val FROM table;
204173
SELECT id, val FROM table1;
205174
SELECT id, val FROM table2;
206175
"""
207176
await cursor.execute(query=yql_text)
208177

178+
assert cursor.rowcount == 12
179+
209180
res = await cursor.fetchall()
210181

211182
assert len(res) == RESULT_SET_COUNT * RESULT_SET_LENGTH
@@ -274,45 +245,19 @@ def test_cursor_fetch_all(self, session_sync: ydb.QuerySession) -> None:
274245

275246
assert cursor.fetchall() == []
276247

277-
def test_cursor_next_set(self, session_sync: ydb.QuerySession) -> None:
278-
with ydb_dbapi.Cursor(session=session_sync) as cursor:
279-
yql_text = """SELECT 1 as val; SELECT 2 as val;"""
280-
cursor.execute(query=yql_text)
281-
282-
res = cursor.fetchall()
283-
assert res is not None
284-
assert len(res) == 1
285-
assert res[0][0] == 1
286-
287-
nextset = cursor.nextset()
288-
assert nextset
289-
290-
res = cursor.fetchall()
291-
assert res is not None
292-
assert len(res) == 1
293-
assert res[0][0] == 2
294-
295-
nextset = cursor.nextset()
296-
assert nextset
297-
298-
assert cursor.fetchall() == []
299-
300-
nextset = cursor.nextset()
301-
assert not nextset
302-
303-
def test_cursor_fetch_one_autoscroll(
248+
def test_cursor_fetch_one_multiple_result_sets(
304249
self, session_sync: ydb.QuerySession
305250
) -> None:
306-
with ydb_dbapi.Cursor(
307-
session=session_sync, auto_scroll_result_sets=True
308-
) as cursor:
251+
with ydb_dbapi.Cursor(session=session_sync) as cursor:
309252
yql_text = """
310253
SELECT id, val FROM table;
311254
SELECT id, val FROM table1;
312255
SELECT id, val FROM table2;
313256
"""
314257
cursor.execute(query=yql_text)
315258

259+
assert cursor.rowcount == 12
260+
316261
for i in range(RESULT_SET_LENGTH * RESULT_SET_COUNT):
317262
res = cursor.fetchone()
318263
assert res is not None
@@ -321,19 +266,19 @@ def test_cursor_fetch_one_autoscroll(
321266
assert cursor.fetchone() is None
322267
assert not cursor.nextset()
323268

324-
def test_cursor_fetch_many_autoscroll(
269+
def test_cursor_fetch_many_multiple_result_sets(
325270
self, session_sync: ydb.QuerySession
326271
) -> None:
327-
with ydb_dbapi.Cursor(
328-
session=session_sync, auto_scroll_result_sets=True
329-
) as cursor:
272+
with ydb_dbapi.Cursor(session=session_sync) as cursor:
330273
yql_text = """
331274
SELECT id, val FROM table;
332275
SELECT id, val FROM table1;
333276
SELECT id, val FROM table2;
334277
"""
335278
cursor.execute(query=yql_text)
336279

280+
assert cursor.rowcount == 12
281+
337282
halfsize = (RESULT_SET_LENGTH * RESULT_SET_COUNT) // 2
338283
for _ in range(2):
339284
res = cursor.fetchmany(size=halfsize)
@@ -343,19 +288,19 @@ def test_cursor_fetch_many_autoscroll(
343288
assert cursor.fetchmany(2) == []
344289
assert not cursor.nextset()
345290

346-
def test_cursor_fetch_all_autoscroll(
291+
def test_cursor_fetch_all_multiple_result_sets(
347292
self, session_sync: ydb.QuerySession
348293
) -> None:
349-
with ydb_dbapi.Cursor(
350-
session=session_sync, auto_scroll_result_sets=True
351-
) as cursor:
294+
with ydb_dbapi.Cursor(session=session_sync) as cursor:
352295
yql_text = """
353296
SELECT id, val FROM table;
354297
SELECT id, val FROM table1;
355298
SELECT id, val FROM table2;
356299
"""
357300
cursor.execute(query=yql_text)
358301

302+
assert cursor.rowcount == 12
303+
359304
res = cursor.fetchall()
360305

361306
assert len(res) == RESULT_SET_COUNT * RESULT_SET_LENGTH

0 commit comments

Comments
 (0)