fix review

Author: vgvoleg

tests/query/test_query_session.py

@@ -89,3 +89,13 @@ def test_basic_execute(self, session: QuerySessionSync):
         assert len(result_sets[0].rows) == 1
         assert len(result_sets[0].columns) == 1
         assert list(result_sets[0].rows[0].values()) == [1]
+
+    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()))
+
+        assert res == [1, 2]

tests/query/test_query_transaction.py

@@ -1,78 +1,78 @@
 import pytest
 
-from ydb.query.transaction import BaseTxContext
+from ydb.query.transaction import BaseQueryTxContext
 from ydb.query.transaction import QueryTxStateEnum
 
 
 class TestQueryTransaction:
-    def test_tx_begin(self, tx: BaseTxContext):
+    def test_tx_begin(self, tx: BaseQueryTxContext):
         assert tx.tx_id is None
 
         tx.begin()
         assert tx.tx_id is not None
 
-    def test_tx_allow_double_commit(self, tx: BaseTxContext):
+    def test_tx_allow_double_commit(self, tx: BaseQueryTxContext):
         tx.begin()
         tx.commit()
         tx.commit()
 
-    def test_tx_allow_double_rollback(self, tx: BaseTxContext):
+    def test_tx_allow_double_rollback(self, tx: BaseQueryTxContext):
         tx.begin()
         tx.rollback()
         tx.rollback()
 
-    def test_tx_commit_before_begin(self, tx: BaseTxContext):
+    def test_tx_commit_before_begin(self, tx: BaseQueryTxContext):
         tx.commit()
         assert tx._tx_state._state == QueryTxStateEnum.COMMITTED
 
-    def test_tx_rollback_before_begin(self, tx: BaseTxContext):
+    def test_tx_rollback_before_begin(self, tx: BaseQueryTxContext):
         tx.rollback()
         assert tx._tx_state._state == QueryTxStateEnum.ROLLBACKED
 
-    def test_tx_first_execute_begins_tx(self, tx: BaseTxContext):
+    def test_tx_first_execute_begins_tx(self, tx: BaseQueryTxContext):
         tx.execute("select 1;")
         tx.commit()
 
-    def test_interactive_tx_commit(self, tx: BaseTxContext):
+    def test_interactive_tx_commit(self, tx: BaseQueryTxContext):
         tx.execute("select 1;", commit_tx=True)
         with pytest.raises(RuntimeError):
             tx.execute("select 1;")
 
-    def test_tx_execute_raises_after_commit(self, tx: BaseTxContext):
+    def test_tx_execute_raises_after_commit(self, tx: BaseQueryTxContext):
         tx.begin()
         tx.commit()
         with pytest.raises(RuntimeError):
             tx.execute("select 1;")
 
-    def test_tx_execute_raises_after_rollback(self, tx: BaseTxContext):
+    def test_tx_execute_raises_after_rollback(self, tx: BaseQueryTxContext):
         tx.begin()
         tx.rollback()
         with pytest.raises(RuntimeError):
             tx.execute("select 1;")
 
-    def test_context_manager_rollbacks_tx(self, tx: BaseTxContext):
+    def test_context_manager_rollbacks_tx(self, tx: BaseQueryTxContext):
         with tx:
             tx.begin()
 
         assert tx._tx_state._state == QueryTxStateEnum.ROLLBACKED
 
-    def test_context_manager_normal_flow(self, tx: BaseTxContext):
+    def test_context_manager_normal_flow(self, tx: BaseQueryTxContext):
         with tx:
             tx.begin()
             tx.execute("select 1;")
             tx.commit()
 
         assert tx._tx_state._state == QueryTxStateEnum.COMMITTED
 
-    def test_context_manager_does_not_hide_exceptions(self, tx: BaseTxContext):
+    def test_context_manager_does_not_hide_exceptions(self, tx: BaseQueryTxContext):
         class CustomException(Exception):
             pass
 
         with pytest.raises(CustomException):
             with tx:
                 raise CustomException()
 
-    def test_execute_as_context_manager(self, tx: BaseTxContext):
+    def test_execute_as_context_manager(self, tx: BaseQueryTxContext):
         tx.begin()
 
         with tx.execute("select 1;") as results:

tox.ini

@@ -75,6 +75,7 @@ builtins = _
 max-line-length = 160
 ignore=E203,W503
 exclude=*_pb2.py,*_grpc.py,.venv,.git,.tox,dist,doc,*egg,ydb/public/api/protos/*,docs/*,ydb/public/api/grpc/*,persqueue/*,client/*,dbapi/*,ydb/default_pem.py,*docs/conf.py
+per-file-ignores = ydb/table.py:F401
 
 [pytest]
 asyncio_mode = auto

ydb/_grpc/grpcwrapper/ydb_query.py

@@ -21,7 +21,7 @@
 
 @dataclass
 class CreateSessionResponse(IFromProto):
-    status: Optional[ServerStatus]
+    status: ServerStatus
     session_id: str
     node_id: int
 
@@ -36,7 +36,7 @@ def from_proto(msg: ydb_query_pb2.CreateSessionResponse) -> "CreateSessionRespon
 
 @dataclass
 class DeleteSessionResponse(IFromProto):
-    status: Optional[ServerStatus]
+    status: ServerStatus
 
     @staticmethod
     def from_proto(msg: ydb_query_pb2.DeleteSessionResponse) -> "DeleteSessionResponse":
@@ -129,10 +129,10 @@ def from_proto(msg: ydb_query_pb2.RollbackTransactionResponse) -> "RollbackTrans
 @dataclass
 class QueryContent(IFromPublic, IToProto):
     text: str
-    syntax: Optional[int] = None
+    syntax: int
 
     @staticmethod
-    def from_public(query: str, syntax: int = None) -> "QueryContent":
+    def from_public(query: str, syntax: int) -> "QueryContent":
         return QueryContent(text=query, syntax=syntax)
 
     def to_proto(self) -> ydb_query_pb2.QueryContent:
@@ -141,9 +141,9 @@ def to_proto(self) -> ydb_query_pb2.QueryContent:
 
 @dataclass
 class TransactionControl(IToProto):
-    begin_tx: Optional[TransactionSettings] = None
-    commit_tx: Optional[bool] = None
-    tx_id: Optional[str] = None
+    begin_tx: Optional[TransactionSettings]
+    commit_tx: Optional[bool]
+    tx_id: Optional[str]
 
     def to_proto(self) -> ydb_query_pb2.TransactionControl:
         if self.tx_id:
@@ -161,11 +161,11 @@ def to_proto(self) -> ydb_query_pb2.TransactionControl:
 class ExecuteQueryRequest(IToProto):
     session_id: str
     query_content: QueryContent
-    tx_control: Optional[TransactionControl] = None
-    concurrent_result_sets: Optional[bool] = False
-    exec_mode: Optional[int] = None
-    parameters: Optional[dict] = None
-    stats_mode: Optional[int] = None
+    tx_control: TransactionControl
+    concurrent_result_sets: bool
+    exec_mode: int
+    parameters: dict
+    stats_mode: int
 
     def to_proto(self) -> ydb_query_pb2.ExecuteQueryRequest:
         tx_control = self.tx_control.to_proto() if self.tx_control is not None else self.tx_control

ydb/query/__init__.py

@@ -5,6 +5,7 @@
     "QueryStaleReadOnly",
     "QuerySessionPool",
     "QueryClientSync",
+    "QuerySessionSync",
 ]
 
 import logging