PYCBC-1646: Migrate transactions to use transaction_context

thejcfactor · thejcfactor · commit 03c270e7382e · 2025-01-27T18:05:45.000Z
Motivation ========== The current path of using the C++ core async_attempt_context adds complexity to both the Python SDK and C++ core (the Python SDK is the only client using the async_attempt_context logic). Migrating to using transaction_context will place the Python SDK inline with other wrappers and also allow some code removal/cleanup in the C++ core. Modification ============ * Remove attempt_context and use transaction_context instead * Remove executing the transaction lambda in the bindings * Rearrange how transaction exceptions are built in the bindings * Update Python (non-binding) transaction code to execute transaction logic Results ======= All tests pass without any modification to tests. Change-Id: Iaca68e5d36507bcc64d2cd34afc8e244dc819fc4 Reviewed-on: https://review.couchbase.org/c/couchbase-python-client/+/221542 Tested-by: Build Bot <build@couchbase.com> Reviewed-by: Dimitris Christodoulou <dimitris.christodoulou@couchbase.com>
diff --git a/acouchbase/transactions/transactions.py b/acouchbase/transactions/transactions.py
@@ -13,13 +13,13 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-import asyncio
 import logging
 from functools import wraps
 from typing import (TYPE_CHECKING,
                     Any,
                     Awaitable,
                     Callable,
+                    Coroutine,
                     Dict,
                     Optional)
 
@@ -111,28 +111,19 @@ def __init__(self,
                  ):
         super().__init__(cluster, config)
 
-    @AsyncWrapper.inject_callbacks(TransactionResult)
-    def run(self,
-            txn_logic,  # type:  Callable[[AttemptContextLogic], None]
-            transaction_options=None,  # type: Optional[TransactionOptions]
-            **kwargs) -> Awaitable[TransactionResult]:
-        def wrapped_logic(c):
-            try:
-                ctx = AttemptContext(c, self._loop, self._transcoder)
-                asyncio.run_coroutine_threadsafe(txn_logic(ctx), self._loop).result()
-                log.debug('wrapped logic completed')
-            except Exception as e:
-                log.debug('wrapped_logic raised %s', e)
-                raise e
-
+    async def run(self,
+                  txn_logic,  # type:  Callable[[AttemptContextLogic], Coroutine[Any, Any, None]]
+                  transaction_options=None,  # type: Optional[TransactionOptions]
+                  **kwargs) -> TransactionResult:
         opts = None
         if transaction_options:
-            opts = transaction_options
+            opts = transaction_options._base
         if 'per_txn_config' in kwargs:
             Supportability.method_param_deprecated('per_txn_config', 'transaction_options')
             opts = kwargs.pop('per_txn_config', None)
 
-        return super().run(wrapped_logic, opts, **kwargs)
+        txn_result = await super().run_async(txn_logic, AttemptContext(self._txns, self._loop, self._transcoder, opts))
+        return TransactionResult(**txn_result)
 
     # TODO: make async?
     def close(self):
@@ -149,12 +140,32 @@ def close(self):
 
 
 class AttemptContext(AttemptContextLogic):
+
     def __init__(self,
                  ctx,    # type: PyCapsuleType
                  loop,    # type: AbstractEventLoop
-                 transcoder  # type: Transcoder
+                 transcoder,  # type: Transcoder
+                 opts    # type: Optional[PyCapsuleType]
                  ):
-        super().__init__(ctx, loop, transcoder)
+        super().__init__(ctx, transcoder, loop, opts)
+
+    @AsyncWrapper.inject_callbacks(None)
+    def _new_attempt(self,
+                     **kwargs  # type: Dict[str, Any]
+                     ) -> Awaitable[None]:
+        return super()._new_attempt_async(**kwargs)
+
+    @AsyncWrapper.inject_callbacks(None)
+    def _rollback(self,
+                  **kwargs  # type: Dict[str, Any]
+                  ) -> Awaitable[None]:
+        return super()._rollback_async(**kwargs)
+
+    @AsyncWrapper.inject_callbacks(TransactionResult)
+    def _commit(self,
+                **kwargs  # type: Dict[str, Any]
+                ) -> Awaitable[TransactionResult]:
+        return super()._commit_async(**kwargs)
 
     @AsyncWrapper.inject_callbacks(TransactionGetResult)
     def _get(self,
diff --git a/couchbase/transactions/logic/attempt_context_logic.py b/couchbase/transactions/logic/attempt_context_logic.py
@@ -14,11 +14,21 @@
 #  limitations under the License.
 
 import logging
-from typing import TYPE_CHECKING, Optional
+from typing import (TYPE_CHECKING,
+                    Any,
+                    Dict,
+                    Optional)
 
-from couchbase.pycbc_core import (transaction_op,
+from couchbase.exceptions import ErrorMapper
+from couchbase.exceptions import exception as CouchbaseBaseException
+from couchbase.options import TransactionOptions
+from couchbase.pycbc_core import (create_new_attempt_context,
+                                  create_transaction_context,
+                                  transaction_commit,
+                                  transaction_op,
                                   transaction_operations,
-                                  transaction_query_op)
+                                  transaction_query_op,
+                                  transaction_rollback)
 
 if TYPE_CHECKING:
     from asyncio import AbstractEventLoop
@@ -31,23 +41,61 @@
 
 class AttemptContextLogic:
     def __init__(self,
-                 ctx,    # type: PyCapsuleType
-                 loop,    # type: Optional[AbstractEventLoop]
+                 txns,    # type: PyCapsuleType
                  transcoder,  # type: Transcoder
-                 ):
-        log.debug('creating new attempt context with context=%s, loop=%s, and transcoder=%s', ctx, loop, transcoder)
-        self._ctx = ctx
+                 loop,    # type: Optional[AbstractEventLoop]
+                 opts    # type: Optional[PyCapsuleType]
+                 ) -> None:
+        if opts is None:
+            opts = TransactionOptions()._base
+        self._txnctx = create_transaction_context(txns=txns, transaction_options=opts)
         self._loop = loop
         self._transcoder = transcoder
 
+    def _handle_exception(self, ex: Any) -> None:
+        if isinstance(ex, Exception):
+            raise ex
+        if isinstance(ex, CouchbaseBaseException):
+            raise ErrorMapper.build_exception(ex)
+
+    def _new_attempt(self) -> None:
+        new_attempt_ctx = create_new_attempt_context(ctx=self._txnctx)
+        self._handle_exception(new_attempt_ctx)
+
+    def _new_attempt_async(self,
+                           **kwargs,  # type: Dict[str, Any]
+                           ) -> None:
+        create_new_attempt_context(ctx=self._txnctx, **kwargs)
+
+    def _rollback(self) -> None:
+        rollback_res = transaction_rollback(ctx=self._txnctx)
+        self._handle_exception(rollback_res)
+
+    def _rollback_async(self,
+                        **kwargs,  # type: Dict[str, Any]
+                        ) -> None:
+        transaction_rollback(ctx=self._txnctx, **kwargs)
+
+    def _commit(self) -> Optional[Dict[str, Any]]:
+        commit_res = transaction_commit(ctx=self._txnctx)
+        self._handle_exception(commit_res)
+        return commit_res
+
+    def _commit_async(self,
+                      **kwargs,  # type: Dict[str, Any]
+                      ) -> Optional[Dict[str, Any]]:
+        return transaction_commit(ctx=self._txnctx, **kwargs)
+
     def get(self, coll, key, **kwargs):
         # make sure we don't pass the transcoder along
         kwargs.pop('transcoder', None)
         kwargs.update(coll._get_connection_args())
         kwargs.pop("conn")
-        kwargs["key"] = key
-        kwargs["ctx"] = self._ctx
-        kwargs["op"] = transaction_operations.GET.value
+        kwargs.update({
+            'key': key,
+            'ctx': self._txnctx,
+            'op': transaction_operations.GET.value
+        })
         log.debug('get calling transaction op with %s', kwargs)
         return transaction_op(**kwargs)
 
@@ -57,7 +105,7 @@ def insert(self, coll, key, value, **kwargs):
         kwargs.pop("conn")
         kwargs.update({
             'key': key,
-            'ctx': self._ctx,
+            'ctx': self._txnctx,
             'op': transaction_operations.INSERT.value,
             'value': transcoder.encode_value(value)
         })
@@ -67,7 +115,7 @@ def insert(self, coll, key, value, **kwargs):
     def replace(self, txn_get_result, value, **kwargs):
         transcoder = kwargs.pop('transcoder', self._transcoder)
         kwargs.update({
-            'ctx': self._ctx,
+            'ctx': self._txnctx,
             'op': transaction_operations.REPLACE.value,
             'value': transcoder.encode_value(value),
             'txn_get_result': txn_get_result._res
@@ -76,13 +124,13 @@ def replace(self, txn_get_result, value, **kwargs):
         return transaction_op(**kwargs)
 
     def remove(self, txn_get_result, **kwargs):
-        kwargs.update({"ctx": self._ctx,
-                       "op": transaction_operations.REMOVE.value,
-                       "txn_get_result": txn_get_result._res})
+        kwargs.update({'ctx': self._txnctx,
+                       'op': transaction_operations.REMOVE.value,
+                       'txn_get_result': txn_get_result._res})
         log.debug('remove calling transaction op with %s', kwargs)
         return transaction_op(**kwargs)
 
     def query(self, query, options, **kwargs):
-        kwargs.update({"ctx": self._ctx, "statement": query, "options": options._base})
+        kwargs.update({'ctx': self._txnctx, 'statement': query, 'options': options._base})
         log.debug('query calling transaction_op with %s', kwargs)
         return transaction_query_op(**kwargs)
diff --git a/couchbase/transactions/logic/transactions_logic.py b/couchbase/transactions/logic/transactions_logic.py
@@ -17,19 +17,18 @@
 from typing import (TYPE_CHECKING,
                     Any,
                     Callable,
-                    Dict,
-                    Optional)
+                    Coroutine)
 
-from couchbase.logic.supportability import Supportability
-from couchbase.pycbc_core import (create_transactions,
-                                  destroy_transactions,
-                                  run_transaction)
+from couchbase.exceptions import TransactionExpired, TransactionFailed
+from couchbase.pycbc_core import create_transactions, destroy_transactions
+from couchbase.transactions.transaction_result import TransactionResult
 from couchbase.transcoder import JSONTranscoder
 
 if TYPE_CHECKING:
+    from acouchbase.transactions import AttemptContext as AsyncAttemptContext
     from couchbase.logic.cluster import ClusterLogic
-    from couchbase.options import TransactionConfig, TransactionOptions
-    from couchbase.transactions.logic.attempt_context_logic import AttemptContextLogic
+    from couchbase.options import TransactionConfig
+    from couchbase.transactions import AttemptContext as BlockingAttemptContext
 
 log = logging.getLogger(__name__)
 
@@ -49,21 +48,54 @@ def __init__(self,
         log.info('created transactions object using config=%s, transcoder=%s', self._config, self._transcoder)
 
     def run(self,
-            logic,                     # type: Callable[[AttemptContextLogic], None]
-            transaction_options=None,  # type: Optional[TransactionOptions]
-            **kwargs                   # type: Optional[Dict[str, Any]]
-            ):
-        if transaction_options:
-            kwargs['transaction_options'] = transaction_options._base
-        if 'per_txn_config' in kwargs:
-            Supportability.method_param_deprecated('per_txn_config', 'transaction_options')
-            kwargs['transaction_options'] = kwargs.pop('per_txn_config', None)
+            logic,          # type: Callable[[BlockingAttemptContext], None]
+            attempt_ctx     # type: BlockingAttemptContext
+            ) -> TransactionResult:
 
-        try:
-            return run_transaction(txns=self._txns, logic=logic, **kwargs)
-        except Exception as e:
-            log.debug('txn_logic.run() got %s:%s, re-raising it', e.__class__.__name__, e)
-            raise e
+        while True:
+            attempt_ctx._new_attempt()
+            try:
+                logic(attempt_ctx)
+            except Exception as ex:
+                attempt_ctx._rollback()
+                if isinstance(ex, TransactionExpired):
+                    raise ex from None
+                raise TransactionFailed(exc_info={'inner_cause': ex}) from None
+
+            try:
+                # calls finalize internally
+                res = attempt_ctx._commit()
+                if not res:
+                    continue
+                return TransactionResult(**res)
+            except Exception:
+                # commit failed, retrying...
+                pass  # nosec
+
+    async def run_async(self,
+                        logic,          # type: Callable[[AsyncAttemptContext], Coroutine[Any, Any, None]]
+                        attempt_ctx     # type: AsyncAttemptContext
+                        ) -> TransactionResult:
+
+        while True:
+            await attempt_ctx._new_attempt()
+            try:
+                await logic(attempt_ctx)
+            except Exception as ex:
+                await attempt_ctx._rollback()
+                if isinstance(ex, TransactionExpired):
+                    raise ex from None
+                raise TransactionFailed(exc_info={'inner_cause': ex}) from None
+
+            try:
+                # calls finalize internally
+                res = await attempt_ctx._commit()
+                if not res:
+                    continue
+                return res
+            except Exception:
+                # commit failed, retrying...
+                pass  # nosec
 
     def close(self, **kwargs):
         log.info('shutting down transactions...')
diff --git a/couchbase/transactions/transactions.py b/couchbase/transactions/transactions.py
@@ -116,22 +116,14 @@ def txn_logic(ctx):
 
         """  # noqa: E501
 
-        def wrapped_txn_logic(c):
-            try:
-                ctx = AttemptContext(c, self._transcoder)
-                return txn_logic(ctx)
-            except Exception as e:
-                log.debug('wrapped_txn_logic got %s:%s, re-raising it', e.__class__.__name__, e)
-                raise e
-
         opts = None
         if transaction_options:
-            opts = transaction_options
+            opts = transaction_options._base
         if 'per_txn_config' in kwargs:
             Supportability.method_param_deprecated('per_txn_config', 'transaction_options')
             opts = kwargs.pop('per_txn_config', None)
 
-        return TransactionResult(**super().run(wrapped_txn_logic, opts))
+        return super().run(txn_logic, AttemptContext(self._txns, self._transcoder, opts))
 
     def close(self):
         super().close()
@@ -141,10 +133,11 @@ def close(self):
 class AttemptContext(AttemptContextLogic):
 
     def __init__(self,
-                 ctx,  # type: PyCapsuleType
-                 transcoder  # type: Transcoder
+                 txns,    # type: PyCapsuleType
+                 transcoder,  # type: Transcoder
+                 opts    # type: Optional[PyCapsuleType]
                  ):
-        super().__init__(ctx, None, transcoder)
+        super().__init__(txns, transcoder, None, opts)
 
     @BlockingWrapper.block(TransactionGetResult)
     def _get(self,
diff --git a/src/client.cxx b/src/client.cxx
@@ -401,10 +401,14 @@ static struct PyMethodDef methods[] = {
     (PyCFunction)pycbc_txns::create_transactions,
     METH_VARARGS | METH_KEYWORDS,
     "Create a transactions object" },
-  { "run_transaction",
-    (PyCFunction)pycbc_txns::run_transactions,
+  { "create_transaction_context",
+    (PyCFunction)pycbc_txns::create_transaction_context,
     METH_VARARGS | METH_KEYWORDS,
-    "Run a transaction" },
+    "Create a transaction context object" },
+  { "create_new_attempt_context",
+    (PyCFunction)pycbc_txns::create_new_attempt_context,
+    METH_VARARGS | METH_KEYWORDS,
+    "Create a new attempt context object" },
   { "transaction_op",
     (PyCFunction)pycbc_txns::transaction_op,
     METH_VARARGS | METH_KEYWORDS,
@@ -413,6 +417,14 @@ static struct PyMethodDef methods[] = {
     (PyCFunction)pycbc_txns::transaction_query_op,
     METH_VARARGS | METH_KEYWORDS,
     "perform a transactional query" },
+  { "transaction_commit",
+    (PyCFunction)pycbc_txns::transaction_commit,
+    METH_VARARGS | METH_KEYWORDS,
+    "Commit a transaction" },
+  { "transaction_rollback",
+    (PyCFunction)pycbc_txns::transaction_rollback,
+    METH_VARARGS | METH_KEYWORDS,
+    "Rollback a transaction" },
   { "destroy_transactions",
     (PyCFunction)pycbc_txns::destroy_transactions,
     METH_VARARGS | METH_KEYWORDS,
diff --git a/src/transactions/transactions.cxx b/src/transactions/transactions.cxx
diff --git a/src/transactions/transactions.hxx b/src/transactions/transactions.hxx
