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

@Test
public void testDropDatabaseDeleteDataWithCatalogId() throws Exception {
  when(glueClient.getDatabase(any(GetDatabaseRequest.class))).thenReturn(
      new GetDatabaseResult().withDatabase(testDb));
  when(glueClient.getTables(any(GetTablesRequest.class))).thenReturn(
      new GetTablesResult().withTableList(ImmutableList.<Table>of()));
  Path dbPath = new Path(testDb.getLocationUri());
  when(wh.deleteDir(dbPath, true)).thenReturn(true);

  metastoreClientDelegateCatalogId.dropDatabase(testDb.getName(), true, false, false);
  ArgumentCaptor<DeleteDatabaseRequest> captor = ArgumentCaptor.forClass(DeleteDatabaseRequest.class);
  verify(glueClient, times(1)).deleteDatabase(captor.capture());
  DeleteDatabaseRequest request = captor.getValue();
  verify(wh, times(1)).deleteDir(dbPath, true);
  assertEquals(CATALOG_ID, request.getCatalogId());
  assertEquals(testDb.getName(), request.getName());
}
 

@Test
public void testGetTableWithCatalogId() throws Exception {
  Table tbl2 = getTestTable();
  List<String> tableNames = ImmutableList.of(testTbl.getName(), tbl2.getName());
  List<Table> tableList = ImmutableList.of(testTbl, tbl2);

  when(glueClient.getTables(new GetTablesRequest().withDatabaseName(testDb.getName()).withExpression("*").withCatalogId(CATALOG_ID)))
    .thenReturn(new GetTablesResult().withTableList(tableList));
  List<String> result = metastoreClientDelegateCatalogId.getTables(testDb.getName(), "*");

  assertThat(result, is(tableNames));
  ArgumentCaptor<GetTablesRequest> captor = ArgumentCaptor.forClass(GetTablesRequest.class);
  verify(glueClient, times(1)).getTables(captor.capture());
  assertEquals(CATALOG_ID, captor.getValue().getCatalogId());
  assertEquals(testDb.getName(), captor.getValue().getDatabaseName());
  assertEquals("*", captor.getValue().getExpression());
}
 

@Test
public void testGetTablesWithPagination() throws Exception {
  Table tbl2 = getTestTable();
  List<String> tableNames = ImmutableList.of(testTbl.getName(), tbl2.getName());
  List<Table> tableList1 = ImmutableList.of(testTbl);
  List<Table> tableList2 = ImmutableList.of(tbl2);

  String nextToken = "1";
  when(glueClient.getTables(any(GetTablesRequest.class)))
    .thenReturn(new GetTablesResult().withTableList(tableList1).withNextToken(nextToken))
    .thenReturn(new GetTablesResult().withTableList(tableList2));
  List<String> result = metastoreClientDelegate.getTables(testDb.getName(), "*");

  verify(glueClient, times(2)).getTables(any(GetTablesRequest.class));
  assertThat(result, is(tableNames));
}
 

@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 tables from AWS Glue DataCatalog with optional filtering for the requested schema (aka database)
 *
 * @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 TableFilter to apply to all tables before adding them to the results list.
 * @return The ListTablesResponse which mostly contains the list of table names.
 */
protected ListTablesResponse doListTables(BlockAllocator blockAllocator, ListTablesRequest request, TableFilter filter)
        throws Exception
{
    GetTablesRequest getTablesRequest = new GetTablesRequest();
    getTablesRequest.setCatalogId(getCatalog(request));
    getTablesRequest.setDatabaseName(request.getSchemaName());

    Set<TableName> tables = new HashSet<>();
    String nextToken = null;
    do {
        getTablesRequest.setNextToken(nextToken);
        GetTablesResult result = awsGlue.getTables(getTablesRequest);

        for (Table next : result.getTableList()) {
            if (filter == null || filter.filter(next)) {
                tables.add(new TableName(request.getSchemaName(), next.getName()));
            }
        }

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

    return new ListTablesResponse(request.getCatalogName(), tables);
}
 

@Override
public List<Table> getTables(String dbname, String tablePattern) {
    List<Table> ret = new ArrayList<>();
    String nextToken = null;
    do {
        GetTablesRequest getTablesRequest = new GetTablesRequest().withDatabaseName(dbname)
                .withExpression(tablePattern).withNextToken(nextToken).withCatalogId(catalogId);
        GetTablesResult result = glueClient.getTables(getTablesRequest);
        ret.addAll(result.getTableList());
        nextToken = result.getNextToken();
    } while (nextToken != null);
    return ret;
}
 

@Test
public void testDropDatabaseDeleteData() throws Exception {
  when(glueClient.getDatabase(any(GetDatabaseRequest.class))).thenReturn(
      new GetDatabaseResult().withDatabase(testDb));
  when(glueClient.getTables(any(GetTablesRequest.class))).thenReturn(
      new GetTablesResult().withTableList(ImmutableList.<Table>of()));
  Path dbPath = new Path(testDb.getLocationUri());
  when(wh.deleteDir(dbPath, true)).thenReturn(true);

  metastoreClientDelegate.dropDatabase(testDb.getName(), true, false, false);
  verify(glueClient, times(1)).deleteDatabase(any(DeleteDatabaseRequest.class));
  verify(wh, times(1)).deleteDir(dbPath, true);
}
 

@Test
public void testDropDatabaseKeepData() throws Exception {
  when(glueClient.getDatabase(any(GetDatabaseRequest.class))).thenReturn(
      new GetDatabaseResult().withDatabase(testDb));
  when(glueClient.getTables(any(GetTablesRequest.class))).thenReturn(
      new GetTablesResult().withTableList(ImmutableList.<Table>of()));
  Path dbPath = new Path(testDb.getLocationUri());
  when(wh.deleteDir(dbPath, true)).thenReturn(true);

  metastoreClientDelegate.dropDatabase(testDb.getName(), false, false, false);
  verify(glueClient, times(1)).deleteDatabase(any(DeleteDatabaseRequest.class));
  verify(wh, never()).deleteDir(dbPath, true);
}
 

@Test
public void testGetTables() throws Exception {
  Table tbl2 = getTestTable();
  List<String> tableNames = ImmutableList.of(testTbl.getName(), tbl2.getName());
  List<Table> tableList = ImmutableList.of(testTbl, tbl2);

  when(glueClient.getTables(new GetTablesRequest().withDatabaseName(testDb.getName()).withExpression("*")))
    .thenReturn(new GetTablesResult().withTableList(tableList));
  List<String> result = metastoreClientDelegate.getTables(testDb.getName(), "*");

  verify(glueClient).getTables(new GetTablesRequest().withDatabaseName(testDb.getName()).withExpression("*"));
  assertThat(result, is(tableNames));
}
 

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

@AfterClass
public static void cleanUp() {
  List<com.amazonaws.services.glue.model.Table> table_to_delete = new ArrayList<>();
  String token = null;
  do {
    GetTablesResult result = glueClient.getTables(new GetTablesRequest().withDatabaseName(
        catalogDB.getName()).withNextToken(token));
    table_to_delete.addAll(result.getTableList());
    token = result.getNextToken();
  } while(token != null);
  for(com.amazonaws.services.glue.model.Table table : table_to_delete) {
    glueClient.deleteTable(new DeleteTableRequest().withDatabaseName(catalogDB.getName()).withName(table.getName()));
  }
  glueClient.deleteDatabase(new DeleteDatabaseRequest().withName(catalogDB.getName()));
}
 

@After
public void cleanUpForClass() {
  String token = null;
  do {
    GetTablesResult result = glueClient.getTables(new GetTablesRequest().withDatabaseName(
        catalogDB.getName()).withNextToken(token));
    token = result.getNextToken();
    for (com.amazonaws.services.glue.model.Table table : result.getTableList()) {
      glueClient.deleteTable(new DeleteTableRequest().withDatabaseName(catalogDB.getName())
              .withName(table.getName()));
    }
  } while (token != null);
}
 

@Test
public void testGetAllTable() throws Exception {
  List<Table> result = ImmutableList.of(HiveToCatalogConverter.convertTable(testTable));
  when(glueClient.getTables(new GetTablesRequest().withDatabaseName(testDB.getName()).withExpression(".*")))
          .thenReturn(new GetTablesResult().withTableList(result));
  List<String> tableList = metastoreClient.getAllTables(testDB.getName());
  verify(glueClient).getTables(new GetTablesRequest().withDatabaseName(testDB.getName()).withExpression(".*"));
  assertEquals(Iterables.getOnlyElement(result).getName(), Iterables.getOnlyElement(tableList));
}
 

@Test
public void doListTablesGlueAndDynamo()
        throws Exception
{
    List<String> tableNames = new ArrayList<>();
    tableNames.add("table1");
    tableNames.add("table2");
    tableNames.add("table3");

    GetTablesResult mockResult = new GetTablesResult();
    List<Table> tableList = new ArrayList<>();
    tableList.add(new Table().withName("table1")
            .withParameters(ImmutableMap.of("classification", "dynamodb"))
            .withStorageDescriptor(new StorageDescriptor()
                    .withLocation("some.location")));
    tableList.add(new Table().withName("table2")
            .withParameters(ImmutableMap.of())
            .withStorageDescriptor(new StorageDescriptor()
                    .withLocation("some.location")
                    .withParameters(ImmutableMap.of("classification", "dynamodb"))));
    tableList.add(new Table().withName("table3")
            .withParameters(ImmutableMap.of())
            .withStorageDescriptor(new StorageDescriptor()
                    .withLocation("arn:aws:dynamodb:us-east-1:012345678910:table/table3")));
    tableList.add(new Table().withName("notADynamoTable").withParameters(ImmutableMap.of()).withStorageDescriptor(
            new StorageDescriptor().withParameters(ImmutableMap.of()).withLocation("some_location")));
    mockResult.setTableList(tableList);
    when(glueClient.getTables(any())).thenReturn(mockResult);

    ListTablesRequest req = new ListTablesRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, DEFAULT_SCHEMA);
    ListTablesResponse res = handler.doListTables(allocator, req);

    logger.info("doListTables - {}", res.getTables());

    List<TableName> expectedTables = tableNames.stream().map(table -> new TableName(DEFAULT_SCHEMA, table)).collect(Collectors.toList());
    expectedTables.add(TEST_TABLE_NAME);
    expectedTables.add(new TableName(DEFAULT_SCHEMA, "test_table2"));
    expectedTables.add(new TableName(DEFAULT_SCHEMA, "test_table3"));
    expectedTables.add(new TableName(DEFAULT_SCHEMA, "test_table4"));

    assertThat(new HashSet<>(res.getTables()), equalTo(new HashSet<>(expectedTables)));
}