Fix name of shared_db in UI, added missing test, fix backwards delete…

… database issue
data-dot-all · Jan 30, 2024 · 5fdb612 · 5fdb612
1 parent d812994
commit 5fdb612
Show file tree

Hide file tree

Showing 4 changed files with 77 additions and 39 deletions.
diff --git a/backend/dataall/modules/dataset_sharing/api/resolvers.py b/backend/dataall/modules/dataset_sharing/api/resolvers.py
@@ -181,18 +181,12 @@ def resolve_group(context: Context, source: ShareObject, **kwargs):
 def resolve_consumption_data(context: Context, source: ShareObject, **kwargs):
     if not source:
         return None
-    with context.engine.scoped_session() as session:
-        ds: Dataset = DatasetRepository.get_dataset_by_uri(session, source.datasetUri)
-        if ds:
-            S3AccessPointName = utils.slugify(
-                source.datasetUri + '-' + source.principalId,
-                max_length=50, lowercase=True, regex_pattern='[^a-zA-Z0-9-]', separator='-'
-            )
-            return {
-                's3AccessPointName': S3AccessPointName,
-                'sharedGlueDatabase': (ds.GlueDatabaseName + '_shared_' + source.shareUri)[:254] if ds else 'Not created',
-                's3bucketName': ds.S3BucketName,
-            }
+    return ShareObjectService.resolve_share_object_consumption_data(
+        uri=source.shareUri,
+        datasetUri=source.datasetUri,
+        principalId=source.principalId,
+        environmentUri=source.environmentUri
+    )
 
 
 def resolve_share_object_statistics(context: Context, source: ShareObject, **kwargs):

diff --git a/backend/dataall/modules/dataset_sharing/services/share_object_service.py b/backend/dataall/modules/dataset_sharing/services/share_object_service.py
@@ -18,6 +18,7 @@
 from dataall.modules.dataset_sharing.services.share_permissions import REJECT_SHARE_OBJECT, APPROVE_SHARE_OBJECT, \
     SUBMIT_SHARE_OBJECT, SHARE_OBJECT_APPROVER, SHARE_OBJECT_REQUESTER, CREATE_SHARE_OBJECT, DELETE_SHARE_OBJECT, \
     GET_SHARE_OBJECT
+from dataall.modules.dataset_sharing.aws.glue_client import GlueClient
 from dataall.modules.datasets_base.db.dataset_repositories import DatasetRepository
 from dataall.modules.datasets_base.db.dataset_models import DatasetTable, Dataset
 from dataall.modules.datasets_base.services.permissions import DATASET_TABLE_READ
@@ -353,6 +354,34 @@ def resolve_share_object_statistics(uri):
             return {'tables': tables, 'locations': locations, 'sharedItems': shared_items, 'revokedItems': revoked_items,
                     'failedItems': failed_items, 'pendingItems': pending_items}
 
+    @staticmethod
+    def resolve_share_object_consumption_data(uri, datasetUri, principalId, environmentUri):
+        with get_context().db_engine.scoped_session() as session:
+            dataset = DatasetRepository.get_dataset_by_uri(session, datasetUri)
+            environment = EnvironmentService.get_environment_by_uri(session, environmentUri)
+            if dataset:
+                S3AccessPointName = utils.slugify(
+                    datasetUri + '-' + principalId,
+                    max_length=50, lowercase=True, regex_pattern='[^a-zA-Z0-9-]', separator='-'
+                )
+                old_shared_db_name = f"{dataset.GlueDatabaseName}_shared_{uri}"[:254]
+                database = GlueClient(
+                    account_id=environment.AwsAccountId,
+                    region=environment.region,
+                    database=old_shared_db_name
+                ).get_glue_database()
+                sharedGlueDatabase = old_shared_db_name if database else f"{dataset.GlueDatabaseName[:247]}_shared"
+                return {
+                    's3AccessPointName': S3AccessPointName,
+                    'sharedGlueDatabase': sharedGlueDatabase,
+                    's3bucketName': dataset.S3BucketName,
+                }
+            return {
+                's3AccessPointName': "Not Created",
+                'sharedGlueDatabase': "Not Created",
+                's3bucketName': dataset.S3BucketName,
+            }
+
     @staticmethod
     def list_shares_in_my_inbox(filter: dict):
         context = get_context()

diff --git a/.../dataall/modules/dataset_sharing/services/share_processors/lakeformation_process_share.py b/.../dataall/modules/dataset_sharing/services/share_processors/lakeformation_process_share.py
@@ -202,30 +202,34 @@ def process_revoked_shares(self) -> bool:
                 self.handle_revoke_failure(table=table, error=e)
 
         try:
-            existing_shared_tables = ShareObjectRepository.check_existing_shared_items_of_type(
-                session=self.session,
-                uri=self.share.shareUri,
-                item_type=ShareableType.Table.value
-            )
-            log.info(
-                f'Still remaining tables shared in this share object = {existing_shared_tables}')
-
-            if not existing_shared_tables and self.revoked_tables:
-                log.info("Revoking permissions to target shared database...")
-                self.revoke_principals_database_permissions_to_shared_database()
-                # TODO Add test for this new method
-
-            existing_shared_tables_environment = ShareObjectRepository.list_dataset_shares_with_existing_shared_items(
-                session=self.session,
-                dataset_uri=self.dataset.datasetUri,
-                environment_uri=self.target_environment.environmentUri,
-                item_type=ShareableType.Table.value
-            )
+            if self.revoked_tables:
+                existing_shared_tables_in_share = ShareObjectRepository.check_existing_shared_items_of_type(
+                    session=self.session,
+                    uri=self.share.shareUri,
+                    item_type=ShareableType.Table.value
+                )
+                log.info(
+                    f'Remaining tables shared in this share object = {existing_shared_tables_in_share}')
+
+                if not existing_shared_tables_in_share:
+                    log.info("Revoking permissions to target shared database...")
+                    self.revoke_principals_database_permissions_to_shared_database()
+
+                    if not self.is_new_share:
+                        log.info("Deleting OLD target shared database...")
+                        self.delete_shared_database_in_target()
+
+                existing_shared_tables_in_environment = ShareObjectRepository.list_dataset_shares_with_existing_shared_items(
+                    session=self.session,
+                    dataset_uri=self.dataset.datasetUri,
+                    environment_uri=self.target_environment.environmentUri,
+                    item_type=ShareableType.Table.value
+                )
 
-            log.info(f'Still remaining tables shared from this dataset to this environment = {existing_shared_tables_environment}')
-            if not existing_shared_tables_environment and self.revoked_tables:
-                log.info("Deleting target shared database...")
-                self.delete_shared_database_in_target()
+                log.info(f'Remaining tables shared from this dataset to this environment = {existing_shared_tables_in_environment}')
+                if self.is_new_share and not existing_shared_tables_in_environment:
+                    log.info("Deleting target shared database...")
+                    self.delete_shared_database_in_target()
         except Exception as e:
             log.error(
                 f'Failed to clean-up database permission or delete shared database {self.shared_db_name} '

diff --git a/tests/modules/datasets/tasks/test_lf_share_manager.py b/tests/modules/datasets/tasks/test_lf_share_manager.py
@@ -345,7 +345,7 @@ def test_grant_pivot_role_all_database_permissions_to_shared_database(
     )
 
 
-def test_grant_principals_permissions_to_shared_database(
+def test_grant_principals_database_permissions_to_shared_database(
         processor_with_mocks,
         dataset1: Dataset
 ):
@@ -475,7 +475,6 @@ def test_revoke_principals_permissions_to_resource_link_table(
         )
 
 
-
 def test_revoke_principals_permissions_to_table_in_target(
         processor_with_mocks,
         table1: DatasetTable,
@@ -494,8 +493,6 @@ def test_revoke_principals_permissions_to_table_in_target(
             permissions=['DESCRIBE', 'SELECT']
         )
 
-
-
 def test_delete_resource_link_table_in_shared_database_true(
         processor_with_mocks,
         table2: DatasetTable
@@ -510,6 +507,20 @@ def test_delete_resource_link_table_in_shared_database_true(
     glue_client.table_exists.assert_called_once()
     glue_client.delete_table.assert_called_once()
 
+def test_revoke_principals_database_permissions_to_shared_database(
+        processor_with_mocks,
+        dataset1: Dataset
+):
+    processor, lf_client, glue_client = processor_with_mocks
+    # When
+    processor.revoke_principals_database_permissions_to_shared_database()
+    # Then
+    lf_client.revoke_permissions_to_database.assert_called_once()
+    lf_client.revoke_permissions_to_database.assert_called_with(
+        principals=processor.principals,
+        database_name=f"{dataset1.GlueDatabaseName}_shared",
+        permissions=['DESCRIBE'],
+    )
 
 def test_delete_shared_database_in_target(
     processor_with_mocks,