Feature/error handling

athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBMetadataHandler.java

@@ -29,6 +29,7 @@
 import com.amazonaws.athena.connector.lambda.domain.TableName;
 import com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet;
 import com.amazonaws.athena.connector.lambda.domain.spill.SpillLocation;
+import com.amazonaws.athena.connector.lambda.exceptions.AthenaConnectorException;
 import com.amazonaws.athena.connector.lambda.handlers.GlueMetadataHandler;
 import com.amazonaws.athena.connector.lambda.metadata.GetDataSourceCapabilitiesRequest;
 import com.amazonaws.athena.connector.lambda.metadata.GetDataSourceCapabilitiesResponse;
@@ -59,6 +60,8 @@
 import com.amazonaws.services.athena.AmazonAthena;
 import com.amazonaws.services.glue.AWSGlue;
 import com.amazonaws.services.glue.model.Database;
+import com.amazonaws.services.glue.model.ErrorDetails;
+import com.amazonaws.services.glue.model.FederationSourceErrorCode;
 import com.amazonaws.services.glue.model.Table;
 import com.amazonaws.services.secretsmanager.AWSSecretsManager;
 import com.amazonaws.util.json.Jackson;
@@ -255,7 +258,7 @@ public ListTablesResponse doListTables(BlockAllocator allocator, ListTablesReque
     public GetTableResponse doGetQueryPassthroughSchema(BlockAllocator allocator, GetTableRequest request) throws Exception
     {
         if (!request.isQueryPassthrough()) {
-            throw new IllegalArgumentException("No Query passed through [{}]" + request);
+            throw new AthenaConnectorException("No Query passed through [{}]" + request, new ErrorDetails().withErrorCode(FederationSourceErrorCode.InvalidInputException.toString()).withErrorMessage("No Query passed through [{}]" + request));
         }
 
         queryPassthrough.verify(request.getQueryPassthroughArguments());
@@ -324,7 +327,7 @@ public void enhancePartitionSchema(SchemaBuilder partitionSchemaBuilder, GetTabl
             table = tableResolver.getTableMetadata(tableName);
         }
         catch (TimeoutException e) {
-            throw new RuntimeException(e);
+            throw new AthenaConnectorException(e.getMessage(), new ErrorDetails().withErrorCode(FederationSourceErrorCode.OperationTimeoutException.toString()).withErrorMessage(e.getMessage()));
         }
         // add table name so we don't have to do case insensitive resolution again
         partitionSchemaBuilder.addMetadata(TABLE_METADATA, table.getName());
@@ -471,7 +474,7 @@ public GetSplitsResponse doGetSplits(BlockAllocator allocator, GetSplitsRequest
         Map<String, String> partitionMetadata = partitions.getSchema().getCustomMetadata();
         String partitionType = partitionMetadata.get(PARTITION_TYPE_METADATA);
         if (partitionType == null) {
-            throw new IllegalStateException(String.format("No metadata %s defined in Schema %s", PARTITION_TYPE_METADATA, partitions.getSchema()));
+            throw new AthenaConnectorException(String.format("No metadata %s defined in Schema %s", PARTITION_TYPE_METADATA, partitions.getSchema()), new ErrorDetails().withErrorCode(FederationSourceErrorCode.InvalidInputException.toString()));
         }
         if (QUERY_PARTITION_TYPE.equals(partitionType)) {
             String hashKeyName = partitionMetadata.get(HASH_KEY_NAME_METADATA);
@@ -527,7 +530,7 @@ else if (SCAN_PARTITION_TYPE.equals(partitionType)) {
             return new GetSplitsResponse(request.getCatalogName(), splits, null);
         }
         else {
-            throw new IllegalStateException("Unexpected partition type " + partitionType);
+            throw new AthenaConnectorException("Unexpected partition type " + partitionType, new ErrorDetails().withErrorCode(FederationSourceErrorCode.InvalidInputException.toString()));
         }
     }
 

athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBRecordHandler.java

@@ -27,6 +27,7 @@
 import com.amazonaws.athena.connector.lambda.data.writers.extractors.Extractor;
 import com.amazonaws.athena.connector.lambda.domain.Split;
 import com.amazonaws.athena.connector.lambda.domain.predicate.Constraints;
+import com.amazonaws.athena.connector.lambda.exceptions.AthenaConnectorException;
 import com.amazonaws.athena.connector.lambda.handlers.RecordHandler;
 import com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest;
 import com.amazonaws.athena.connectors.dynamodb.credentials.CrossAccountCredentialsProviderV2;
@@ -36,6 +37,8 @@
 import com.amazonaws.athena.connectors.dynamodb.util.DDBRecordMetadata;
 import com.amazonaws.athena.connectors.dynamodb.util.DDBTypeUtils;
 import com.amazonaws.services.athena.AmazonAthena;
+import com.amazonaws.services.glue.model.ErrorDetails;
+import com.amazonaws.services.glue.model.FederationSourceErrorCode;
 import com.amazonaws.services.s3.AmazonS3;
 import com.amazonaws.services.secretsmanager.AWSSecretsManager;
 import com.amazonaws.util.json.Jackson;
@@ -206,7 +209,7 @@ protected void readWithConstraint(BlockSpiller spiller, ReadRecordsRequest recor
     private void handleQueryPassthroughPartiQLQuery(BlockSpiller spiller, ReadRecordsRequest recordsRequest, QueryStatusChecker queryStatusChecker)
     {
         if (!recordsRequest.getConstraints().isQueryPassThrough()) {
-            throw new RuntimeException("Attempting to readConstraints with Query Passthrough without PartiQL Query");
+            throw new AthenaConnectorException("Attempting to readConstraints with Query Passthrough without PartiQL Query", new ErrorDetails().withErrorCode(FederationSourceErrorCode.InvalidInputException.toString()));
         }
         queryPassthrough.verify(recordsRequest.getConstraints().getQueryPassthroughArguments());
 
@@ -325,7 +328,7 @@ private QueryRequest buildQueryRequest(Split split, String tableName, Schema sch
                 expressionAttributeValues.putAll(EnhancedDocument.fromJson(split.getProperty(EXPRESSION_VALUES_METADATA)).toMap());
             }
             catch (IOException e) {
-                throw new RuntimeException(e);
+                throw new AthenaConnectorException(e.getMessage(), new ErrorDetails().withErrorCode(FederationSourceErrorCode.InternalServiceException.toString()));
             }
         }
 
@@ -392,7 +395,7 @@ private ScanRequest buildScanRequest(Split split, String tableName, Schema schem
                 expressionAttributeValues.putAll(EnhancedDocument.fromJson(split.getProperty(EXPRESSION_VALUES_METADATA)).toMap());
             }
             catch (IOException e) {
-                throw new RuntimeException(e);
+                throw new AthenaConnectorException(e.getMessage(), new ErrorDetails().withErrorCode(FederationSourceErrorCode.InternalServiceException.toString()));
             }
         }
 
@@ -465,7 +468,7 @@ public Map<String, AttributeValue> next()
                     }
                 }
                 catch (TimeoutException | ExecutionException e) {
-                    throw new RuntimeException(e);
+                    throw new AthenaConnectorException(e.getMessage(), new ErrorDetails().withErrorCode(FederationSourceErrorCode.OperationTimeoutException.toString()));
                 }
                 currentPageIterator.set(iterator);
                 if (iterator.hasNext()) {

athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/qpt/DDBQueryPassthrough.java

@@ -19,7 +19,10 @@
  */
 package com.amazonaws.athena.connectors.dynamodb.qpt;
 
+import com.amazonaws.athena.connector.lambda.exceptions.AthenaConnectorException;
 import com.amazonaws.athena.connector.lambda.metadata.optimizations.querypassthrough.QueryPassthroughSignature;
+import com.amazonaws.services.glue.model.ErrorDetails;
+import com.amazonaws.services.glue.model.FederationSourceErrorCode;
 import com.google.common.collect.ImmutableSet;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -77,7 +80,7 @@ public void customConnectorVerifications(Map<String, String> engineQptArguments)
 
         // Immediately check if the statement starts with "SELECT"
         if (!upperCaseStatement.startsWith("SELECT")) {
-            throw new UnsupportedOperationException("Statement does not start with SELECT.");
+            throw new AthenaConnectorException("Statement does not start with SELECT.", new ErrorDetails().withErrorCode(FederationSourceErrorCode.OperationNotSupportedException.toString()));
         }
 
         // List of disallowed keywords
@@ -86,7 +89,7 @@ public void customConnectorVerifications(Map<String, String> engineQptArguments)
         // Check if the statement contains any disallowed keywords
         for (String keyword : disallowedKeywords) {
             if (upperCaseStatement.contains(keyword)) {
-                throw new UnsupportedOperationException("Unaccepted operation; only SELECT statements are allowed. Found: " + keyword);
+                throw new AthenaConnectorException("Unaccepted operation; only SELECT statements are allowed. Found: " + keyword, new ErrorDetails().withErrorCode(FederationSourceErrorCode.OperationNotSupportedException.toString()));
             }
         }
     }

athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/resolver/DynamoDBFieldResolver.java

@@ -20,8 +20,11 @@
 package com.amazonaws.athena.connectors.dynamodb.resolver;
 
 import com.amazonaws.athena.connector.lambda.data.FieldResolver;
+import com.amazonaws.athena.connector.lambda.exceptions.AthenaConnectorException;
 import com.amazonaws.athena.connectors.dynamodb.util.DDBRecordMetadata;
 import com.amazonaws.athena.connectors.dynamodb.util.DDBTypeUtils;
+import com.amazonaws.services.glue.model.ErrorDetails;
+import com.amazonaws.services.glue.model.FederationSourceErrorCode;
 import org.apache.arrow.vector.types.Types;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.slf4j.Logger;
@@ -86,8 +89,8 @@ public Object getFieldValue(Field field, Object originalValue)
                 return DDBTypeUtils.coerceValueToExpectedType(fieldValue, field, fieldType, metadata);
         }
 
-        throw new RuntimeException("Invalid field value encountered in DB record for field: " + field +
-                ",value: " + fieldValue);
+        throw new AthenaConnectorException("Invalid field value encountered in DB record for field: " + field +
+                ",value: " + fieldValue, new ErrorDetails().withErrorCode(FederationSourceErrorCode.InvalidInputException.toString()));
     }
 
     // Return the field value of a map key

athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/resolver/DynamoDBTableResolver.java

@@ -20,9 +20,12 @@
 package com.amazonaws.athena.connectors.dynamodb.resolver;
 
 import com.amazonaws.athena.connector.lambda.ThrottlingInvoker;
+import com.amazonaws.athena.connector.lambda.exceptions.AthenaConnectorException;
 import com.amazonaws.athena.connectors.dynamodb.model.DynamoDBPaginatedTables;
 import com.amazonaws.athena.connectors.dynamodb.model.DynamoDBTable;
 import com.amazonaws.athena.connectors.dynamodb.util.DDBTableUtils;
+import com.amazonaws.services.glue.model.ErrorDetails;
+import com.amazonaws.services.glue.model.FederationSourceErrorCode;
 import com.google.common.collect.ArrayListMultimap;
 import com.google.common.collect.Multimap;
 import org.apache.arrow.vector.types.pojo.Schema;
@@ -118,7 +121,7 @@ public Schema getTableSchema(String tableName)
                 return DDBTableUtils.peekTableForSchema(caseInsensitiveMatch.get(), invoker, ddbClient);
             }
             else {
-                throw e;
+                throw new AthenaConnectorException(e.getMessage(), new ErrorDetails().withErrorCode(FederationSourceErrorCode.EntityNotFoundException.toString()));
             }
         }
     }

athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/util/DDBPredicateUtils.java

@@ -24,8 +24,11 @@
 import com.amazonaws.athena.connector.lambda.domain.predicate.Range;
 import com.amazonaws.athena.connector.lambda.domain.predicate.SortedRangeSet;
 import com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet;
+import com.amazonaws.athena.connector.lambda.exceptions.AthenaConnectorException;
 import com.amazonaws.athena.connectors.dynamodb.model.DynamoDBIndex;
 import com.amazonaws.athena.connectors.dynamodb.model.DynamoDBTable;
+import com.amazonaws.services.glue.model.ErrorDetails;
+import com.amazonaws.services.glue.model.FederationSourceErrorCode;
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
@@ -189,15 +192,15 @@ private static void validateColumnRange(Range range)
                 case EXACTLY:
                     break;
                 case BELOW:
-                    throw new IllegalArgumentException("Low marker should never use BELOW bound");
+                    throw new AthenaConnectorException("Low marker should never use BELOW bound", new ErrorDetails().withErrorCode(FederationSourceErrorCode.InvalidInputException.toString()));
                 default:
                     throw new AssertionError("Unhandled lower bound: " + range.getLow().getBound());
             }
         }
         if (!range.getHigh().isUpperUnbounded()) {
             switch (range.getHigh().getBound()) {
                 case ABOVE:
-                    throw new IllegalArgumentException("High marker should never use ABOVE bound");
+                    throw new AthenaConnectorException("High marker should never use ABOVE bound", new ErrorDetails().withErrorCode(FederationSourceErrorCode.InvalidInputException.toString()));
                 case EXACTLY:
                     break;
                 case BELOW:

athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/util/DDBTableUtils.java

@@ -21,8 +21,11 @@
 
 import com.amazonaws.athena.connector.lambda.ThrottlingInvoker;
 import com.amazonaws.athena.connector.lambda.data.SchemaBuilder;
+import com.amazonaws.athena.connector.lambda.exceptions.AthenaConnectorException;
 import com.amazonaws.athena.connectors.dynamodb.model.DynamoDBIndex;
 import com.amazonaws.athena.connectors.dynamodb.model.DynamoDBTable;
+import com.amazonaws.services.glue.model.ErrorDetails;
+import com.amazonaws.services.glue.model.FederationSourceErrorCode;
 import com.google.common.collect.ImmutableList;
 import org.apache.arrow.vector.types.pojo.Schema;
 import org.slf4j.Logger;
@@ -167,7 +170,7 @@ public static Schema peekTableForSchema(String tableName, ThrottlingInvoker invo
                 logger.warn("Failed to retrieve table schema due to KMS issue, empty schema for table: {}. Error Message: {}", tableName, runtimeException.getMessage());
             }
             else {
-                throw runtimeException;
+                throw new AthenaConnectorException(runtimeException.getMessage(), new ErrorDetails().withErrorCode(FederationSourceErrorCode.OperationTimeoutException.toString()));
             }
         }
         return schemaBuilder.build();

athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/util/DDBTypeUtils.java

@@ -30,7 +30,10 @@
 import com.amazonaws.athena.connector.lambda.data.writers.holders.NullableDecimalHolder;
 import com.amazonaws.athena.connector.lambda.data.writers.holders.NullableVarBinaryHolder;
 import com.amazonaws.athena.connector.lambda.domain.predicate.ConstraintProjector;
+import com.amazonaws.athena.connector.lambda.exceptions.AthenaConnectorException;
 import com.amazonaws.athena.connectors.dynamodb.resolver.DynamoDBFieldResolver;
+import com.amazonaws.services.glue.model.ErrorDetails;
+import com.amazonaws.services.glue.model.FederationSourceErrorCode;
 import org.apache.arrow.vector.FieldVector;
 import org.apache.arrow.vector.holders.NullableBitHolder;
 import org.apache.arrow.vector.types.Types;
@@ -188,7 +191,7 @@ else if (enhancedAttributeValue.isMap()) {
         }
 
         String attributeTypeName = (value == null || value.getClass() == null) ? "null" : enhancedAttributeValue.type().name();
-        throw new RuntimeException("Unknown Attribute Value Type[" + attributeTypeName + "] for field[" + key + "]");
+        throw new AthenaConnectorException("Unknown Attribute Value Type[" + attributeTypeName + "] for field[" + key + "]", new ErrorDetails().withErrorCode(FederationSourceErrorCode.InvalidInputException.toString()));
     }
 
     /**
@@ -262,7 +265,7 @@ public static Field getArrowFieldFromDDBType(String attributeName, String attrib
             case MAP:
                 return new Field(attributeName, FieldType.nullable(Types.MinorType.STRUCT.getType()), null);
             default:
-                throw new RuntimeException("Unknown type[" + attributeType + "] for field[" + attributeName + "]");
+                throw new AthenaConnectorException("Unknown type[" + attributeType + "] for field[" + attributeName + "]", new ErrorDetails().withErrorCode(FederationSourceErrorCode.InvalidInputException.toString()));
         }
     }
 
@@ -382,7 +385,7 @@ public static List<Object> coerceListToExpectedType(Object value, Field field, D
 
         if (!(value instanceof Collection)) {
             if (value instanceof Map) {
-                throw new RuntimeException("Unexpected type (Map) encountered for: " + childField.getName());
+                throw new AthenaConnectorException("Unexpected type (Map) encountered for: " + childField.getName(), new ErrorDetails().withErrorCode(FederationSourceErrorCode.InvalidInputException.toString()));
             }
             return Collections.singletonList(coerceValueToExpectedType(value, childField, fieldType, recordMetadata));
         }
@@ -618,7 +621,7 @@ else if (value instanceof Map<?, ?>) {
             return handleMapType((Map<String, Object>) value);
         }
         else {
-            throw new UnsupportedOperationException("Unsupported value type: " + value.getClass());
+            throw new AthenaConnectorException("Unsupported value type: " + value.getClass(), new ErrorDetails().withErrorCode(FederationSourceErrorCode.InvalidInputException.toString()));
         }
     }
 
@@ -632,7 +635,7 @@ public static AttributeValue jsonToAttributeValue(String jsonString, String key)
     {
         EnhancedDocument enhancedDocument = EnhancedDocument.fromJson(jsonString);
         if (!enhancedDocument.isPresent(key)) {
-            throw new RuntimeException("Unknown attribute Key");
+            throw new AthenaConnectorException("Unknown attribute Key", new ErrorDetails().withErrorCode(FederationSourceErrorCode.InvalidInputException.toString()));
         }
         return enhancedDocument.toMap().get(key);
     }
@@ -655,7 +658,7 @@ else if (firstElement instanceof Number) {
         } // Add other types if needed
 
         // Fallback for unsupported set types
-        throw new UnsupportedOperationException("Unsupported Set element type: " + firstElement.getClass());
+        throw new AthenaConnectorException("Unsupported Set element type: " + firstElement.getClass(), new ErrorDetails().withErrorCode(FederationSourceErrorCode.InvalidInputException.toString()));
     }
 
     private static AttributeValue handleListType(List<?> value)

athena-elasticsearch/src/main/java/com/amazonaws/athena/connectors/elasticsearch/AWSRequestSigningApacheInterceptor.java

@@ -20,9 +20,12 @@
 package com.amazonaws.athena.connectors.elasticsearch;
 
 import com.amazonaws.DefaultRequest;
+import com.amazonaws.athena.connector.lambda.exceptions.AthenaConnectorException;
 import com.amazonaws.auth.AWSCredentialsProvider;
 import com.amazonaws.auth.Signer;
 import com.amazonaws.http.HttpMethodName;
+import com.amazonaws.services.glue.model.ErrorDetails;
+import com.amazonaws.services.glue.model.FederationSourceErrorCode;
 import org.apache.http.Header;
 import org.apache.http.HttpEntityEnclosingRequest;
 import org.apache.http.HttpException;
@@ -95,7 +98,7 @@ public void process(final HttpRequest request, final HttpContext context)
             uriBuilder = new URIBuilder(request.getRequestLine().getUri());
         }
         catch (URISyntaxException e) {
-            throw new IOException("Invalid URI", e);
+            throw new AthenaConnectorException("Invalid URI", new ErrorDetails().withErrorCode(FederationSourceErrorCode.InternalServiceException.toString()).withErrorMessage(e.getMessage()));
         }
 
         // Copy Apache HttpRequest to AWS DefaultRequest
@@ -112,7 +115,7 @@ public void process(final HttpRequest request, final HttpContext context)
             signableRequest.setResourcePath(uriBuilder.build().getRawPath());
         }
         catch (URISyntaxException e) {
-            throw new IOException("Invalid URI", e);
+            throw new AthenaConnectorException("Invalid URI", new ErrorDetails().withErrorCode(FederationSourceErrorCode.InternalServiceException.toString()).withErrorMessage(e.getMessage()));
         }
 
         if (request instanceof HttpEntityEnclosingRequest) {