response iterator as context manager

Author: vgvoleg

examples/query-service/basic_example.py

@@ -38,39 +38,40 @@ def callee(session):
 
         print("=" * 50)
         print("INSERT WITH COMMIT TX")
-        tx = session.transaction()
 
-        tx.begin()
+        with session.transaction() as tx:
+            tx.begin()
 
-        tx.execute("""INSERT INTO example (key, value) VALUES (1, "onepieceisreal");""")
+            with tx.execute("""INSERT INTO example (key, value) VALUES (1, "onepieceisreal");""") as results:
+                pass
 
-        for result_set in tx.execute("""SELECT COUNT(*) as rows_count FROM example;"""):
-            print(f"rows: {str(result_set.rows)}")
+            with tx.execute("""SELECT COUNT(*) as rows_count FROM example;""") as results:
+                for result_set in results:
+                    print(f"rows: {str(result_set.rows)}")
 
-        tx.commit()
+            tx.commit()
 
-        print("=" * 50)
-        print("AFTER COMMIT TX")
+            print("=" * 50)
+            print("AFTER COMMIT TX")
 
-        for result_set in session.execute("""SELECT COUNT(*) as rows_count FROM example;"""):
-            print(f"rows: {str(result_set.rows)}")
+            for result_set in session.execute("""SELECT COUNT(*) as rows_count FROM example;"""):
+                print(f"rows: {str(result_set.rows)}")
 
         print("=" * 50)
         print("INSERT WITH ROLLBACK TX")
 
-        tx = session.transaction()
+        with session.transaction() as tx:
+            tx.begin()
 
-        tx.begin()
+            tx.execute("""INSERT INTO example (key, value) VALUES (2, "onepieceisreal");""")
 
-        tx.execute("""INSERT INTO example (key, value) VALUES (2, "onepieceisreal");""")
+            for result_set in tx.execute("""SELECT COUNT(*) as rows_count FROM example;"""):
+                print(f"rows: {str(result_set.rows)}")
 
-        for result_set in tx.execute("""SELECT COUNT(*) as rows_count FROM example;"""):
-            print(f"rows: {str(result_set.rows)}")
-
-        tx.rollback()
+            tx.rollback()
 
-        print("=" * 50)
-        print("AFTER ROLLBACK TX")
+            print("=" * 50)
+            print("AFTER ROLLBACK TX")
 
         for result_set in session.execute("""SELECT COUNT(*) as rows_count FROM example;"""):
             print(f"rows: {str(result_set.rows)}")

tests/query/test_query_transaction.py

@@ -68,3 +68,11 @@ def test_context_manager_does_not_hide_exceptions(self, tx: BaseTxContext):
         with pytest.raises(Exception):
             with tx:
                 raise Exception()
+
+    def test_execute_as_context_manager(self, tx: BaseTxContext):
+        tx.begin()
+
+        with tx.execute("select 1;") as results:
+            res = [result_set for result_set in results]
+
+        assert len(res) == 1

ydb/query/base.py

@@ -16,6 +16,7 @@
 )
 from .. import convert
 from .. import issues
+from .. import _utilities
 
 
 class QueryClientSettings:
@@ -215,3 +216,12 @@ def decorator(rpc_state, response_pb, session_state, *args, **kwargs):
             raise
 
     return decorator
+
+
+class SyncResponseContextIterator(_utilities.SyncResponseIterator):
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        for _ in self.it:
+            pass

ydb/query/transaction.py

@@ -6,7 +6,6 @@
 from .. import (
     _apis,
     issues,
-    _utilities,
 )
 from .._grpc.grpcwrapper import ydb_query as _ydb_query
 from .._grpc.grpcwrapper import ydb_query_public_types as _ydb_query_public
@@ -340,7 +339,7 @@ def execute(
             parameters=parameters,
             concurrent_result_sets=concurrent_result_sets,
         )
-        self._prev_stream = _utilities.SyncResponseIterator(
+        self._prev_stream = base.SyncResponseContextIterator(
             stream_it,
             lambda resp: base.wrap_execute_query_response(
                 rpc_state=None,