GenAI metadata generation - improvements

Author: dlpzx

backend/dataall/modules/s3_datasets/api/dataset/resolvers.py

@@ -175,6 +175,7 @@ def generate_metadata(
     tableSampleData: dict = {},
 ):
     RequestValidator.validate_uri(param_name='resourceUri', param_value=resourceUri)
+    RequestValidator.validate_table_sample_data(tableSampleData)
     if any(metadata_type not in [item.value for item in MetadataGenerationTypes] for metadata_type in metadataTypes):
         raise InvalidInput(
             'metadataType',
@@ -264,3 +265,92 @@ def validate_import_request(data):
         RequestValidator.validate_creation_request(data)
         if not data.get('bucketName'):
             raise RequiredParameter('bucketName')
+
+    @staticmethod
+    def validate_table_sample_data(table_sample_data: dict):
+        """
+        Validates tableSampleData parameter structure to match readTableSampleData API output.
+        Expected structure:
+        {
+            "fields": [JSON string objects with "name" property],
+            "rows": [JSON array strings matching field count]
+        }
+        """
+        if not table_sample_data:
+            return  # Empty dict is allowed as default parameter
+
+        if not isinstance(table_sample_data, dict):
+            raise InvalidInput(
+                param_name='tableSampleData',
+                param_value=str(table_sample_data),
+                constraint='must be a dictionary object',
+            )
+
+        # Validate fields array
+        fields = table_sample_data.get('fields')
+        if fields is not None:
+            if not isinstance(fields, list):
+                raise InvalidInput(
+                    param_name='tableSampleData.fields', param_value=str(fields), constraint='must be an array'
+                )
+
+            for i, field in enumerate(fields):
+                if not isinstance(field, str):
+                    raise InvalidInput(
+                        param_name=f'tableSampleData.fields[{i}]',
+                        param_value=str(field),
+                        constraint='must be a JSON string',
+                    )
+
+                try:
+                    import json
+
+                    field_obj = json.loads(field)
+                    if not isinstance(field_obj, dict) or 'name' not in field_obj:
+                        raise InvalidInput(
+                            param_name=f'tableSampleData.fields[{i}]',
+                            param_value=field,
+                            constraint='must be a JSON object with "name" property',
+                        )
+                except json.JSONDecodeError:
+                    raise InvalidInput(
+                        param_name=f'tableSampleData.fields[{i}]', param_value=field, constraint='must be valid JSON'
+                    )
+
+        # Validate rows array
+        rows = table_sample_data.get('rows')
+        if rows is not None:
+            if not isinstance(rows, list):
+                raise InvalidInput(
+                    param_name='tableSampleData.rows', param_value=str(rows), constraint='must be an array'
+                )
+
+            expected_field_count = len(fields) if fields else 0
+
+            for i, row in enumerate(rows):
+                if not isinstance(row, str):
+                    raise InvalidInput(
+                        param_name=f'tableSampleData.rows[{i}]',
+                        param_value=str(row),
+                        constraint='must be a JSON string',
+                    )
+
+                try:
+                    import json
+
+                    row_array = json.loads(row)
+                    if not isinstance(row_array, list):
+                        raise InvalidInput(
+                            param_name=f'tableSampleData.rows[{i}]', param_value=row, constraint='must be a JSON array'
+                        )
+
+                    if expected_field_count > 0 and len(row_array) != expected_field_count:
+                        raise InvalidInput(
+                            param_name=f'tableSampleData.rows[{i}]',
+                            param_value=row,
+                            constraint=f'must contain {expected_field_count} elements to match fields count',
+                        )
+                except json.JSONDecodeError:
+                    raise InvalidInput(
+                        param_name=f'tableSampleData.rows[{i}]', param_value=row, constraint='must be valid JSON'
+                    )

backend/dataall/modules/s3_datasets/api/table/queries.py

@@ -43,8 +43,8 @@
     type=gql.Ref('QueryPreviewResult'),
 )
 
-listTableSampleData = gql.QueryField(
-    name='listTableSampleData',
+readTableSampleData = gql.QueryField(
+    name='readTableSampleData',
     args=[gql.Argument(name='tableUri', type=gql.NonNullableType(gql.String))],
     type=gql.Ref('QueryPreviewResult'),
     resolver=read_sample_data,

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

@@ -591,6 +591,7 @@ def list_dataset_tables_folders(uri, filter):
             return DatasetRepository.paginated_dataset_tables_folders(session, 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(

frontend/src/modules/S3_Datasets/components/ReviewMetadataComponent.js

@@ -18,7 +18,7 @@ import { useClient } from 'services';
 import { updateDatasetTable } from 'modules/Tables/services';
 import { updateDatasetStorageLocation } from 'modules/Folders/services';
 import {
-  listTableSampleData,
+  readTableSampleData,
   updateDataset,
   generateMetadataBedrock
 } from '../services';
@@ -55,12 +55,12 @@ export const ReviewMetadataComponent = (props) => {
     try {
       setLoadingSampleData(true);
       const response = await client.query(
-        listTableSampleData({
+        readTableSampleData({
           tableUri: table.targetUri
         })
       );
       if (!response.errors) {
-        openSampleDataPopup(response.data.listTableSampleData);
+        openSampleDataPopup(response.data.readTableSampleData);
         setTargetUri(table.targetUri);
         enqueueSnackbar('Successfully read sample data', {
           variant: 'success'
@@ -152,7 +152,7 @@ export const ReviewMetadataComponent = (props) => {
           metadataTypes: Object.entries(selectedMetadataTypes)
             .filter(([key, value]) => value === true)
             .map(([key]) => key),
-          sampleData: sampleDataWithoutTypename
+          tableSampleData: sampleDataWithoutTypename
         })
       );
 

frontend/src/modules/S3_Datasets/services/index.js

@@ -1,4 +1,3 @@
-export * from './batchUpdateTableColumnDescriptions';
 export * from './createDataset';
 export * from './deleteDataset';
 export * from './generateDatasetAccessToken';
@@ -7,7 +6,7 @@ export * from './getDatasetPresignedUrl';
 export * from './importDataset';
 export * from './listDatasetStorageLocations';
 export * from './listDatasetTablesFolders';
-export * from './listTableSampleData';
+export * from './readTableSampleData';
 export * from './startGlueCrawler';
 export * from './syncTables';
 export * from './updateDataset';

frontend/src/modules/S3_Datasets/services/listTableSampleData.js

@@ -0,0 +1,15 @@
+import { gql } from 'apollo-boost';
+
+export const readTableSampleData = ({ tableUri }) => ({
+  variables: {
+    tableUri
+  },
+  query: gql`
+    query readTableSampleData($tableUri: String!) {
+      readTableSampleData(tableUri: $tableUri) {
+        fields
+        rows
+      }
+    }
+  `
+});