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

The following examples show how to use com.amazonaws.services.glue.model.StorageDescriptor. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example #1
Source File: AWSCatalogMetastoreClientTest.java    From aws-glue-data-catalog-client-for-apache-hive-metastore with Apache License 2.0 6 votes vote down vote up
@Test(expected=InvalidOperationException.class)
public void testRenamePartitionForUnknownTable() throws Exception {
  String dbName = testDB.getName();
  Table externalTable = getTestTable();
  externalTable.setTableType(TableType.EXTERNAL_TABLE.name());   
  
  StorageDescriptor sd = HiveToCatalogConverter.convertStorageDescriptor(testPartition.getSd());
  
  Partition oldPartition = new Partition()
      .withDatabaseName(dbName).withTableName(externalTable.getName())
      .withValues(Lists.newArrayList("oldval")).withStorageDescriptor(sd);
  
  Partition newPartition = new Partition()
                                    .withDatabaseName(dbName).withTableName(externalTable.getName())
                                    .withValues(Lists.newArrayList("newval")).withStorageDescriptor(sd);    
      
  when(glueClient.getDatabase(any(GetDatabaseRequest.class)))
    .thenReturn(new GetDatabaseResult().withDatabase(HiveToCatalogConverter.convertDatabase(testDB)));
  doThrow(EntityNotFoundException.class).when(glueClient).getTable(any(GetTableRequest.class));
  
  when(glueClient.getPartition(any(GetPartitionRequest.class)))
    .thenReturn(new GetPartitionResult().withPartition(oldPartition));     
  
  metastoreClient.renamePartition(dbName, externalTable.getName(), oldPartition.getValues(),
      CatalogToHiveConverter.convertPartition(newPartition));
}
 
Example #2
Source File: AWSCatalogMetastoreClientTest.java    From aws-glue-data-catalog-client-for-apache-hive-metastore with Apache License 2.0 6 votes vote down vote up
@Test(expected=InvalidOperationException.class)
public void testRenamePartitionForUnknownPartition() throws Exception {
  String dbName = testDB.getName();
  Table externalTable = getTestTable();
  externalTable.setTableType(TableType.EXTERNAL_TABLE.name());   
  
  StorageDescriptor sd = HiveToCatalogConverter.convertStorageDescriptor(testPartition.getSd());
  
  Partition oldPartition = new Partition()
      .withDatabaseName(dbName).withTableName(externalTable.getName())
      .withValues(Lists.newArrayList("oldval")).withStorageDescriptor(sd);
  
  Partition newPartition = new Partition()
                                    .withDatabaseName(dbName).withTableName(externalTable.getName())
                                    .withValues(Lists.newArrayList("newval")).withStorageDescriptor(sd);

  when(glueClient.getDatabase(any(GetDatabaseRequest.class)))
    .thenReturn(new GetDatabaseResult().withDatabase(HiveToCatalogConverter.convertDatabase(testDB)));
  when(glueClient.getTable(any(GetTableRequest.class)))
  .thenReturn(new GetTableResult().withTable(externalTable));
  doThrow(EntityNotFoundException.class).when(glueClient).getPartition(any(GetPartitionRequest.class));
  
  metastoreClient.renamePartition(dbName, externalTable.getName(), oldPartition.getValues(),
      CatalogToHiveConverter.convertPartition(newPartition));
}
 
Example #3
Source File: AWSCatalogMetastoreClientTest.java    From aws-glue-data-catalog-client-for-apache-hive-metastore with Apache License 2.0 6 votes vote down vote up
@Test(expected=InvalidOperationException.class)
public void testRenamePartitionForInvalidSD() throws Exception {
  String dbName = testDB.getName();
  Table externalTable = getTestTable();
  externalTable.setTableType(TableType.EXTERNAL_TABLE.name());   
  
  StorageDescriptor sd = HiveToCatalogConverter.convertStorageDescriptor(testPartition.getSd());
  
  Partition oldPartition = new Partition()
      .withDatabaseName(dbName).withTableName(externalTable.getName())
      .withValues(Lists.newArrayList("oldval")).withStorageDescriptor(null);
  
  Partition newPartition = new Partition()
                                    .withDatabaseName(dbName).withTableName(externalTable.getName())
                                    .withValues(Lists.newArrayList("newval")).withStorageDescriptor(sd);

  when(glueClient.getDatabase(any(GetDatabaseRequest.class)))
    .thenReturn(new GetDatabaseResult().withDatabase(HiveToCatalogConverter.convertDatabase(testDB)));
  when(glueClient.getTable(any(GetTableRequest.class)))
  .thenReturn(new GetTableResult().withTable(externalTable));
  when(glueClient.getPartition(any(GetPartitionRequest.class)))
    .thenReturn(new GetPartitionResult().withPartition(oldPartition));

  metastoreClient.renamePartition(dbName, externalTable.getName(), oldPartition.getValues(),
      CatalogToHiveConverter.convertPartition(newPartition));
}
 
Example #4
Source File: GlueToPrestoConverter.java    From presto with Apache License 2.0 6 votes vote down vote up
public static Partition convertPartition(com.amazonaws.services.glue.model.Partition gluePartition, Map<String, String> tableParameters)
{
    requireNonNull(gluePartition.getStorageDescriptor(), "Partition StorageDescriptor is null");
    StorageDescriptor sd = gluePartition.getStorageDescriptor();

    Partition.Builder partitionBuilder = Partition.builder()
            .setDatabaseName(gluePartition.getDatabaseName())
            .setTableName(gluePartition.getTableName())
            .setValues(gluePartition.getValues())
            .setColumns(sd.getColumns().stream()
                    .map(GlueToPrestoConverter::convertColumn)
                    .collect(toImmutableList()))
            .setParameters(firstNonNull(gluePartition.getParameters(), ImmutableMap.of()));

    setStorageBuilder(sd, partitionBuilder.getStorageBuilder(), tableParameters);
    return partitionBuilder.build();
}
 
Example #5
Source File: TestObjects.java    From aws-glue-data-catalog-client-for-apache-hive-metastore with Apache License 2.0 6 votes vote down vote up
public static StorageDescriptor getTestStorageDescriptor() {
  StorageDescriptor sd = new StorageDescriptor();
  List<String> cols = new ArrayList<>();
  cols.add("sampleCols");
  sd.setBucketColumns(cols);
  sd.setColumns(getTestFieldList());
  sd.setParameters(new HashMap<String, String>());
  sd.setSerdeInfo(getTestSerdeInfo());
  sd.setSkewedInfo(getSkewedInfo());
  sd.setSortColumns(new ArrayList<Order>());
  sd.setInputFormat("inputFormat");
  sd.setOutputFormat("outputFormat");
  sd.setLocation("/test-table");
  sd.withSortColumns(new Order().withColumn("foo").withSortOrder(1));
  sd.setCompressed(false);
  sd.setStoredAsSubDirectories(false);
  sd.setNumberOfBuckets(0);
  return sd;
}
 
Example #6
Source File: DynamoDBMetadataHandlerTest.java    From aws-athena-query-federation with Apache License 2.0 5 votes vote down vote up
@Test
public void validateSourceTableNamePropagation()
        throws Exception
{
    List<Column> columns = new ArrayList<>();
    columns.add(new Column().withName("col1").withType("int"));
    columns.add(new Column().withName("col2").withType("bigint"));
    columns.add(new Column().withName("col3").withType("string"));

    Map<String, String> param = ImmutableMap.of(
            SOURCE_TABLE_PROPERTY, TEST_TABLE,
            COLUMN_NAME_MAPPING_PROPERTY, "col1=Col1 , col2=Col2 ,col3=Col3",
            DATETIME_FORMAT_MAPPING_PROPERTY, "col1=datetime1,col3=datetime3 ");
    Table table = new Table()
            .withParameters(param)
            .withPartitionKeys()
            .withStorageDescriptor(new StorageDescriptor().withColumns(columns));
    GetTableResult mockResult = new GetTableResult().withTable(table);
    when(glueClient.getTable(any())).thenReturn(mockResult);

    TableName tableName = new TableName(DEFAULT_SCHEMA, "glueTableForTestTable");
    GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName);
    GetTableResponse getTableResponse = handler.doGetTable(allocator, getTableRequest);
    logger.info("validateSourceTableNamePropagation: GetTableResponse[{}]", getTableResponse);
    Map<String, String> customMetadata = getTableResponse.getSchema().getCustomMetadata();
    assertThat(customMetadata.get(SOURCE_TABLE_PROPERTY), equalTo(TEST_TABLE));
    assertThat(customMetadata.get(DATETIME_FORMAT_MAPPING_PROPERTY_NORMALIZED), equalTo("Col1=datetime1,Col3=datetime3"));

    GetTableLayoutRequest getTableLayoutRequest = new GetTableLayoutRequest(TEST_IDENTITY,
            TEST_QUERY_ID,
            TEST_CATALOG_NAME,
            tableName,
            new Constraints(ImmutableMap.of()),
            getTableResponse.getSchema(),
            Collections.EMPTY_SET);

    GetTableLayoutResponse getTableLayoutResponse = handler.doGetTableLayout(allocator, getTableLayoutRequest);
    logger.info("validateSourceTableNamePropagation: GetTableLayoutResponse[{}]", getTableLayoutResponse);
    assertThat(getTableLayoutResponse.getPartitions().getSchema().getCustomMetadata().get(TABLE_METADATA), equalTo(TEST_TABLE));
}
 
Example #7
Source File: TestGlueToPrestoConverter.java    From presto with Apache License 2.0 5 votes vote down vote up
private static void assertStorage(Storage actual, StorageDescriptor expected)
{
    assertEquals(actual.getLocation(), expected.getLocation());
    assertEquals(actual.getStorageFormat().getSerDe(), expected.getSerdeInfo().getSerializationLibrary());
    assertEquals(actual.getStorageFormat().getInputFormat(), expected.getInputFormat());
    assertEquals(actual.getStorageFormat().getOutputFormat(), expected.getOutputFormat());
    if (!isNullOrEmpty(expected.getBucketColumns())) {
        HiveBucketProperty bucketProperty = actual.getBucketProperty().get();
        assertEquals(bucketProperty.getBucketedBy(), expected.getBucketColumns());
        assertEquals(bucketProperty.getBucketCount(), expected.getNumberOfBuckets().intValue());
    }
}
 
Example #8
Source File: TestingMetastoreObjects.java    From presto with Apache License 2.0 5 votes vote down vote up
public static StorageDescriptor getGlueTestStorageDescriptor()
{
    return new StorageDescriptor()
            .withBucketColumns(ImmutableList.of("test-bucket-col"))
            .withColumns(ImmutableList.of(getGlueTestColumn()))
            .withParameters(ImmutableMap.of())
            .withSerdeInfo(new SerDeInfo()
                    .withSerializationLibrary("SerdeLib")
                    .withParameters(ImmutableMap.of()))
            .withInputFormat("InputFormat")
            .withOutputFormat("OutputFormat")
            .withLocation("/test-tbl")
            .withNumberOfBuckets(1);
}
 
Example #9
Source File: TestGlueInputConverter.java    From presto with Apache License 2.0 5 votes vote down vote up
private static void assertStorage(StorageDescriptor actual, Storage expected)
{
    assertEquals(actual.getLocation(), expected.getLocation());
    assertEquals(actual.getSerdeInfo().getSerializationLibrary(), expected.getStorageFormat().getSerDe());
    assertEquals(actual.getInputFormat(), expected.getStorageFormat().getInputFormat());
    assertEquals(actual.getOutputFormat(), expected.getStorageFormat().getOutputFormat());

    if (expected.getBucketProperty().isPresent()) {
        HiveBucketProperty bucketProperty = expected.getBucketProperty().get();
        assertEquals(actual.getBucketColumns(), bucketProperty.getBucketedBy());
        assertEquals(actual.getNumberOfBuckets().intValue(), bucketProperty.getBucketCount());
    }
}
 
Example #10
Source File: GlueToPrestoConverter.java    From presto with Apache License 2.0 5 votes vote down vote up
private static void setStorageBuilder(StorageDescriptor sd, Storage.Builder storageBuilder, Map<String, String> tableParameters)
{
    requireNonNull(sd.getSerdeInfo(), "StorageDescriptor SerDeInfo is null");
    SerDeInfo serdeInfo = sd.getSerdeInfo();

    Optional<HiveBucketProperty> bucketProperty = Optional.empty();
    if (sd.getNumberOfBuckets() > 0) {
        if (isNullOrEmpty(sd.getBucketColumns())) {
            throw new PrestoException(HIVE_INVALID_METADATA, "Table/partition metadata has 'numBuckets' set, but 'bucketCols' is not set");
        }
        List<SortingColumn> sortedBy = ImmutableList.of();
        if (!isNullOrEmpty(sd.getSortColumns())) {
            sortedBy = sd.getSortColumns().stream()
                    .map(column -> new SortingColumn(
                            column.getColumn(),
                            Order.fromMetastoreApiOrder(column.getSortOrder(), "unknown")))
                    .collect(toImmutableList());
        }
        BucketingVersion bucketingVersion = HiveBucketing.getBucketingVersion(tableParameters);
        bucketProperty = Optional.of(new HiveBucketProperty(sd.getBucketColumns(), bucketingVersion, sd.getNumberOfBuckets(), sortedBy));
    }

    storageBuilder.setStorageFormat(StorageFormat.createNullable(serdeInfo.getSerializationLibrary(), sd.getInputFormat(), sd.getOutputFormat()))
            .setLocation(nullToEmpty(sd.getLocation()))
            .setBucketProperty(bucketProperty)
            .setSkewed(sd.getSkewedInfo() != null && !isNullOrEmpty(sd.getSkewedInfo().getSkewedColumnNames()))
            .setSerdeParameters(firstNonNull(serdeInfo.getParameters(), ImmutableMap.of()))
            .build();
}
 
Example #11
Source File: GlueToPrestoConverter.java    From presto with Apache License 2.0 5 votes vote down vote up
public static Table convertTable(com.amazonaws.services.glue.model.Table glueTable, String dbName)
{
    requireNonNull(glueTable.getStorageDescriptor(), "Table StorageDescriptor is null");
    Map<String, String> tableParameters = firstNonNull(glueTable.getParameters(), ImmutableMap.of());
    StorageDescriptor sd = glueTable.getStorageDescriptor();

    Table.Builder tableBuilder = Table.builder()
            .setDatabaseName(dbName)
            .setTableName(glueTable.getName())
            .setOwner(nullToEmpty(glueTable.getOwner()))
            // Athena treats missing table type as EXTERNAL_TABLE.
            .setTableType(firstNonNull(glueTable.getTableType(), EXTERNAL_TABLE.name()))
            .setDataColumns(sd.getColumns().stream()
                    .map(GlueToPrestoConverter::convertColumn)
                    .collect(toImmutableList()))
            .setParameters(tableParameters)
            .setViewOriginalText(Optional.ofNullable(glueTable.getViewOriginalText()))
            .setViewExpandedText(Optional.ofNullable(glueTable.getViewExpandedText()));

    if (glueTable.getPartitionKeys() != null) {
        tableBuilder.setPartitionColumns(glueTable.getPartitionKeys().stream()
                .map(GlueToPrestoConverter::convertColumn)
                .collect(toImmutableList()));
    }
    else {
        tableBuilder.setPartitionColumns(new ArrayList<>());
    }

    setStorageBuilder(sd, tableBuilder.getStorageBuilder(), tableParameters);
    return tableBuilder.build();
}
 
Example #12
Source File: GlueInputConverter.java    From presto with Apache License 2.0 5 votes vote down vote up
private static StorageDescriptor convertStorage(Storage storage, List<Column> columns)
{
    if (storage.isSkewed()) {
        throw new IllegalArgumentException("Writing to skewed table/partition is not supported");
    }
    SerDeInfo serdeInfo = new SerDeInfo()
            .withSerializationLibrary(storage.getStorageFormat().getSerDeNullable())
            .withParameters(storage.getSerdeParameters());

    StorageDescriptor sd = new StorageDescriptor();
    sd.setLocation(storage.getLocation());
    sd.setColumns(columns.stream().map(GlueInputConverter::convertColumn).collect(toImmutableList()));
    sd.setSerdeInfo(serdeInfo);
    sd.setInputFormat(storage.getStorageFormat().getInputFormatNullable());
    sd.setOutputFormat(storage.getStorageFormat().getOutputFormatNullable());
    sd.setParameters(ImmutableMap.of());

    Optional<HiveBucketProperty> bucketProperty = storage.getBucketProperty();
    if (bucketProperty.isPresent()) {
        sd.setNumberOfBuckets(bucketProperty.get().getBucketCount());
        sd.setBucketColumns(bucketProperty.get().getBucketedBy());
        if (!bucketProperty.get().getSortedBy().isEmpty()) {
            sd.setSortColumns(bucketProperty.get().getSortedBy().stream()
                    .map(column -> new Order().withColumn(column.getColumnName()).withSortOrder(column.getOrder().getHiveOrder()))
                    .collect(toImmutableList()));
        }
    }

    return sd;
}
 
Example #13
Source File: GlueMetadataHandlerTest.java    From aws-athena-query-federation with Apache License 2.0 5 votes vote down vote up
@Test
public void populateSourceTableFromLocation() {
    Map<String, String> params = new HashMap<>();
    StorageDescriptor storageDescriptor = new StorageDescriptor().withLocation("arn:aws:dynamodb:us-east-1:012345678910:table/My-Table");
    Table table = new Table().withParameters(params).withStorageDescriptor(storageDescriptor);
    SchemaBuilder schemaBuilder = new SchemaBuilder();
    populateSourceTableNameIfAvailable(table, schemaBuilder);
    Schema schema = schemaBuilder.build();
    assertEquals("My-Table", getSourceTableName(schema));
}
 
Example #14
Source File: AWSCatalogMetastoreClientTest.java    From aws-glue-data-catalog-client-for-apache-hive-metastore with Apache License 2.0 5 votes vote down vote up
@Test
public void testRenamePartitionForExternalTable() throws Exception {
  String dbName = testDB.getName();
  Table externalTable = getTestTable();
  externalTable.setTableType(TableType.EXTERNAL_TABLE.name());   
  
  StorageDescriptor sd = HiveToCatalogConverter.convertStorageDescriptor(testPartition.getSd());
  
  Partition oldPartition = new Partition()
      .withDatabaseName(dbName).withTableName(externalTable.getName())
      .withValues(Lists.newArrayList("oldval")).withStorageDescriptor(sd);
  
  Partition newPartition = new Partition()
                                    .withDatabaseName(dbName).withTableName(externalTable.getName())
                                    .withValues(Lists.newArrayList("newval")).withStorageDescriptor(sd);

  when(glueClient.getDatabase(any(GetDatabaseRequest.class)))
    .thenReturn(new GetDatabaseResult().withDatabase(HiveToCatalogConverter.convertDatabase(testDB)));
  when(glueClient.getTable(any(GetTableRequest.class)))
  .thenReturn(new GetTableResult().withTable(externalTable));
  when(glueClient.getPartition(any(GetPartitionRequest.class)))
    .thenReturn(new GetPartitionResult().withPartition(oldPartition));

  metastoreClient.renamePartition(dbName, externalTable.getName(), oldPartition.getValues(),
      CatalogToHiveConverter.convertPartition(newPartition));
  
  // Verify catalog service is called and no interactions with wh
  verify(glueClient, times(1)).updatePartition(any(UpdatePartitionRequest.class));
  verifyNoMoreInteractions(wh);
}
 
Example #15
Source File: DynamoDBMetadataHandlerTest.java    From aws-athena-query-federation with Apache License 2.0 4 votes vote down vote up
@Test
public void doGetTableLayoutScanWithTypeOverride()
        throws Exception
{
    List<Column> columns = new ArrayList<>();
    columns.add(new Column().withName("col1").withType("int"));
    columns.add(new Column().withName("col2").withType("timestamptz"));
    columns.add(new Column().withName("col3").withType("string"));

    Map<String, String> param = ImmutableMap.of(
            SOURCE_TABLE_PROPERTY, TEST_TABLE,
            COLUMN_NAME_MAPPING_PROPERTY, "col1=Col1",
            DATETIME_FORMAT_MAPPING_PROPERTY, "col1=datetime1,col3=datetime3 ");
    Table table = new Table()
            .withParameters(param)
            .withPartitionKeys()
            .withStorageDescriptor(new StorageDescriptor().withColumns(columns));
    GetTableResult mockResult = new GetTableResult().withTable(table);
    when(glueClient.getTable(any())).thenReturn(mockResult);

    TableName tableName = new TableName(DEFAULT_SCHEMA, "glueTableForTestTable");
    GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName);
    GetTableResponse getTableResponse = handler.doGetTable(allocator, getTableRequest);
    logger.info("validateSourceTableNamePropagation: GetTableResponse[{}]", getTableResponse);
    Map<String, String> customMetadata = getTableResponse.getSchema().getCustomMetadata();
    assertThat(customMetadata.get(SOURCE_TABLE_PROPERTY), equalTo(TEST_TABLE));
    assertThat(customMetadata.get(DATETIME_FORMAT_MAPPING_PROPERTY_NORMALIZED), equalTo("Col1=datetime1,col3=datetime3"));

    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put("col3",
            EquatableValueSet.newBuilder(allocator, new ArrowType.Bool(), true, true)
                    .add(true).build());
    constraintsMap.put("col2",
            EquatableValueSet.newBuilder(allocator, new ArrowType.Bool(), true, true)
                    .add(true).build());

    GetTableLayoutRequest getTableLayoutRequest = new GetTableLayoutRequest(TEST_IDENTITY,
            TEST_QUERY_ID,
            TEST_CATALOG_NAME,
            tableName,
            new Constraints(constraintsMap),
            getTableResponse.getSchema(),
            Collections.EMPTY_SET);


    GetTableLayoutResponse res = handler.doGetTableLayout(allocator, getTableLayoutRequest);

    logger.info("doGetTableLayoutScanWithTypeOverride schema - {}", res.getPartitions().getSchema());
    logger.info("doGetTableLayoutScanWithTypeOverride partitions - {}", res.getPartitions());

    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(PARTITION_TYPE_METADATA), equalTo(SCAN_PARTITION_TYPE));
    // no hash key constraints, so look for segment count column
    assertThat(res.getPartitions().getSchema().findField(SEGMENT_COUNT_METADATA) != null, is(true));
    assertThat(res.getPartitions().getRowCount(), equalTo(1));

    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(NON_KEY_FILTER_METADATA), equalTo("(#col3 = :v0 OR attribute_not_exists(#col3) OR #col3 = :v1)"));

    ImmutableMap<String, String> expressionNames = ImmutableMap.of("#col3", "col3", "#col2", "col2");
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_NAMES_METADATA), equalTo(Jackson.toJsonString(expressionNames)));

    ImmutableMap<String, AttributeValue> expressionValues = ImmutableMap.of(":v0", ItemUtils.toAttributeValue(true), ":v1", ItemUtils.toAttributeValue(null));
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_VALUES_METADATA), equalTo(Jackson.toJsonString(expressionValues)));
}
 
Example #16
Source File: DynamoDBMetadataHandlerTest.java    From aws-athena-query-federation with Apache License 2.0 4 votes vote down vote up
@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)));
}
 
Example #17
Source File: GlueMetadataHandlerTest.java    From aws-athena-query-federation with Apache License 2.0 4 votes vote down vote up
@Test
public void doGetTable()
        throws Exception
{
    String sourceTable = "My-Table";

    Map<String, String> expectedParams = new HashMap<>();
    expectedParams.put(SOURCE_TABLE_PROPERTY, sourceTable);
    expectedParams.put(COLUMN_NAME_MAPPING_PROPERTY, "col2=Col2,col3=Col3, col4=Col4");
    expectedParams.put(DATETIME_FORMAT_MAPPING_PROPERTY, "col2=someformat2, col1=someformat1 ");

    List<Column> columns = new ArrayList<>();
    columns.add(new Column().withName("col1").withType("int").withComment("comment"));
    columns.add(new Column().withName("col2").withType("bigint").withComment("comment"));
    columns.add(new Column().withName("col3").withType("string").withComment("comment"));
    columns.add(new Column().withName("col4").withType("timestamp").withComment("comment"));
    columns.add(new Column().withName("col5").withType("date").withComment("comment"));
    columns.add(new Column().withName("col6").withType("timestamptz").withComment("comment"));
    columns.add(new Column().withName("col7").withType("timestamptz").withComment("comment"));

    Table mockTable = mock(Table.class);
    StorageDescriptor mockSd = mock(StorageDescriptor.class);

    when(mockTable.getName()).thenReturn(table);
    when(mockTable.getStorageDescriptor()).thenReturn(mockSd);
    when(mockTable.getParameters()).thenReturn(expectedParams);
    when(mockSd.getColumns()).thenReturn(columns);

    when(mockGlue.getTable(any(com.amazonaws.services.glue.model.GetTableRequest.class)))
            .thenAnswer((InvocationOnMock invocationOnMock) ->
            {
                com.amazonaws.services.glue.model.GetTableRequest request =
                        (com.amazonaws.services.glue.model.GetTableRequest) invocationOnMock.getArguments()[0];

                assertEquals(accountId, request.getCatalogId());
                assertEquals(schema, request.getDatabaseName());
                assertEquals(table, request.getName());

                GetTableResult mockResult = mock(GetTableResult.class);
                when(mockResult.getTable()).thenReturn(mockTable);
                return mockResult;
            });

    GetTableRequest req = new GetTableRequest(IdentityUtil.fakeIdentity(), queryId, catalog, new TableName(schema, table));
    GetTableResponse res = handler.doGetTable(allocator, req);

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

    assertTrue(res.getSchema().getFields().size() == 7);
    assertTrue(res.getSchema().getCustomMetadata().size() > 0);
    assertTrue(res.getSchema().getCustomMetadata().containsKey(DATETIME_FORMAT_MAPPING_PROPERTY));
    assertEquals(res.getSchema().getCustomMetadata().get(DATETIME_FORMAT_MAPPING_PROPERTY_NORMALIZED), "Col2=someformat2,col1=someformat1");
    assertEquals(sourceTable, getSourceTableName(res.getSchema()));

    //Verify column name mapping works
    assertNotNull(res.getSchema().findField("col1"));
    assertNotNull(res.getSchema().findField("Col2"));
    assertNotNull(res.getSchema().findField("Col3"));
    assertNotNull(res.getSchema().findField("Col4"));
    assertNotNull(res.getSchema().findField("col5"));
    assertNotNull(res.getSchema().findField("col6"));
    assertNotNull(res.getSchema().findField("col7"));

    //Verify types
    assertTrue(Types.getMinorTypeForArrowType(res.getSchema().findField("col1").getType()).equals(Types.MinorType.INT));
    assertTrue(Types.getMinorTypeForArrowType(res.getSchema().findField("Col2").getType()).equals(Types.MinorType.BIGINT));
    assertTrue(Types.getMinorTypeForArrowType(res.getSchema().findField("Col3").getType()).equals(Types.MinorType.VARCHAR));
    assertTrue(Types.getMinorTypeForArrowType(res.getSchema().findField("Col4").getType()).equals(Types.MinorType.DATEMILLI));
    assertTrue(Types.getMinorTypeForArrowType(res.getSchema().findField("col5").getType()).equals(Types.MinorType.DATEDAY));
    assertTrue(Types.getMinorTypeForArrowType(res.getSchema().findField("col6").getType()).equals(Types.MinorType.TIMESTAMPMILLITZ));
    assertTrue(Types.getMinorTypeForArrowType(res.getSchema().findField("col7").getType()).equals(Types.MinorType.TIMESTAMPMILLITZ));
}