Skip to content

Commit 49e7471

Browse files
committed
review fixes
1 parent 2336624 commit 49e7471

File tree

3 files changed

+17
-10
lines changed

3 files changed

+17
-10
lines changed

examples/query-service/basic_example.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ def main():
1818

1919
print("=" * 50)
2020
print("DELETE TABLE IF EXISTS")
21-
pool.execute_with_retries("drop table if exists example")
21+
pool.execute_with_retries("DROP TABLE IF EXISTS example")
2222

2323
print("=" * 50)
2424
print("CREATE TABLE")
@@ -28,11 +28,11 @@ def main():
2828

2929
def callee(session):
3030
print("=" * 50)
31-
with session.execute("delete from example"):
31+
with session.execute("DELETE FROM example"):
3232
pass
3333

3434
print("BEFORE ACTION")
35-
with session.execute("SELECT COUNT(*) as rows_count FROM example") as results:
35+
with session.execute("SELECT COUNT(*) AS rows_count FROM example") as results:
3636
for result_set in results:
3737
print(f"rows: {str(result_set.rows)}")
3838

@@ -45,7 +45,7 @@ def callee(session):
4545
with tx.execute("INSERT INTO example (key, value) VALUES (1, 'onepieceisreal')"):
4646
pass
4747

48-
with tx.execute("SELECT COUNT(*) as rows_count FROM example") as results:
48+
with tx.execute("SELECT COUNT(*) AS rows_count FROM example") as results:
4949
for result_set in results:
5050
print(f"rows: {str(result_set.rows)}")
5151

@@ -54,7 +54,7 @@ def callee(session):
5454
print("=" * 50)
5555
print("AFTER COMMIT TX")
5656

57-
with session.execute("SELECT COUNT(*) as rows_count FROM example") as results:
57+
with session.execute("SELECT COUNT(*) AS rows_count FROM example") as results:
5858
for result_set in results:
5959
print(f"rows: {str(result_set.rows)}")
6060

@@ -67,7 +67,7 @@ def callee(session):
6767
with tx.execute("INSERT INTO example (key, value) VALUES (2, 'onepieceisreal')"):
6868
pass
6969

70-
with tx.execute("SELECT COUNT(*) as rows_count FROM example") as results:
70+
with tx.execute("SELECT COUNT(*) AS rows_count FROM example") as results:
7171
for result_set in results:
7272
print(f"rows: {str(result_set.rows)}")
7373

@@ -76,7 +76,7 @@ def callee(session):
7676
print("=" * 50)
7777
print("AFTER ROLLBACK TX")
7878

79-
with session.execute("SELECT COUNT(*) as rows_count FROM example") as results:
79+
with session.execute("SELECT COUNT(*) AS rows_count FROM example") as results:
8080
for result_set in results:
8181
print(f"rows: {str(result_set.rows)}")
8282

tests/query/test_query_session.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -93,9 +93,10 @@ def test_basic_execute(self, session: QuerySessionSync):
9393
def test_two_results(self, session: QuerySessionSync):
9494
session.create()
9595
res = []
96+
9697
with session.execute("select 1; select 2") as results:
9798
for result_set in results:
9899
if len(result_set.rows) > 0:
99-
res.extend(list(result_set.rows[0].values()))
100+
res.append(list(result_set.rows[0].values()))
100101

101-
assert res == [1, 2]
102+
assert res == [[1], [2]]

ydb/query/pool.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
from typing import (
33
Callable,
44
Optional,
5+
List,
56
)
67

78
from . import base
@@ -12,6 +13,7 @@
1213
RetrySettings,
1314
retry_operation_sync,
1415
)
16+
from .. import convert
1517

1618
logger = logging.getLogger(__name__)
1719

@@ -49,8 +51,12 @@ def wrapped_callee():
4951

5052
return retry_operation_sync(wrapped_callee, retry_settings)
5153

52-
def execute_with_retries(self, query: str, retry_settings: Optional[RetrySettings] = None, *args, **kwargs):
54+
def execute_with_retries(
55+
self, query: str, retry_settings: Optional[RetrySettings] = None, *args, **kwargs
56+
) -> List[convert.ResultSet]:
5357
"""Special interface to execute a one-shot queries in a safe, retriable way.
58+
Note: this method loads all data from stream before return, do not use this
59+
method with huge read queries.
5460
5561
:param query: A query, yql or sql text.
5662
:param retry_settings: RetrySettings object.

0 commit comments

Comments
 (0)