tx state handler

Author: vgvoleg

ydb/_grpc/grpcwrapper/ydb_query.py

@@ -116,6 +116,29 @@ def from_proto(msg: ydb_query_pb2.BeginTransactionResponse) -> "BeginTransaction
             tx_meta=TransactionMeta.from_proto(msg.tx_meta),
         )
 
+
+@dataclass
+class CommitTransactionResponse(IFromProto):
+    status: Optional[ServerStatus]
+
+    @staticmethod
+    def from_proto(msg: ydb_query_pb2.CommitTransactionResponse) -> "CommitTransactionResponse":
+        return CommitTransactionResponse(
+            status=ServerStatus(msg.status, msg.issues),
+        )
+
+
+@dataclass
+class RollbackTransactionResponse(IFromProto):
+    status: Optional[ServerStatus]
+
+    @staticmethod
+    def from_proto(msg: ydb_query_pb2.RollbackTransactionResponse) -> "RollbackTransactionResponse":
+        return RollbackTransactionResponse(
+            status=ServerStatus(msg.status, msg.issues),
+        )
+
+
 @dataclass
 class QueryContent(IFromPublic, IToProto):
     text: str

ydb/query/base.py

@@ -1,4 +1,6 @@
 import abc
+import enum
+import functools
 
 from typing import (
     Optional,
@@ -13,10 +15,51 @@
     QuerySerializableReadWrite
 )
 from .. import convert
+from .. import issues
 
 class QueryClientSettings: ...
 
 
+class QuerySessionStateEnum(enum.Enum):
+    NOT_INITIALIZED = "NOT_INITIALIZED"
+    CREATED = "CREATED"
+    CLOSED = "CLOSED"
+
+
+class QuerySessionStateHelper(abc.ABC):
+    _VALID_TRANSITIONS = {
+        QuerySessionStateEnum.NOT_INITIALIZED: [QuerySessionStateEnum.CREATED],
+        QuerySessionStateEnum.CREATED: [QuerySessionStateEnum.CLOSED],
+        QuerySessionStateEnum.CLOSED: []
+    }
+
+    @classmethod
+    def valid_transition(cls, before: QuerySessionStateEnum, after: QuerySessionStateEnum) -> bool:
+        return after in cls._VALID_TRANSITIONS[before]
+
+
+class QueryTxStateEnum(enum.Enum):
+    NOT_INITIALIZED = "NOT_INITIALIZED"
+    BEGINED = "BEGINED"
+    COMMITTED = "COMMITTED"
+    ROLLBACKED = "ROLLBACKED"
+    DEAD = "DEAD"
+
+
+class QueryTxStateHelper(abc.ABC):
+    _VALID_TRANSITIONS = {
+        QueryTxStateEnum.NOT_INITIALIZED: [QueryTxStateEnum.BEGINED, QueryTxStateEnum.DEAD],
+        QueryTxStateEnum.BEGINED: [QueryTxStateEnum.COMMITTED, QueryTxStateEnum.ROLLBACKED, QueryTxStateEnum.DEAD],
+        QueryTxStateEnum.COMMITTED: [],
+        QueryTxStateEnum.ROLLBACKED: [],
+        QueryTxStateEnum.DEAD: [],
+    }
+
+    @classmethod
+    def valid_transition(cls, before: QueryTxStateEnum, after: QueryTxStateEnum) -> bool:
+        return after in cls._VALID_TRANSITIONS[before]
+
+
 class QuerySessionState:
     _session_id: Optional[str]
     _node_id: Optional[int]
@@ -99,15 +142,15 @@ def tx_id(self):
         pass
 
     @abc.abstractmethod
-    def begin():
+    def begin(settings: QueryClientSettings = None):
         pass
 
     @abc.abstractmethod
-    def commit():
+    def commit(settings: QueryClientSettings = None):
         pass
 
     @abc.abstractmethod
-    def rollback():
+    def rollback(settings: QueryClientSettings = None):
         pass
 
     @abc.abstractmethod
@@ -142,9 +185,26 @@ def create_execute_query_request(query: str, session_id: str, commit_tx: bool):
 
 def wrap_execute_query_response(rpc_state, response_pb):
 
-    # print("RESP:")
-    # print(f"meta: {response_pb.tx_meta}")
-    # print(response_pb)
+    return convert.ResultSet.from_message(response_pb.result_set)
 
+X_YDB_SERVER_HINTS = "x-ydb-server-hints"
+X_YDB_SESSION_CLOSE = "session-close"
 
-    return convert.ResultSet.from_message(response_pb.result_set)
+
+def _check_session_is_closing(rpc_state, session_state):
+    metadata = rpc_state.trailing_metadata()
+    if X_YDB_SESSION_CLOSE in metadata.get(X_YDB_SERVER_HINTS, []):
+        session_state.set_closing()
+
+
+def bad_session_handler(func):
+    @functools.wraps(func)
+    def decorator(rpc_state, response_pb, session_state, *args, **kwargs):
+        try:
+            _check_session_is_closing(rpc_state, session_state)
+            return func(rpc_state, response_pb, session_state, *args, **kwargs)
+        except issues.BadSession:
+            session_state.reset()
+            raise
+
+    return decorator

ydb/query/transaction.py

@@ -1,5 +1,7 @@
 import abc
 import logging
+import enum
+import functools
 
 from .. import (
     _apis,
@@ -9,8 +11,6 @@
 from .._grpc.grpcwrapper import ydb_query as _ydb_query
 from .._grpc.grpcwrapper import ydb_query_public_types as _ydb_query_public
 
-from .._tx_ctx_impl import TxState, reset_tx_id_handler
-from .._session_impl import bad_session_handler
 from . import base
 
 logger = logging.getLogger(__name__)
@@ -30,6 +30,38 @@
 #     return tx_mode
 
 
+def reset_tx_id_handler(func):
+    @functools.wraps(func)
+    def decorator(rpc_state, response_pb, session_state, tx_state, *args, **kwargs):
+        try:
+            return func(rpc_state, response_pb, session_state, tx_state, *args, **kwargs)
+        except issues.Error:
+            tx_state.change_state(base.QueryTxStateEnum.DEAD)
+            tx_state.tx_id = None
+            raise
+
+    return decorator
+
+
+class QueryTxState:
+    def __init__(self, tx_mode: base.BaseQueryTxMode):
+        """
+        Holds transaction context manager info
+        :param tx_mode: A mode of transaction
+        """
+        self.tx_id = None
+        self.tx_mode = tx_mode
+        self._state = base.QueryTxStateEnum.NOT_INITIALIZED
+
+    def check_invalid_transition(self, target: base.QueryTxStateEnum):
+        if not base.QueryTxStateHelper.is_valid_transition(self._state, target):
+            raise RuntimeError(f"Transaction could not be moved from {self._state.value} to {target.value}")
+
+    def change_state(self, target: base.QueryTxStateEnum):
+        self.check_invalid_transition(target)
+        self._state = target
+
+
 def _construct_tx_settings(tx_state):
     tx_settings = _ydb_query.TransactionSettings.from_public(tx_state.tx_mode)
     return tx_settings
@@ -41,8 +73,6 @@ def _create_begin_transaction_request(session_state, tx_state):
         tx_settings=_construct_tx_settings(tx_state),
     ).to_proto()
 
-    print(request)
-
     return request
 
 
@@ -59,32 +89,35 @@ def _create_rollback_transaction_request(session_state, tx_state):
     return request
 
 
-@bad_session_handler
+@base.bad_session_handler
 def wrap_tx_begin_response(rpc_state, response_pb, session_state, tx_state, tx):
-    # session_state.complete_query()
-    # issues._process_response(response_pb.operation)
-    print("wrap result")
     message = _ydb_query.BeginTransactionResponse.from_proto(response_pb)
-
+    issues._process_response(message.status)
+    tx_state.change_state(base.QueryTxStateEnum.BEGINED)
     tx_state.tx_id = message.tx_meta.tx_id
     return tx
 
 
-@bad_session_handler
+@base.bad_session_handler
 @reset_tx_id_handler
-def wrap_result_on_rollback_or_commit_tx(rpc_state, response_pb, session_state, tx_state, tx):
+def wrap_tx_commit_response(rpc_state, response_pb, session_state, tx_state, tx):
+    message = _ydb_query.CommitTransactionResponse(response_pb)
+    issues._process_response(message.status)
+    tx_state.tx_id = None
+    tx_state.change_state(base.QueryTxStateEnum.COMMITTED)
+    return tx
 
-    # issues._process_response(response_pb.operation)
-    # transaction successfully committed or rolled back
+@base.bad_session_handler
+@reset_tx_id_handler
+def wrap_tx_rollback_response(rpc_state, response_pb, session_state, tx_state, tx):
+    message = _ydb_query.RollbackTransactionResponse(response_pb)
+    issues._process_response(message.status)
     tx_state.tx_id = None
+    tx_state.change_state(base.QueryTxStateEnum.ROLLBACKED)
     return tx
 
 
 class BaseTxContext(base.IQueryTxContext):
-
-    _COMMIT = "commit"
-    _ROLLBACK = "rollback"
-
     def __init__(self, driver, session_state, session, tx_mode=None):
         """
         An object that provides a simple transaction context manager that allows statements execution
@@ -106,7 +139,7 @@ def __init__(self, driver, session_state, session, tx_mode=None):
         self._driver = driver
         if tx_mode is None:
             tx_mode = _ydb_query_public.QuerySerializableReadWrite()
-        self._tx_state = TxState(tx_mode)
+        self._tx_state = QueryTxState(tx_mode)
         self._session_state = session_state
         self.session = session
         self._finished = ""
@@ -163,17 +196,15 @@ def begin(self, settings=None):
 
         :return: An open transaction
         """
-        if self._tx_state.tx_id is not None:
-            return self
-
-        print('try to begin tx')
+        self._tx_state.check_invalid_transition(base.QueryTxStateEnum.BEGINED)
 
         return self._driver(
             _create_begin_transaction_request(self._session_state, self._tx_state),
             _apis.QueryService.Stub,
             _apis.QueryService.BeginTransaction,
-            wrap_result=wrap_tx_begin_response,
-            wrap_args=(self._session_state, self._tx_state, self),
+            wrap_tx_begin_response,
+            settings,
+            (self._session_state, self._tx_state, self),
         )
 
     def commit(self, settings=None):
@@ -186,22 +217,28 @@ def commit(self, settings=None):
         :return: A committed transaction or exception if commit is failed
         """
 
-        self._set_finish(self._COMMIT)
-
-        if self._tx_state.tx_id is None and not self._tx_state.dead:
-            return self
+        self._tx_state.check_invalid_transition(base.QueryTxStateEnum.COMMITTED)
 
         return self._driver(
             _create_commit_transaction_request(self._session_state, self._tx_state),
             _apis.QueryService.Stub,
             _apis.QueryService.CommitTransaction,
-            wrap_result_on_rollback_or_commit_tx,
+            wrap_tx_commit_response,
             settings,
             (self._session_state, self._tx_state, self),
         )
 
     def rollback(self, settings=None):
-        pass
+        self._tx_state.check_invalid_transition(base.QueryTxStateEnum.ROLLBACKED)
+
+        return self._driver(
+            _create_rollback_transaction_request(self._session_state, self._tx_state),
+            _apis.QueryService.Stub,
+            _apis.QueryService.RollbackTransaction,
+            wrap_tx_rollback_response,
+            settings,
+            (self._session_state, self._tx_state, self),
+        )
 
     def execute(self, query, parameters=None, commit_tx=False, settings=None):
         pass