Search in sources :

Example 16 with LocalKeyFactory

use of com.amazonaws.athena.connector.lambda.security.LocalKeyFactory in project aws-athena-query-federation by awslabs.

the class AwsCmdbMetadataHandlerTest method setUp.

@Before
public void setUp() throws Exception {
    blockAllocator = new BlockAllocatorImpl();
    Map<TableName, TableProvider> tableProviderMap = new HashMap<>();
    tableProviderMap.putIfAbsent(new TableName("schema1", "table1"), mockTableProvider1);
    tableProviderMap.putIfAbsent(new TableName("schema1", "table2"), mockTableProvider2);
    tableProviderMap.putIfAbsent(new TableName("schema2", "table1"), mockTableProvider3);
    when(mockTableProviderFactory.getTableProviders()).thenReturn(tableProviderMap);
    Map<String, List<TableName>> schemas = new HashMap<>();
    schemas.put("schema1", new ArrayList<>());
    schemas.put("schema2", new ArrayList<>());
    schemas.get("schema1").add(new TableName("schema1", "table1"));
    schemas.get("schema1").add(new TableName("schema1", "table2"));
    schemas.get("schema2").add(new TableName("schema2", "table1"));
    when(mockTableProviderFactory.getSchemas()).thenReturn(schemas);
    handler = new AwsCmdbMetadataHandler(mockTableProviderFactory, new LocalKeyFactory(), mockSecretsManager, mockAthena, bucket, prefix);
    verify(mockTableProviderFactory, times(1)).getTableProviders();
    verify(mockTableProviderFactory, times(1)).getSchemas();
    verifyNoMoreInteractions(mockTableProviderFactory);
}
Also used : TableName(com.amazonaws.athena.connector.lambda.domain.TableName) BlockAllocatorImpl(com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl) HashMap(java.util.HashMap) ArrayList(java.util.ArrayList) List(java.util.List) TableProvider(com.amazonaws.athena.connectors.aws.cmdb.tables.TableProvider) LocalKeyFactory(com.amazonaws.athena.connector.lambda.security.LocalKeyFactory) Before(org.junit.Before)

Example 17 with LocalKeyFactory

use of com.amazonaws.athena.connector.lambda.security.LocalKeyFactory in project aws-athena-query-federation by awslabs.

the class MetricsMetadataHandlerTest method setUp.

@Before
public void setUp() throws Exception {
    handler = new MetricsMetadataHandler(mockMetrics, new LocalKeyFactory(), mockSecretsManager, mockAthena, "spillBucket", "spillPrefix");
    allocator = new BlockAllocatorImpl();
}
Also used : BlockAllocatorImpl(com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl) LocalKeyFactory(com.amazonaws.athena.connector.lambda.security.LocalKeyFactory) Before(org.junit.Before)

Example 18 with LocalKeyFactory

use of com.amazonaws.athena.connector.lambda.security.LocalKeyFactory in project aws-athena-query-federation by awslabs.

the class DynamoDBRecordHandlerTest method setup.

@Before
public void setup() {
    logger.info("{}: enter", testName.getMethodName());
    allocator = new BlockAllocatorImpl();
    handler = new DynamoDBRecordHandler(ddbClient, mock(AmazonS3.class), mock(AWSSecretsManager.class), mock(AmazonAthena.class), "source_type");
    metadataHandler = new DynamoDBMetadataHandler(new LocalKeyFactory(), secretsManager, athena, "spillBucket", "spillPrefix", ddbClient, glueClient);
}
Also used : BlockAllocatorImpl(com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl) LocalKeyFactory(com.amazonaws.athena.connector.lambda.security.LocalKeyFactory) Before(org.junit.Before)

Example 19 with LocalKeyFactory

use of com.amazonaws.athena.connector.lambda.security.LocalKeyFactory in project aws-athena-query-federation by awslabs.

the class DocDBRecordHandlerTest method setUp.

@Before
public void setUp() {
    logger.info("{}: enter", testName.getMethodName());
    schemaForRead = SchemaBuilder.newBuilder().addField("col1", new ArrowType.Int(32, true)).addField("col2", new ArrowType.Utf8()).addField("col3", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)).addField("int", Types.MinorType.INT.getType()).addField("tinyint", Types.MinorType.TINYINT.getType()).addField("smallint", Types.MinorType.SMALLINT.getType()).addField("bigint", Types.MinorType.BIGINT.getType()).addField("uint1", Types.MinorType.UINT1.getType()).addField("uint2", Types.MinorType.UINT2.getType()).addField("uint4", Types.MinorType.UINT4.getType()).addField("uint8", Types.MinorType.UINT8.getType()).addField("float4", Types.MinorType.FLOAT4.getType()).addField("float8", Types.MinorType.FLOAT8.getType()).addField("bit", Types.MinorType.BIT.getType()).addField("varchar", Types.MinorType.VARCHAR.getType()).addField("varbinary", Types.MinorType.VARBINARY.getType()).addField("decimal", new ArrowType.Decimal(10, 2)).addField("decimalLong", new ArrowType.Decimal(36, 2)).addField("unsupported", Types.MinorType.VARCHAR.getType()).addStructField("struct").addChildField("struct", "struct_string", Types.MinorType.VARCHAR.getType()).addChildField("struct", "struct_int", Types.MinorType.INT.getType()).addListField("list", Types.MinorType.VARCHAR.getType()).build();
    when(connectionFactory.getOrCreateConn(anyString())).thenReturn(mockClient);
    allocator = new BlockAllocatorImpl();
    amazonS3 = mock(AmazonS3.class);
    mockDatabase = mock(MongoDatabase.class);
    mockCollection = mock(MongoCollection.class);
    mockIterable = mock(FindIterable.class);
    when(mockClient.getDatabase(eq(DEFAULT_SCHEMA))).thenReturn(mockDatabase);
    when(mockDatabase.getCollection(eq(TEST_TABLE))).thenReturn(mockCollection);
    when(amazonS3.putObject(anyObject(), anyObject(), anyObject(), anyObject())).thenAnswer((InvocationOnMock invocationOnMock) -> {
        InputStream inputStream = (InputStream) invocationOnMock.getArguments()[2];
        ByteHolder byteHolder = new ByteHolder();
        byteHolder.setBytes(ByteStreams.toByteArray(inputStream));
        synchronized (mockS3Storage) {
            mockS3Storage.add(byteHolder);
            logger.info("puObject: total size " + mockS3Storage.size());
        }
        return mock(PutObjectResult.class);
    });
    when(amazonS3.getObject(anyString(), anyString())).thenAnswer((InvocationOnMock invocationOnMock) -> {
        S3Object mockObject = mock(S3Object.class);
        ByteHolder byteHolder;
        synchronized (mockS3Storage) {
            byteHolder = mockS3Storage.get(0);
            mockS3Storage.remove(0);
            logger.info("getObject: total size " + mockS3Storage.size());
        }
        when(mockObject.getObjectContent()).thenReturn(new S3ObjectInputStream(new ByteArrayInputStream(byteHolder.getBytes()), null));
        return mockObject;
    });
    handler = new DocDBRecordHandler(amazonS3, mockSecretsManager, mockAthena, connectionFactory);
    spillReader = new S3BlockSpillReader(amazonS3, allocator);
    mdHandler = new DocDBMetadataHandler(awsGlue, connectionFactory, new LocalKeyFactory(), secretsManager, mockAthena, "spillBucket", "spillPrefix");
}
Also used : AmazonS3(com.amazonaws.services.s3.AmazonS3) ByteArrayInputStream(java.io.ByteArrayInputStream) S3ObjectInputStream(com.amazonaws.services.s3.model.S3ObjectInputStream) InputStream(java.io.InputStream) ArrowType(org.apache.arrow.vector.types.pojo.ArrowType) S3ObjectInputStream(com.amazonaws.services.s3.model.S3ObjectInputStream) S3BlockSpillReader(com.amazonaws.athena.connector.lambda.data.S3BlockSpillReader) FindIterable(com.mongodb.client.FindIterable) LocalKeyFactory(com.amazonaws.athena.connector.lambda.security.LocalKeyFactory) MongoCollection(com.mongodb.client.MongoCollection) BlockAllocatorImpl(com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl) ByteArrayInputStream(java.io.ByteArrayInputStream) InvocationOnMock(org.mockito.invocation.InvocationOnMock) S3Object(com.amazonaws.services.s3.model.S3Object) MongoDatabase(com.mongodb.client.MongoDatabase) Before(org.junit.Before)

Aggregations

LocalKeyFactory (com.amazonaws.athena.connector.lambda.security.LocalKeyFactory)19 Before (org.junit.Before)15 BlockAllocatorImpl (com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl)14 TableName (com.amazonaws.athena.connector.lambda.domain.TableName)5 Test (org.junit.Test)4 InvocationOnMock (org.mockito.invocation.InvocationOnMock)4 QueryStatusChecker (com.amazonaws.athena.connector.lambda.QueryStatusChecker)2 Constraints (com.amazonaws.athena.connector.lambda.domain.predicate.Constraints)2 AmazonS3 (com.amazonaws.services.s3.AmazonS3)2 GetSecretValueRequest (com.amazonaws.services.secretsmanager.model.GetSecretValueRequest)2 GetSecretValueResult (com.amazonaws.services.secretsmanager.model.GetSecretValueResult)2 Field (org.apache.arrow.vector.types.pojo.Field)2 Schema (org.apache.arrow.vector.types.pojo.Schema)2 Mockito.anyString (org.mockito.Mockito.anyString)2 Block (com.amazonaws.athena.connector.lambda.data.Block)1 BlockAllocator (com.amazonaws.athena.connector.lambda.data.BlockAllocator)1 BlockWriter (com.amazonaws.athena.connector.lambda.data.BlockWriter)1 S3BlockSpillReader (com.amazonaws.athena.connector.lambda.data.S3BlockSpillReader)1 GetSplitsRequest (com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest)1 GetSplitsResponse (com.amazonaws.athena.connector.lambda.metadata.GetSplitsResponse)1