Convert Dataset Lock Mechanism to Generic Resource Lock

backend/dataall/base/db/exceptions.py

@@ -169,3 +169,15 @@ def __init__(self, action, message):
 
     def __str__(self):
         return f'{self.message}'
+
+
+class ResourceLockTimeout(Exception):
+    def __init__(self, action, message):
+        self.action = action
+        self.message = f"""
+                    An error occurred (ResourceLockTimeout) when calling {self.action} operation:
+                    {message}
+                """
+
+    def __str__(self):
+        return f'{self.message}'

backend/dataall/core/__init__.py

@@ -1,11 +1,3 @@
 """The package contains the core functionality that is required by data.all to work correctly"""
 
-from dataall.core import (
-    permissions,
-    stacks,
-    groups,
-    environment,
-    organizations,
-    tasks,
-    vpc,
-)
+from dataall.core import permissions, stacks, groups, environment, organizations, tasks, vpc, resource_lock

backend/dataall/core/resource_lock/db/resource_lock_models.py

@@ -0,0 +1,25 @@
+from typing import Optional
+from sqlalchemy import Column, String, Boolean
+
+from dataall.base.db import Base
+
+
+class ResourceLock(Base):
+    __tablename__ = 'resource_lock'
+
+    resourceUri = Column(String, nullable=False, primary_key=True)
+    resourceType = Column(String, nullable=False, primary_key=True)
+    acquiredByUri = Column(String, nullable=True)
+    acquiredByType = Column(String, nullable=True)
+
+    def __init__(
+        self,
+        resourceUri: str,
+        resourceType: str,
+        acquiredByUri: Optional[str] = None,
+        acquiredByType: Optional[str] = None,
+    ):
+        self.resourceUri = resourceUri
+        self.resourceType = resourceType
+        self.acquiredByUri = acquiredByUri
+        self.acquiredByType = acquiredByType

backend/dataall/core/resource_lock/db/resource_lock_repositories.py

@@ -0,0 +1,135 @@
+import logging
+
+from dataall.core.resource_lock.db.resource_lock_models import ResourceLock
+from sqlalchemy import and_, or_
+from sqlalchemy.orm import Session
+from time import sleep
+from typing import List, Tuple
+from contextlib import contextmanager
+from dataall.base.db.exceptions import ResourceLockTimeout
+
+log = logging.getLogger(__name__)
+
+MAX_RETRIES = 10
+RETRY_INTERVAL = 60
+
+
+class ResourceLockRepository:
+    @staticmethod
+    def acquire_locks(resources, session, acquired_by_uri, acquired_by_type):
+        """
+        Attempts to acquire/create one or more locks on the resources identified by resourceUri and resourceType.
+
+        Args:
+            resources: List of resource tuples (resourceUri, resourceType) to acquire locks for.
+            session (sqlalchemy.orm.Session): The SQLAlchemy session object used for interacting with the database.
+            acquired_by_uri: The ID of the resource that is attempting to acquire the lock.
+            acquired_by_type: The resource type that is attempting to acquire the lock.qu
+
+        Returns:
+            bool: True if the lock is successfully acquired, False otherwise.
+        """
+        try:
+            # Execute the query to get the ResourceLock object
+            filter_conditions = [
+                and_(
+                    ResourceLock.resourceUri == resource[0],
+                    ResourceLock.resourceType == resource[1],
+                )
+                for resource in resources
+            ]
+
+            if not session.query(ResourceLock).filter(or_(*filter_conditions)).first():
+                records = []
+                for resource in resources:
+                    records.append(
+                        ResourceLock(
+                            resourceUri=resource[0],
+                            resourceType=resource[1],
+                            acquiredByUri=acquired_by_uri,
+                            acquiredByType=acquired_by_type,
+                        )
+                    )
+                session.add_all(records)
+                session.commit()
+                return True
+            else:
+                log.info(
+                    'Not all ResourceLocks were found. One or more ResourceLocks may be acquired by another resource...'
+                )
+                return False
+        except Exception as e:
+            session.expunge_all()
+            session.rollback()
+            log.error('Error occurred while acquiring lock:', e)
+            return False
+
+    @staticmethod
+    def release_lock(session, resource_uri, resource_type, share_uri):
+        """
+        Releases/delete the lock on the resource identified by resource_uri, resource_type.
+
+        Args:
+            session (sqlalchemy.orm.Session): The SQLAlchemy session object used for interacting with the database.
+            resource_uri: The ID of the resource that owns the lock.
+            resource_type: The type of the resource that owns the lock.
+            share_uri: The ID of the share that is attempting to release the lock.
+
+        Returns:
+            bool: True if the lock is successfully released, False otherwise.
+        """
+        try:
+            log.info(f'Releasing lock for resource: {resource_uri=}, {resource_type=}')
+
+            resource_lock = (
+                session.query(ResourceLock)
+                .filter(
+                    and_(
+                        ResourceLock.resourceUri == resource_uri,
+                        ResourceLock.resourceType == resource_type,
+                        ResourceLock.acquiredByUri == share_uri,
+                    )
+                )
+                .with_for_update()
+                .first()
+            )
+
+            if resource_lock:
+                session.delete(resource_lock)
+                session.commit()
+                return True
+            else:
+                log.info(f'ResourceLock not found for resource: {resource_uri=}, {resource_type=}')
+                return False
+
+        except Exception as e:
+            session.expunge_all()
+            session.rollback()
+            log.error('Error occurred while releasing lock:', e)
+            return False
+
+    @staticmethod
+    @contextmanager
+    def acquire_lock_with_retry(
+        resources: List[Tuple[str, str]], session: Session, acquired_by_uri: str, acquired_by_type: str
+    ):
+        retries_remaining = MAX_RETRIES
+        log.info(f'Attempting to acquire lock for resources {resources} by share {acquired_by_uri}...')
+        while not (
+            lock_acquired := ResourceLockRepository.acquire_locks(resources, session, acquired_by_uri, acquired_by_type)
+        ):
+            log.info(
+                f'Lock for one or more resources {resources} already acquired. Retrying in {RETRY_INTERVAL} seconds...'
+            )
+            sleep(RETRY_INTERVAL)
+            retries_remaining -= 1
+            if retries_remaining <= 0:
+                raise ResourceLockTimeout(
+                    'process shares',
+                    f'Failed to acquire lock for one or more of {resources=}',
+                )
+        try:
+            yield lock_acquired
+        finally:
+            for resource in resources:
+                ResourceLockRepository.release_lock(session, resource[0], resource[1], acquired_by_uri)

backend/dataall/modules/datasets_base/db/dataset_models.py

@@ -40,15 +40,3 @@ def uri(cls):
 
 
 DatasetBase.__name__ = 'Dataset'
-
-
-class DatasetLock(Base):
-    __tablename__ = 'dataset_lock'
-    datasetUri = Column(String, ForeignKey('dataset.datasetUri'), nullable=False, primary_key=True)
-    isLocked = Column(Boolean, default=False)
-    acquiredBy = Column(String, nullable=True)
-
-    def __init__(self, datasetUri, isLocked=False, acquiredBy=None):
-        self.datasetUri = datasetUri
-        self.isLocked = isLocked
-        self.acquiredBy = acquiredBy

backend/dataall/modules/datasets_base/db/dataset_repositories.py

@@ -5,26 +5,14 @@
 from dataall.base.db import paginate
 from dataall.base.db.exceptions import ObjectNotFound
 from dataall.core.activity.db.activity_models import Activity
-from dataall.modules.datasets_base.db.dataset_models import DatasetBase, DatasetLock
+from dataall.modules.datasets_base.db.dataset_models import DatasetBase
 
 logger = logging.getLogger(__name__)
 
 
 class DatasetBaseRepository:
     """DAO layer for GENERIC Datasets"""
 
-    @staticmethod
-    def create_dataset_lock(session, dataset: DatasetBase):
-        dataset_lock = DatasetLock(datasetUri=dataset.datasetUri, isLocked=False, acquiredBy='')
-        session.add(dataset_lock)
-        session.commit()
-
-    @staticmethod
-    def delete_dataset_lock(session, dataset: DatasetBase):
-        dataset_lock = session.query(DatasetLock).filter(DatasetLock.datasetUri == dataset.datasetUri).first()
-        session.delete(dataset_lock)
-        session.commit()
-
     @staticmethod
     def update_dataset_activity(session, dataset: DatasetBase, username):
         activity = Activity(

backend/dataall/modules/s3_datasets/services/dataset_service.py

@@ -2,6 +2,7 @@
 import json
 import logging
 from typing import List
+from dataall.core.resource_lock.db.resource_lock_repositories import ResourceLockRepository
 from dataall.base.aws.quicksight import QuicksightClient
 from dataall.base.db import exceptions
 from dataall.base.utils.naming_convention import NamingConventionPattern
@@ -164,8 +165,6 @@ def create_dataset(uri, admin_group, data: dict):
                 DatasetService.check_imported_resources(dataset)
 
             dataset = DatasetRepository.create_dataset(session=session, env=environment, dataset=dataset, data=data)
-            DatasetBaseRepository.create_dataset_lock(session=session, dataset=dataset)
-
             DatasetBucketRepository.create_dataset_bucket(session, dataset, data)
 
             ResourcePolicyService.attach_resource_policy(
@@ -411,7 +410,6 @@ def delete_dataset(uri: str, delete_from_aws: bool = False):
                 ResourcePolicyService.delete_resource_policy(session=session, resource_uri=uri, group=env.SamlGroupName)
             if dataset.stewards:
                 ResourcePolicyService.delete_resource_policy(session=session, resource_uri=uri, group=dataset.stewards)
-            DatasetBaseRepository.delete_dataset_lock(session=session, dataset=dataset)
             DatasetRepository.delete_dataset(session, dataset)
 
         if delete_from_aws:

backend/dataall/modules/shares_base/services/share_exceptions.py

@@ -19,8 +19,3 @@ def __init__(self, action, message):
 class PrincipalRoleNotFound(BaseShareException):
     def __init__(self, action, message):
         super().__init__('PrincipalRoleNotFound', action, message)
-
-
-class DatasetLockTimeout(BaseShareException):
-    def __init__(self, action, message):
-        super().__init__('DatasetLockTimeout', action, message)