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/environment/services/environment_service.py

@@ -19,6 +19,7 @@
     CREDENTIALS_ENVIRONMENT,
 )
 from dataall.core.permissions.services.resource_policy_service import ResourcePolicyService
+from dataall.core.resource_lock.db.resource_lock_repositories import ResourceLockRepository
 from dataall.core.permissions.services.tenant_policy_service import TenantPolicyService
 from dataall.core.activity.db.activity_models import Activity
 from dataall.core.environment.db.environment_models import EnvironmentParameter, ConsumptionRole
@@ -295,6 +296,12 @@ def create_environment(uri, data=None):
             )
             session.commit()
 
+            ResourceLockRepository.create_resource_lock(
+                session=session,
+                resource_uri=f'{env_group.groupUri}-{env_group.environmentUri}',
+                resource_type=env_group.__tablename__,
+            )
+
             activity = Activity(
                 action='ENVIRONMENT:CREATE',
                 label='ENVIRONMENT:CREATE',
@@ -413,6 +420,13 @@ def invite_group(uri, data=None) -> (Environment, EnvironmentGroup):
             )
             session.add(environment_group)
             session.commit()
+
+            ResourceLockRepository.create_resource_lock(
+                session=session,
+                resource_uri=f'{environment_group.groupUri}-{environment_group.environmentUri}',
+                resource_type=environment_group.__tablename__,
+            )
+
             ResourcePolicyService.attach_resource_policy(
                 session=session,
                 group=group,
@@ -466,6 +480,9 @@ def remove_group(uri, group):
             ).delete_all_policies()
 
             if group_membership:
+                ResourceLockRepository.delete_resource_lock(
+                    session=session, resource_uri=f'{group}-{environment.environmentUri}'
+                )
                 session.delete(group_membership)
                 session.commit()
 
@@ -592,6 +609,12 @@ def add_consumption_role(uri, data=None) -> (Environment, EnvironmentGroup):
             session.add(consumption_role)
             session.commit()
 
+            ResourceLockRepository.create_resource_lock(
+                session=session,
+                resource_uri=consumption_role.consumptionRoleUri,
+                resource_type=consumption_role.__tablename__,
+            )
+
             ResourcePolicyService.attach_resource_policy(
                 session=session,
                 group=group,
@@ -631,6 +654,7 @@ def remove_consumption_role(uri, env_uri):
                     resource_uri=consumption_role.consumptionRoleUri,
                     resource_type=ConsumptionRole.__name__,
                 )
+                ResourceLockRepository.delete_resource_lock(session=session, resource_uri=uri)
 
                 session.delete(consumption_role)
                 session.commit()
@@ -881,6 +905,9 @@ def delete_environment(uri):
                 EnvironmentParameterRepository(session).delete_params(environment.environmentUri)
 
                 for group in env_groups:
+                    ResourceLockRepository.delete_resource_lock(
+                        session=session, resource_uri=f'{group.groupUri}-{environment.environmentUri}'
+                    )
                     session.delete(group)
 
                     ResourcePolicyService.delete_resource_policy(
@@ -890,6 +917,7 @@ def delete_environment(uri):
                     )
 
                 for role in env_roles:
+                    ResourceLockRepository.delete_resource_lock(session=session, resource_uri=role.consumptionRoleUri)
                     session.delete(role)
 
                 return session.delete(environment), environment

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

@@ -0,0 +1,28 @@
+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)
+    isLocked = Column(Boolean, default=False)
+    acquiredByUri = Column(String, nullable=True)
+    acquiredByType = Column(String, nullable=True)
+
+    def __init__(
+        self,
+        resourceUri: str,
+        resourceType: str,
+        isLocked: bool = False,
+        acquiredByUri: Optional[str] = None,
+        acquiredByType: Optional[str] = None,
+    ):
+        self.resourceUri = resourceUri
+        self.resourceType = resourceType
+        self.isLocked = isLocked
+        self.acquiredByUri = acquiredByUri
+        self.acquiredByType = acquiredByType

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

@@ -0,0 +1,122 @@
+import logging
+
+from dataall.core.resource_lock.db.resource_lock_models import ResourceLock
+from sqlalchemy import and_, or_
+
+log = logging.getLogger(__name__)
+
+
+class ResourceLockRepository:
+    @staticmethod
+    def create_resource_lock(
+        session, resource_uri, resource_type, is_locked=False, acquired_by_uri=None, acquired_by_type=None
+    ):
+        resource_lock = ResourceLock(
+            resourceUri=resource_uri,
+            resourceType=resource_type,
+            isLocked=is_locked,
+            acquiredByUri=acquired_by_uri,
+            acquiredByType=acquired_by_type,
+        )
+        session.add(resource_lock)
+        session.commit()
+
+    @staticmethod
+    def delete_resource_lock(session, resource_uri):
+        resource_lock = session.query(ResourceLock).filter(ResourceLock.resourceUri == resource_uri).first()
+        session.delete(resource_lock)
+        session.commit()
+
+    @staticmethod
+    def acquire_locks(resources, session, acquired_by_uri, acquired_by_type):
+        """
+        Attempts to acquire 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],
+                    ~ResourceLock.isLocked,
+                )
+                for resource in resources
+            ]
+            resource_locks = session.query(ResourceLock).filter(or_(*filter_conditions)).with_for_update().all()
+
+            # Ensure lock record found for each resource
+            if len(resource_locks) == len(resources):
+                # Update the attributes of the ResourceLock object
+                for resource_lock in resource_locks:
+                    resource_lock.isLocked = True
+                    resource_lock.acquiredByUri = acquired_by_uri
+                    resource_lock.acquiredByType = acquired_by_type
+                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 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.isLocked,
+                        ResourceLock.acquiredByUri == share_uri,
+                    )
+                )
+                .with_for_update()
+                .first()
+            )
+
+            if resource_lock:
+                resource_lock.isLocked = False
+                resource_lock.acquiredByUri = ''
+                resource_lock.acquiredByType = ''
+
+                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

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,9 @@ 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)
-
+            ResourceLockRepository.create_resource_lock(
+                session=session, resource_uri=dataset.datasetUri, resource_type=dataset.__tablename__
+            )
             DatasetBucketRepository.create_dataset_bucket(session, dataset, data)
 
             ResourcePolicyService.attach_resource_policy(
@@ -411,7 +413,7 @@ 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)
+            ResourceLockRepository.delete_resource_lock(session=session, resource_uri=dataset.datasetUri)
             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)