Search in sources :

Example 1 with ListTablesResponse

use of com.amazonaws.athena.connector.lambda.metadata.ListTablesResponse 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 ListTablesResponse

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

the class GlueMetadataHandlerTest method doListTablesWithLargePageSize.

@Test
public void doListTablesWithLargePageSize() throws Exception {
    ListTablesRequest req = new ListTablesRequest(IdentityUtil.fakeIdentity(), queryId, catalog, schema, null, GET_TABLES_REQUEST_MAX_RESULTS + 50);
    logger.info("Request - {}", req);
    ListTablesResponse actualResponse = handler.doListTables(allocator, req);
    logger.info("Response - {}", actualResponse);
    assertEquals("Lists do not match.", fullListResponse, actualResponse);
}
Also used : ListTablesRequest(com.amazonaws.athena.connector.lambda.metadata.ListTablesRequest) ListTablesResponse(com.amazonaws.athena.connector.lambda.metadata.ListTablesResponse) Test(org.junit.Test)

Example 3 with ListTablesResponse

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

the class GlueMetadataHandlerTest method doListTablesWithPagination.

@Test
public void doListTablesWithPagination() throws Exception {
    logger.info("First paginated request");
    ListTablesRequest req = new ListTablesRequest(IdentityUtil.fakeIdentity(), queryId, catalog, schema, null, 3);
    logger.info("Request - {}", req);
    ListTablesResponse expectedResponse = new ListTablesResponse(req.getCatalogName(), new ImmutableList.Builder<TableName>().add(new TableName(req.getSchemaName(), "table1")).add(new TableName(req.getSchemaName(), "table2")).add(new TableName(req.getSchemaName(), "table3")).build(), "table4");
    ListTablesResponse actualResponse = handler.doListTables(allocator, req);
    logger.info("Response - {}", actualResponse);
    assertEquals("Lists do not match.", expectedResponse, actualResponse);
    logger.info("Second paginated request");
    req = new ListTablesRequest(IdentityUtil.fakeIdentity(), queryId, catalog, schema, actualResponse.getNextToken(), 3);
    logger.info("Request - {}", req);
    expectedResponse = new ListTablesResponse(req.getCatalogName(), new ImmutableList.Builder<TableName>().add(new TableName(req.getSchemaName(), "table4")).add(new TableName(req.getSchemaName(), "table5")).build(), null);
    actualResponse = handler.doListTables(allocator, req);
    logger.info("Response - {}", actualResponse);
    assertEquals("Lists do not match.", expectedResponse, actualResponse);
}
Also used : GlueMetadataHandler.getSourceTableName(com.amazonaws.athena.connector.lambda.handlers.GlueMetadataHandler.getSourceTableName) TableName(com.amazonaws.athena.connector.lambda.domain.TableName) ImmutableList(com.google.common.collect.ImmutableList) ListTablesRequest(com.amazonaws.athena.connector.lambda.metadata.ListTablesRequest) ListTablesResponse(com.amazonaws.athena.connector.lambda.metadata.ListTablesResponse) Test(org.junit.Test)

Example 4 with ListTablesResponse

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

the class GlueMetadataHandlerTest method doListTablesWithUnlimitedPageSize.

@Test
public void doListTablesWithUnlimitedPageSize() throws Exception {
    ListTablesRequest req = new ListTablesRequest(IdentityUtil.fakeIdentity(), queryId, catalog, schema, null, UNLIMITED_PAGE_SIZE_VALUE);
    logger.info("Request - {}", req);
    ListTablesResponse actualResponse = handler.doListTables(allocator, req);
    logger.info("Response - {}", actualResponse);
    assertEquals("Lists do not match.", fullListResponse, actualResponse);
}
Also used : ListTablesRequest(com.amazonaws.athena.connector.lambda.metadata.ListTablesRequest) ListTablesResponse(com.amazonaws.athena.connector.lambda.metadata.ListTablesResponse) Test(org.junit.Test)

Example 5 with ListTablesResponse

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

the class BigQueryMetadataHandler method doListTables.

@Override
public ListTablesResponse doListTables(BlockAllocator blockAllocator, ListTablesRequest listTablesRequest) {
    try {
        logger.info("doListTables called with request {}:{}", listTablesRequest.getCatalogName(), listTablesRequest.getSchemaName());
        // Get the project name, dataset name, and dataset id. Google BigQuery is case sensitive.
        final String projectName = BigQueryUtils.getProjectName(listTablesRequest);
        final String datasetName = fixCaseForDatasetName(projectName, listTablesRequest.getSchemaName(), bigQuery);
        final DatasetId datasetId = DatasetId.of(projectName, datasetName);
        Page<Table> response = bigQuery.listTables(datasetId, BigQuery.TableListOption.pageSize(100));
        List<TableName> tables = new ArrayList<>();
        do {
            for (Table table : response.iterateAll()) {
                if (tables.size() > BigQueryConstants.MAX_RESULTS) {
                    throw new BigQueryExceptions.TooManyTablesException();
                }
                tables.add(new TableName(listTablesRequest.getSchemaName(), table.getTableId().getTable().toLowerCase()));
            }
        } while (response.hasNextPage());
        logger.info("Found {} table(s)!", tables.size());
        return new ListTablesResponse(listTablesRequest.getCatalogName(), tables, null);
    } catch (Exception e) {
        logger.error("Error:", e);
    }
    return null;
}
Also used : TableName(com.amazonaws.athena.connector.lambda.domain.TableName) BigQueryUtils.fixCaseForTableName(com.amazonaws.athena.connectors.google.bigquery.BigQueryUtils.fixCaseForTableName) Table(com.google.cloud.bigquery.Table) ArrayList(java.util.ArrayList) ListTablesResponse(com.amazonaws.athena.connector.lambda.metadata.ListTablesResponse) IOException(java.io.IOException) DatasetId(com.google.cloud.bigquery.DatasetId)

Aggregations

ListTablesResponse (com.amazonaws.athena.connector.lambda.metadata.ListTablesResponse)31 TableName (com.amazonaws.athena.connector.lambda.domain.TableName)24 ListTablesRequest (com.amazonaws.athena.connector.lambda.metadata.ListTablesRequest)18 Test (org.junit.Test)17 ArrayList (java.util.ArrayList)13 Table (com.amazonaws.services.glue.model.Table)5 SchemaBuilder (com.amazonaws.athena.connector.lambda.data.SchemaBuilder)4 GetSplitsResponse (com.amazonaws.athena.connector.lambda.metadata.GetSplitsResponse)4 GetTableLayoutResponse (com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutResponse)4 GetTableResponse (com.amazonaws.athena.connector.lambda.metadata.GetTableResponse)4 ListSchemasResponse (com.amazonaws.athena.connector.lambda.metadata.ListSchemasResponse)4 GetTablesResult (com.amazonaws.services.glue.model.GetTablesResult)4 ImmutableList (com.google.common.collect.ImmutableList)4 BlockAllocatorImpl (com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl)3 Before (org.junit.Before)3 BlockAllocator (com.amazonaws.athena.connector.lambda.data.BlockAllocator)2 Split (com.amazonaws.athena.connector.lambda.domain.Split)2 Constraints (com.amazonaws.athena.connector.lambda.domain.predicate.Constraints)2 GetSplitsRequest (com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest)2 GetTableLayoutRequest (com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest)2