use of com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory in project aws-athena-query-federation by awslabs.
the class HiveMetadataHandlerTest 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.hiveMetadataHandler.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());
HiveMetadataHandler implalaMetadataHandler = new HiveMetadataHandler(databaseConnectionConfig, this.secretsManager, this.athena, jdbcConnectionFactory);
implalaMetadataHandler.doGetTableLayout(Mockito.mock(BlockAllocator.class), getTableLayoutRequest);
}
use of com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory 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);
}
use of com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory 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);
}
use of com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory in project aws-athena-query-federation by awslabs.
the class PostGreSqlRecordHandlerTest 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 PostGreSqlQueryStringBuilder("\"");
final DatabaseConnectionConfig databaseConnectionConfig = new DatabaseConnectionConfig("testCatalog", POSTGRES_NAME, "postgres://jdbc:postgresql://hostname/user=A&password=B");
this.postGreSqlRecordHandler = new PostGreSqlRecordHandler(databaseConnectionConfig, amazonS3, secretsManager, athena, jdbcConnectionFactory, jdbcSplitQueryBuilder);
}
use of com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory in project aws-athena-query-federation by awslabs.
the class MySqlRecordHandlerTest 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 MySqlQueryStringBuilder("`");
final DatabaseConnectionConfig databaseConnectionConfig = new DatabaseConnectionConfig("testCatalog", MYSQL_NAME, "mysql://jdbc:mysql://hostname/user=A&password=B");
this.mySqlRecordHandler = new MySqlRecordHandler(databaseConnectionConfig, amazonS3, secretsManager, athena, jdbcConnectionFactory, jdbcSplitQueryBuilder);
}
Aggregations