Examples with GetTableRequest com.amazonaws.athena.connector.lambda.metadata.GetTableRequest used on opensource projects

Example 21 with GetTableRequest

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

the class CloudwatchMetadataHandlerTest method doGetTable.

@Test
public void doGetTable() {
    logger.info("doGetTable - enter");
    String expectedSchema = "schema-20";
    when(mockAwsLogs.describeLogStreams(any(DescribeLogStreamsRequest.class))).thenAnswer((InvocationOnMock invocationOnMock) -> {
        DescribeLogStreamsRequest request = (DescribeLogStreamsRequest) invocationOnMock.getArguments()[0];
        assertTrue(request.getLogGroupName().equals(expectedSchema));
        DescribeLogStreamsResult result = new DescribeLogStreamsResult();
        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<LogStream> logStreams = new ArrayList<>();
        if (request.getNextToken() == null || Integer.valueOf(request.getNextToken()) < 3) {
            for (int i = 0; i < 10; i++) {
                LogStream nextLogStream = new LogStream();
                nextLogStream.setLogStreamName("table-" + String.valueOf(i));
                logStreams.add(nextLogStream);
            }
        }
        result.withLogStreams(logStreams);
        if (nextToken != null) {
            result.setNextToken(String.valueOf(nextToken));
        }
        return result;
    });
    GetTableRequest req = new GetTableRequest(identity, "queryId", "default", new TableName(expectedSchema, "table-9"));
    GetTableResponse res = handler.doGetTable(allocator, req);
    logger.info("doGetTable - {} {}", res.getTableName(), res.getSchema());
    assertEquals(new TableName(expectedSchema, "table-9"), res.getTableName());
    assertTrue(res.getSchema() != null);
    verify(mockAwsLogs, times(1)).describeLogStreams(any(DescribeLogStreamsRequest.class));
    logger.info("doGetTable - exit");
}

Example 22 with GetTableRequest

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

the class DocDBMetadataHandlerTest method doGetTable.

/**
 * TODO: Add more types.
 */
@Test
public void doGetTable() throws Exception {
    List<Document> documents = new ArrayList<>();
    Document doc1 = new Document();
    documents.add(doc1);
    doc1.put("stringCol", "stringVal");
    doc1.put("intCol", 1);
    doc1.put("doubleCol", 2.2D);
    doc1.put("longCol", 100L);
    doc1.put("unsupported", new UnsupportedType());
    Document doc2 = new Document();
    documents.add(doc2);
    doc2.put("stringCol2", "stringVal");
    doc2.put("intCol2", 1);
    doc2.put("doubleCol2", 2.2D);
    doc2.put("longCol2", 100L);
    Document doc3 = new Document();
    documents.add(doc3);
    doc3.put("stringCol", "stringVal");
    doc3.put("intCol2", 1);
    doc3.put("doubleCol", 2.2D);
    doc3.put("longCol2", 100L);
    MongoDatabase mockDatabase = mock(MongoDatabase.class);
    MongoCollection mockCollection = mock(MongoCollection.class);
    FindIterable mockIterable = mock(FindIterable.class);
    when(mockClient.getDatabase(eq(DEFAULT_SCHEMA))).thenReturn(mockDatabase);
    when(mockDatabase.getCollection(eq(TEST_TABLE))).thenReturn(mockCollection);
    when(mockCollection.find()).thenReturn(mockIterable);
    when(mockIterable.limit(anyInt())).thenReturn(mockIterable);
    when(mockIterable.maxScan(anyInt())).thenReturn(mockIterable);
    when(mockIterable.batchSize(anyInt())).thenReturn(mockIterable);
    when(mockIterable.iterator()).thenReturn(new StubbingCursor(documents.iterator()));
    GetTableRequest req = new GetTableRequest(IDENTITY, QUERY_ID, DEFAULT_CATALOG, TABLE_NAME);
    GetTableResponse res = handler.doGetTable(allocator, req);
    logger.info("doGetTable - {}", res);
    assertEquals(9, res.getSchema().getFields().size());
    Field stringCol = res.getSchema().findField("stringCol");
    assertEquals(Types.MinorType.VARCHAR, Types.getMinorTypeForArrowType(stringCol.getType()));
    Field stringCol2 = res.getSchema().findField("stringCol2");
    assertEquals(Types.MinorType.VARCHAR, Types.getMinorTypeForArrowType(stringCol2.getType()));
    Field intCol = res.getSchema().findField("intCol");
    assertEquals(Types.MinorType.INT, Types.getMinorTypeForArrowType(intCol.getType()));
    Field intCol2 = res.getSchema().findField("intCol2");
    assertEquals(Types.MinorType.INT, Types.getMinorTypeForArrowType(intCol2.getType()));
    Field doubleCol = res.getSchema().findField("doubleCol");
    assertEquals(Types.MinorType.FLOAT8, Types.getMinorTypeForArrowType(doubleCol.getType()));
    Field doubleCol2 = res.getSchema().findField("doubleCol2");
    assertEquals(Types.MinorType.FLOAT8, Types.getMinorTypeForArrowType(doubleCol2.getType()));
    Field longCol = res.getSchema().findField("longCol");
    assertEquals(Types.MinorType.BIGINT, Types.getMinorTypeForArrowType(longCol.getType()));
    Field longCol2 = res.getSchema().findField("longCol2");
    assertEquals(Types.MinorType.BIGINT, Types.getMinorTypeForArrowType(longCol2.getType()));
    Field unsupported = res.getSchema().findField("unsupported");
    assertEquals(Types.MinorType.VARCHAR, Types.getMinorTypeForArrowType(unsupported.getType()));
}

Example 23 with GetTableRequest

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

the class DocDBRecordHandlerTest method nestedStructTest.

@Test
public void nestedStructTest() throws Exception {
    List<Document> documents = new ArrayList<>();
    Document result = new Document();
    documents.add(result);
    Document listStruct1 = new Document();
    listStruct1.put("SomeSubStruct", "someSubStruct1");
    List<Document> subList = new ArrayList<>();
    Document listSubStruct1 = new Document();
    listSubStruct1.put("SomeSubSubStruct", "someSubSubStruct");
    subList.add(listSubStruct1);
    listStruct1.put("SomeSubList", subList);
    Document listStruct2 = new Document();
    listStruct2.put("SomeSubStruct1", "someSubStruct2");
    List<Document> list = new ArrayList<>();
    list.add(listStruct1);
    list.add(listStruct1);
    Document structWithList = new Document();
    structWithList.put("SomeList", list);
    Document structWithNullList = new Document();
    structWithNullList.put("SomeNullList", null);
    Document simpleSubStruct = new Document();
    simpleSubStruct.put("SomeSimpleSubStruct", "someSimpleSubStruct");
    structWithList.put("SimpleSubStruct", simpleSubStruct);
    structWithList.put("SimpleSubStructNullList", structWithNullList);
    result.put("ComplexStruct", structWithList);
    Document simpleStruct = new Document();
    simpleStruct.put("SomeSimpleStruct", "someSimpleStruct");
    result.put("SimpleStruct", simpleStruct);
    when(mockCollection.find()).thenReturn(mockIterable);
    when(mockIterable.limit(anyInt())).thenReturn(mockIterable);
    when(mockIterable.maxScan(anyInt())).thenReturn(mockIterable);
    when(mockIterable.batchSize(anyInt())).thenReturn(mockIterable);
    when(mockIterable.iterator()).thenReturn(new StubbingCursor(documents.iterator()));
    GetTableRequest req = new GetTableRequest(IDENTITY, QUERY_ID, DEFAULT_CATALOG, TABLE_NAME);
    GetTableResponse res = mdHandler.doGetTable(allocator, req);
    logger.info("doGetTable - {}", res);
    when(mockCollection.find(any(Document.class))).thenAnswer((InvocationOnMock invocationOnMock) -> {
        logger.info("doReadRecordsNoSpill: query[{}]", invocationOnMock.getArguments()[0]);
        return mockIterable;
    });
    when(mockIterable.projection(any(Document.class))).thenAnswer((InvocationOnMock invocationOnMock) -> {
        logger.info("doReadRecordsNoSpill: projection[{}]", invocationOnMock.getArguments()[0]);
        return mockIterable;
    });
    when(mockIterable.batchSize(anyInt())).thenReturn(mockIterable);
    when(mockIterable.iterator()).thenReturn(new StubbingCursor(documents.iterator()));
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    S3SpillLocation splitLoc = S3SpillLocation.newBuilder().withBucket(UUID.randomUUID().toString()).withSplitId(UUID.randomUUID().toString()).withQueryId(UUID.randomUUID().toString()).withIsDirectory(true).build();
    ReadRecordsRequest request = new ReadRecordsRequest(IDENTITY, DEFAULT_CATALOG, "queryId-" + System.currentTimeMillis(), TABLE_NAME, res.getSchema(), Split.newBuilder(splitLoc, keyFactory.create()).add(DOCDB_CONN_STR, CONNECTION_STRING).build(), new Constraints(constraintsMap), // 100GB don't expect this to spill
    100_000_000_000L, 100_000_000_000L);
    RecordResponse rawResponse = handler.doReadRecords(allocator, request);
    assertTrue(rawResponse instanceof ReadRecordsResponse);
    ReadRecordsResponse response = (ReadRecordsResponse) rawResponse;
    logger.info("doReadRecordsNoSpill: rows[{}]", response.getRecordCount());
    logger.info("doReadRecordsNoSpill: {}", BlockUtils.rowToString(response.getRecords(), 0));
    assertTrue(response.getRecordCount() == 1);
    String expectedString = "[ComplexStruct : {[SomeList : {{[SomeSubStruct : someSubStruct1]," + "[SomeSubList : {{[SomeSubSubStruct : someSubSubStruct]}}]}," + "{[SomeSubStruct : someSubStruct1],[SomeSubList : {{[SomeSubSubStruct : someSubSubStruct]}}]}}]," + "[SimpleSubStruct : {[SomeSimpleSubStruct : someSimpleSubStruct]}]," + "[SimpleSubStructNullList : {[SomeNullList : null]}]}], [SimpleStruct : {[SomeSimpleStruct : someSimpleStruct]}]";
    assertEquals(expectedString, BlockUtils.rowToString(response.getRecords(), 0));
}

Example 24 with GetTableRequest

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

the class MetricsMetadataHandlerTest method doGetMetricSamplesTable.

@Test
public void doGetMetricSamplesTable() {
    logger.info("doGetMetricSamplesTable - enter");
    GetTableRequest metricsTableReq = new GetTableRequest(identity, "queryId", "default", new TableName(defaultSchema, "metric_samples"));
    GetTableResponse metricsTableRes = handler.doGetTable(allocator, metricsTableReq);
    logger.info("doGetMetricSamplesTable - {} {}", metricsTableRes.getTableName(), metricsTableRes.getSchema());
    assertEquals(new TableName(defaultSchema, "metric_samples"), metricsTableRes.getTableName());
    assertNotNull(metricsTableRes.getSchema());
    assertEquals(9, metricsTableRes.getSchema().getFields().size());
    logger.info("doGetMetricSamplesTable - exit");
}

Example 25 with GetTableRequest

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

the class DataLakeGen2MetadataHandlerTest method doGetTable.

@Test
public void doGetTable() throws Exception {
    BlockAllocator blockAllocator = new BlockAllocatorImpl();
    String[] schema = { "DATA_TYPE", "COLUMN_SIZE", "COLUMN_NAME", "DECIMAL_DIGITS", "NUM_PREC_RADIX" };
    Object[][] values = { { Types.INTEGER, 12, "testCol1", 0, 0 }, { Types.VARCHAR, 25, "testCol2", 0, 0 }, { Types.TIMESTAMP, 93, "testCol3", 0, 0 }, { Types.TIMESTAMP_WITH_TIMEZONE, 93, "testCol4", 0, 0 } };
    AtomicInteger rowNumber = new AtomicInteger(-1);
    ResultSet resultSet = mockResultSet(schema, values, rowNumber);
    SchemaBuilder expectedSchemaBuilder = SchemaBuilder.newBuilder();
    expectedSchemaBuilder.addField(FieldBuilder.newBuilder("testCol1", org.apache.arrow.vector.types.Types.MinorType.INT.getType()).build());
    expectedSchemaBuilder.addField(FieldBuilder.newBuilder("testCol2", org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build());
    expectedSchemaBuilder.addField(FieldBuilder.newBuilder("testCol3", org.apache.arrow.vector.types.Types.MinorType.DATEMILLI.getType()).build());
    expectedSchemaBuilder.addField(FieldBuilder.newBuilder("testCol4", org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build());
    PARTITION_SCHEMA.getFields().forEach(expectedSchemaBuilder::addField);
    Schema expected = expectedSchemaBuilder.build();
    TableName inputTableName = new TableName("TESTSCHEMA", "TESTTABLE");
    Mockito.when(connection.getMetaData().getColumns("testCatalog", inputTableName.getSchemaName(), inputTableName.getTableName(), null)).thenReturn(resultSet);
    Mockito.when(connection.getCatalog()).thenReturn("testCatalog");
    GetTableResponse getTableResponse = this.dataLakeGen2MetadataHandler.doGetTable(blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName));
    Assert.assertEquals(expected, getTableResponse.getSchema());
    Assert.assertEquals(inputTableName, getTableResponse.getTableName());
    Assert.assertEquals("testCatalog", getTableResponse.getCatalogName());
}