Search in sources :

Example 1 with GetTableResponse

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

the class CompositeHandlerTest method setUp.

@Before
public void setUp() throws Exception {
    logger.info("{}: enter", testName.getMethodName());
    allocator = new BlockAllocatorImpl();
    objectMapper = ObjectMapperFactory.create(allocator);
    mockMetadataHandler = mock(MetadataHandler.class);
    mockRecordHandler = mock(RecordHandler.class);
    schemaForRead = SchemaBuilder.newBuilder().addField("col1", new ArrowType.Int(32, true)).build();
    when(mockMetadataHandler.doGetTableLayout(any(BlockAllocatorImpl.class), any(GetTableLayoutRequest.class))).thenReturn(new GetTableLayoutResponse("catalog", new TableName("schema", "table"), BlockUtils.newBlock(allocator, "col1", Types.MinorType.BIGINT.getType(), 1L)));
    when(mockMetadataHandler.doListTables(any(BlockAllocatorImpl.class), any(ListTablesRequest.class))).thenReturn(new ListTablesResponse("catalog", Collections.singletonList(new TableName("schema", "table")), null));
    when(mockMetadataHandler.doGetTable(any(BlockAllocatorImpl.class), any(GetTableRequest.class))).thenReturn(new GetTableResponse("catalog", new TableName("schema", "table"), SchemaBuilder.newBuilder().addStringField("col1").build()));
    when(mockMetadataHandler.doListSchemaNames(any(BlockAllocatorImpl.class), any(ListSchemasRequest.class))).thenReturn(new ListSchemasResponse("catalog", Collections.singleton("schema1")));
    when(mockMetadataHandler.doGetSplits(any(BlockAllocatorImpl.class), any(GetSplitsRequest.class))).thenReturn(new GetSplitsResponse("catalog", Split.newBuilder(null, null).build()));
    when(mockMetadataHandler.doPing(any(PingRequest.class))).thenReturn(new PingResponse("catalog", "queryId", "type", 23, 2));
    when(mockRecordHandler.doReadRecords(any(BlockAllocatorImpl.class), any(ReadRecordsRequest.class))).thenReturn(new ReadRecordsResponse("catalog", BlockUtils.newEmptyBlock(allocator, "col", new ArrowType.Int(32, true))));
    compositeHandler = new CompositeHandler(mockMetadataHandler, mockRecordHandler);
}
Also used : PingRequest(com.amazonaws.athena.connector.lambda.request.PingRequest) GetSplitsRequest(com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest) ReadRecordsResponse(com.amazonaws.athena.connector.lambda.records.ReadRecordsResponse) ArrowType(org.apache.arrow.vector.types.pojo.ArrowType) PingResponse(com.amazonaws.athena.connector.lambda.request.PingResponse) TableName(com.amazonaws.athena.connector.lambda.domain.TableName) GetTableRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableRequest) ListSchemasRequest(com.amazonaws.athena.connector.lambda.metadata.ListSchemasRequest) ReadRecordsRequest(com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest) GetTableLayoutResponse(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutResponse) BlockAllocatorImpl(com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl) GetTableResponse(com.amazonaws.athena.connector.lambda.metadata.GetTableResponse) GetSplitsResponse(com.amazonaws.athena.connector.lambda.metadata.GetSplitsResponse) GetTableLayoutRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest) ListTablesRequest(com.amazonaws.athena.connector.lambda.metadata.ListTablesRequest) ListTablesResponse(com.amazonaws.athena.connector.lambda.metadata.ListTablesResponse) ListSchemasResponse(com.amazonaws.athena.connector.lambda.metadata.ListSchemasResponse) Before(org.junit.Before)

Example 2 with GetTableResponse

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

the class ExampleMetadataHandlerTest method doGetTable.

@Test
public void doGetTable() {
    logger.info("doGetTable - enter");
    GetTableRequest req = new GetTableRequest(IdentityUtil.fakeIdentity(), "queryId", "default", new TableName("custom_source", "fake_table"));
    ObjectMapperUtil.assertSerialization(req);
    GetTableResponse res = metadataHandler.doGetTable(allocator, req);
    ObjectMapperUtil.assertSerialization(res);
    assertTrue(res.getSchema().getFields().size() > 0);
    assertTrue(res.getSchema().getCustomMetadata().size() > 0);
    logger.info("doGetTable - {}", res);
    logger.info("doGetTable - exit");
}
Also used : GetTableRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableRequest) TableName(com.amazonaws.athena.connector.lambda.domain.TableName) GetTableResponse(com.amazonaws.athena.connector.lambda.metadata.GetTableResponse) Test(org.junit.Test)

Example 3 with GetTableResponse

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

the class GetTableResponseSerDeTest method deserialize.

@Test
public void deserialize() throws IOException {
    logger.info("deserialize: enter");
    InputStream input = new ByteArrayInputStream(expectedSerDeText.getBytes());
    GetTableResponse actual = (GetTableResponse) mapper.readValue(input, FederationResponse.class);
    logger.info("deserialize: deserialized[{}]", actual);
    assertEquals(expected, actual);
    logger.info("deserialize: exit");
}
Also used : ByteArrayInputStream(java.io.ByteArrayInputStream) GetTableResponse(com.amazonaws.athena.connector.lambda.metadata.GetTableResponse) ByteArrayInputStream(java.io.ByteArrayInputStream) InputStream(java.io.InputStream) FederationResponse(com.amazonaws.athena.connector.lambda.request.FederationResponse) TypedSerDeTest(com.amazonaws.athena.connector.lambda.serde.TypedSerDeTest) Test(org.junit.Test)

Example 4 with GetTableResponse

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

the class BigQueryMetadataHandler method doGetTable.

@Override
public GetTableResponse doGetTable(BlockAllocator blockAllocator, GetTableRequest getTableRequest) {
    try {
        logger.info("doGetTable called with request {}:{}", BigQueryUtils.getProjectName(getTableRequest), getTableRequest.getTableName());
        final Schema tableSchema = getSchema(BigQueryUtils.getProjectName(getTableRequest), getTableRequest.getTableName().getSchemaName(), getTableRequest.getTableName().getTableName());
        return new GetTableResponse(BigQueryUtils.getProjectName(getTableRequest).toLowerCase(), getTableRequest.getTableName(), tableSchema);
    } catch (Exception e) {
        logger.error("Error: ", e);
    }
    return null;
}
Also used : GetTableResponse(com.amazonaws.athena.connector.lambda.metadata.GetTableResponse) Schema(org.apache.arrow.vector.types.pojo.Schema) IOException(java.io.IOException)

Example 5 with GetTableResponse

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

the class HbaseMetadataHandler method doGetTable.

/**
 * If Glue is enabled as a source of supplemental metadata we look up the requested Schema/Table in Glue and
 * filters out any results that don't have the HBASE_METADATA_FLAG set. If no matching results were found in Glue,
 * then we resort to inferring the schema of the HBase table using HbaseSchemaUtils.inferSchema(...). If there
 * is no such table in HBase the operation will fail.
 *
 * @see GlueMetadataHandler
 */
@Override
public GetTableResponse doGetTable(BlockAllocator blockAllocator, GetTableRequest request) throws Exception {
    logger.info("doGetTable: enter", request.getTableName());
    Schema origSchema = null;
    try {
        if (awsGlue != null) {
            origSchema = super.doGetTable(blockAllocator, request, TABLE_FILTER).getSchema();
        }
    } catch (RuntimeException ex) {
        logger.warn("doGetTable: Unable to retrieve table[{}:{}] from AWSGlue.", request.getTableName().getSchemaName(), request.getTableName().getTableName(), ex);
    }
    if (origSchema == null) {
        origSchema = HbaseSchemaUtils.inferSchema(getOrCreateConn(request), request.getTableName(), NUM_ROWS_TO_SCAN);
    }
    SchemaBuilder schemaBuilder = SchemaBuilder.newBuilder();
    origSchema.getFields().forEach((Field field) -> schemaBuilder.addField(field.getName(), field.getType(), field.getChildren()));
    origSchema.getCustomMetadata().entrySet().forEach((Map.Entry<String, String> meta) -> schemaBuilder.addMetadata(meta.getKey(), meta.getValue()));
    schemaBuilder.addField(HbaseSchemaUtils.ROW_COLUMN_NAME, Types.MinorType.VARCHAR.getType());
    Schema schema = schemaBuilder.build();
    logger.info("doGetTable: return {}", schema);
    return new GetTableResponse(request.getCatalogName(), request.getTableName(), schema);
}
Also used : Field(org.apache.arrow.vector.types.pojo.Field) GetTableResponse(com.amazonaws.athena.connector.lambda.metadata.GetTableResponse) Schema(org.apache.arrow.vector.types.pojo.Schema) SchemaBuilder(com.amazonaws.athena.connector.lambda.data.SchemaBuilder)

Aggregations

GetTableResponse (com.amazonaws.athena.connector.lambda.metadata.GetTableResponse)56 TableName (com.amazonaws.athena.connector.lambda.domain.TableName)37 GetTableRequest (com.amazonaws.athena.connector.lambda.metadata.GetTableRequest)34 Test (org.junit.Test)32 Schema (org.apache.arrow.vector.types.pojo.Schema)25 Field (org.apache.arrow.vector.types.pojo.Field)16 ArrayList (java.util.ArrayList)14 SchemaBuilder (com.amazonaws.athena.connector.lambda.data.SchemaBuilder)12 GetTableResult (com.amazonaws.services.glue.model.GetTableResult)11 Column (com.amazonaws.services.glue.model.Column)10 StorageDescriptor (com.amazonaws.services.glue.model.StorageDescriptor)9 Constraints (com.amazonaws.athena.connector.lambda.domain.predicate.Constraints)8 Table (com.amazonaws.services.glue.model.Table)8 InvocationOnMock (org.mockito.invocation.InvocationOnMock)8 BlockAllocatorImpl (com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl)7 HashMap (java.util.HashMap)7 GetTableLayoutResponse (com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutResponse)6 ReadRecordsRequest (com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest)6 ReadRecordsResponse (com.amazonaws.athena.connector.lambda.records.ReadRecordsResponse)6 Connection (java.sql.Connection)6