Search in sources :

Example 6 with JdbcConnectionFactory

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

the class TeradataMetadataHandlerTest 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.teradataMetadataHandler.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());
    TeradataMetadataHandler teradataMetadataHandler = new TeradataMetadataHandler(databaseConnectionConfig, this.secretsManager, this.athena, jdbcConnectionFactory);
    teradataMetadataHandler.doGetTableLayout(Mockito.mock(BlockAllocator.class), getTableLayoutRequest);
}
Also used : JdbcConnectionFactory(com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory) TableName(com.amazonaws.athena.connector.lambda.domain.TableName) Constraints(com.amazonaws.athena.connector.lambda.domain.predicate.Constraints) Schema(org.apache.arrow.vector.types.pojo.Schema) JdbcCredentialProvider(com.amazonaws.athena.connectors.jdbc.connection.JdbcCredentialProvider) Test(org.junit.Test)

Example 7 with JdbcConnectionFactory

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

the class TeradataRecordHandlerTest 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 TeradataQueryStringBuilder("`");
    final DatabaseConnectionConfig databaseConnectionConfig = new DatabaseConnectionConfig("testCatalog", TeradataConstants.TERADATA_NAME, "teradata://jdbc:teradata://115.113.87.100/TMODE=ANSI,CHARSET=UTF8,DATABASE=TEST,USER=DBC,PASSWORD=DBC");
    this.teradataRecordHandler = new TeradataRecordHandler(databaseConnectionConfig, amazonS3, secretsManager, athena, jdbcConnectionFactory, jdbcSplitQueryBuilder);
}
Also used : JdbcConnectionFactory(com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory) AmazonS3(com.amazonaws.services.s3.AmazonS3) AWSSecretsManager(com.amazonaws.services.secretsmanager.AWSSecretsManager) Connection(java.sql.Connection) DatabaseConnectionConfig(com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig) JdbcCredentialProvider(com.amazonaws.athena.connectors.jdbc.connection.JdbcCredentialProvider) AmazonAthena(com.amazonaws.services.athena.AmazonAthena) Before(org.junit.Before)

Example 8 with JdbcConnectionFactory

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

the class SynapseMetadataHandlerTest 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.synapseMetadataHandler.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());
    SynapseMetadataHandler synapseMetadataHandler = new SynapseMetadataHandler(databaseConnectionConfig, this.secretsManager, this.athena, jdbcConnectionFactory);
    synapseMetadataHandler.doGetTableLayout(Mockito.mock(BlockAllocator.class), getTableLayoutRequest);
}
Also used : JdbcConnectionFactory(com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory) TableName(com.amazonaws.athena.connector.lambda.domain.TableName) Constraints(com.amazonaws.athena.connector.lambda.domain.predicate.Constraints) SQLException(java.sql.SQLException) Schema(org.apache.arrow.vector.types.pojo.Schema) GetTableLayoutRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest) BlockAllocator(com.amazonaws.athena.connector.lambda.data.BlockAllocator) Connection(java.sql.Connection) JdbcCredentialProvider(com.amazonaws.athena.connectors.jdbc.connection.JdbcCredentialProvider) Test(org.junit.Test)

Example 9 with JdbcConnectionFactory

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);
}
Also used : JdbcConnectionFactory(com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory) TableName(com.amazonaws.athena.connector.lambda.domain.TableName) Constraints(com.amazonaws.athena.connector.lambda.domain.predicate.Constraints) Schema(org.apache.arrow.vector.types.pojo.Schema) JdbcCredentialProvider(com.amazonaws.athena.connectors.jdbc.connection.JdbcCredentialProvider) Test(org.junit.Test)

Example 10 with JdbcConnectionFactory

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

the class ImpalaRecordHandlerTest 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 ImpalaQueryStringBuilder("`");
    final DatabaseConnectionConfig databaseConnectionConfig = new DatabaseConnectionConfig("testCatalog", ImpalaConstants.IMPALA_NAME, "impala://jdbc:impala://54.89.6.2:10000/authena;{testSecret}", "testSecret");
    this.impalaRecordHandler = new ImpalaRecordHandler(databaseConnectionConfig, amazonS3, secretsManager, athena, jdbcConnectionFactory, jdbcSplitQueryBuilder);
}
Also used : JdbcConnectionFactory(com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory) AmazonS3(com.amazonaws.services.s3.AmazonS3) GetSecretValueResult(com.amazonaws.services.secretsmanager.model.GetSecretValueResult) AWSSecretsManager(com.amazonaws.services.secretsmanager.AWSSecretsManager) Connection(java.sql.Connection) GetSecretValueRequest(com.amazonaws.services.secretsmanager.model.GetSecretValueRequest) DatabaseConnectionConfig(com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig) JdbcCredentialProvider(com.amazonaws.athena.connectors.jdbc.connection.JdbcCredentialProvider) AmazonAthena(com.amazonaws.services.athena.AmazonAthena) Before(org.junit.Before)

Aggregations

JdbcConnectionFactory (com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory)28 JdbcCredentialProvider (com.amazonaws.athena.connectors.jdbc.connection.JdbcCredentialProvider)27 Connection (java.sql.Connection)22 DatabaseConnectionConfig (com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig)15 AmazonAthena (com.amazonaws.services.athena.AmazonAthena)14 AWSSecretsManager (com.amazonaws.services.secretsmanager.AWSSecretsManager)14 Schema (org.apache.arrow.vector.types.pojo.Schema)14 Before (org.junit.Before)14 TableName (com.amazonaws.athena.connector.lambda.domain.TableName)13 Constraints (com.amazonaws.athena.connector.lambda.domain.predicate.Constraints)13 AmazonS3 (com.amazonaws.services.s3.AmazonS3)13 Test (org.junit.Test)13 BlockAllocator (com.amazonaws.athena.connector.lambda.data.BlockAllocator)9 SQLException (java.sql.SQLException)8 GetTableLayoutRequest (com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest)7 GetSecretValueRequest (com.amazonaws.services.secretsmanager.model.GetSecretValueRequest)4 GetSecretValueResult (com.amazonaws.services.secretsmanager.model.GetSecretValueResult)4 QueryStatusChecker (com.amazonaws.athena.connector.lambda.QueryStatusChecker)1 BlockWriter (com.amazonaws.athena.connector.lambda.data.BlockWriter)1 GetSplitsRequest (com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest)1