wow tx begin works

Author: vgvoleg

examples/query-service/basic_example.py

@@ -21,9 +21,9 @@ def main():
 
     it = session.execute("select 1; select 2;", commit_tx=False)
     for result_set in it:
-        pass
-        # print(f"columns: {str(result_set.columns)}")
-        # print(f"rows: {str(result_set.rows)}")
+        # pass
+        print(f"columns: {str(result_set.columns)}")
+        print(f"rows: {str(result_set.rows)}")
 
 
 

tests/query/test_query_transaction.py

@@ -10,8 +10,8 @@ def test_transaction_begin(self, driver_sync):
 
         tx = session.transaction()
 
-        assert tx._tx_state.tx_id == None
+        assert tx.tx_id == None
 
         tx.begin()
 
-        assert tx._tx_state.tx_id != None
+        assert tx.tx_id != None

ydb/query/base.py

@@ -17,9 +17,6 @@
 class QueryClientSettings: ...
 
 
-class IQueryTxContext: ...
-
-
 class QuerySessionState:
     _session_id: Optional[str]
     _node_id: Optional[int]
@@ -60,6 +57,7 @@ def set_attached(self, attached: bool) -> None:
 
 
 class IQuerySession(abc.ABC):
+    @abc.abstractmethod
     def __init__(self, driver: SupportedDriverType, settings: QueryClientSettings = None):
         pass
 
@@ -72,7 +70,48 @@ def delete(self) -> None:
         pass
 
     @abc.abstractmethod
-    def transaction(self, tx_mode: BaseQueryTxMode) -> IQueryTxContext:
+    def transaction(self, tx_mode: BaseQueryTxMode) -> "IQueryTxContext":
+        pass
+
+
+class IQueryTxContext(abc.ABC):
+
+    @abc.abstractmethod
+    def __init__(self, driver: SupportedDriverType, session_state: QuerySessionState, session: IQuerySession, tx_mode: BaseQueryTxMode = None):
+        pass
+
+    @abc.abstractmethod
+    def __enter__(self):
+        pass
+
+    @abc.abstractmethod
+    def __exit__(self, *args, **kwargs):
+        pass
+
+    @property
+    @abc.abstractmethod
+    def session_id(self):
+        pass
+
+    @property
+    @abc.abstractmethod
+    def tx_id(self):
+        pass
+
+    @abc.abstractmethod
+    def begin():
+        pass
+
+    @abc.abstractmethod
+    def commit():
+        pass
+
+    @abc.abstractmethod
+    def rollback():
+        pass
+
+    @abc.abstractmethod
+    def execute(query: str):
         pass
 
 
@@ -103,6 +142,9 @@ def create_execute_query_request(query: str, session_id: str, commit_tx: bool):
 
 def wrap_execute_query_response(rpc_state, response_pb):
 
-    print(response_pb)
+    # print("RESP:")
+    # print(f"meta: {response_pb.tx_meta}")
+    # print(response_pb)
+
 
     return convert.ResultSet.from_message(response_pb.result_set)

ydb/query/session.py

@@ -77,6 +77,7 @@ def _execute_call(self, query: str, commit_tx: bool):
             session_id=self._state.session_id,
             commit_tx=commit_tx
         )
+        print(request)
         return self._driver(
             request,
             _apis.QueryService.Stub,

ydb/query/transaction.py

@@ -11,27 +11,23 @@
 
 from .._tx_ctx_impl import TxState, reset_tx_id_handler
 from .._session_impl import bad_session_handler
-from ..table import (
-    AbstractTransactionModeBuilder,
-    ITxContext,
-    SerializableReadWrite
-)
+from . import base
 
 logger = logging.getLogger(__name__)
 
-def patch_table_service_tx_mode_to_query_service(tx_mode: AbstractTransactionModeBuilder):
-    if tx_mode.name == 'snapshot_read_only':
-        tx_mode = _ydb_query_public.QuerySnapshotReadOnly()
-    elif tx_mode.name == 'serializable_read_write':
-        tx_mode = _ydb_query_public.QuerySerializableReadWrite()
-    elif tx_mode.name =='online_read_only':
-        tx_mode = _ydb_query_public.QueryOnlineReadOnly()
-    elif tx_mode.name == 'stale_read_only':
-        tx_mode = _ydb_query_public.QueryStaleReadOnly()
-    else:
-        raise issues.YDBInvalidArgumentError(f'Unknown transaction mode: {tx_mode.name}')
+# def patch_table_service_tx_mode_to_query_service(tx_mode: AbstractTransactionModeBuilder):
+#     if tx_mode.name == 'snapshot_read_only':
+#         tx_mode = _ydb_query_public.QuerySnapshotReadOnly()
+#     elif tx_mode.name == 'serializable_read_write':
+#         tx_mode = _ydb_query_public.QuerySerializableReadWrite()
+#     elif tx_mode.name =='online_read_only':
+#         tx_mode = _ydb_query_public.QueryOnlineReadOnly()
+#     elif tx_mode.name == 'stale_read_only':
+#         tx_mode = _ydb_query_public.QueryStaleReadOnly()
+#     else:
+#         raise issues.YDBInvalidArgumentError(f'Unknown transaction mode: {tx_mode.name}')
 
-    return tx_mode
+#     return tx_mode
 
 
 def _construct_tx_settings(tx_state):
@@ -69,7 +65,8 @@ def wrap_tx_begin_response(rpc_state, response_pb, session_state, tx_state, tx):
     # issues._process_response(response_pb.operation)
     print("wrap result")
     message = _ydb_query.BeginTransactionResponse.from_proto(response_pb)
-    tx_state.tx_id = message.tx_meta.id
+
+    tx_state.tx_id = message.tx_meta.tx_id
     return tx
 
 
@@ -83,7 +80,7 @@ def wrap_result_on_rollback_or_commit_tx(rpc_state, response_pb, session_state,
     return tx
 
 
-class BaseTxContext(ITxContext):
+class BaseTxContext(base.IQueryTxContext):
 
     _COMMIT = "commit"
     _ROLLBACK = "rollback"
@@ -108,9 +105,7 @@ def __init__(self, driver, session_state, session, tx_mode=None):
         """
         self._driver = driver
         if tx_mode is None:
-            tx_mode = patch_table_service_tx_mode_to_query_service(SerializableReadWrite())
-        else:
-            tx_mode = patch_table_service_tx_mode_to_query_service(tx_mode)
+            tx_mode = _ydb_query_public.QuerySerializableReadWrite()
         self._tx_state = TxState(tx_mode)
         self._session_state = session_state
         self.session = session