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

Example 6 with S3BlockSpillReader

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

the class CloudwatchRecordHandlerTest method setUp.

@Before
public void setUp() throws Exception {
    schemaForRead = CloudwatchMetadataHandler.CLOUDWATCH_SCHEMA;
    mockS3Storage = new ArrayList<>();
    allocator = new BlockAllocatorImpl();
    handler = new CloudwatchRecordHandler(mockS3, mockSecretsManager, mockAthena, mockAwsLogs);
    spillReader = new S3BlockSpillReader(mockS3, allocator);
    when(mockS3.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(mockS3.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(mockAwsLogs.getLogEvents(any(GetLogEventsRequest.class))).thenAnswer((InvocationOnMock invocationOnMock) -> {
        GetLogEventsRequest request = (GetLogEventsRequest) invocationOnMock.getArguments()[0];
        // Check that predicate pushdown was propagated to cloudwatch
        assertNotNull(request.getStartTime());
        assertNotNull(request.getEndTime());
        GetLogEventsResult result = new GetLogEventsResult();
        Integer nextToken;
        if (request.getNextToken() == null) {
            nextToken = 1;
        } else if (Integer.valueOf(request.getNextToken()) < 3) {
            nextToken = Integer.valueOf(request.getNextToken()) + 1;
        } else {
            nextToken = null;
        }
        List<OutputLogEvent> logEvents = new ArrayList<>();
        if (request.getNextToken() == null || Integer.valueOf(request.getNextToken()) < 3) {
            long continuation = request.getNextToken() == null ? 0 : Integer.valueOf(request.getNextToken());
            for (int i = 0; i < 100_000; i++) {
                OutputLogEvent outputLogEvent = new OutputLogEvent();
                outputLogEvent.setMessage("message-" + (continuation * i));
                outputLogEvent.setTimestamp(i * 100L);
                logEvents.add(outputLogEvent);
            }
        }
        result.withEvents(logEvents);
        if (nextToken != null) {
            result.setNextForwardToken(String.valueOf(nextToken));
        }
        return result;
    });
}

Example 7 with S3BlockSpillReader

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

the class ElasticsearchRecordHandlerTest method setUp.

@Before
public void setUp() throws IOException {
    logger.info("setUpBefore - enter");
    Map<String, Object> document1 = new ObjectMapper().readValue("{\n" + "  \"mytext\" : \"My favorite Sci-Fi movie is Interstellar.\",\n" + "  \"mykeyword\" : \"I love keywords.\",\n" + "  \"mylong\" : [\n" + "    11,\n" + "    12,\n" + "    13\n" + "  ],\n" + "  \"myinteger\" : 666115,\n" + "  \"myshort\" : 1972,\n" + "  \"mybyte\" : [\n" + "    5,\n" + "    6\n" + "  ],\n" + "  \"mydouble\" : 47.5,\n" + "  \"myscaled\" : 0.666,\n" + "  \"myfloat\" : 5.6,\n" + "  \"myhalf\" : 6.2,\n" + "  \"mydatemilli\" : \"2020-05-15T06:49:30\",\n" + "  \"mydatenano\" : \"2020-05-15T06:50:01.45678\",\n" + "  \"myboolean\" : true,\n" + "  \"mybinary\" : \"U29tZSBiaW5hcnkgYmxvYg==\",\n" + "  \"mynested\" : {\n" + "    \"l1long\" : 357345987,\n" + "    \"l1date\" : \"2020-05-15T06:57:44.123\",\n" + "    \"l1nested\" : {\n" + "      \"l2short\" : [\n" + "        1,\n" + "        2,\n" + "        3,\n" + "        4,\n" + "        5,\n" + "        6,\n" + "        7,\n" + "        8,\n" + "        9,\n" + "        10\n" + "      ],\n" + "      \"l2binary\" : \"U29tZSBiaW5hcnkgYmxvYg==\"\n" + "    }\n" + "  }\n" + "}\n", HashMap.class);
    Map<String, Object> document2 = new ObjectMapper().readValue("{\n" + "  \"mytext\" : \"My favorite TV comedy is Seinfeld.\",\n" + "  \"mykeyword\" : \"I hate key-values.\",\n" + "  \"mylong\" : [\n" + "    \"14.5\",\n" + "    null,\n" + "    16\n" + "  ],\n" + "  \"myinteger\" : \"732765666.5\",\n" + "  \"myshort\" : \"1971.1\",\n" + "  \"mybyte\" : \"7.2\",\n" + "  \"mydouble\" : \"27.6\",\n" + "  \"myscaled\" : \"0.999\",\n" + "  \"myfloat\" : \"7.8\",\n" + "  \"myhalf\" : \"7.3\",\n" + "  \"mydatemilli\" : null,\n" + "  \"mydatenano\" : 1589525370001,\n" + "  \"myboolean\" : \"false\",\n" + "  \"mybinary\" : \"U29tZSBiaW5hcnkgYmxvYg==\",\n" + "  \"mynested\" : {\n" + "    \"l1long\" : \"7322775555\",\n" + "    \"l1date\" : \"2020-05-15T06:57:44.7765+05:00\",\n" + "    \"l1nested\" : {\n" + "      \"l2short\" : [\n" + "        11,\n" + "        12,\n" + "        13,\n" + "        \"14.3\",\n" + "        15,\n" + "        16.5,\n" + "        null,\n" + "        18,\n" + "        \"19.4\",\n" + "        20\n" + "      ],\n" + "      \"l2binary\" : \"U29tZSBiaW5hcnkgYmxvYg==\"\n" + "    }\n" + "  }\n" + "}\n", HashMap.class);
    mapping = SchemaBuilder.newBuilder().addField("mytext", Types.MinorType.VARCHAR.getType()).addField("mykeyword", Types.MinorType.VARCHAR.getType()).addField(new Field("mylong", FieldType.nullable(Types.MinorType.LIST.getType()), Collections.singletonList(new Field("mylong", FieldType.nullable(Types.MinorType.BIGINT.getType()), null)))).addField("myinteger", Types.MinorType.INT.getType()).addField("myshort", Types.MinorType.SMALLINT.getType()).addField("mybyte", Types.MinorType.TINYINT.getType()).addField("mydouble", Types.MinorType.FLOAT8.getType()).addField(new Field("myscaled", new FieldType(true, Types.MinorType.BIGINT.getType(), null, ImmutableMap.of("scaling_factor", "10.51")), null)).addField("myfloat", Types.MinorType.FLOAT4.getType()).addField("myhalf", Types.MinorType.FLOAT4.getType()).addField("mydatemilli", Types.MinorType.DATEMILLI.getType()).addField(new Field("mydatenano", FieldType.nullable(Types.MinorType.LIST.getType()), Collections.singletonList(new Field("mydatenano", FieldType.nullable(Types.MinorType.DATEMILLI.getType()), null)))).addField("myboolean", Types.MinorType.BIT.getType()).addField("mybinary", Types.MinorType.VARCHAR.getType()).addField("mynested", Types.MinorType.STRUCT.getType(), ImmutableList.of(new Field("l1long", FieldType.nullable(Types.MinorType.BIGINT.getType()), null), new Field("l1date", FieldType.nullable(Types.MinorType.DATEMILLI.getType()), null), new Field("l1nested", FieldType.nullable(Types.MinorType.STRUCT.getType()), ImmutableList.of(new Field("l2short", FieldType.nullable(Types.MinorType.LIST.getType()), Collections.singletonList(new Field("l2short", FieldType.nullable(Types.MinorType.SMALLINT.getType()), null))), new Field("l2binary", FieldType.nullable(Types.MinorType.VARCHAR.getType()), null))))).build();
    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));
        synchronized (mockS3Storage) {
            mockS3Storage.add(byteHolder);
            logger.info("puObject: total size " + mockS3Storage.size());
        }
        return mock(PutObjectResult.class);
    });
    when(amazonS3.getObject(Matchers.anyString(), Matchers.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;
    });
    spillReader = new S3BlockSpillReader(amazonS3, allocator);
    split = Split.newBuilder(makeSpillLocation(), null).add("movies", "https://search-movies-ne3fcqzfipy6jcrew2wca6kyqu.us-east-1.es.amazonaws.com").build();
    when(clientFactory.getOrCreateClient(anyString())).thenReturn(mockClient);
    when(mockClient.getDocuments(any())).thenReturn(mockResponse);
    when(mockClient.getDocument(any())).thenReturn(document1, document2);
    handler = new ElasticsearchRecordHandler(amazonS3, awsSecretsManager, athena, clientFactory, 720);
    logger.info("setUpBefore - exit");
}

Example 8 with S3BlockSpillReader

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

the class ExampleRecordHandlerTest method setUp.

@Before
public void setUp() {
    logger.info("setUpBefore - enter");
    schemaForRead = SchemaBuilder.newBuilder().addField("year", new ArrowType.Int(32, true)).addField("month", new ArrowType.Int(32, true)).addField("day", new ArrowType.Int(32, true)).addField("col2", new ArrowType.Utf8()).addField("col3", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)).addField("int", Types.MinorType.INT.getType()).addField("tinyint", Types.MinorType.TINYINT.getType()).addField("smallint", Types.MinorType.SMALLINT.getType()).addField("bigint", Types.MinorType.BIGINT.getType()).addField("float4", Types.MinorType.FLOAT4.getType()).addField("float8", Types.MinorType.FLOAT8.getType()).addField("bit", Types.MinorType.BIT.getType()).addField("varchar", Types.MinorType.VARCHAR.getType()).addField("varbinary", Types.MinorType.VARBINARY.getType()).addField("datemilli", Types.MinorType.DATEMILLI.getType()).addField("dateday", Types.MinorType.DATEDAY.getType()).addField("decimal", new ArrowType.Decimal(10, 2)).addField("decimalLong", // Example of a List of Structs
    new ArrowType.Decimal(36, 2)).addField(FieldBuilder.newBuilder("list", new ArrowType.List()).addField(FieldBuilder.newBuilder("innerStruct", Types.MinorType.STRUCT.getType()).addStringField("varchar").addBigIntField("bigint").build()).build()).addField(FieldBuilder.newBuilder("outerlist", new ArrowType.List()).addListField("innerList", Types.MinorType.VARCHAR.getType()).build()).addField(FieldBuilder.newBuilder("simplemap", new ArrowType.Map(false)).addField("entries", Types.MinorType.STRUCT.getType(), false, Arrays.asList(FieldBuilder.newBuilder("key", Types.MinorType.VARCHAR.getType(), false).build(), FieldBuilder.newBuilder("value", Types.MinorType.INT.getType()).build())).build()).addMetadata("partitionCols", "year,month,day").build();
    allocator = new BlockAllocatorImpl();
    amazonS3 = mock(AmazonS3.class);
    awsSecretsManager = mock(AWSSecretsManager.class);
    athena = mock(AmazonAthena.class);
    when(amazonS3.putObject(anyObject(), anyObject(), anyObject(), anyObject())).thenAnswer(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocationOnMock) throws Throwable {
            InputStream inputStream = (InputStream) invocationOnMock.getArguments()[2];
            ByteHolder byteHolder = new ByteHolder();
            byteHolder.setBytes(ByteStreams.toByteArray(inputStream));
            mockS3Storage.add(byteHolder);
            return mock(PutObjectResult.class);
        }
    });
    when(amazonS3.getObject(anyString(), anyString())).thenAnswer(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocationOnMock) throws Throwable {
            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;
        }
    });
    recordService = new LocalHandler(allocator, amazonS3, awsSecretsManager, athena);
    spillReader = new S3BlockSpillReader(amazonS3, allocator);
    logger.info("setUpBefore - exit");
}

Example 9 with S3BlockSpillReader

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

the class HbaseRecordHandlerTest method setUp.

@Before
public void setUp() throws IOException {
    logger.info("{}: enter", testName.getMethodName());
    when(mockConnFactory.getOrCreateConn(anyString())).thenReturn(mockClient);
    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;
    });
    schemaForRead = TestUtils.makeSchema().addStringField(HbaseSchemaUtils.ROW_COLUMN_NAME).build();
    handler = new HbaseRecordHandler(amazonS3, mockSecretsManager, mockAthena, mockConnFactory);
    spillReader = new S3BlockSpillReader(amazonS3, allocator);
}

Example 10 with S3BlockSpillReader

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

the class ExampleRecordHandlerTest method setUp.

@Before
public void setUp() {
    logger.info("{}: enter", testName.getMethodName());
    schemaForRead = SchemaBuilder.newBuilder().addIntField("year").addIntField("month").addIntField("day").addStringField("account_id").addStringField("encrypted_payload").addStructField("transaction").addChildField("transaction", "id", Types.MinorType.INT.getType()).addChildField("transaction", "completed", Types.MinorType.BIT.getType()).addMetadata("partitionCols", "year,month,day").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.getObject(anyString(), anyString())).thenAnswer(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocationOnMock) throws Throwable {
            S3Object mockObject = mock(S3Object.class);
            when(mockObject.getObjectContent()).thenReturn(new S3ObjectInputStream(new ByteArrayInputStream(getFakeObject()), null));
            return mockObject;
        }
    });
    handler = new ExampleRecordHandler(amazonS3, awsSecretsManager, athena);
    spillReader = new S3BlockSpillReader(amazonS3, allocator);
}