uow: import from invenio_db

invenio_records_resources/services/uow.py

@@ -8,6 +8,8 @@
 
 """Unit of work.
 
+Classes were moved to invenio-db.
+
 Used to group multiple service operations into a single atomic unit. The Unit
 of Work maintains a list of operations and coordinates the commit, indexing and
 task execution.
@@ -101,68 +103,25 @@ def on_commit(self, uow):
             # ... executed after the database transaction commit ...
 """
 
-
-from functools import wraps
-
 from celery import current_app
-from invenio_db import db
+
+# backwards compatible imports
+from invenio_db.uow import (
+    ModelCommitOp,
+    ModelDeleteOp,
+    Operation,
+    UnitOfWork,
+    unit_of_work,
+)
 
 from ..tasks import send_change_notifications
 
+__all__ = ["ModelCommitOp", "ModelDeleteOp", "Operation", "UnitOfWork", "unit_of_work"]
+
 
 #
 # Unit of work operations
 #
-class Operation:
-    """Base class for unit of work operations."""
-
-    def on_register(self, uow):
-        """Called upon operation registration."""
-        pass
-
-    def on_commit(self, uow):
-        """Called in the commit phase (after the transaction is committed)."""
-        pass
-
-    def on_post_commit(self, uow):
-        """Called right after the commit phase."""
-        pass
-
-    def on_rollback(self, uow):
-        """Called in the rollback phase (after the transaction rollback)."""
-        pass
-
-    def on_post_rollback(self, uow):
-        """Called right after the rollback phase."""
-        pass
-
-
-class ModelCommitOp(Operation):
-    """SQLAlchemy model add/update operation."""
-
-    def __init__(self, model):
-        """Initialize the commit operation."""
-        super().__init__()
-        self._model = model
-
-    def on_register(self, uow):
-        """Add model to db session."""
-        uow.session.add(self._model)
-
-
-class ModelDeleteOp(Operation):
-    """SQLAlchemy model delete operation."""
-
-    def __init__(self, model):
-        """Initialize the set delete operation."""
-        super().__init__()
-        self._model = model
-
-    def on_register(self, uow):
-        """Delete model."""
-        uow.session.delete(self._model)
-
-
 class RecordCommitOp(Operation):
     """Record commit operation with indexing."""
 
@@ -325,104 +284,3 @@ def on_post_commit(self, uow):
             self._record_type,
             [(r.pid.pid_value, str(r.id), r.revision_id) for r in self._records],
         )
-
-
-#
-# Unit of work context manager
-#
-class UnitOfWork:
-    """Unit of work context manager.
-
-    Note, the unit of work does not currently take duplication of work into
-    account. Thus, you can e.g. add two record commit operations of the same
-    record which will then index the record twice, even though only one time
-    is needed.
-    """
-
-    def __init__(self, session=None):
-        """Initialize unit of work context."""
-        self._session = session or db.session
-        self._operations = []
-        self._dirty = False
-
-    def __enter__(self):
-        """Entering the context."""
-        return self
-
-    def __exit__(self, exc_type, *args):
-        """Rollback on exception."""
-        if exc_type is not None:
-            self.rollback()
-            self._mark_dirty()
-
-    @property
-    def session(self):
-        """The SQLAlchemy database session associated with this UoW."""
-        return self._session
-
-    def _mark_dirty(self):
-        """Mark the unit of work as dirty."""
-        if self._dirty:
-            raise RuntimeError("The unit of work is already committed or rolledback.")
-        self._dirty = True
-
-    def commit(self):
-        """Commit the unit of work."""
-        self.session.commit()
-        # Run commit operations
-        for op in self._operations:
-            op.on_commit(self)
-        # Run post commit operations
-        for op in self._operations:
-            op.on_post_commit(self)
-        self._mark_dirty()
-
-    def rollback(self):
-        """Rollback the database session."""
-        self.session.rollback()
-        # Run rollback operations
-        for op in self._operations:
-            op.on_rollback(self)
-        # Run post rollback operations
-        for op in self._operations:
-            op.on_post_rollback(self)
-
-    def register(self, op):
-        """Register an operation."""
-        # Run on register
-        op.on_register(self)
-        # Append to list of operations.
-        self._operations.append(op)
-
-
-def unit_of_work(**kwargs):
-    """Decorator to auto-inject a unit of work if not provided.
-
-    If no unit of work is provided, this decorator will create a new unit of
-    work and commit it after the function has been executed.
-
-    .. code-block:: python
-
-        @unit_of_work()
-        def aservice_method(self, ...., uow=None):
-            # ...
-            uow.register(...)
-
-    """
-
-    def decorator(f):
-        @wraps(f)
-        def inner(self, *args, **kwargs):
-            if "uow" not in kwargs:
-                # Migration path - start a UoW and commit
-                with UnitOfWork(db.session) as uow:
-                    kwargs["uow"] = uow
-                    res = f(self, *args, **kwargs)
-                    uow.commit()
-                    return res
-            else:
-                return f(self, *args, **kwargs)
-
-        return inner
-
-    return decorator

setup.cfg

@@ -32,6 +32,7 @@ install_requires =
     flask-resources>=1.0.0,<2.0.0
     invenio-accounts>=5.0.0,<6.0.0
     invenio-base>=1.3.0,<2.0.0
+    invenio-db>=1.2.0,<2.0.0
     invenio-files-rest>=2.0.0,<3.0.0
     invenio-i18n>=2.0.0,<3.0.0
     invenio-indexer>=2.1.0,<3.0.0