apply resource_thresholds

Author: petrkalos

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

@@ -5,8 +5,7 @@
 from dataall.core.permissions.services.resource_policy_service import ResourcePolicyService
 from dataall.core.permissions.services.tenant_policy_service import TenantPolicyService
 
-##TODO Uncomment the following to use the ResourceThresholdService once https://github.com/data-dot-all/dataall/pull/1653 is merged
-##from dataall.core.resource_threshold.services.resource_threshold_service import ResourceThresholdService
+from dataall.core.resource_threshold.services.resource_threshold_service import ResourceThresholdService
 from dataall.modules.catalog.db.glossary_repositories import GlossaryRepository
 from dataall.base.db.exceptions import ResourceAlreadyExists
 from dataall.modules.s3_datasets.services.dataset_service import DatasetService
@@ -153,10 +152,9 @@ def get_folder_restricted_information(uri: str, folder: DatasetStorageLocation):
 
     @staticmethod
     @ResourcePolicyService.has_resource_permission(UPDATE_DATASET_FOLDER, parent_resource=_get_dataset_uri)
-    ##TODO Uncomment the following to use the ResourceThresholdService once https://github.com/data-dot-all/dataall/pull/1653 is merged
-    # @ResourceThresholdService.check_invocation_count(
-    #     'metadata', 'modules.s3_datasets.features.generate_metadata_ai.max_count_per_day'
-    # )
+    @ResourceThresholdService.check_invocation_count(
+        'metadata', 'modules.s3_datasets.features.generate_metadata_ai.max_count_per_day'
+    )
     def generate_metadata_for_folder(uri, metadata_types):
         context = get_context()
         with context.db_engine.scoped_session() as session:

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

@@ -20,8 +20,7 @@
 from dataall.core.stacks.db.stack_models import Stack
 from dataall.core.tasks.db.task_models import Task
 
-##TODO
-##from dataall.core.resource_threshold.services.resource_threshold_service import ResourceThresholdService
+from dataall.core.resource_threshold.services.resource_threshold_service import ResourceThresholdService
 from dataall.modules.catalog.db.glossary_repositories import GlossaryRepository
 from dataall.modules.s3_datasets.db.dataset_bucket_repositories import DatasetBucketRepository
 from dataall.modules.shares_base.db.share_object_repositories import ShareObjectRepository
@@ -593,10 +592,9 @@ def list_dataset_tables_folders(uri, filter):
     @staticmethod
     @TenantPolicyService.has_tenant_permission(MANAGE_DATASETS)
     @ResourcePolicyService.has_resource_permission(UPDATE_DATASET)
-    ##TODO Uncomment the following to use the ResourceThresholdService once https://github.com/data-dot-all/dataall/pull/1653 is merged
-    # @ResourceThresholdService.check_invocation_count(
-    #     'metadata', 'modules.s3_datasets.features.generate_metadata_ai.max_count_per_day'
-    # )
+    @ResourceThresholdService.check_invocation_count(
+        'metadata', 'modules.s3_datasets.features.generate_metadata_ai.max_count_per_day'
+    )
     def generate_metadata_for_dataset(uri, metadata_types):
         context = get_context()
         with context.db_engine.scoped_session() as session:

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

@@ -3,8 +3,7 @@
 from dataall.core.permissions.services.resource_policy_service import ResourcePolicyService
 from dataall.core.permissions.services.tenant_policy_service import TenantPolicyService
 
-##TODO
-##from dataall.core.resource_threshold.services.resource_threshold_service import ResourceThresholdService
+from dataall.core.resource_threshold.services.resource_threshold_service import ResourceThresholdService
 from dataall.modules.catalog.db.glossary_repositories import GlossaryRepository
 from dataall.core.environment.services.environment_service import EnvironmentService
 from dataall.modules.s3_datasets.aws.athena_table_client import AthenaTableClient
@@ -203,10 +202,9 @@ def _delete_dataset_table_read_permission(session, table_uri):
 
     @staticmethod
     @ResourcePolicyService.has_resource_permission(UPDATE_DATASET_TABLE, parent_resource=_get_dataset_uri)
-    ##TODO Uncomment the following to use the ResourceThresholdService once https://github.com/data-dot-all/dataall/pull/1653 is merged
-    # @ResourceThresholdService.check_invocation_count(
-    #     'metadata', 'modules.s3_datasets.features.generate_metadata_ai.max_count_per_day'
-    # )
+    @ResourceThresholdService.check_invocation_count(
+        'metadata', 'modules.s3_datasets.features.generate_metadata_ai.max_count_per_day'
+    )
     def generate_metadata_for_table(uri, metadata_types, sample_data):
         metadataTypesForTable = [MetadataGenerationTypes.Description.value, MetadataGenerationTypes.Tag.value]
         table_metadata_types = [item for item in metadata_types if item in metadataTypesForTable]

backend/migrations/versions/2258cd8d6e9f_resource_threshholds_added.py

@@ -12,7 +12,7 @@
 
 # revision identifiers, used by Alembic.
 revision = '2258cd8d6e9f'
-down_revision = '5a798acc6282'
+down_revision = 'ba2da94739ab'
 branch_labels = None
 depends_on = None