Examples with FederationThrottleException com.amazonaws.athena.connector.lambda.exceptions.FederationThrottleException used on opensource projects

Example 1 with FederationThrottleException

use of com.amazonaws.athena.connector.lambda.exceptions.FederationThrottleException in project aws-athena-query-federation by awslabs.

the class ThrottlingInvokerTest method invokeWithThrottle.

@Test
public void invokeWithThrottle() throws TimeoutException {
    ThrottlingInvoker invoker = ThrottlingInvoker.newBuilder().withDecrease(0.8).withIncrease(1).withInitialDelayMs(10).withMaxDelayMs(200).withFilter((Exception ex) -> ex instanceof FederationThrottleException).build();
    for (int i = 0; i < 5; i++) {
        // Make a call and validate that the state didn't change
        final AtomicLong count = new AtomicLong(0);
        final int val = i;
        long result = invoker.invoke(() -> {
            if (count.incrementAndGet() < 4) {
                throw new FederationThrottleException();
            }
            return val;
        }, 10_000);
        assertEquals(val, result);
        assertEquals(4, count.get());
        assertEquals(ThrottlingInvoker.State.AVOIDANCE, invoker.getState());
        assertTrue(invoker.getDelay() > 0);
    }
    assertEquals(199, invoker.getDelay());
}

Example 2 with FederationThrottleException

use of com.amazonaws.athena.connector.lambda.exceptions.FederationThrottleException in project aws-athena-query-federation by awslabs.

the class ExampleRecordHandler method readWithConstraint.

/**
 * Here we generate our simulated row data. A real connector would instead connect to the actual source and read
 * the data corresponding to the requested split.
 *
 * @param spiller A BlockSpiller that should be used to write the row data associated with this Split.
 * The BlockSpiller automatically handles applying constraints, chunking the response, encrypting, and spilling to S3.
 * @param request The ReadRecordsRequest containing the split and other details about what to read.
 * @param queryStatusChecker A QueryStatusChecker that you can use to stop doing work for a query that has already terminated
 */
@Override
protected void readWithConstraint(BlockSpiller spiller, ReadRecordsRequest request, QueryStatusChecker queryStatusChecker) {
    long startTime = System.currentTimeMillis();
    /**
     * It is important to try and throw any throttling events before writing data since Athena may not be able to
     * continue the query, due to consistency errors, if you throttle after writing data.
     */
    if (simulateThrottle > 0 && count++ % simulateThrottle == 0) {
        logger.info("readWithConstraint: throwing throttle Exception!");
        throw new FederationThrottleException("Please slow down for this simulated throttling event");
    }
    logCaller(request);
    Set<String> partitionCols = new HashSet<>();
    String partitionColsMetadata = request.getSchema().getCustomMetadata().get("partitionCols");
    if (partitionColsMetadata != null) {
        partitionCols.addAll(Arrays.asList(partitionColsMetadata.split(",")));
    }
    int year = Integer.valueOf(request.getSplit().getProperty("year"));
    int month = Integer.valueOf(request.getSplit().getProperty("month"));
    int day = Integer.valueOf(request.getSplit().getProperty("day"));
    final RowContext rowContext = new RowContext(year, month, day);
    GeneratedRowWriter.RowWriterBuilder builder = GeneratedRowWriter.newBuilder(request.getConstraints());
    for (Field next : request.getSchema().getFields()) {
        Extractor extractor = makeExtractor(next, rowContext);
        if (extractor != null) {
            builder.withExtractor(next.getName(), extractor);
        } else {
            builder.withFieldWriterFactory(next.getName(), makeFactory(next, rowContext));
        }
    }
    GeneratedRowWriter rowWriter = builder.build();
    for (int i = 0; i < numRowsPerSplit; i++) {
        rowContext.seed = i;
        rowContext.negative = i % 2 == 0;
        if (!queryStatusChecker.isQueryRunning()) {
            return;
        }
        spiller.writeRows((Block block, int rowNum) -> rowWriter.writeRow(block, rowNum, rowContext) ? 1 : 0);
    }
    logger.info("readWithConstraint: Completed generating rows in {} ms", System.currentTimeMillis() - startTime);
}

Example 3 with FederationThrottleException

use of com.amazonaws.athena.connector.lambda.exceptions.FederationThrottleException in project aws-athena-query-federation by awslabs.

the class ExampleMetadataHandler method doGetSplits.

/**
 * For each partition we generate a pre-determined number of splits based on the NUM_PARTS_PER_SPLIT setting. This
 * method also demonstrates how to handle calls for batches of partitions and also leverage this API's ability
 * to paginated. A connector for a real data source would likely query that source's metadata to determine if/how
 * to split up the read operations for a particular partition.
 *
 * @param allocator Tool for creating and managing Apache Arrow Blocks.
 * @param request Provides details of the catalog, database, table, andpartition(s) being queried as well as
 * any filter predicate.
 * @return A GetSplitsResponse which contains a list of splits as an optional continuation token if we were not
 * able to generate all splits for the partitions in this batch.
 */
@Override
public GetSplitsResponse doGetSplits(BlockAllocator allocator, GetSplitsRequest request) {
    logCaller(request);
    logger.info("doGetSplits: spill location " + makeSpillLocation(request));
    /**
     * It is important to try and throw any throttling events before writing data since Athena may not be able to
     * continue the query, due to consistency errors, if you throttle after writing data.
     */
    if (simulateThrottle > 0 && count++ % simulateThrottle == 0) {
        logger.info("readWithConstraint: throwing throttle Exception!");
        throw new FederationThrottleException("Please slow down for this simulated throttling event");
    }
    ContinuationToken requestToken = ContinuationToken.decode(request.getContinuationToken());
    int partitionContd = requestToken.getPartition();
    int partContd = requestToken.getPart();
    Set<Split> splits = new HashSet<>();
    Block partitions = request.getPartitions();
    for (int curPartition = partitionContd; curPartition < partitions.getRowCount(); curPartition++) {
        // We use the makeEncryptionKey() method from our parent class to make an EncryptionKey
        EncryptionKey encryptionKey = makeEncryptionKey();
        // We prepare to read our custom metadata fields from the partition so that we can pass this info to the split(s)
        FieldReader locationReader = partitions.getFieldReader(SplitProperties.LOCATION.getId());
        locationReader.setPosition(curPartition);
        FieldReader storageClassReader = partitions.getFieldReader(SplitProperties.SERDE.getId());
        storageClassReader.setPosition(curPartition);
        // table scan operations (aka splits)
        for (int curPart = partContd; curPart < NUM_PARTS_PER_SPLIT; curPart++) {
            if (splits.size() >= MAX_SPLITS_PER_REQUEST) {
                // a continuation token.
                return new GetSplitsResponse(request.getCatalogName(), splits, ContinuationToken.encode(curPartition, curPart));
            }
            // We use makeSpillLocation(...) from our parent class to get a unique SpillLocation for each split
            Split.Builder splitBuilder = Split.newBuilder(makeSpillLocation(request), encryptionEnabled ? encryptionKey : null).add(SplitProperties.LOCATION.getId(), String.valueOf(locationReader.readText())).add(SplitProperties.SERDE.getId(), String.valueOf(storageClassReader.readText())).add(SplitProperties.SPLIT_PART.getId(), String.valueOf(curPart));
            // will likely vary. Our example only supports a limited number of partition column types.
            for (String next : request.getPartitionCols()) {
                FieldReader reader = partitions.getFieldReader(next);
                reader.setPosition(curPartition);
                switch(reader.getMinorType()) {
                    case UINT2:
                        splitBuilder.add(next, Integer.valueOf(reader.readCharacter()).toString());
                        break;
                    case UINT4:
                    case INT:
                        splitBuilder.add(next, String.valueOf(reader.readInteger()));
                        break;
                    case UINT8:
                    case BIGINT:
                        splitBuilder.add(next, String.valueOf(reader.readLong()));
                        break;
                    default:
                        throw new RuntimeException("Unsupported partition column type. " + reader.getMinorType());
                }
            }
            splits.add(splitBuilder.build());
        }
        // part continuation only applies within a partition so we complete that partial partition and move on
        // to the next one.
        partContd = 0;
    }
    return new GetSplitsResponse(request.getCatalogName(), splits, null);
}