Search in sources :

Example 1 with SnowflakeMetadataHandler

use of com.amazonaws.athena.connectors.snowflake.SnowflakeMetadataHandler in project aws-athena-query-federation by awslabs.

the class SnowflakeMetadataHandlerTest method setup.

@Before
public void setup() {
    environmentVariables.set("pagecount", "300000");
    environmentVariables.set("partitionlimit", "300000");
    this.jdbcConnectionFactory = Mockito.mock(JdbcConnectionFactory.class, Mockito.RETURNS_DEEP_STUBS);
    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\"}"));
    this.snowflakeMetadataHandler = new SnowflakeMetadataHandler(databaseConnectionConfig, this.secretsManager, this.athena, this.jdbcConnectionFactory);
    this.federatedIdentity = Mockito.mock(FederatedIdentity.class);
    this.blockAllocator = Mockito.mock(BlockAllocator.class);
}
Also used : JdbcConnectionFactory(com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory) GetSecretValueResult(com.amazonaws.services.secretsmanager.model.GetSecretValueResult) FederatedIdentity(com.amazonaws.athena.connector.lambda.security.FederatedIdentity) AWSSecretsManager(com.amazonaws.services.secretsmanager.AWSSecretsManager) GetSecretValueRequest(com.amazonaws.services.secretsmanager.model.GetSecretValueRequest) JdbcCredentialProvider(com.amazonaws.athena.connectors.jdbc.connection.JdbcCredentialProvider) AmazonAthena(com.amazonaws.services.athena.AmazonAthena) SnowflakeMetadataHandler(com.amazonaws.athena.connectors.snowflake.SnowflakeMetadataHandler) Before(org.junit.Before)

Example 2 with SnowflakeMetadataHandler

use of com.amazonaws.athena.connectors.snowflake.SnowflakeMetadataHandler in project aws-athena-query-federation by awslabs.

the class SnowflakeMetadataHandlerTest 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.snowflakeMetadataHandler.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());
    SnowflakeMetadataHandler snowflakeMetadataHandler = new SnowflakeMetadataHandler(databaseConnectionConfig, this.secretsManager, this.athena, jdbcConnectionFactory);
    snowflakeMetadataHandler.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) SnowflakeMetadataHandler(com.amazonaws.athena.connectors.snowflake.SnowflakeMetadataHandler) Test(org.junit.Test)

Aggregations

JdbcConnectionFactory (com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory)2 JdbcCredentialProvider (com.amazonaws.athena.connectors.jdbc.connection.JdbcCredentialProvider)2 SnowflakeMetadataHandler (com.amazonaws.athena.connectors.snowflake.SnowflakeMetadataHandler)2 TableName (com.amazonaws.athena.connector.lambda.domain.TableName)1 Constraints (com.amazonaws.athena.connector.lambda.domain.predicate.Constraints)1 FederatedIdentity (com.amazonaws.athena.connector.lambda.security.FederatedIdentity)1 AmazonAthena (com.amazonaws.services.athena.AmazonAthena)1 AWSSecretsManager (com.amazonaws.services.secretsmanager.AWSSecretsManager)1 GetSecretValueRequest (com.amazonaws.services.secretsmanager.model.GetSecretValueRequest)1 GetSecretValueResult (com.amazonaws.services.secretsmanager.model.GetSecretValueResult)1 Schema (org.apache.arrow.vector.types.pojo.Schema)1 Before (org.junit.Before)1 Test (org.junit.Test)1