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);
}
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));
}
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");
}
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");
}
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;
}
Aggregations