transaction identity

Author: vgvoleg

examples/topic/topic_transactions_example.py

@@ -5,13 +5,13 @@ def writer_example(driver: ydb.Driver, topic: str):
     session_pool = ydb.QuerySessionPool(driver)
 
     def callee(tx: ydb.QueryTxContext):
-        tx_writer: ydb.TopicTxWriter = driver.topic_client.tx_writer(tx, topic) # <=======
+        tx_writer: ydb.TopicTxWriter = driver.topic_client.tx_writer(tx, topic)  # <=======
         # дефолт - без дедупликации, без ретраев и без producer_id.
 
         with tx.execute(query="select 1") as result_sets:
             messages = [result_set.rows[0] for result_set in result_sets]
 
-        tx_writer.write(messages) # вне зависимости от состояния вышестоящего стрима поведение должно быть одинаковое
+        tx_writer.write(messages)  # вне зависимости от состояния вышестоящего стрима поведение должно быть одинаковое
 
     session_pool.retry_tx_sync(callee)
 
@@ -20,7 +20,7 @@ def reader_example(driver: ydb.Driver, reader: ydb.TopicReader):
     session_pool = ydb.QuerySessionPool(driver)
 
     def callee(tx: ydb.QueryTxContext):
-        batch = reader.receive_batch_with_tx(tx, max_messages=5) # <=======
+        batch = reader.receive_batch_with_tx(tx, max_messages=5)  # <=======
 
         with tx.execute(query="INSERT INTO max_values(val) VALUES ($val)", parameters={"$val": max(batch)}) as _:
             pass

tests/query/test_query_transaction.py

@@ -92,3 +92,15 @@ def test_execute_two_results(self, tx: QueryTxContext):
 
         assert res == [[1], [2]]
         assert counter == 2
+
+    def test_tx_identity_before_begin_raises(self, tx: QueryTxContext):
+        with pytest.raises(RuntimeError):
+            tx._tx_identity()
+
+    def test_tx_identity_after_begin_works(self, tx: QueryTxContext):
+        tx.begin()
+
+        identity = tx._tx_identity()
+
+        assert identity.tx_id == tx.tx_id
+        assert identity.session_id == tx.session_id

ydb/_grpc/grpcwrapper/ydb_query.py

@@ -6,8 +6,10 @@
 # Workaround for good IDE and universal for runtime
 if typing.TYPE_CHECKING:
     from ..v4.protos import ydb_query_pb2
+    from ..v4.protos import ydb_topic_pb2
 else:
     from ..common.protos import ydb_query_pb2
+    from ..common.protos import ydb_topic_pb2
 
 from . import ydb_query_public_types as public_types
 
@@ -81,6 +83,18 @@ def to_proto(self) -> ydb_query_pb2.TransactionSettings:
             return ydb_query_pb2.TransactionSettings(stale_read_only=self.tx_mode.to_proto())
 
 
+@dataclass
+class TransactionIdentity(IToProto):
+    tx_id: str
+    session_id: str
+
+    def to_proto(self) -> ydb_topic_pb2.TransactionIdentity:
+        return ydb_topic_pb2.TransactionIdentity(
+            id=self.tx_id,
+            session=self.session_id,
+        )
+
+
 @dataclass
 class BeginTransactionRequest(IToProto):
     session_id: str

ydb/query/transaction.py

@@ -215,6 +215,11 @@ def tx_id(self) -> Optional[str]:
         """
         return self._tx_state.tx_id
 
+    def _tx_identity(self) -> _ydb_query.TransactionIdentity:
+        if not self.tx_id:
+            raise RuntimeError("Unable to get tx identity without started tx.")
+        return _ydb_query.TransactionIdentity(self.tx_id, self.session_id)
+
     def _begin_call(self, settings: Optional[BaseRequestSettings]) -> "BaseQueryTxContext":
         self._tx_state._check_invalid_transition(QueryTxStateEnum.BEGINED)
 

ydb/topic.py

@@ -294,7 +294,7 @@ def tx_writer(
         # If max_worker in the executor is 1 - then encoders will be called from the thread without parallel.
         encoder_executor: Optional[concurrent.futures.Executor] = None,
     ) -> TopicTxWriterAsyncIO:
-        
+        pass
 
     def close(self):
         if self._closed: