com.amazonaws.athena.connectors.aws.cmdb.AwsCmdbMetadataHandler Maven / Gradle / Ivy

Go to download
/*-
 * #%L
 * athena-aws-cmdb
 * %%
 * Copyright (C) 2019 Amazon Web Services
 * %%
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 * #L%
 */
package com.amazonaws.athena.connectors.aws.cmdb;

import com.amazonaws.athena.connector.lambda.QueryStatusChecker;
import com.amazonaws.athena.connector.lambda.data.BlockAllocator;
import com.amazonaws.athena.connector.lambda.data.BlockWriter;
import com.amazonaws.athena.connector.lambda.data.SchemaBuilder;
import com.amazonaws.athena.connector.lambda.domain.Split;
import com.amazonaws.athena.connector.lambda.domain.TableName;
import com.amazonaws.athena.connector.lambda.domain.spill.SpillLocation;
import com.amazonaws.athena.connector.lambda.handlers.MetadataHandler;
import com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest;
import com.amazonaws.athena.connector.lambda.metadata.GetSplitsResponse;
import com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest;
import com.amazonaws.athena.connector.lambda.metadata.GetTableRequest;
import com.amazonaws.athena.connector.lambda.metadata.GetTableResponse;
import com.amazonaws.athena.connector.lambda.metadata.ListSchemasRequest;
import com.amazonaws.athena.connector.lambda.metadata.ListSchemasResponse;
import com.amazonaws.athena.connector.lambda.metadata.ListTablesRequest;
import com.amazonaws.athena.connector.lambda.metadata.ListTablesResponse;
import com.amazonaws.athena.connector.lambda.security.EncryptionKey;
import com.amazonaws.athena.connector.lambda.security.EncryptionKeyFactory;
import com.amazonaws.athena.connectors.aws.cmdb.tables.TableProvider;
import com.amazonaws.services.athena.AmazonAthena;
import com.amazonaws.services.secretsmanager.AWSSecretsManager;
import org.apache.arrow.util.VisibleForTesting;

import java.util.List;
import java.util.Map;

/**
 * Handles metadata requests for the Athena AWS CMDB Connector.
 * 
 * For more detail, please see the module's README.md, some notable characteristics of this class include:
 * 
 * 1. Maps AWS Resources to SQL tables using a set of TableProviders constructed from a TableProviderFactory.
 * 2. This class is largely a mux that delegates requests to the appropriate TableProvider based on the
 * requested TableName.
 * 3. Provides a schema and table list by scanning all loaded TableProviders.
 */
public class AwsCmdbMetadataHandler
        extends MetadataHandler
{
    private static final String SOURCE_TYPE = "cmdb";
    //Map of schema name to list of TableNames generated by scanning all loaded TableProviders.
    private Map> schemas;
    //Map of available fully qualified TableNames to their respective TableProviders.
    private Map tableProviders;

    public AwsCmdbMetadataHandler(java.util.Map configOptions)
    {
        super(SOURCE_TYPE, configOptions);
        TableProviderFactory tableProviderFactory = new TableProviderFactory(configOptions);
        schemas = tableProviderFactory.getSchemas();
        tableProviders = tableProviderFactory.getTableProviders();
    }

    @VisibleForTesting
    protected AwsCmdbMetadataHandler(
        TableProviderFactory tableProviderFactory,
        EncryptionKeyFactory keyFactory,
        AWSSecretsManager secretsManager,
        AmazonAthena athena,
        String spillBucket,
        String spillPrefix,
        java.util.Map configOptions)
    {
        super(keyFactory, secretsManager, athena, SOURCE_TYPE, spillBucket, spillPrefix, configOptions);
        schemas = tableProviderFactory.getSchemas();
        tableProviders = tableProviderFactory.getTableProviders();
    }

    /**
     * Returns the list of supported schemas discovered from the loaded TableProvider scan.
     *
     * @see MetadataHandler
     */
    @Override
    public ListSchemasResponse doListSchemaNames(BlockAllocator blockAllocator, ListSchemasRequest listSchemasRequest)
    {
        return new ListSchemasResponse(listSchemasRequest.getCatalogName(), schemas.keySet());
    }

    /**
     * Returns the list of supported tables on the requested schema discovered from the loaded TableProvider scan.
     *
     * @see MetadataHandler
     */
    @Override
    public ListTablesResponse doListTables(BlockAllocator blockAllocator, ListTablesRequest listTablesRequest)
    {
        return new ListTablesResponse(listTablesRequest.getCatalogName(),
                schemas.get(listTablesRequest.getSchemaName()), null);
    }

    /**
     * Delegates to the TableProvider that is registered for the requested table.
     *
     * @see MetadataHandler
     */
    @Override
    public GetTableResponse doGetTable(BlockAllocator blockAllocator, GetTableRequest getTableRequest)
    {
        TableProvider tableProvider = tableProviders.get(getTableRequest.getTableName());
        if (tableProvider == null) {
            throw new RuntimeException("Unknown table " + getTableRequest.getTableName());
        }
        return tableProvider.getTable(blockAllocator, getTableRequest);
    }

    /**
     * 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);
    }

    /**
     * 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);
    }

    /**
     * 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);
    }
}