com.amazonaws.services.glue.model.GetDatabasesRequest Java Examples

@Test
public void testGetTableMetaNullEmptyTableType() throws Exception {
  List<Table> tables = Lists.newArrayList(testTbl);
  List<String> tableTypes = null;

  when(glueClient.getDatabases(any(GetDatabasesRequest.class))).thenReturn(
    new GetDatabasesResult().withDatabaseList(testDb));
  when(glueClient.getTables(any(GetTablesRequest.class))).thenReturn(
    new GetTablesResult().withTableList(tables));

  List<TableMeta> tableMetaResult = metastoreClientDelegate.getTableMeta(testDb.getName(), testTbl.getName(), tableTypes);
  assertEquals(CatalogToHiveConverter.convertTableMeta(testTbl, testDb.getName()), Iterables.getOnlyElement(tableMetaResult));

  tableTypes = Lists.newArrayList();
  tableMetaResult = metastoreClientDelegate.getTableMeta(testDb.getName(), testTbl.getName(), tableTypes);
  assertEquals(CatalogToHiveConverter.convertTableMeta(testTbl, testDb.getName()), Iterables.getOnlyElement(tableMetaResult));
}
 

/**
 * Returns a list of schemas (aka databases) from AWS Glue DataCatalog with optional filtering.
 *
 * @param blockAllocator Tool for creating and managing Apache Arrow Blocks.
 * @param request Provides details on who made the request and which Athena catalog they are querying.
 * @param filter The DatabaseFilter to apply to all schemas (aka databases) before adding them to the results list.
 * @return The ListSchemasResponse which mostly contains the list of schemas (aka databases).
 */
protected ListSchemasResponse doListSchemaNames(BlockAllocator blockAllocator, ListSchemasRequest request, DatabaseFilter filter)
        throws Exception
{
    GetDatabasesRequest getDatabasesRequest = new GetDatabasesRequest();
    getDatabasesRequest.setCatalogId(getCatalog(request));

    List<String> schemas = new ArrayList<>();
    String nextToken = null;
    do {
        getDatabasesRequest.setNextToken(nextToken);
        GetDatabasesResult result = awsGlue.getDatabases(getDatabasesRequest);

        for (Database next : result.getDatabaseList()) {
            if (filter == null || filter.filter(next)) {
                schemas.add(next.getName());
            }
        }

        nextToken = result.getNextToken();
    }
    while (nextToken != null);

    return new ListSchemasResponse(request.getCatalogName(), schemas);
}
 

@Override
public List<String> getAllDatabases()
{
    try {
        return stats.getGetAllDatabases().call(() -> {
            List<String> databaseNames = new ArrayList<>();
            String nextToken = null;

            do {
                GetDatabasesResult result = glueClient.getDatabases(new GetDatabasesRequest().withCatalogId(catalogId).withNextToken(nextToken));
                nextToken = result.getNextToken();
                result.getDatabaseList().forEach(database -> databaseNames.add(database.getName()));
            }
            while (nextToken != null);

            return databaseNames;
        });
    }
    catch (AmazonServiceException e) {
        throw new PrestoException(HIVE_METASTORE_ERROR, e);
    }
}
 

@Override
public List<Database> getAllDatabases() {
    List<Database> ret = Lists.newArrayList();
    String nextToken = null;
    do {
        GetDatabasesRequest getDatabasesRequest = new GetDatabasesRequest().withNextToken(nextToken).withCatalogId(
                catalogId);
        GetDatabasesResult result = glueClient.getDatabases(getDatabasesRequest);
        nextToken = result.getNextToken();
        ret.addAll(result.getDatabaseList());
    } while (nextToken != null);
    return ret;
}
 

@Test
public void testGetDatabases() throws Exception {
  when(glueClient.getDatabases(any(GetDatabasesRequest.class))).thenReturn(
      new GetDatabasesResult().withDatabaseList(testDb));

  List<String> dbs = metastoreClientDelegate.getDatabases("*");
  assertEquals(testDb.getName(), Iterables.getOnlyElement(dbs));
}
 

@Test
public void testGetDatabasesWithCatalogId() throws Exception {
  when(glueClient.getDatabases(any(GetDatabasesRequest.class))).thenReturn(
      new GetDatabasesResult().withDatabaseList(testDb));

  List<String> dbs = metastoreClientDelegateCatalogId.getDatabases("*");
  ArgumentCaptor<GetDatabasesRequest> captor = ArgumentCaptor.forClass(GetDatabasesRequest.class);
  verify(glueClient, times(1)).getDatabases(captor.capture());
  assertEquals(CATALOG_ID, captor.getValue().getCatalogId());
  assertEquals(testDb.getName(), Iterables.getOnlyElement(dbs));
}
 

@Test
public void testGetDatabasesNullPattern() throws Exception {
  when(glueClient.getDatabases(any(GetDatabasesRequest.class))).thenReturn(
      new GetDatabasesResult().withDatabaseList(testDb));

  List<String> dbs = metastoreClientDelegate.getDatabases(null);
  assertEquals(testDb.getName(), Iterables.getOnlyElement(dbs));
}
 

@Test
public void testGetAllDatabases() throws Exception {
  when(glueClient.getDatabases(any(GetDatabasesRequest.class))).thenReturn(
      new GetDatabasesResult().withDatabaseList(getTestDatabase()));
  metastoreClientDelegate.getDatabases("*");
  // Ensure this gets invoked
  verify(glueClient, atLeastOnce()).getDatabases(any(GetDatabasesRequest.class));
}
 

@Test
public void testGetAllDatabasesPaginated() throws Exception {
  when(glueClient.getDatabases(any(GetDatabasesRequest.class)))
      .thenReturn(new GetDatabasesResult().withDatabaseList(testDb).withNextToken("token"))
      .thenReturn(new GetDatabasesResult().withDatabaseList(getTestDatabase()));
  List<String> databases = metastoreClientDelegate.getDatabases(".*");

  assertEquals(2, databases.size());
  verify(glueClient, times(2)).getDatabases(any(GetDatabasesRequest.class));
}
 

@Test
public void testGetTableMeta() throws Exception {
  List<Table> tables = Lists.newArrayList(testTbl);
  List<String> tableTypes = Lists.newArrayList(TableType.MANAGED_TABLE.name());

  when(glueClient.getDatabases(any(GetDatabasesRequest.class))).thenReturn(
      new GetDatabasesResult().withDatabaseList(testDb));
  when(glueClient.getTables(any(GetTablesRequest.class))).thenReturn(
      new GetTablesResult().withTableList(tables));

  List<TableMeta> tableMetaResult = metastoreClientDelegate.getTableMeta(testDb.getName(), testTbl.getName(), tableTypes);
  assertEquals(CatalogToHiveConverter.convertTableMeta(testTbl, testDb.getName()), Iterables.getOnlyElement(tableMetaResult));
}
 

@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));
}