Examples with JdbcConnectionFactory com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory used on opensource projects

Example 1 with JdbcConnectionFactory

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

the class HiveRecordHandlerTest method setup.

@Before
public void setup() {
    this.amazonS3 = Mockito.mock(AmazonS3.class);
    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\"}"));
    this.connection = Mockito.mock(Connection.class);
    this.jdbcConnectionFactory = Mockito.mock(JdbcConnectionFactory.class);
    Mockito.when(this.jdbcConnectionFactory.getConnection(Mockito.mock(JdbcCredentialProvider.class))).thenReturn(this.connection);
    jdbcSplitQueryBuilder = new HiveQueryStringBuilder("`");
    final DatabaseConnectionConfig databaseConnectionConfig = new DatabaseConnectionConfig("testCatalog", HiveConstants.HIVE_NAME, "hive2://jdbc:hive2://54.89.6.2:10000/authena;AuthMech=3;UID=hive;PWD=hive");
    this.hiveRecordHandler = new HiveRecordHandler(databaseConnectionConfig, amazonS3, secretsManager, athena, jdbcConnectionFactory, jdbcSplitQueryBuilder);
}

Example 2 with JdbcConnectionFactory

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

the class PostGreSqlMetadataHandlerTest method doGetTableLayoutWithSQLException.

@Test(expected = RuntimeException.class)
public void doGetTableLayoutWithSQLException() throws Exception {
    Constraints constraints = Mockito.mock(Constraints.class);
    TableName tableName = new TableName("testSchema", "testTable");
    Schema partitionSchema = this.postGreSqlMetadataHandler.getPartitionSchema("testCatalogName");
    Set<String> partitionCols = partitionSchema.getFields().stream().map(Field::getName).collect(Collectors.toSet());
    GetTableLayoutRequest getTableLayoutRequest = new GetTableLayoutRequest(this.federatedIdentity, "testQueryId", "testCatalogName", tableName, constraints, partitionSchema, partitionCols);
    Connection connection = Mockito.mock(Connection.class, Mockito.RETURNS_DEEP_STUBS);
    JdbcConnectionFactory jdbcConnectionFactory = Mockito.mock(JdbcConnectionFactory.class);
    Mockito.when(jdbcConnectionFactory.getConnection(Mockito.any(JdbcCredentialProvider.class))).thenReturn(connection);
    Mockito.when(connection.getMetaData().getSearchStringEscape()).thenThrow(new SQLException());
    PostGreSqlMetadataHandler postGreSqlMetadataHandler = new PostGreSqlMetadataHandler(databaseConnectionConfig, this.secretsManager, this.athena, jdbcConnectionFactory);
    postGreSqlMetadataHandler.doGetTableLayout(Mockito.mock(BlockAllocator.class), getTableLayoutRequest);
}

Example 3 with JdbcConnectionFactory

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

the class OracleMetadataHandlerTest method doGetTableLayoutWithSQLException.

@Test(expected = RuntimeException.class)
public void doGetTableLayoutWithSQLException() throws Exception {
    Constraints constraints = Mockito.mock(Constraints.class);
    TableName tableName = new TableName("testSchema", "testTable");
    Schema partitionSchema = this.oracleMetadataHandler.getPartitionSchema("testCatalogName");
    Set<String> partitionCols = partitionSchema.getFields().stream().map(Field::getName).collect(Collectors.toSet());
    GetTableLayoutRequest getTableLayoutRequest = new GetTableLayoutRequest(this.federatedIdentity, "testQueryId", "testCatalogName", tableName, constraints, partitionSchema, partitionCols);
    Connection connection = Mockito.mock(Connection.class, Mockito.RETURNS_DEEP_STUBS);
    JdbcConnectionFactory jdbcConnectionFactory = Mockito.mock(JdbcConnectionFactory.class);
    Mockito.when(jdbcConnectionFactory.getConnection(Mockito.any(JdbcCredentialProvider.class))).thenReturn(connection);
    Mockito.when(connection.getMetaData().getSearchStringEscape()).thenThrow(new SQLException());
    OracleMetadataHandler oracleMetadataHandler = new OracleMetadataHandler(databaseConnectionConfig, this.secretsManager, this.athena, jdbcConnectionFactory);
    oracleMetadataHandler.doGetTableLayout(Mockito.mock(BlockAllocator.class), getTableLayoutRequest);
}

Example 4 with JdbcConnectionFactory

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

the class RedshiftMetadataHandlerTest method doGetTableLayoutWithSQLException.

@Test(expected = RuntimeException.class)
public void doGetTableLayoutWithSQLException() throws Exception {
    Constraints constraints = Mockito.mock(Constraints.class);
    TableName tableName = new TableName("testSchema", "testTable");
    Schema partitionSchema = this.redshiftMetadataHandler.getPartitionSchema("testCatalogName");
    Set<String> partitionCols = partitionSchema.getFields().stream().map(Field::getName).collect(Collectors.toSet());
    GetTableLayoutRequest getTableLayoutRequest = new GetTableLayoutRequest(this.federatedIdentity, "testQueryId", "testCatalogName", tableName, constraints, partitionSchema, partitionCols);
    Connection connection = Mockito.mock(Connection.class, Mockito.RETURNS_DEEP_STUBS);
    JdbcConnectionFactory jdbcConnectionFactory = Mockito.mock(JdbcConnectionFactory.class);
    Mockito.when(jdbcConnectionFactory.getConnection(Mockito.any(JdbcCredentialProvider.class))).thenReturn(connection);
    Mockito.when(connection.getMetaData().getSearchStringEscape()).thenThrow(new SQLException());
    RedshiftMetadataHandler redshiftMetadataHandler = new RedshiftMetadataHandler(databaseConnectionConfig, this.secretsManager, this.athena, jdbcConnectionFactory);
    redshiftMetadataHandler.doGetTableLayout(Mockito.mock(BlockAllocator.class), getTableLayoutRequest);
}

Example 5 with JdbcConnectionFactory

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

the class SnowflakeRecordHandlerTest method setup.

@Before
public void setup() {
    this.amazonS3 = Mockito.mock(AmazonS3.class);
    this.secretsManager = Mockito.mock(AWSSecretsManager.class);
    this.athena = Mockito.mock(AmazonAthena.class);
    this.connection = Mockito.mock(Connection.class);
    this.jdbcConnectionFactory = Mockito.mock(JdbcConnectionFactory.class);
    Mockito.when(this.jdbcConnectionFactory.getConnection(Mockito.mock(JdbcCredentialProvider.class))).thenReturn(this.connection);
    jdbcSplitQueryBuilder = new SnowflakeQueryStringBuilder("`");
    final DatabaseConnectionConfig databaseConnectionConfig = new DatabaseConnectionConfig("testCatalog", SnowflakeConstants.SNOWFLAKE_NAME, "snowflake://jdbc:snowflake://hostname/?warehouse=warehousename&db=dbname&schema=schemaname&user=xxx&password=xxx");
    this.snowflakeRecordHandler = new SnowflakeRecordHandler(databaseConnectionConfig, amazonS3, secretsManager, athena, jdbcConnectionFactory, jdbcSplitQueryBuilder);
}