Examples with EncryptionKey com.amazonaws.athena.connector.lambda.security.EncryptionKey used on opensource projects

Example 6 with EncryptionKey

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

the class S3BlockSpiller method write.

/**
 * Writes (aka spills) a Block.
 */
protected SpillLocation write(Block block) {
    try {
        S3SpillLocation spillLocation = makeSpillLocation();
        EncryptionKey encryptionKey = spillConfig.getEncryptionKey();
        logger.info("write: Started encrypting block for write to {}", spillLocation);
        byte[] bytes = blockCrypto.encrypt(encryptionKey, block);
        totalBytesSpilled.addAndGet(bytes.length);
        logger.info("write: Started spilling block of size {} bytes", bytes.length);
        amazonS3.putObject(spillLocation.getBucket(), spillLocation.getKey(), new ByteArrayInputStream(bytes), new ObjectMetadata());
        logger.info("write: Completed spilling block of size {} bytes", bytes.length);
        return spillLocation;
    } catch (RuntimeException ex) {
        asyncException.compareAndSet(null, ex);
        logger.warn("write: Encountered error while writing block.", ex);
        throw ex;
    }
}

Example 7 with EncryptionKey

use of com.amazonaws.athena.connector.lambda.security.EncryptionKey in project foundry-athena-query-federation-connector by palantir.

the class FoundryMetadataHandler method doGetSplits.

@Override
public GetSplitsResponse doGetSplits(BlockAllocator _allocator, GetSplitsRequest request) {
    log.debug("Getting splits with constraints: {}", request.getConstraints());
    // can be shared across
    EncryptionKey encryptionKey = makeEncryptionKey();
    return splitsFetcher.getSplits(request, () -> makeSpillLocation(request), encryptionKey);
}

Example 8 with EncryptionKey

use of com.amazonaws.athena.connector.lambda.security.EncryptionKey in project foundry-athena-query-federation-connector by palantir.

the class S3Spiller method write.

private SpillLocation write(Block block) {
    try {
        S3SpillLocation spillLocation = makeSpillLocation();
        EncryptionKey encryptionKey = spillConfig.getEncryptionKey();
        log.info("write: Started encrypting block for write to {}", spillLocation);
        byte[] bytes = blockCrypto.encrypt(encryptionKey, block);
        totalBytesSpilled.addAndGet(bytes.length);
        log.info("write: Started spilling block of size {} bytes", bytes.length);
        ObjectMetadata objectMetadata = new ObjectMetadata();
        objectMetadata.setContentLength(bytes.length);
        amazonS3.putObject(spillLocation.getBucket(), spillLocation.getKey(), new ByteArrayInputStream(bytes), objectMetadata);
        log.info("write: Completed spilling block of size {} bytes", bytes.length);
        return spillLocation;
    } catch (RuntimeException ex) {
        asyncException.compareAndSet(null, ex);
        log.warn("write: Encountered error while writing block.", ex);
        throw ex;
    }
}

Example 9 with EncryptionKey

use of com.amazonaws.athena.connector.lambda.security.EncryptionKey in project foundry-athena-query-federation-connector by palantir.

the class SplitsFetcher method getSplits.

GetSplitsResponse getSplits(GetSplitsRequest request, SpillLocationFactory spillLocationFactory, EncryptionKey encryptionKey) {
    CatalogLocator locator = FoundryAthenaObjectMapper.objectMapper().convertValue(request.getSchema().getCustomMetadata(), CatalogLocator.class);
    Optional<Filter> filter;
    if (request.getConstraints().getSummary().isEmpty()) {
        filter = Optional.empty();
    } else {
        // we just push down all constraints which will include those for any partition columns
        filter = Optional.of(Filter.and(AndFilter.of(request.getConstraints().getSummary().entrySet().stream().map(entry -> ConstraintConverter.convert(entry.getKey(), entry.getValue())).collect(Collectors.toList()))));
    }
    Set<Split> splits = new HashSet<>();
    Optional<String> pageToken = Optional.empty();
    while (true) {
        GetSlicesResponse response = metadataService.getSlices(authProvider.getAuthHeader(), GetSlicesRequest.builder().locator(locator).filter(filter).nextPageToken(pageToken).build());
        splits.addAll(response.getSlices().stream().map(slice -> slices.toSplit(spillLocationFactory.makeSpillLocation(), encryptionKey, slice)).collect(Collectors.toSet()));
        if (response.getNextPageToken().isPresent()) {
            pageToken = response.getNextPageToken();
        } else {
            log.debug("finished planning splits. number of splits: {}", splits.size());
            return new GetSplitsResponse(request.getCatalogName(), splits);
        }
    }
}

Example 10 with EncryptionKey

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

the class ExampleRecordHandlerTest method doReadRecordsSpill.

@Test
public void doReadRecordsSpill() throws Exception {
    logger.info("doReadRecordsSpill: enter");
    for (int i = 0; i < 2; i++) {
        EncryptionKey encryptionKey = (i % 2 == 0) ? keyFactory.create() : null;
        logger.info("doReadRecordsSpill: Using encryptionKey[" + encryptionKey + "]");
        Map<String, ValueSet> constraintsMap = new HashMap<>();
        constraintsMap.put("col3", SortedRangeSet.copyOf(Types.MinorType.FLOAT8.getType(), ImmutableList.of(Range.greaterThan(allocator, Types.MinorType.FLOAT8.getType(), -10000D)), false));
        constraintsMap.put("unknown", EquatableValueSet.newBuilder(allocator, Types.MinorType.FLOAT8.getType(), false, true).add(1.1D).build());
        constraintsMap.put("unknown2", new AllOrNoneValueSet(Types.MinorType.FLOAT8.getType(), false, true));
        ReadRecordsRequest request = new ReadRecordsRequest(IdentityUtil.fakeIdentity(), "catalog", "queryId-" + System.currentTimeMillis(), new TableName("schema", "table"), schemaForRead, Split.newBuilder(makeSpillLocation(), encryptionKey).add("year", "10").add("month", "10").add("day", "10").build(), new Constraints(constraintsMap), // ~1.5MB so we should see some spill
        1_600_000L, 1000L);
        ObjectMapperUtil.assertSerialization(request);
        RecordResponse rawResponse = recordService.readRecords(request);
        ObjectMapperUtil.assertSerialization(rawResponse);
        assertTrue(rawResponse instanceof RemoteReadRecordsResponse);
        try (RemoteReadRecordsResponse response = (RemoteReadRecordsResponse) rawResponse) {
            logger.info("doReadRecordsSpill: remoteBlocks[{}]", response.getRemoteBlocks().size());
            assertTrue(response.getNumberBlocks() > 1);
            int blockNum = 0;
            for (SpillLocation next : response.getRemoteBlocks()) {
                S3SpillLocation spillLocation = (S3SpillLocation) next;
                try (Block block = spillReader.read(spillLocation, response.getEncryptionKey(), response.getSchema())) {
                    logger.info("doReadRecordsSpill: blockNum[{}] and recordCount[{}]", blockNum++, block.getRowCount());
                    // assertTrue(++blockNum < response.getRemoteBlocks().size() && block.getRowCount() > 10_000);
                    logger.info("doReadRecordsSpill: {}", BlockUtils.rowToString(block, 0));
                    assertNotNull(BlockUtils.rowToString(block, 0));
                }
            }
        }
    }
    logger.info("doReadRecordsSpill: exit");
}