Examples with S3BlockSpillReader com.amazonaws.athena.connector.lambda.data.S3BlockSpillReader used on opensource projects

Example 1 with S3BlockSpillReader

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

the class NeptuneRecordHandlerTest method setUp.

@Before
public void setUp() {
    logger.info("{}: enter", testName.getMethodName());
    schemaPGVertexForRead = SchemaBuilder.newBuilder().addMetadata("componenttype", "vertex").addStringField("id").addIntField("property1").addStringField("property2").addFloat8Field("property3").addBitField("property4").addBigIntField("property5").addFloat4Field("property6").build();
    schemaPGEdgeForRead = SchemaBuilder.newBuilder().addMetadata("componenttype", "edge").addStringField("in").addStringField("out").addStringField("id").addIntField("property1").addStringField("property2").addFloat8Field("property3").addBitField("property4").addBigIntField("property5").addFloat4Field("property6").build();
    allocator = new BlockAllocatorImpl();
    amazonS3 = mock(AmazonS3.class);
    awsSecretsManager = mock(AWSSecretsManager.class);
    athena = mock(AmazonAthena.class);
    when(amazonS3.doesObjectExist(anyString(), anyString())).thenReturn(true);
    when(amazonS3.putObject(anyObject(), anyObject(), anyObject(), anyObject())).thenAnswer((InvocationOnMock invocationOnMock) -> {
        InputStream inputStream = (InputStream) invocationOnMock.getArguments()[2];
        ByteHolder byteHolder = new ByteHolder();
        byteHolder.setBytes(ByteStreams.toByteArray(inputStream));
        synchronized (mockS3Storage) {
            mockS3Storage.add(byteHolder);
            logger.info("puObject: total size " + mockS3Storage.size());
        }
        return mock(PutObjectResult.class);
    });
    when(amazonS3.getObject(anyString(), anyString())).thenAnswer((InvocationOnMock invocationOnMock) -> {
        S3Object mockObject = mock(S3Object.class);
        ByteHolder byteHolder;
        synchronized (mockS3Storage) {
            byteHolder = mockS3Storage.get(0);
            mockS3Storage.remove(0);
            logger.info("getObject: total size " + mockS3Storage.size());
        }
        when(mockObject.getObjectContent()).thenReturn(new S3ObjectInputStream(new ByteArrayInputStream(byteHolder.getBytes()), null));
        return mockObject;
    });
    handler = new NeptuneRecordHandler(amazonS3, awsSecretsManager, athena, neptuneConnection);
    spillReader = new S3BlockSpillReader(amazonS3, allocator);
}

Example 2 with S3BlockSpillReader

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

the class RedisRecordHandlerTest method setUp.

@Before
public void setUp() {
    logger.info("{}: enter", testName.getMethodName());
    when(mockFactory.getOrCreateConn(eq(decodedEndpoint), anyBoolean(), anyBoolean(), anyString())).thenReturn(mockConnection);
    when(mockConnection.sync()).thenReturn(mockSyncCommands);
    allocator = new BlockAllocatorImpl();
    amazonS3 = mock(AmazonS3.class);
    when(amazonS3.putObject(anyObject(), anyObject(), anyObject(), anyObject())).thenAnswer((InvocationOnMock invocationOnMock) -> {
        InputStream inputStream = (InputStream) invocationOnMock.getArguments()[2];
        ByteHolder byteHolder = new ByteHolder();
        byteHolder.setBytes(ByteStreams.toByteArray(inputStream));
        synchronized (mockS3Storage) {
            mockS3Storage.add(byteHolder);
            logger.info("puObject: total size " + mockS3Storage.size());
        }
        return mock(PutObjectResult.class);
    });
    when(amazonS3.getObject(anyString(), anyString())).thenAnswer((InvocationOnMock invocationOnMock) -> {
        S3Object mockObject = mock(S3Object.class);
        ByteHolder byteHolder;
        synchronized (mockS3Storage) {
            byteHolder = mockS3Storage.get(0);
            mockS3Storage.remove(0);
            logger.info("getObject: total size " + mockS3Storage.size());
        }
        when(mockObject.getObjectContent()).thenReturn(new S3ObjectInputStream(new ByteArrayInputStream(byteHolder.getBytes()), null));
        return mockObject;
    });
    when(mockSecretsManager.getSecretValue(any(GetSecretValueRequest.class))).thenAnswer((InvocationOnMock invocation) -> {
        GetSecretValueRequest request = invocation.getArgumentAt(0, GetSecretValueRequest.class);
        if ("endpoint".equalsIgnoreCase(request.getSecretId())) {
            return new GetSecretValueResult().withSecretString(decodedEndpoint);
        }
        throw new RuntimeException("Unknown secret " + request.getSecretId());
    });
    handler = new RedisRecordHandler(amazonS3, mockSecretsManager, mockAthena, mockFactory);
    spillReader = new S3BlockSpillReader(amazonS3, allocator);
    logger.info("setUpBefore - exit");
}

Example 3 with S3BlockSpillReader

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

the class TPCDSRecordHandlerTest method setUp.

@Before
public void setUp() throws Exception {
    for (Table next : Table.getBaseTables()) {
        if (next.getName().equals("customer")) {
            table = next;
        }
    }
    SchemaBuilder schemaBuilder = SchemaBuilder.newBuilder();
    for (Column nextCol : table.getColumns()) {
        schemaBuilder.addField(TPCDSUtils.convertColumn(nextCol));
    }
    schemaForRead = schemaBuilder.build();
    mockS3Storage = new ArrayList<>();
    allocator = new BlockAllocatorImpl();
    handler = new TPCDSRecordHandler(mockS3, mockSecretsManager, mockAthena);
    spillReader = new S3BlockSpillReader(mockS3, allocator);
    when(mockS3.putObject(anyObject(), anyObject(), anyObject(), anyObject())).thenAnswer((InvocationOnMock invocationOnMock) -> {
        synchronized (mockS3Storage) {
            InputStream inputStream = (InputStream) invocationOnMock.getArguments()[2];
            ByteHolder byteHolder = new ByteHolder();
            byteHolder.setBytes(ByteStreams.toByteArray(inputStream));
            mockS3Storage.add(byteHolder);
            return mock(PutObjectResult.class);
        }
    });
    when(mockS3.getObject(anyString(), anyString())).thenAnswer((InvocationOnMock invocationOnMock) -> {
        synchronized (mockS3Storage) {
            S3Object mockObject = mock(S3Object.class);
            ByteHolder byteHolder = mockS3Storage.get(0);
            mockS3Storage.remove(0);
            when(mockObject.getObjectContent()).thenReturn(new S3ObjectInputStream(new ByteArrayInputStream(byteHolder.getBytes()), null));
            return mockObject;
        }
    });
}

Example 4 with S3BlockSpillReader

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

the class TimestreamRecordHandlerTest method setUp.

@Before
public void setUp() throws IOException {
    logger.info("{}: enter", testName.getMethodName());
    allocator = new BlockAllocatorImpl();
    amazonS3 = mock(AmazonS3.class);
    when(amazonS3.putObject(anyObject(), anyObject(), anyObject(), anyObject())).thenAnswer((InvocationOnMock invocationOnMock) -> {
        InputStream inputStream = (InputStream) invocationOnMock.getArguments()[2];
        ByteHolder byteHolder = new ByteHolder();
        byteHolder.setBytes(ByteStreams.toByteArray(inputStream));
        synchronized (mockS3Storage) {
            mockS3Storage.add(byteHolder);
        }
        return mock(PutObjectResult.class);
    });
    when(amazonS3.getObject(anyString(), anyString())).thenAnswer((InvocationOnMock invocationOnMock) -> {
        S3Object mockObject = mock(S3Object.class);
        ByteHolder byteHolder;
        synchronized (mockS3Storage) {
            byteHolder = mockS3Storage.get(0);
            mockS3Storage.remove(0);
        }
        when(mockObject.getObjectContent()).thenReturn(new S3ObjectInputStream(new ByteArrayInputStream(byteHolder.getBytes()), null));
        return mockObject;
    });
    schemaForRead = SchemaBuilder.newBuilder().addField("measure_name", Types.MinorType.VARCHAR.getType()).addField("measure_value::double", Types.MinorType.FLOAT8.getType()).addField("az", Types.MinorType.VARCHAR.getType()).addField("time", Types.MinorType.DATEMILLI.getType()).addField("hostname", Types.MinorType.VARCHAR.getType()).addField("region", Types.MinorType.VARCHAR.getType()).build();
    handler = new TimestreamRecordHandler(amazonS3, mockSecretsManager, mockAthena, mockClient);
    spillReader = new S3BlockSpillReader(amazonS3, allocator);
}

Example 5 with S3BlockSpillReader

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

the class AbstractTableProviderTest method setUp.

@Before
public void setUp() {
    allocator = new BlockAllocatorImpl();
    when(amazonS3.putObject(anyObject(), anyObject(), anyObject(), anyObject())).thenAnswer((InvocationOnMock invocationOnMock) -> {
        InputStream inputStream = (InputStream) invocationOnMock.getArguments()[2];
        ByteHolder byteHolder = new ByteHolder();
        byteHolder.setBytes(ByteStreams.toByteArray(inputStream));
        mockS3Store.add(byteHolder);
        return mock(PutObjectResult.class);
    });
    when(amazonS3.getObject(anyString(), anyString())).thenAnswer((InvocationOnMock invocationOnMock) -> {
        S3Object mockObject = mock(S3Object.class);
        ByteHolder byteHolder = mockS3Store.get(0);
        mockS3Store.remove(0);
        when(mockObject.getObjectContent()).thenReturn(new S3ObjectInputStream(new ByteArrayInputStream(byteHolder.getBytes()), null));
        return mockObject;
    });
    blockSpillReader = new S3BlockSpillReader(amazonS3, allocator);
    provider = setUpSource();
    when(queryStatusChecker.isQueryRunning()).thenReturn(true);
}