review fixes

Author: vgvoleg

examples/query-service/basic_example.py

@@ -18,7 +18,7 @@ def main():
 
     print("=" * 50)
     print("DELETE TABLE IF EXISTS")
-    pool.execute_with_retries("drop table if exists example")
+    pool.execute_with_retries("DROP TABLE IF EXISTS example")
 
     print("=" * 50)
     print("CREATE TABLE")
@@ -28,11 +28,11 @@ def main():
 
     def callee(session):
         print("=" * 50)
-        with session.execute("delete from example"):
+        with session.execute("DELETE FROM example"):
             pass
 
         print("BEFORE ACTION")
-        with session.execute("SELECT COUNT(*) as rows_count FROM example") as results:
+        with session.execute("SELECT COUNT(*) AS rows_count FROM example") as results:
             for result_set in results:
                 print(f"rows: {str(result_set.rows)}")
 
@@ -45,7 +45,7 @@ def callee(session):
             with tx.execute("INSERT INTO example (key, value) VALUES (1, 'onepieceisreal')"):
                 pass
 
-            with tx.execute("SELECT COUNT(*) as rows_count FROM example") as results:
+            with tx.execute("SELECT COUNT(*) AS rows_count FROM example") as results:
                 for result_set in results:
                     print(f"rows: {str(result_set.rows)}")
 
@@ -54,7 +54,7 @@ def callee(session):
         print("=" * 50)
         print("AFTER COMMIT TX")
 
-        with session.execute("SELECT COUNT(*) as rows_count FROM example") as results:
+        with session.execute("SELECT COUNT(*) AS rows_count FROM example") as results:
             for result_set in results:
                 print(f"rows: {str(result_set.rows)}")
 
@@ -67,7 +67,7 @@ def callee(session):
             with tx.execute("INSERT INTO example (key, value) VALUES (2, 'onepieceisreal')"):
                 pass
 
-            with tx.execute("SELECT COUNT(*) as rows_count FROM example") as results:
+            with tx.execute("SELECT COUNT(*) AS rows_count FROM example") as results:
                 for result_set in results:
                     print(f"rows: {str(result_set.rows)}")
 
@@ -76,7 +76,7 @@ def callee(session):
         print("=" * 50)
         print("AFTER ROLLBACK TX")
 
-        with session.execute("SELECT COUNT(*) as rows_count FROM example") as results:
+        with session.execute("SELECT COUNT(*) AS rows_count FROM example") as results:
             for result_set in results:
                 print(f"rows: {str(result_set.rows)}")
 

tests/query/test_query_session.py

@@ -93,9 +93,10 @@ def test_basic_execute(self, session: QuerySessionSync):
     def test_two_results(self, session: QuerySessionSync):
         session.create()
         res = []
+
         with session.execute("select 1; select 2") as results:
             for result_set in results:
                 if len(result_set.rows) > 0:
-                    res.extend(list(result_set.rows[0].values()))
+                    res.append(list(result_set.rows[0].values()))
 
-        assert res == [1, 2]
+        assert res == [[1], [2]]

ydb/query/pool.py

@@ -2,6 +2,7 @@
 from typing import (
     Callable,
     Optional,
+    List,
 )
 
 from . import base
@@ -12,6 +13,7 @@
     RetrySettings,
     retry_operation_sync,
 )
+from .. import convert
 
 logger = logging.getLogger(__name__)
 
@@ -49,8 +51,12 @@ def wrapped_callee():
 
         return retry_operation_sync(wrapped_callee, retry_settings)
 
-    def execute_with_retries(self, query: str, retry_settings: Optional[RetrySettings] = None, *args, **kwargs):
+    def execute_with_retries(
+        self, query: str, retry_settings: Optional[RetrySettings] = None, *args, **kwargs
+    ) -> List[convert.ResultSet]:
         """Special interface to execute a one-shot queries in a safe, retriable way.
+        Note: this method loads all data from stream before return, do not use this
+        method with huge read queries.
 
         :param query: A query, yql or sql text.
         :param retry_settings: RetrySettings object.