GenAI metadata generation - resolve conflicts+add guardrails

Author: dlpzx

backend/dataall/modules/s3_datasets/aws/s3_dataset_client.py

@@ -81,7 +81,7 @@ def list_bucket_files(self, bucket_name, prefix):
                 Bucket=bucket_name,
                 Prefix=prefix,
                 ExpectedBucketOwner=dataset.AwsAccountId,
-                MaxKeys=1000,
+                MaxKeys=100,
             )
             return response.get('Contents', [])
         except ClientError as e:

backend/dataall/modules/s3_datasets/db/dataset_column_repositories.py

@@ -48,6 +48,8 @@ def paginate_active_columns_for_table(session, table_uri: str, filter: dict):
         return paginate(query=q, page=filter.get('page', 1), page_size=filter.get('pageSize', 10)).to_dict()
 
     @staticmethod
-    def list_active_columns_for_table(session, table_uri: str):
+    def list_active_columns_for_table(session, table_uri: str, limit=None):
         q = DatasetColumnRepository.query_active_columns_for_table(session, table_uri)
+        if limit:
+            q = q.limit(limit)
         return q.all()

backend/dataall/modules/s3_datasets/db/dataset_location_repositories.py

@@ -97,9 +97,12 @@ def delete_dataset_locations(session, dataset_uri) -> bool:
         return True
 
     @staticmethod
-    def get_dataset_folders(session, dataset_uri):
+    def get_dataset_folders(session, dataset_uri, limit=None):
         """return the dataset folders"""
-        return session.query(DatasetStorageLocation).filter(DatasetStorageLocation.datasetUri == dataset_uri).all()
+        query = session.query(DatasetStorageLocation).filter(DatasetStorageLocation.datasetUri == dataset_uri)
+        if limit:
+            query = query.limit(limit)
+        return query.all()
 
     @staticmethod
     def paginated_dataset_locations(session, uri, data=None) -> dict:

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

@@ -186,9 +186,12 @@ def update_glue_database_status(session, dataset_uri):
         dataset.glueDatabaseCreated = True
 
     @staticmethod
-    def get_dataset_tables(session, dataset_uri):
+    def get_dataset_tables(session, dataset_uri, limit=None):
         """return the dataset tables"""
-        return session.query(DatasetTable).filter(DatasetTable.datasetUri == dataset_uri).all()
+        query = session.query(DatasetTable).filter(DatasetTable.datasetUri == dataset_uri)
+        if limit:
+            query = query.limit(limit)
+        return query.all()
 
     @staticmethod
     def delete_dataset(session, dataset) -> bool:

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

@@ -601,8 +601,8 @@ def generate_metadata_for_dataset(uri, metadata_types):
         context = get_context()
         with context.db_engine.scoped_session() as session:
             dataset = DatasetBaseRepository.get_dataset_by_uri(session, uri)
-            tables = DatasetRepository.get_dataset_tables(session, dataset.datasetUri)
-            folders = DatasetLocationRepository.get_dataset_folders(session, dataset.datasetUri)
+            tables = DatasetRepository.get_dataset_tables(session, dataset.datasetUri, limit=50)
+            folders = DatasetLocationRepository.get_dataset_folders(session, dataset.datasetUri, limit=50)
             metadata = BedrockClient().invoke_model_dataset_metadata(
                 metadata_types=metadata_types,
                 dataset=dataset,

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

@@ -210,7 +210,7 @@ def generate_metadata_for_table(uri, metadata_types, sample_data):
         context = get_context()
         with context.db_engine.scoped_session() as session:
             table = DatasetTableRepository.get_dataset_table_by_uri(session, uri)
-            table_columns = DatasetColumnRepository.list_active_columns_for_table(session, table.tableUri)
+            table_columns = DatasetColumnRepository.list_active_columns_for_table(session, table.tableUri, limit=50)
             metadata = BedrockClient().invoke_model_table_metadata(
                 table=table, columns=table_columns, metadata_types=metadata_types, sample_data=sample_data
             )