Search in sources :

Example 1 with ListSchemasResponse

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

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

the class GlueMetadataHandlerTest method doListSchemaNames.

@Test
public void doListSchemaNames() throws Exception {
    List<Database> databases = new ArrayList<>();
    databases.add(new Database().withName("db1"));
    databases.add(new Database().withName("db2"));
    when(mockGlue.getDatabases(any(GetDatabasesRequest.class))).thenAnswer((InvocationOnMock invocationOnMock) -> {
        GetDatabasesRequest request = (GetDatabasesRequest) invocationOnMock.getArguments()[0];
        assertEquals(accountId, request.getCatalogId());
        GetDatabasesResult mockResult = mock(GetDatabasesResult.class);
        if (request.getNextToken() == null) {
            when(mockResult.getDatabaseList()).thenReturn(databases);
            when(mockResult.getNextToken()).thenReturn("next");
        } else {
            // only return real info on 1st call
            when(mockResult.getDatabaseList()).thenReturn(new ArrayList<>());
            when(mockResult.getNextToken()).thenReturn(null);
        }
        return mockResult;
    });
    ListSchemasRequest req = new ListSchemasRequest(IdentityUtil.fakeIdentity(), queryId, catalog);
    ListSchemasResponse res = handler.doListSchemaNames(allocator, req);
    logger.info("doListSchemas - {}", res.getSchemas());
    assertEquals(databases.stream().map(next -> next.getName()).collect(Collectors.toList()), new ArrayList<>(res.getSchemas()));
    verify(mockGlue, times(2)).getDatabases(any(GetDatabasesRequest.class));
}
Also used : ListSchemasRequest(com.amazonaws.athena.connector.lambda.metadata.ListSchemasRequest) InvocationOnMock(org.mockito.invocation.InvocationOnMock) Database(com.amazonaws.services.glue.model.Database) ArrayList(java.util.ArrayList) GetDatabasesResult(com.amazonaws.services.glue.model.GetDatabasesResult) GetDatabasesRequest(com.amazonaws.services.glue.model.GetDatabasesRequest) ListSchemasResponse(com.amazonaws.athena.connector.lambda.metadata.ListSchemasResponse) Test(org.junit.Test)

Example 3 with ListSchemasResponse

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

the class ExampleMetadataHandlerTest method doListSchemas.

@Test
public void doListSchemas() {
    logger.info("doListSchemas - enter");
    ListSchemasRequest req = new ListSchemasRequest(IdentityUtil.fakeIdentity(), "queryId", "default");
    ObjectMapperUtil.assertSerialization(req);
    ListSchemasResponse res = metadataHandler.doListSchemaNames(allocator, req);
    ObjectMapperUtil.assertSerialization(res);
    logger.info("doListSchemas - {}", res.getSchemas());
    assertFalse(res.getSchemas().isEmpty());
    logger.info("doListSchemas - exit");
}
Also used : ListSchemasRequest(com.amazonaws.athena.connector.lambda.metadata.ListSchemasRequest) ListSchemasResponse(com.amazonaws.athena.connector.lambda.metadata.ListSchemasResponse) Test(org.junit.Test)

Example 4 with ListSchemasResponse

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

the class HbaseMetadataHandlerTest method doListSchemaNames.

@Test
public void doListSchemaNames() throws IOException {
    NamespaceDescriptor[] schemaNames = { NamespaceDescriptor.create("schema1").build(), NamespaceDescriptor.create("schema2").build(), NamespaceDescriptor.create("schema3").build() };
    when(mockClient.listNamespaceDescriptors()).thenReturn(schemaNames);
    ListSchemasRequest req = new ListSchemasRequest(IDENTITY, QUERY_ID, DEFAULT_CATALOG);
    ListSchemasResponse res = handler.doListSchemaNames(allocator, req);
    logger.info("doListSchemas - {}", res.getSchemas());
    Set<String> expectedSchemaName = new HashSet<>();
    expectedSchemaName.add("schema1");
    expectedSchemaName.add("schema2");
    expectedSchemaName.add("schema3");
    assertEquals(expectedSchemaName, new HashSet<>(res.getSchemas()));
    logger.info("doListSchemaNames: exit");
}
Also used : ListSchemasRequest(com.amazonaws.athena.connector.lambda.metadata.ListSchemasRequest) NamespaceDescriptor(org.apache.hadoop.hbase.NamespaceDescriptor) ListSchemasResponse(com.amazonaws.athena.connector.lambda.metadata.ListSchemasResponse) Matchers.anyString(org.mockito.Matchers.anyString) HashSet(java.util.HashSet) Test(org.junit.Test)

Example 5 with ListSchemasResponse

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

the class BigQueryMetadataHandler method doListSchemaNames.

@Override
public ListSchemasResponse doListSchemaNames(BlockAllocator blockAllocator, ListSchemasRequest listSchemasRequest) {
    try {
        logger.info("doListSchemaNames called with Catalog: {}", listSchemasRequest.getCatalogName());
        final List<String> schemas = new ArrayList<>();
        final String projectName = BigQueryUtils.getProjectName(listSchemasRequest);
        Page<Dataset> response = bigQuery.listDatasets(projectName, BigQuery.DatasetListOption.pageSize(100));
        if (response == null) {
            logger.info("Dataset does not contain any models: {}");
        } else {
            do {
                for (Dataset dataset : response.iterateAll()) {
                    if (schemas.size() > BigQueryConstants.MAX_RESULTS) {
                        throw new BigQueryExceptions.TooManyTablesException();
                    }
                    schemas.add(dataset.getDatasetId().getDataset().toLowerCase());
                    logger.debug("Found Dataset: {}", dataset.getDatasetId().getDataset());
                }
            } while (response.hasNextPage());
        }
        logger.info("Found {} schemas!", schemas.size());
        return new ListSchemasResponse(listSchemasRequest.getCatalogName(), schemas);
    } catch (Exception e) {
        logger.error("Error: ", e);
    }
    return null;
}
Also used : Dataset(com.google.cloud.bigquery.Dataset) ArrayList(java.util.ArrayList) ListSchemasResponse(com.amazonaws.athena.connector.lambda.metadata.ListSchemasResponse) IOException(java.io.IOException)

Aggregations

ListSchemasResponse (com.amazonaws.athena.connector.lambda.metadata.ListSchemasResponse)26 ListSchemasRequest (com.amazonaws.athena.connector.lambda.metadata.ListSchemasRequest)16 Test (org.junit.Test)15 ArrayList (java.util.ArrayList)11 GetSplitsResponse (com.amazonaws.athena.connector.lambda.metadata.GetSplitsResponse)3 GetTableLayoutResponse (com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutResponse)3 GetTableResponse (com.amazonaws.athena.connector.lambda.metadata.GetTableResponse)3 ListTablesResponse (com.amazonaws.athena.connector.lambda.metadata.ListTablesResponse)3 Database (com.amazonaws.services.glue.model.Database)3 GetDatabasesResult (com.amazonaws.services.glue.model.GetDatabasesResult)3 BlockAllocatorImpl (com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl)2 TableName (com.amazonaws.athena.connector.lambda.domain.TableName)2 ReadRecordsResponse (com.amazonaws.athena.connector.lambda.records.ReadRecordsResponse)2 GetDatabasesRequest (com.amazonaws.services.glue.model.GetDatabasesRequest)2 DescribeLogGroupsRequest (com.amazonaws.services.logs.model.DescribeLogGroupsRequest)2 DescribeLogGroupsResult (com.amazonaws.services.logs.model.DescribeLogGroupsResult)2 Database (com.amazonaws.services.timestreamwrite.model.Database)2 ListDatabasesRequest (com.amazonaws.services.timestreamwrite.model.ListDatabasesRequest)2 ListDatabasesResult (com.amazonaws.services.timestreamwrite.model.ListDatabasesResult)2 InvocationOnMock (org.mockito.invocation.InvocationOnMock)2