Search in sources :

Example 26 with DatabaseConnectionConfig

use of com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig in project aws-athena-query-federation by awslabs.

the class HiveCompositeHandlerTest method HiveCompositeHandlerTest.

@Test
public void HiveCompositeHandlerTest() {
    Exception ex = null;
    try {
        DatabaseConnectionConfig databaseConnectionConfig = new DatabaseConnectionConfig("testCatalog1", HiveConstants.HIVE_NAME, "hdphive2://jdbc:hive2://54.89.6.2:10000/authena;AuthMech=3;${testSecret}", "testSecret");
        PowerMockito.mockStatic(JDBCUtil.class);
        JDBCUtil tested = PowerMockito.mock(JDBCUtil.class);
        PowerMockito.when(tested.getSingleDatabaseConfigFromEnv(HiveConstants.HIVE_NAME)).thenReturn(databaseConnectionConfig);
        new HiveCompositeHandler();
    } catch (Exception e) {
        ex = e;
    }
    Assert.assertEquals(null, ex);
}
Also used : JDBCUtil(com.amazonaws.athena.connectors.jdbc.manager.JDBCUtil) DatabaseConnectionConfig(com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig) PrepareForTest(org.powermock.core.classloader.annotations.PrepareForTest) Test(org.junit.Test)

Example 27 with DatabaseConnectionConfig

use of com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig in project aws-athena-query-federation by awslabs.

the class JdbcTableUtils method getDbConnection.

/**
 * Gets a JDBC DB connection.
 * @return Connection object.
 * @param databaseConnectionInfo
 */
protected Connection getDbConnection(DatabaseConnectionInfo databaseConnectionInfo) {
    DatabaseConnectionConfig connectionConfig = getDbConfig();
    JdbcConnectionFactory connectionFactory = new GenericJdbcConnectionFactory(connectionConfig, properties, databaseConnectionInfo);
    return connectionFactory.getConnection(null);
}
Also used : GenericJdbcConnectionFactory(com.amazonaws.athena.connectors.jdbc.connection.GenericJdbcConnectionFactory) JdbcConnectionFactory(com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory) DatabaseConnectionConfig(com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig) GenericJdbcConnectionFactory(com.amazonaws.athena.connectors.jdbc.connection.GenericJdbcConnectionFactory)

Example 28 with DatabaseConnectionConfig

use of com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig in project aws-athena-query-federation by awslabs.

the class JdbcMetadataHandlerTest method setup.

@Before
public void setup() {
    this.jdbcConnectionFactory = Mockito.mock(JdbcConnectionFactory.class);
    this.connection = Mockito.mock(Connection.class, Mockito.RETURNS_DEEP_STUBS);
    Mockito.when(this.jdbcConnectionFactory.getConnection(Mockito.any(JdbcCredentialProvider.class))).thenReturn(this.connection);
    this.secretsManager = Mockito.mock(AWSSecretsManager.class);
    this.athena = Mockito.mock(AmazonAthena.class);
    Mockito.when(this.secretsManager.getSecretValue(Mockito.eq(new GetSecretValueRequest().withSecretId("testSecret")))).thenReturn(new GetSecretValueResult().withSecretString("{\"username\": \"testUser\", \"password\": \"testPassword\"}"));
    DatabaseConnectionConfig databaseConnectionConfig = new DatabaseConnectionConfig("testCatalog", "fakedatabase", "fakedatabase://jdbc:fakedatabase://hostname/${testSecret}", "testSecret");
    this.jdbcMetadataHandler = new JdbcMetadataHandler(databaseConnectionConfig, this.secretsManager, this.athena, jdbcConnectionFactory) {

        @Override
        public Schema getPartitionSchema(final String catalogName) {
            return PARTITION_SCHEMA;
        }

        @Override
        public void getPartitions(final BlockWriter blockWriter, final GetTableLayoutRequest getTableLayoutRequest, QueryStatusChecker queryStatusChecker) {
        }

        @Override
        public GetSplitsResponse doGetSplits(BlockAllocator blockAllocator, GetSplitsRequest getSplitsRequest) {
            return null;
        }
    };
    this.federatedIdentity = Mockito.mock(FederatedIdentity.class);
    this.blockAllocator = Mockito.mock(BlockAllocator.class);
}
Also used : JdbcConnectionFactory(com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory) GetSplitsRequest(com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest) AWSSecretsManager(com.amazonaws.services.secretsmanager.AWSSecretsManager) Schema(org.apache.arrow.vector.types.pojo.Schema) Connection(java.sql.Connection) DatabaseConnectionConfig(com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig) GetSecretValueResult(com.amazonaws.services.secretsmanager.model.GetSecretValueResult) QueryStatusChecker(com.amazonaws.athena.connector.lambda.QueryStatusChecker) FederatedIdentity(com.amazonaws.athena.connector.lambda.security.FederatedIdentity) GetSplitsResponse(com.amazonaws.athena.connector.lambda.metadata.GetSplitsResponse) GetTableLayoutRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest) BlockAllocator(com.amazonaws.athena.connector.lambda.data.BlockAllocator) GetSecretValueRequest(com.amazonaws.services.secretsmanager.model.GetSecretValueRequest) BlockWriter(com.amazonaws.athena.connector.lambda.data.BlockWriter) JdbcCredentialProvider(com.amazonaws.athena.connectors.jdbc.connection.JdbcCredentialProvider) AmazonAthena(com.amazonaws.services.athena.AmazonAthena) Before(org.junit.Before)

Example 29 with DatabaseConnectionConfig

use of com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig in project aws-athena-query-federation by awslabs.

the class MultiplexingJdbcMetadataHandlerTest method setup.

@Before
public void setup() {
    // this.allocator = Mockito.mock(BlockAllocator.class);
    this.allocator = new BlockAllocatorImpl();
    // Mockito.when(this.allocator.createBlock(Mockito.any(Schema.class))).thenReturn(Mockito.mock(Block.class));
    this.fakeDatabaseHandler = Mockito.mock(JdbcMetadataHandler.class);
    this.metadataHandlerMap = Collections.singletonMap("fakedatabase", this.fakeDatabaseHandler);
    this.secretsManager = Mockito.mock(AWSSecretsManager.class);
    this.athena = Mockito.mock(AmazonAthena.class);
    this.queryStatusChecker = Mockito.mock(QueryStatusChecker.class);
    this.jdbcConnectionFactory = Mockito.mock(JdbcConnectionFactory.class);
    DatabaseConnectionConfig databaseConnectionConfig = new DatabaseConnectionConfig("testCatalog", "fakedatabase", "fakedatabase://jdbc:fakedatabase://hostname/${testSecret}", "testSecret");
    this.jdbcMetadataHandler = new MultiplexingJdbcMetadataHandler(this.secretsManager, this.athena, this.jdbcConnectionFactory, this.metadataHandlerMap, databaseConnectionConfig);
}
Also used : JdbcConnectionFactory(com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory) QueryStatusChecker(com.amazonaws.athena.connector.lambda.QueryStatusChecker) BlockAllocatorImpl(com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl) AWSSecretsManager(com.amazonaws.services.secretsmanager.AWSSecretsManager) DatabaseConnectionConfig(com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig) JdbcMetadataHandler(com.amazonaws.athena.connectors.jdbc.manager.JdbcMetadataHandler) AmazonAthena(com.amazonaws.services.athena.AmazonAthena) Before(org.junit.Before)

Example 30 with DatabaseConnectionConfig

use of com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig in project aws-athena-query-federation by awslabs.

the class JDBCUtil method createJdbcMetadataHandlerMap.

/**
 * Creates a map of Catalog to respective metadata handler to be used by Multiplexer.
 *
 * @param properties system properties.
 * @param metadataHandlerFactory factory for creating the appropriate metadata handler for the database type
 * @return Map of String -> {@link JdbcMetadataHandler}
 */
public static Map<String, JdbcMetadataHandler> createJdbcMetadataHandlerMap(final Map<String, String> properties, JdbcMetadataHandlerFactory metadataHandlerFactory) {
    ImmutableMap.Builder<String, JdbcMetadataHandler> metadataHandlerMap = ImmutableMap.builder();
    final String functionName = Validate.notBlank(properties.get(LAMBDA_FUNCTION_NAME_PROPERTY), "Lambda function name not present in environment.");
    List<DatabaseConnectionConfig> databaseConnectionConfigs = new DatabaseConnectionConfigBuilder().engine(metadataHandlerFactory.getEngine()).properties(properties).build();
    if (databaseConnectionConfigs.isEmpty()) {
        throw new RuntimeException("At least one connection string required.");
    }
    boolean defaultPresent = false;
    for (DatabaseConnectionConfig databaseConnectionConfig : databaseConnectionConfigs) {
        JdbcMetadataHandler jdbcMetadataHandler = metadataHandlerFactory.createJdbcMetadataHandler(databaseConnectionConfig);
        metadataHandlerMap.put(databaseConnectionConfig.getCatalog(), jdbcMetadataHandler);
        if (DatabaseConnectionConfigBuilder.DEFAULT_CONNECTION_STRING_PROPERTY.equals(databaseConnectionConfig.getCatalog())) {
            metadataHandlerMap.put(DEFAULT_CATALOG_PREFIX + functionName, jdbcMetadataHandler);
            defaultPresent = true;
        }
    }
    if (!defaultPresent) {
        throw new RuntimeException("Must provide connection parameters for default database instance " + DatabaseConnectionConfigBuilder.DEFAULT_CONNECTION_STRING_PROPERTY);
    }
    return metadataHandlerMap.build();
}
Also used : DatabaseConnectionConfigBuilder(com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfigBuilder) DatabaseConnectionConfig(com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig) ImmutableMap(com.google.common.collect.ImmutableMap)

Aggregations

DatabaseConnectionConfig (com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig)62 JdbcConnectionFactory (com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory)44 AmazonAthena (com.amazonaws.services.athena.AmazonAthena)43 AWSSecretsManager (com.amazonaws.services.secretsmanager.AWSSecretsManager)43 Before (org.junit.Before)43 QueryStatusChecker (com.amazonaws.athena.connector.lambda.QueryStatusChecker)30 AmazonS3 (com.amazonaws.services.s3.AmazonS3)28 Test (org.junit.Test)16 JdbcCredentialProvider (com.amazonaws.athena.connectors.jdbc.connection.JdbcCredentialProvider)15 Connection (java.sql.Connection)15 BlockAllocatorImpl (com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl)14 JDBCUtil (com.amazonaws.athena.connectors.jdbc.manager.JDBCUtil)10 PrepareForTest (org.powermock.core.classloader.annotations.PrepareForTest)10 HashMap (java.util.HashMap)6 GetSecretValueRequest (com.amazonaws.services.secretsmanager.model.GetSecretValueRequest)5 GetSecretValueResult (com.amazonaws.services.secretsmanager.model.GetSecretValueResult)5 JdbcMetadataHandler (com.amazonaws.athena.connectors.jdbc.manager.JdbcMetadataHandler)4 JdbcRecordHandler (com.amazonaws.athena.connectors.jdbc.manager.JdbcRecordHandler)4 SQLException (java.sql.SQLException)4 FederatedIdentity (com.amazonaws.athena.connector.lambda.security.FederatedIdentity)2