Examples with TableProvider com.amazonaws.athena.connectors.aws.cmdb.tables.TableProvider used on opensource projects

Example 1 with TableProvider

use of com.amazonaws.athena.connectors.aws.cmdb.tables.TableProvider in project aws-athena-query-federation by awslabs.

the class AwsCmdbMetadataHandler method getPartitions.

/**
 * Delegates to the TableProvider that is registered for the requested table.
 *
 * @see MetadataHandler
 */
@Override
public void getPartitions(BlockWriter blockWriter, GetTableLayoutRequest request, QueryStatusChecker queryStatusChecker) throws Exception {
    TableProvider tableProvider = tableProviders.get(request.getTableName());
    if (tableProvider == null) {
        throw new RuntimeException("Unknown table " + request.getTableName());
    }
    tableProvider.getPartitions(blockWriter, request);
}

Example 2 with TableProvider

use of com.amazonaws.athena.connectors.aws.cmdb.tables.TableProvider in project aws-athena-query-federation by awslabs.

the class AwsCmdbMetadataHandler method doGetSplits.

/**
 * Delegates to the TableProvider that is registered for the requested table.
 *
 * @see MetadataHandler
 */
@Override
public GetSplitsResponse doGetSplits(BlockAllocator blockAllocator, GetSplitsRequest getSplitsRequest) {
    TableProvider tableProvider = tableProviders.get(getSplitsRequest.getTableName());
    if (tableProvider == null) {
        throw new RuntimeException("Unknown table " + getSplitsRequest.getTableName());
    }
    // Every split needs a unique spill location.
    SpillLocation spillLocation = makeSpillLocation(getSplitsRequest);
    EncryptionKey encryptionKey = makeEncryptionKey();
    Split split = Split.newBuilder(spillLocation, encryptionKey).build();
    return new GetSplitsResponse(getSplitsRequest.getCatalogName(), split);
}

Example 3 with TableProvider

use of com.amazonaws.athena.connectors.aws.cmdb.tables.TableProvider in project aws-athena-query-federation by awslabs.

the class AwsCmdbMetadataHandler method enhancePartitionSchema.

/**
 * Delegates to the TableProvider that is registered for the requested table.
 *
 * @see MetadataHandler
 */
@Override
public void enhancePartitionSchema(SchemaBuilder partitionSchemaBuilder, GetTableLayoutRequest request) {
    TableProvider tableProvider = tableProviders.get(request.getTableName());
    if (tableProvider == null) {
        throw new RuntimeException("Unknown table " + request.getTableName());
    }
    tableProvider.enhancePartitionSchema(partitionSchemaBuilder, request);
}

Example 4 with TableProvider

use of com.amazonaws.athena.connectors.aws.cmdb.tables.TableProvider in project aws-athena-query-federation by awslabs.

the class AwsCmdbMetadataHandlerTest method setUp.

@Before
public void setUp() throws Exception {
    blockAllocator = new BlockAllocatorImpl();
    Map<TableName, TableProvider> tableProviderMap = new HashMap<>();
    tableProviderMap.putIfAbsent(new TableName("schema1", "table1"), mockTableProvider1);
    tableProviderMap.putIfAbsent(new TableName("schema1", "table2"), mockTableProvider2);
    tableProviderMap.putIfAbsent(new TableName("schema2", "table1"), mockTableProvider3);
    when(mockTableProviderFactory.getTableProviders()).thenReturn(tableProviderMap);
    Map<String, List<TableName>> schemas = new HashMap<>();
    schemas.put("schema1", new ArrayList<>());
    schemas.put("schema2", new ArrayList<>());
    schemas.get("schema1").add(new TableName("schema1", "table1"));
    schemas.get("schema1").add(new TableName("schema1", "table2"));
    schemas.get("schema2").add(new TableName("schema2", "table1"));
    when(mockTableProviderFactory.getSchemas()).thenReturn(schemas);
    handler = new AwsCmdbMetadataHandler(mockTableProviderFactory, new LocalKeyFactory(), mockSecretsManager, mockAthena, bucket, prefix);
    verify(mockTableProviderFactory, times(1)).getTableProviders();
    verify(mockTableProviderFactory, times(1)).getSchemas();
    verifyNoMoreInteractions(mockTableProviderFactory);
}

Example 5 with TableProvider

use of com.amazonaws.athena.connectors.aws.cmdb.tables.TableProvider in project aws-athena-query-federation by awslabs.

the class AwsCmdbRecordHandler method readWithConstraint.

/**
 * Delegates to the TableProvider that is registered for the requested table.
 *
 * @see RecordHandler
 */
@Override
protected void readWithConstraint(BlockSpiller blockSpiller, ReadRecordsRequest readRecordsRequest, QueryStatusChecker queryStatusChecker) {
    TableProvider tableProvider = tableProviders.get(readRecordsRequest.getTableName());
    tableProvider.readWithConstraint(blockSpiller, readRecordsRequest, queryStatusChecker);
}