diff --git a/.gitignore b/.gitignore index d3c5ae7f08..37d101daaa 100644 --- a/.gitignore +++ b/.gitignore @@ -1,9 +1,10 @@ */target */dependency-reduced-pom.xml +dependency-reduced-pom.xml */*.xml.unformatted .idea/ /target/ -*/*.iml +**/*.iml .classpath .factorypath .project diff --git a/athena-aws-cmdb/src/test/java/com/amazonaws/athena/connectors/aws/cmdb/AwsCmdbMetadataHandlerTest.java b/athena-aws-cmdb/src/test/java/com/amazonaws/athena/connectors/aws/cmdb/AwsCmdbMetadataHandlerTest.java index 565a7939e1..ba8f6f815e 100644 --- a/athena-aws-cmdb/src/test/java/com/amazonaws/athena/connectors/aws/cmdb/AwsCmdbMetadataHandlerTest.java +++ b/athena-aws-cmdb/src/test/java/com/amazonaws/athena/connectors/aws/cmdb/AwsCmdbMetadataHandlerTest.java @@ -22,10 +22,8 @@ import com.amazonaws.athena.connector.lambda.data.Block; import com.amazonaws.athena.connector.lambda.data.BlockAllocator; import com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl; -import com.amazonaws.athena.connector.lambda.data.BlockWriter; import com.amazonaws.athena.connector.lambda.data.SchemaBuilder; import com.amazonaws.athena.connector.lambda.domain.TableName; -import com.amazonaws.athena.connector.lambda.domain.predicate.ConstraintEvaluator; import com.amazonaws.athena.connector.lambda.domain.predicate.Constraints; import com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest; import com.amazonaws.athena.connector.lambda.metadata.GetSplitsResponse; @@ -168,7 +166,7 @@ public void doListTables() @Test public void doGetTable() { - GetTableRequest request = new GetTableRequest(identity, queryId, catalog, new TableName("schema1", "table1")); + GetTableRequest request = new GetTableRequest(identity, queryId, catalog, new TableName("schema1", "table1"), Collections.emptyMap()); when(mockTableProvider1.getTable(eq(blockAllocator), eq(request))).thenReturn(mock(GetTableResponse.class)); GetTableResponse response = handler.doGetTable(blockAllocator, request); diff --git a/athena-aws-cmdb/src/test/java/com/amazonaws/athena/connectors/aws/cmdb/tables/AbstractTableProviderTest.java b/athena-aws-cmdb/src/test/java/com/amazonaws/athena/connectors/aws/cmdb/tables/AbstractTableProviderTest.java index a3f1c2857e..f4d6ba505a 100644 --- a/athena-aws-cmdb/src/test/java/com/amazonaws/athena/connectors/aws/cmdb/tables/AbstractTableProviderTest.java +++ b/athena-aws-cmdb/src/test/java/com/amazonaws/athena/connectors/aws/cmdb/tables/AbstractTableProviderTest.java @@ -178,7 +178,7 @@ public void getTableName() @Test public void readTableTest() { - GetTableRequest request = new GetTableRequest(identity, expectedQuery, expectedCatalog, expectedTableName); + GetTableRequest request = new GetTableRequest(identity, expectedQuery, expectedCatalog, expectedTableName, Collections.emptyMap()); GetTableResponse response = provider.getTable(allocator, request); assertTrue(response.getSchema().getFields().size() > 1); diff --git a/athena-cloudera-hive/src/test/java/com/amazonaws/athena/connectors/cloudera/HiveMetadataHandlerTest.java b/athena-cloudera-hive/src/test/java/com/amazonaws/athena/connectors/cloudera/HiveMetadataHandlerTest.java index 24bde7a335..b1520d4a3d 100644 --- a/athena-cloudera-hive/src/test/java/com/amazonaws/athena/connectors/cloudera/HiveMetadataHandlerTest.java +++ b/athena-cloudera-hive/src/test/java/com/amazonaws/athena/connectors/cloudera/HiveMetadataHandlerTest.java @@ -28,7 +28,6 @@ import com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig; import com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory; import com.amazonaws.athena.connectors.jdbc.connection.JdbcCredentialProvider; -import com.amazonaws.athena.connectors.jdbc.manager.JDBCUtil; import com.amazonaws.services.athena.AmazonAthena; import com.amazonaws.services.secretsmanager.AWSSecretsManager; import com.amazonaws.services.secretsmanager.model.GetSecretValueRequest; @@ -296,7 +295,7 @@ public void doGetTable() Mockito.when(this.connection.getMetaData().getColumns("testCatalog", inputTableName.getSchemaName(), inputTableName.getTableName(), null)).thenReturn(resultSet1); Mockito.when(this.connection.getCatalog()).thenReturn("testCatalog"); GetTableResponse getTableResponse = this.hiveMetadataHandler.doGetTable( - this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); Assert.assertEquals(inputTableName, getTableResponse.getTableName()); Assert.assertEquals("testCatalog", getTableResponse.getCatalogName()); } @@ -305,7 +304,7 @@ public void doGetTable() public void doGetTableNoColumns() throws Exception { TableName inputTableName = new TableName("testSchema", "testTable"); - this.hiveMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.hiveMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } @Test(expected = SQLException.class) @@ -315,6 +314,6 @@ public void doGetTableSQLException() TableName inputTableName = new TableName("testSchema", "testTable"); Mockito.when(this.connection.getMetaData().getColumns(nullable(String.class), nullable(String.class), nullable(String.class), nullable(String.class))) .thenThrow(new SQLException()); - this.hiveMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.hiveMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } } diff --git a/athena-cloudera-impala/src/test/java/com/amazonaws/athena/connectors/cloudera/ImpalaMetadataHandlerTest.java b/athena-cloudera-impala/src/test/java/com/amazonaws/athena/connectors/cloudera/ImpalaMetadataHandlerTest.java index ecfef549c5..09746df6da 100644 --- a/athena-cloudera-impala/src/test/java/com/amazonaws/athena/connectors/cloudera/ImpalaMetadataHandlerTest.java +++ b/athena-cloudera-impala/src/test/java/com/amazonaws/athena/connectors/cloudera/ImpalaMetadataHandlerTest.java @@ -278,7 +278,7 @@ public void doGetTable() Mockito.when(this.connection.getMetaData().getColumns("testCatalog", inputTableName.getSchemaName(), inputTableName.getTableName(), null)).thenReturn(resultSet1); Mockito.when(this.connection.getCatalog()).thenReturn("testCatalog"); GetTableResponse getTableResponse = this.impalaMetadataHandler.doGetTable( - this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); Assert.assertEquals(inputTableName, getTableResponse.getTableName()); Assert.assertEquals("testCatalog", getTableResponse.getCatalogName()); } @@ -286,7 +286,7 @@ public void doGetTable() public void doGetTableNoColumns() throws Exception { TableName inputTableName = new TableName("testSchema", "testTable"); - this.impalaMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.impalaMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } @Test(expected = SQLException.class) @@ -296,6 +296,6 @@ public void doGetTableSQLException() TableName inputTableName = new TableName("testSchema", "testTable"); Mockito.when(this.connection.getMetaData().getColumns(nullable(String.class), nullable(String.class), nullable(String.class), nullable(String.class))) .thenThrow(new SQLException()); - this.impalaMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.impalaMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } } diff --git a/athena-cloudwatch-metrics/src/test/java/com/amazonaws/athena/connectors/cloudwatch/metrics/MetricsMetadataHandlerTest.java b/athena-cloudwatch-metrics/src/test/java/com/amazonaws/athena/connectors/cloudwatch/metrics/MetricsMetadataHandlerTest.java index 7b9d1f0e26..2e1992da0e 100644 --- a/athena-cloudwatch-metrics/src/test/java/com/amazonaws/athena/connectors/cloudwatch/metrics/MetricsMetadataHandlerTest.java +++ b/athena-cloudwatch-metrics/src/test/java/com/amazonaws/athena/connectors/cloudwatch/metrics/MetricsMetadataHandlerTest.java @@ -151,7 +151,7 @@ public void doGetMetricsTable() { logger.info("doGetMetricsTable - enter"); - GetTableRequest metricsTableReq = new GetTableRequest(identity, "queryId", "default", new TableName(defaultSchema, "metrics")); + GetTableRequest metricsTableReq = new GetTableRequest(identity, "queryId", "default", new TableName(defaultSchema, "metrics"), Collections.emptyMap()); GetTableResponse metricsTableRes = handler.doGetTable(allocator, metricsTableReq); logger.info("doGetMetricsTable - {} {}", metricsTableRes.getTableName(), metricsTableRes.getSchema()); @@ -170,7 +170,7 @@ public void doGetMetricSamplesTable() GetTableRequest metricsTableReq = new GetTableRequest(identity, "queryId", "default", - new TableName(defaultSchema, "metric_samples")); + new TableName(defaultSchema, "metric_samples"), Collections.emptyMap()); GetTableResponse metricsTableRes = handler.doGetTable(allocator, metricsTableReq); logger.info("doGetMetricSamplesTable - {} {}", metricsTableRes.getTableName(), metricsTableRes.getSchema()); diff --git a/athena-cloudwatch/src/test/java/com/amazonaws/athena/connectors/cloudwatch/CloudwatchMetadataHandlerTest.java b/athena-cloudwatch/src/test/java/com/amazonaws/athena/connectors/cloudwatch/CloudwatchMetadataHandlerTest.java index 41e37078ff..22a876dbae 100644 --- a/athena-cloudwatch/src/test/java/com/amazonaws/athena/connectors/cloudwatch/CloudwatchMetadataHandlerTest.java +++ b/athena-cloudwatch/src/test/java/com/amazonaws/athena/connectors/cloudwatch/CloudwatchMetadataHandlerTest.java @@ -269,7 +269,7 @@ else if (Integer.valueOf(request.getNextToken()) < 3) { return result; }); - GetTableRequest req = new GetTableRequest(identity, "queryId", "default", new TableName(expectedSchema, "table-9")); + GetTableRequest req = new GetTableRequest(identity, "queryId", "default", new TableName(expectedSchema, "table-9"), Collections.emptyMap()); GetTableResponse res = handler.doGetTable(allocator, req); logger.info("doGetTable - {} {}", res.getTableName(), res.getSchema()); diff --git a/athena-datalakegen2/src/test/java/com/amazonaws/athena/connectors/datalakegen2/DataLakeGen2MetadataHandlerTest.java b/athena-datalakegen2/src/test/java/com/amazonaws/athena/connectors/datalakegen2/DataLakeGen2MetadataHandlerTest.java index 144aad4648..c37359bab8 100644 --- a/athena-datalakegen2/src/test/java/com/amazonaws/athena/connectors/datalakegen2/DataLakeGen2MetadataHandlerTest.java +++ b/athena-datalakegen2/src/test/java/com/amazonaws/athena/connectors/datalakegen2/DataLakeGen2MetadataHandlerTest.java @@ -199,7 +199,7 @@ public void doGetTable() Mockito.when(connection.getMetaData().getColumns("testCatalog", inputTableName.getSchemaName(), inputTableName.getTableName(), null)).thenReturn(resultSet); Mockito.when(connection.getCatalog()).thenReturn("testCatalog"); GetTableResponse getTableResponse = this.dataLakeGen2MetadataHandler.doGetTable( - blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); Assert.assertEquals(expected, getTableResponse.getSchema()); Assert.assertEquals(inputTableName, getTableResponse.getTableName()); diff --git a/athena-db2-as400/src/test/java/com/amazonaws/athena/connectors/db2as400/Db2As400MetadataHandlerTest.java b/athena-db2-as400/src/test/java/com/amazonaws/athena/connectors/db2as400/Db2As400MetadataHandlerTest.java index ed088c68e4..5f16236d1a 100644 --- a/athena-db2-as400/src/test/java/com/amazonaws/athena/connectors/db2as400/Db2As400MetadataHandlerTest.java +++ b/athena-db2-as400/src/test/java/com/amazonaws/athena/connectors/db2as400/Db2As400MetadataHandlerTest.java @@ -222,7 +222,7 @@ public void doGetTable() TableName inputTableName = new TableName("TESTSCHEMA", "TESTTABLE"); GetTableResponse getTableResponse = this.db2As400MetadataHandler.doGetTable( - this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); Assert.assertEquals(expected, getTableResponse.getSchema()); Assert.assertEquals(new TableName(schemaName, tableName), getTableResponse.getTableName()); Assert.assertEquals("testCatalog", getTableResponse.getCatalogName()); @@ -248,7 +248,7 @@ public void doGetTableCaseSensitivity() TableName inputTableName = new TableName(schemaName, tableName); Mockito.when(this.connection.getMetaData().getColumns(nullable(String.class), nullable(String.class), nullable(String.class), nullable(String.class))) .thenThrow(new SQLException()); - this.db2As400MetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.db2As400MetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } @Test(expected = SQLException.class) @@ -270,7 +270,7 @@ public void doGetTableCaseSensitivity2() TableName inputTableName = new TableName(schemaName, tableName); Mockito.when(this.connection.getMetaData().getColumns(nullable(String.class), nullable(String.class), nullable(String.class), nullable(String.class))) .thenThrow(new SQLException()); - this.db2As400MetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.db2As400MetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } @Test(expected = SQLException.class) @@ -292,7 +292,7 @@ public void doGetTableCaseSensitivity3() TableName inputTableName = new TableName(schemaName, tableName); Mockito.when(this.connection.getMetaData().getColumns(nullable(String.class), nullable(String.class), nullable(String.class), nullable(String.class))) .thenThrow(new SQLException()); - this.db2As400MetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.db2As400MetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } @Test(expected = SQLException.class) @@ -314,7 +314,7 @@ public void doGetTableCaseSensitivity4() TableName inputTableName = new TableName(schemaName, tableName); Mockito.when(this.connection.getMetaData().getColumns(nullable(String.class), nullable(String.class), nullable(String.class), nullable(String.class))) .thenThrow(new SQLException()); - this.db2As400MetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.db2As400MetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } @Test diff --git a/athena-db2/src/test/java/com/amazonaws/athena/connectors/db2/Db2MetadataHandlerTest.java b/athena-db2/src/test/java/com/amazonaws/athena/connectors/db2/Db2MetadataHandlerTest.java index c64b00ff7e..02ff20fa93 100644 --- a/athena-db2/src/test/java/com/amazonaws/athena/connectors/db2/Db2MetadataHandlerTest.java +++ b/athena-db2/src/test/java/com/amazonaws/athena/connectors/db2/Db2MetadataHandlerTest.java @@ -223,7 +223,7 @@ public void doGetTable() TableName inputTableName = new TableName("TESTSCHEMA", "TESTTABLE"); GetTableResponse getTableResponse = this.db2MetadataHandler.doGetTable( - this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); Assert.assertEquals(expected, getTableResponse.getSchema()); Assert.assertEquals(new TableName(schemaName, tableName), getTableResponse.getTableName()); Assert.assertEquals("testCatalog", getTableResponse.getCatalogName()); @@ -249,7 +249,7 @@ public void doGetTableCaseSensitivity() TableName inputTableName = new TableName(schemaName, tableName); Mockito.when(this.connection.getMetaData().getColumns(nullable(String.class), nullable(String.class), nullable(String.class), nullable(String.class))) .thenThrow(new SQLException()); - this.db2MetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.db2MetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } @Test(expected = SQLException.class) @@ -271,7 +271,7 @@ public void doGetTableCaseSensitivity2() TableName inputTableName = new TableName(schemaName, tableName); Mockito.when(this.connection.getMetaData().getColumns(nullable(String.class), nullable(String.class), nullable(String.class), nullable(String.class))) .thenThrow(new SQLException()); - this.db2MetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.db2MetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } @Test(expected = SQLException.class) @@ -293,7 +293,7 @@ public void doGetTableCaseSensitivity3() TableName inputTableName = new TableName(schemaName, tableName); Mockito.when(this.connection.getMetaData().getColumns(nullable(String.class), nullable(String.class), nullable(String.class), nullable(String.class))) .thenThrow(new SQLException()); - this.db2MetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.db2MetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } @Test(expected = SQLException.class) @@ -315,7 +315,7 @@ public void doGetTableCaseSensitivity4() TableName inputTableName = new TableName(schemaName, tableName); Mockito.when(this.connection.getMetaData().getColumns(nullable(String.class), nullable(String.class), nullable(String.class), nullable(String.class))) .thenThrow(new SQLException()); - this.db2MetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.db2MetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } @Test diff --git a/athena-docdb/src/test/java/com/amazonaws/athena/connectors/docdb/DocDBMetadataHandlerTest.java b/athena-docdb/src/test/java/com/amazonaws/athena/connectors/docdb/DocDBMetadataHandlerTest.java index dd16740da2..bc46eaf1be 100644 --- a/athena-docdb/src/test/java/com/amazonaws/athena/connectors/docdb/DocDBMetadataHandlerTest.java +++ b/athena-docdb/src/test/java/com/amazonaws/athena/connectors/docdb/DocDBMetadataHandlerTest.java @@ -219,7 +219,7 @@ public void doGetTable() when(mockIterable.batchSize(anyInt())).thenReturn(mockIterable); when(mockIterable.iterator()).thenReturn(new StubbingCursor(documents.iterator())); - GetTableRequest req = new GetTableRequest(IDENTITY, QUERY_ID, DEFAULT_CATALOG, TABLE_NAME); + GetTableRequest req = new GetTableRequest(IDENTITY, QUERY_ID, DEFAULT_CATALOG, TABLE_NAME, Collections.emptyMap()); GetTableResponse res = handler.doGetTable(allocator, req); logger.info("doGetTable - {}", res); @@ -307,7 +307,7 @@ connectionFactory, new LocalKeyFactory(), secretsManager, mockAthena, when(mockIterable.iterator()).thenReturn(new StubbingCursor(documents.iterator())); TableName tableNameInput = new TableName("DEfault", TEST_TABLE.toUpperCase()); - GetTableRequest req = new GetTableRequest(IDENTITY, QUERY_ID, DEFAULT_CATALOG, tableNameInput); + GetTableRequest req = new GetTableRequest(IDENTITY, QUERY_ID, DEFAULT_CATALOG, tableNameInput, Collections.emptyMap()); GetTableResponse res = caseInsensitiveHandler.doGetTable(allocator, req); assertEquals(DEFAULT_SCHEMA, res.getTableName().getSchemaName()); @@ -359,7 +359,7 @@ connectionFactory, new LocalKeyFactory(), secretsManager, mockAthena, when(mockListDatabaseNamesIterable.spliterator()).thenReturn(ImmutableList.of(DEFAULT_SCHEMA, DEFAULT_SCHEMA.toUpperCase()).spliterator()); TableName tableNameInput = new TableName("deFAULT", TEST_TABLE.toUpperCase()); - GetTableRequest req = new GetTableRequest(IDENTITY, QUERY_ID, DEFAULT_CATALOG, tableNameInput); + GetTableRequest req = new GetTableRequest(IDENTITY, QUERY_ID, DEFAULT_CATALOG, tableNameInput, Collections.emptyMap()); try { GetTableResponse res = caseInsensitiveHandler.doGetTable(allocator, req); fail("doGetTableCaseInsensitiveMatchMultipleMatch should failed"); @@ -411,7 +411,7 @@ public void doGetTableCaseInsensitiveMatchNotEnable() when(mockIterable.iterator()).thenReturn(new StubbingCursor(documents.iterator())); TableName tableNameInput = new TableName(mixedCaseSchemaName, mixedCaseTableName); - GetTableRequest req = new GetTableRequest(IDENTITY, QUERY_ID, DEFAULT_CATALOG, tableNameInput); + GetTableRequest req = new GetTableRequest(IDENTITY, QUERY_ID, DEFAULT_CATALOG, tableNameInput, Collections.emptyMap()); GetTableResponse res = handler.doGetTable(allocator, req); assertEquals(mixedCaseSchemaName, res.getTableName().getSchemaName()); diff --git a/athena-docdb/src/test/java/com/amazonaws/athena/connectors/docdb/DocDBRecordHandlerTest.java b/athena-docdb/src/test/java/com/amazonaws/athena/connectors/docdb/DocDBRecordHandlerTest.java index c0f7333f7f..18a1947c79 100644 --- a/athena-docdb/src/test/java/com/amazonaws/athena/connectors/docdb/DocDBRecordHandlerTest.java +++ b/athena-docdb/src/test/java/com/amazonaws/athena/connectors/docdb/DocDBRecordHandlerTest.java @@ -69,7 +69,6 @@ import org.mockito.Mockito; import org.mockito.invocation.InvocationOnMock; import org.mockito.junit.MockitoJUnitRunner; -import org.mockito.stubbing.Answer; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -391,7 +390,7 @@ public void nestedStructTest() when(mockIterable.batchSize(anyInt())).thenReturn(mockIterable); when(mockIterable.iterator()).thenReturn(new StubbingCursor(documents.iterator())); - GetTableRequest req = new GetTableRequest(IDENTITY, QUERY_ID, DEFAULT_CATALOG, TABLE_NAME); + GetTableRequest req = new GetTableRequest(IDENTITY, QUERY_ID, DEFAULT_CATALOG, TABLE_NAME, Collections.emptyMap()); GetTableResponse res = mdHandler.doGetTable(allocator, req); logger.info("doGetTable - {}", res); diff --git a/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBMetadataHandlerTest.java b/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBMetadataHandlerTest.java index c210c7d8d9..5228664d75 100644 --- a/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBMetadataHandlerTest.java +++ b/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBMetadataHandlerTest.java @@ -255,7 +255,7 @@ public void doGetTable() { when(glueClient.getTable(any())).thenThrow(new AmazonServiceException("")); - GetTableRequest req = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, TEST_TABLE_NAME); + GetTableRequest req = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, TEST_TABLE_NAME, Collections.emptyMap()); GetTableResponse res = handler.doGetTable(allocator, req); logger.info("doGetTable - {}", res.getSchema()); @@ -271,7 +271,7 @@ public void doGetEmptyTable() { when(glueClient.getTable(any())).thenThrow(new AmazonServiceException("")); - GetTableRequest req = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, TEST_TABLE_2_NAME); + GetTableRequest req = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, TEST_TABLE_2_NAME, Collections.emptyMap()); GetTableResponse res = handler.doGetTable(allocator, req); logger.info("doGetEmptyTable - {}", res.getSchema()); @@ -286,7 +286,7 @@ public void testCaseInsensitiveResolve() { when(glueClient.getTable(any())).thenThrow(new AmazonServiceException("")); - GetTableRequest req = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, TEST_TABLE_2_NAME); + GetTableRequest req = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, TEST_TABLE_2_NAME, Collections.emptyMap()); GetTableResponse res = handler.doGetTable(allocator, req); logger.info("doGetTable - {}", res.getSchema()); @@ -402,7 +402,7 @@ public void doGetTableLayoutQueryIndex() TEST_QUERY_ID, TEST_CATALOG_NAME, TEST_TABLE_NAME, - new Constraints(constraintsMap), + new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), SchemaBuilder.newBuilder().build(), Collections.EMPTY_SET)); // Verify that both bounds are present for col_6 which is not a sort key @@ -435,7 +435,7 @@ public void doGetTableLayoutQueryIndex() TEST_QUERY_ID, TEST_CATALOG_NAME, TEST_TABLE_NAME, - new Constraints(constraintsMap), + new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), SchemaBuilder.newBuilder().build(), Collections.EMPTY_SET)); // Verify that both bounds are present for col_6 which is not a sort key @@ -509,7 +509,7 @@ public void doGetSplitsScan() TEST_QUERY_ID, TEST_CATALOG_NAME, TEST_TABLE_NAME, - new Constraints(ImmutableMap.of()), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), SchemaBuilder.newBuilder().build(), Collections.EMPTY_SET)); @@ -519,7 +519,7 @@ public void doGetSplitsScan() TEST_TABLE_NAME, layoutResponse.getPartitions(), ImmutableList.of(), - new Constraints(new HashMap<>()), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), null); logger.info("doGetSplits: req[{}]", req); @@ -563,7 +563,7 @@ public void doGetSplitsQuery() TEST_TABLE_NAME, layoutResponse.getPartitions(), ImmutableList.of("col_0"), - new Constraints(new HashMap<>()), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), null); logger.info("doGetSplits: req[{}]", req); @@ -608,7 +608,7 @@ public void validateSourceTableNamePropagation() when(glueClient.getTable(any())).thenReturn(mockResult); TableName tableName = new TableName(DEFAULT_SCHEMA, "glueTableForTestTable"); - GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName); + GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName, Collections.emptyMap()); GetTableResponse getTableResponse = handler.doGetTable(allocator, getTableRequest); logger.info("validateSourceTableNamePropagation: GetTableResponse[{}]", getTableResponse); Map customMetadata = getTableResponse.getSchema().getCustomMetadata(); @@ -619,7 +619,7 @@ public void validateSourceTableNamePropagation() TEST_QUERY_ID, TEST_CATALOG_NAME, tableName, - new Constraints(ImmutableMap.of()), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), getTableResponse.getSchema(), Collections.EMPTY_SET); @@ -649,7 +649,7 @@ public void doGetTableLayoutScanWithTypeOverride() when(glueClient.getTable(any())).thenReturn(mockResult); TableName tableName = new TableName(DEFAULT_SCHEMA, "glueTableForTestTable"); - GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName); + GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName, Collections.emptyMap()); GetTableResponse getTableResponse = handler.doGetTable(allocator, getTableRequest); logger.info("validateSourceTableNamePropagation: GetTableResponse[{}]", getTableResponse); Map customMetadata = getTableResponse.getSchema().getCustomMetadata(); diff --git a/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBRecordHandlerTest.java b/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBRecordHandlerTest.java index 8f226734ce..a601c7728a 100644 --- a/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBRecordHandlerTest.java +++ b/athena-dynamodb/src/test/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBRecordHandlerTest.java @@ -72,11 +72,13 @@ import java.time.LocalDateTime; import java.time.ZonedDateTime; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import java.util.Map; import java.util.TimeZone; import java.util.UUID; +import static com.amazonaws.athena.connector.lambda.domain.predicate.Constraints.DEFAULT_NO_LIMIT; import static com.amazonaws.athena.connector.lambda.handlers.GlueMetadataHandler.COLUMN_NAME_MAPPING_PROPERTY; import static com.amazonaws.athena.connector.lambda.handlers.GlueMetadataHandler.DATETIME_FORMAT_MAPPING_PROPERTY; import static com.amazonaws.athena.connector.lambda.handlers.GlueMetadataHandler.SOURCE_TABLE_PROPERTY; @@ -162,7 +164,7 @@ public void testReadScanSplit() TEST_TABLE_NAME, schema, split, - new Constraints(ImmutableMap.of()), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), 100_000_000_000L, // too big to spill 100_000_000_000L); @@ -263,7 +265,7 @@ public void testReadScanSplitFiltered() TEST_TABLE_NAME, schema, split, - new Constraints(ImmutableMap.of()), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), 100_000_000_000L, // too big to spill 100_000_000_000L); @@ -300,7 +302,7 @@ public void testReadQuerySplit() TEST_TABLE_NAME, schema, split, - new Constraints(ImmutableMap.of()), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), 100_000_000_000L, // too big to spill 100_000_000_000L); @@ -374,7 +376,7 @@ public void testZeroRowQuery() TEST_TABLE_NAME, schema, split, - new Constraints(ImmutableMap.of()), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), 100_000_000_000L, // too big to spill 100_000_000_000L); @@ -415,7 +417,7 @@ public void testDateTimeSupportFromGlueTable() throws Exception when(glueClient.getTable(any())).thenReturn(mockResult); TableName tableName = new TableName(DEFAULT_SCHEMA, TEST_TABLE3); - GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName); + GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName, Collections.emptyMap()); GetTableResponse getTableResponse = metadataHandler.doGetTable(allocator, getTableRequest); logger.info("testDateTimeSupportFromGlueTable: GetTableResponse[{}]", getTableResponse); logger.info("testDateTimeSupportFromGlueTable: GetTableResponse Schema[{}]", getTableResponse.getSchema()); @@ -435,7 +437,7 @@ public void testDateTimeSupportFromGlueTable() throws Exception TEST_TABLE_3_NAME, schema3, split, - new Constraints(ImmutableMap.of()), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), 100_000_000_000L, // too big to spill 100_000_000_000L); @@ -472,7 +474,7 @@ public void testStructWithNullFromGlueTable() throws Exception when(glueClient.getTable(any())).thenReturn(mockResult); TableName tableName = new TableName(DEFAULT_SCHEMA, TEST_TABLE4); - GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName); + GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName, Collections.emptyMap()); GetTableResponse getTableResponse = metadataHandler.doGetTable(allocator, getTableRequest); logger.info("testStructWithNullFromGlueTable: GetTableResponse[{}]", getTableResponse); logger.info("testStructWithNullFromGlueTable: GetTableResponse Schema[{}]", getTableResponse.getSchema()); @@ -499,7 +501,7 @@ public void testStructWithNullFromGlueTable() throws Exception TEST_TABLE_4_NAME, schema4, split, - new Constraints(ImmutableMap.of()), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), 100_000_000_000L, // too big to spill 100_000_000_000L); @@ -519,7 +521,7 @@ public void testStructWithNullFromDdbTable() throws Exception when(glueClient.getTable(any())).thenThrow(new EntityNotFoundException("")); TableName tableName = new TableName(DEFAULT_SCHEMA, TEST_TABLE4); - GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName); + GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName, Collections.emptyMap()); GetTableResponse getTableResponse = metadataHandler.doGetTable(allocator, getTableRequest); logger.info("testStructWithNullFromGlueTable: GetTableResponse[{}]", getTableResponse); logger.info("testStructWithNullFromGlueTable: GetTableResponse Schema[{}]", getTableResponse.getSchema()); @@ -544,7 +546,7 @@ public void testStructWithNullFromDdbTable() throws Exception TEST_TABLE_4_NAME, schema4, split, - new Constraints(ImmutableMap.of()), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), 100_000_000_000L, // too big to spill 100_000_000_000L); @@ -582,7 +584,7 @@ public void testMapWithSchemaFromGlueTable() throws Exception when(glueClient.getTable(any())).thenReturn(mockResult); TableName tableName = new TableName(DEFAULT_SCHEMA, TEST_TABLE5); - GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName); + GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName, Collections.emptyMap()); GetTableResponse getTableResponse = metadataHandler.doGetTable(allocator, getTableRequest); logger.info("testMapWithSchemaFromGlueTable: GetTableResponse[{}]", getTableResponse); logger.info("testMapWithSchemaFromGlueTable: GetTableResponse Schema[{}]", getTableResponse.getSchema()); @@ -602,7 +604,7 @@ public void testMapWithSchemaFromGlueTable() throws Exception TEST_TABLE_5_NAME, schema5, split, - new Constraints(ImmutableMap.of()), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), 100_000_000_000L, // too big to spill 100_000_000_000L); @@ -635,7 +637,7 @@ public void testStructWithSchemaFromGlueTable() throws Exception when(glueClient.getTable(any())).thenReturn(mockResult); TableName tableName = new TableName(DEFAULT_SCHEMA, TEST_TABLE6); - GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName); + GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName, Collections.emptyMap()); GetTableResponse getTableResponse = metadataHandler.doGetTable(allocator, getTableRequest); logger.info("testStructWithSchemaFromGlueTable: GetTableResponse[{}]", getTableResponse); logger.info("testStructWithSchemaFromGlueTable: GetTableResponse Schema[{}]", getTableResponse.getSchema()); @@ -655,7 +657,7 @@ public void testStructWithSchemaFromGlueTable() throws Exception TEST_TABLE_6_NAME, schema, split, - new Constraints(ImmutableMap.of()), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), 100_000_000_000L, // too big to spill 100_000_000_000L); @@ -690,7 +692,7 @@ public void testListWithSchemaFromGlueTable() throws Exception when(glueClient.getTable(any())).thenReturn(mockResult); TableName tableName = new TableName(DEFAULT_SCHEMA, TEST_TABLE7); - GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName); + GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName, Collections.emptyMap()); GetTableResponse getTableResponse = metadataHandler.doGetTable(allocator, getTableRequest); logger.info("testListWithSchemaFromGlueTable: GetTableResponse[{}]", getTableResponse); logger.info("testListWithSchemaFromGlueTable: GetTableResponse Schema[{}]", getTableResponse.getSchema()); @@ -710,7 +712,7 @@ public void testListWithSchemaFromGlueTable() throws Exception TEST_TABLE_7_NAME, schema, split, - new Constraints(ImmutableMap.of()), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), 100_000_000_000L, // too big to spill 100_000_000_000L); @@ -770,7 +772,7 @@ public void testNumMapWithSchemaFromGlueTable() throws Exception when(glueClient.getTable(any())).thenReturn(mockResult); TableName tableName = new TableName(DEFAULT_SCHEMA, TEST_TABLE8); - GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName); + GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName, Collections.emptyMap()); GetTableResponse getTableResponse = metadataHandler.doGetTable(allocator, getTableRequest); logger.info("testNumMapWithSchemaFromGlueTable: GetTableResponse[{}]", getTableResponse); logger.info("testNumMapWithSchemaFromGlueTable: GetTableResponse Schema[{}]", getTableResponse.getSchema()); @@ -790,7 +792,7 @@ public void testNumMapWithSchemaFromGlueTable() throws Exception TEST_TABLE_8_NAME, schema, split, - new Constraints(ImmutableMap.of()), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), 100_000_000_000L, // too big to spill 100_000_000_000L); @@ -834,7 +836,7 @@ public void testNumStructWithSchemaFromGlueTable() throws Exception when(glueClient.getTable(any())).thenReturn(mockResult); TableName tableName = new TableName(DEFAULT_SCHEMA, TEST_TABLE8); - GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName); + GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName, Collections.emptyMap()); GetTableResponse getTableResponse = metadataHandler.doGetTable(allocator, getTableRequest); logger.info("testNumStructWithSchemaFromGlueTable: GetTableResponse[{}]", getTableResponse); logger.info("testNumStructWithSchemaFromGlueTable: GetTableResponse Schema[{}]", getTableResponse.getSchema()); @@ -854,7 +856,7 @@ public void testNumStructWithSchemaFromGlueTable() throws Exception TEST_TABLE_8_NAME, schema, split, - new Constraints(ImmutableMap.of()), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), 100_000_000_000L, // too big to spill 100_000_000_000L); diff --git a/athena-elasticsearch/src/test/java/com/amazonaws/athena/connectors/elasticsearch/ElasticsearchMetadataHandlerTest.java b/athena-elasticsearch/src/test/java/com/amazonaws/athena/connectors/elasticsearch/ElasticsearchMetadataHandlerTest.java index d51136df80..1f9dc2d48b 100644 --- a/athena-elasticsearch/src/test/java/com/amazonaws/athena/connectors/elasticsearch/ElasticsearchMetadataHandlerTest.java +++ b/athena-elasticsearch/src/test/java/com/amazonaws/athena/connectors/elasticsearch/ElasticsearchMetadataHandlerTest.java @@ -59,7 +59,6 @@ import java.util.ArrayList; import java.util.Collection; import java.util.Collections; -import java.util.HashMap; import java.util.HashSet; import java.util.Iterator; import java.util.LinkedHashMap; @@ -67,7 +66,10 @@ import java.util.Set; import static com.amazonaws.athena.connector.lambda.metadata.ListTablesRequest.UNLIMITED_PAGE_SIZE_VALUE; -import static org.junit.Assert.*; +import static com.amazonaws.athena.connector.lambda.domain.predicate.Constraints.DEFAULT_NO_LIMIT; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.ArgumentMatchers.nullable; @@ -385,7 +387,7 @@ public void doGetTable() handler = new ElasticsearchMetadataHandler(awsGlue, new LocalKeyFactory(), awsSecretsManager, amazonAthena, "spill-bucket", "spill-prefix", domainMapProvider, clientFactory, 10, com.google.common.collect.ImmutableMap.of()); GetTableRequest req = new GetTableRequest(fakeIdentity(), "queryId", "elasticsearch", - new TableName("movies", "mishmash")); + new TableName("movies", "mishmash"), Collections.emptyMap()); GetTableResponse res = handler.doGetTable(allocator, req); Schema realMapping = res.getSchema(); @@ -421,7 +423,7 @@ public void doGetSplits() new TableName("movies", index), partitions, partitionCols, - new Constraints(new HashMap<>()), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), null); GetSplitsRequest req = new GetSplitsRequest(originalReq, continuationToken); diff --git a/athena-elasticsearch/src/test/java/com/amazonaws/athena/connectors/elasticsearch/ElasticsearchQueryUtilsTest.java b/athena-elasticsearch/src/test/java/com/amazonaws/athena/connectors/elasticsearch/ElasticsearchQueryUtilsTest.java index faadad0d39..f26de59abd 100644 --- a/athena-elasticsearch/src/test/java/com/amazonaws/athena/connectors/elasticsearch/ElasticsearchQueryUtilsTest.java +++ b/athena-elasticsearch/src/test/java/com/amazonaws/athena/connectors/elasticsearch/ElasticsearchQueryUtilsTest.java @@ -44,6 +44,7 @@ import java.util.Map; import static org.junit.Assert.assertEquals; +import static com.amazonaws.athena.connector.lambda.domain.predicate.Constraints.DEFAULT_NO_LIMIT; /** * This class is used to test the ElasticsearchQueryUtils class. @@ -124,7 +125,7 @@ public void getRangePredicateTest() Range.equal(allocator, Types.MinorType.INT.getType(), 1996), Range.greaterThanOrEqual(allocator, Types.MinorType.INT.getType(), 2010)), false)); - Constraints constraints = new Constraints(constraintsMap); + Constraints constraints = new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()); String expectedPredicate = "(_exists_:year) AND year:([* TO 1950} OR {1955 TO 1972] OR [2010 TO *] OR 1952 OR 1996)"; // Get the actual predicate and compare to the expected one. @@ -144,7 +145,7 @@ public void getWhitelistedEquitableValuesPredicate() constraintsMap.put("age", EquatableValueSet.newBuilder(allocator, Types.MinorType.INT.getType(), true, true).addAll(ImmutableList.of(20, 25, 30, 35)).build()); - Constraints constraints = new Constraints(constraintsMap); + Constraints constraints = new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()); String expectedPredicate = "age:(20 OR 25 OR 30 OR 35)"; // Get the actual predicate and compare to the expected one. @@ -164,7 +165,7 @@ public void getExclusiveEquitableValuesPredicate() constraintsMap.put("age", EquatableValueSet.newBuilder(allocator, Types.MinorType.INT.getType(), false, true).addAll(ImmutableList.of(20, 25, 30, 35)).build()); - Constraints constraints = new Constraints(constraintsMap); + Constraints constraints = new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()); String expectedPredicate = "NOT age:(20 OR 25 OR 30 OR 35)"; // Get the actual predicate and compare to the expected one. @@ -183,7 +184,7 @@ public void getAllValuePredicate() logger.info("getAllValuePredicate - enter"); constraintsMap.put("number", new AllOrNoneValueSet(Types.MinorType.INT.getType(), true, true)); - Constraints constraints = new Constraints(constraintsMap); + Constraints constraints = new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()); String expectedPredicate = "(_exists_:number)"; // Get the actual predicate and compare to the expected one. @@ -202,7 +203,7 @@ public void getNoneValuePredicate() logger.info("getNoneValuePredicate - enter"); constraintsMap.put("number", new AllOrNoneValueSet(Types.MinorType.INT.getType(), false, false)); - Constraints constraints = new Constraints(constraintsMap); + Constraints constraints = new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()); String expectedPredicate = "(NOT _exists_:number)"; // Get the actual predicate and compare to the expected one. diff --git a/athena-example/src/test/java/com/amazonaws/athena/connectors/example/ExampleMetadataHandlerTest.java b/athena-example/src/test/java/com/amazonaws/athena/connectors/example/ExampleMetadataHandlerTest.java index fdcdb9a308..5f8ff32501 100644 --- a/athena-example/src/test/java/com/amazonaws/athena/connectors/example/ExampleMetadataHandlerTest.java +++ b/athena-example/src/test/java/com/amazonaws/athena/connectors/example/ExampleMetadataHandlerTest.java @@ -187,7 +187,7 @@ public void doGetTable() logger.info("doGetTable - enter"); GetTableRequest req = new GetTableRequest(fakeIdentity(), "queryId", "default", - new TableName("schema1", "table1")); + new TableName("schema1", "table1"), Collections.emptyMap()); GetTableResponse res = handler.doGetTable(allocator, req); assertTrue(res.getSchema().getFields().size() > 0); assertTrue(res.getSchema().getCustomMetadata().size() > 0); diff --git a/athena-federation-integ-test/pom.xml b/athena-federation-integ-test/pom.xml index f10fc09e36..215ccf9bd8 100644 --- a/athena-federation-integ-test/pom.xml +++ b/athena-federation-integ-test/pom.xml @@ -198,7 +198,8 @@ org.apache.commons commons-lang3 - 3.14.0 + + 3.12.0 diff --git a/athena-federation-sdk-tools/src/main/java/com/amazonaws/athena/connector/validation/LambdaMetadataProvider.java b/athena-federation-sdk-tools/src/main/java/com/amazonaws/athena/connector/validation/LambdaMetadataProvider.java index 43f07ee57d..d3b2a93614 100644 --- a/athena-federation-sdk-tools/src/main/java/com/amazonaws/athena/connector/validation/LambdaMetadataProvider.java +++ b/athena-federation-sdk-tools/src/main/java/com/amazonaws/athena/connector/validation/LambdaMetadataProvider.java @@ -37,6 +37,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.Collections; import java.util.List; import java.util.Set; @@ -133,7 +134,7 @@ public static GetTableResponse getTable(String catalog, log.info("Submitting GetTableRequest with ID " + queryId); try (GetTableRequest request = - new GetTableRequest(identity, queryId, catalog, tableName)) { + new GetTableRequest(identity, queryId, catalog, tableName, Collections.emptyMap())) { log.info("Submitting request: {}", request); GetTableResponse response = (GetTableResponse) getService(metadataFunction, identity, catalog).call(request); log.info("Received response: {}", response); diff --git a/athena-federation-sdk/pom.xml b/athena-federation-sdk/pom.xml index e1085b30ae..83be1478b3 100644 --- a/athena-federation-sdk/pom.xml +++ b/athena-federation-sdk/pom.xml @@ -287,7 +287,8 @@ org.apache.commons commons-lang3 - 3.14.0 + + 3.12.0 net.jqwik @@ -457,7 +458,25 @@ - + + + org.apache.maven.plugins + maven-source-plugin + ${mvn.source.plugin.version} + + + attach-sources + package + + jar-no-fork + + + + + true + false + + diff --git a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/domain/Split.java b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/domain/Split.java index 2011f2d6e0..29ffe18e38 100644 --- a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/domain/Split.java +++ b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/domain/Split.java @@ -239,6 +239,17 @@ public Builder withEncryptionKey(EncryptionKey key) return this; } + /** + * Apply a map of properties to the split + * @param inputProperties + * @return Builder + */ + public Builder applyProperties(Map inputProperties) + { + properties.putAll(inputProperties); + return this; + } + /** * Builds the Split * diff --git a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/domain/predicate/Constraints.java b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/domain/predicate/Constraints.java index 42c7434743..9fd52f95f5 100644 --- a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/domain/predicate/Constraints.java +++ b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/domain/predicate/Constraints.java @@ -25,7 +25,6 @@ import com.fasterxml.jackson.annotation.JsonProperty; import com.google.common.base.Objects; -import java.util.Collections; import java.util.List; import java.util.Map; @@ -49,23 +48,40 @@ public class Constraints private List expression; private final List orderByClause; private long limit; + private Map queryPassthroughArguments; @Deprecated - public Constraints(Map summary) + public Constraints(@JsonProperty("summary") Map summary, + @JsonProperty("expression") List expression, + @JsonProperty("orderByClause") List orderByClause, + @JsonProperty("limit") long limit) { - this(summary, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT); + this.summary = summary; + this.expression = expression; + this.orderByClause = orderByClause; + this.limit = limit; } + /** + * + * @param summary + * @param expression + * @param orderByClause + * @param limit + * @param queryPassthroughArguments + */ @JsonCreator public Constraints(@JsonProperty("summary") Map summary, @JsonProperty("expression") List expression, @JsonProperty("orderByClause") List orderByClause, - @JsonProperty("limit") long limit) + @JsonProperty("limit") long limit, + @JsonProperty("queryPassthroughArguments") Map queryPassthroughArguments) { this.summary = summary; this.expression = expression; this.orderByClause = orderByClause; this.limit = limit; + this.queryPassthroughArguments = queryPassthroughArguments; } /** @@ -103,6 +119,16 @@ public boolean hasNonEmptyOrderByClause() return this.orderByClause != null && this.orderByClause.size() > 0; } + public Map getQueryPassthroughArguments() + { + return this.queryPassthroughArguments; + } + + public boolean isQueryPassThrough() + { + return this.queryPassthroughArguments != null && !this.queryPassthroughArguments.isEmpty(); + } + @Override public boolean equals(Object o) { @@ -118,7 +144,8 @@ public boolean equals(Object o) return Objects.equal(this.summary, that.summary) && Objects.equal(this.expression, that.expression) && Objects.equal(this.orderByClause, that.orderByClause) && - Objects.equal(this.limit, that.limit); + Objects.equal(this.limit, that.limit) && + Objects.equal(this.getQueryPassthroughArguments(), that.getQueryPassthroughArguments()); } @Override @@ -129,6 +156,7 @@ public String toString() "expression=" + expression + "orderByClause=" + orderByClause + "limit=" + limit + + "queryPassthroughArguments=" + queryPassthroughArguments + '}'; } diff --git a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/handlers/CompositeHandler.java b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/handlers/CompositeHandler.java index e25e0d5df7..78c4f51ba6 100644 --- a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/handlers/CompositeHandler.java +++ b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/handlers/CompositeHandler.java @@ -102,8 +102,9 @@ public final void handleRequest(InputStream inputStream, OutputStream outputStre try { rawReq = objectMapper.readValue(allInputBytes, FederationRequest.class); } - catch (IllegalStateException e) { // if client has not upgraded to our latest, fallback to v2 - objectMapper = VersionedObjectMapperFactory.create(allocator, 2); + catch (IllegalStateException e) { // if client has not upgraded to our latest, fallback to v4 + logger.warn("Client's SerDe is not upgraded to latest version, defaulting to V4:", e); + objectMapper = VersionedObjectMapperFactory.create(allocator, 4); rawReq = objectMapper.readValue(allInputBytes, FederationRequest.class); } if (rawReq instanceof MetadataRequest) { diff --git a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/handlers/MetadataHandler.java b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/handlers/MetadataHandler.java index 7fa137e77e..868cff1857 100644 --- a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/handlers/MetadataHandler.java +++ b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/handlers/MetadataHandler.java @@ -270,7 +270,7 @@ protected final void doHandleRequest(BlockAllocator allocator, } return; case GET_TABLE: - try (GetTableResponse response = doGetTable(allocator, (GetTableRequest) req)) { + try (GetTableResponse response = resolveDoGetTableImplementation(allocator, (GetTableRequest) req)) { logger.info("doHandleRequest: response[{}]", response); assertNotNull(response); assertTypes(response); @@ -326,6 +326,32 @@ public abstract ListSchemasResponse doListSchemaNames(final BlockAllocator alloc public abstract ListTablesResponse doListTables(final BlockAllocator allocator, final ListTablesRequest request) throws Exception; + private GetTableResponse resolveDoGetTableImplementation(final BlockAllocator allocator, final GetTableRequest request) + throws Exception + { + logger.info("resolveDoGetTableImplementation: resolving implementation - isQueryPassthrough[{}]", request.isQueryPassthrough()); + if (request.isQueryPassthrough()) { + return doGetQueryPassthroughSchema(allocator, request); + } + return doGetTable(allocator, request); + } + + /** + * Used to get definition (field names, types, descriptions, etc...) of a Query PassThrough. + * + * @param allocator Tool for creating and managing Apache Arrow Blocks. + * @param request Provides details on who made the request and which Athena catalog, database, and table they are querying. + * @return A GetTableResponse which primarily contains: + * 1. An Apache Arrow Schema object describing the table's columns, types, and descriptions. + * 2. A Set of partition column names (or empty if the table isn't partitioned). + */ + public GetTableResponse doGetQueryPassthroughSchema(final BlockAllocator allocator, final GetTableRequest request) + throws Exception + { + //todo; maybe we need a better name for this method, + throw new UnsupportedOperationException("Not implemented"); + } + /** * Used to get definition (field names, types, descriptions, etc...) of a Table. * diff --git a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/handlers/SerDeVersion.java b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/handlers/SerDeVersion.java index e31111dae2..31ba79b3af 100644 --- a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/handlers/SerDeVersion.java +++ b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/handlers/SerDeVersion.java @@ -30,5 +30,5 @@ public class SerDeVersion { private SerDeVersion() {} - public static final int SERDE_VERSION = 4; + public static final int SERDE_VERSION = 5; } diff --git a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/metadata/GetTableRequest.java b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/metadata/GetTableRequest.java index 691ae281c3..f7ada19d7d 100644 --- a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/metadata/GetTableRequest.java +++ b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/metadata/GetTableRequest.java @@ -26,6 +26,8 @@ import com.fasterxml.jackson.annotation.JsonProperty; import com.google.common.base.Objects; +import java.util.Map; + import static java.util.Objects.requireNonNull; /** @@ -35,24 +37,30 @@ public class GetTableRequest extends MetadataRequest { private final TableName tableName; + private Map queryPassthroughArguments; /** * Constructs a new GetTableRequest object. * - * @param identity The identity of the caller. - * @param queryId The ID of the query requesting metadata. - * @param catalogName The catalog name that the table belongs to. - * @param tableName The name of the table metadata is being requested for. + * @param identity The identity of the caller. + * @param queryId The ID of the query requesting metadata. + * @param catalogName The catalog name that the table belongs to. + * @param tableName The name of the table metadata is being requested for. + * @param queryPassthroughArguments */ @JsonCreator public GetTableRequest(@JsonProperty("identity") FederatedIdentity identity, - @JsonProperty("queryId") String queryId, - @JsonProperty("catalogName") String catalogName, - @JsonProperty("tableName") TableName tableName) + @JsonProperty("queryId") String queryId, + @JsonProperty("catalogName") String catalogName, + @JsonProperty("tableName") TableName tableName, + @JsonProperty("queryPassthroughArguments") Map queryPassthroughArguments) { super(identity, MetadataRequestType.GET_TABLE, queryId, catalogName); requireNonNull(tableName, "tableName is null"); + requireNonNull(queryPassthroughArguments, "queryPassthroughArguments is null"); + this.tableName = tableName; + this.queryPassthroughArguments = queryPassthroughArguments; } /** @@ -72,12 +80,22 @@ public void close() //No Op } + public Map getQueryPassthroughArguments() + { + return this.queryPassthroughArguments; + } + + public boolean isQueryPassthrough() + { + return this.queryPassthroughArguments != null && !this.queryPassthroughArguments.isEmpty(); + } @Override public String toString() { return "GetTableRequest{" + "queryId=" + getQueryId() + ", tableName=" + tableName + + "queryPassthroughArguments=" + queryPassthroughArguments + '}'; } @@ -95,12 +113,13 @@ public boolean equals(Object o) return Objects.equal(this.tableName, that.tableName) && Objects.equal(this.getRequestType(), that.getRequestType()) && - Objects.equal(this.getCatalogName(), that.getCatalogName()); + Objects.equal(this.getCatalogName(), that.getCatalogName()) && + Objects.equal(this.getQueryPassthroughArguments(), that.getQueryPassthroughArguments()); } @Override public int hashCode() { - return Objects.hashCode(tableName, getRequestType(), getCatalogName()); + return Objects.hashCode(tableName, getRequestType(), getCatalogName(), getQueryPassthroughArguments()); } } diff --git a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/metadata/MetadataRequestType.java b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/metadata/MetadataRequestType.java index 4e75062f73..0f98f065f5 100644 --- a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/metadata/MetadataRequestType.java +++ b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/metadata/MetadataRequestType.java @@ -9,9 +9,9 @@ * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at - * + * * http://www.apache.org/licenses/LICENSE-2.0 - * + * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. diff --git a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/metadata/optimizations/querypassthrough/QueryPassthrough.java b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/metadata/optimizations/querypassthrough/QueryPassthrough.java new file mode 100644 index 0000000000..f5289b651a --- /dev/null +++ b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/metadata/optimizations/querypassthrough/QueryPassthrough.java @@ -0,0 +1,64 @@ +/*- + * #%L + * Amazon Athena Query Federation SDK + * %% + * Copyright (C) 2019 - 2023 Amazon Web Services + * %% + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * #L% + */ +package com.amazonaws.athena.connector.lambda.metadata.optimizations.querypassthrough; + +import com.amazonaws.athena.connector.lambda.metadata.optimizations.OptimizationSubType; + +import java.util.Arrays; +import java.util.List; + +/** + * This class describes Query Passthrough Signature; + * Schema; is where the function will reside in the catalog attaching this namespace + * Name; is the table function name of the QPT; + * Arguments; list of all arguments that this QPT is expecting to have + * + */ +public enum QueryPassthrough { + QUERY_PASSTHROUGH_SCHEMA("query_passthrough_schema"), + QUERY_PASSTHROUGH_NAME("query_passthrough_name"), + QUERY_PASSTHROUGH_ARGUMENTS("query_passthrough_arguments"); + private final String value; + + QueryPassthrough(String value) + { + this.value = value; + } + + public String getValue() + { + return value; + } + + public final OptimizationSubType withSchema(String schema) + { + return new OptimizationSubType(getValue(), Arrays.asList(schema)); + } + + public final OptimizationSubType withName(String name) + { + return new OptimizationSubType(getValue(), Arrays.asList(name)); + } + + public final OptimizationSubType withArguments(List arguments) + { + return new OptimizationSubType(getValue(), arguments); + } +} diff --git a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/metadata/optimizations/querypassthrough/QueryPassthroughSignature.java b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/metadata/optimizations/querypassthrough/QueryPassthroughSignature.java new file mode 100644 index 0000000000..1b5472cc23 --- /dev/null +++ b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/metadata/optimizations/querypassthrough/QueryPassthroughSignature.java @@ -0,0 +1,162 @@ +/*- + * #%L + * Amazon Athena Query Federation SDK + * %% + * Copyright (C) 2019 - 2023 Amazon Web Services + * %% + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * #L% + */ +package com.amazonaws.athena.connector.lambda.metadata.optimizations.querypassthrough; + +import com.amazonaws.athena.connector.lambda.metadata.optimizations.OptimizationSubType; +import com.google.common.collect.ImmutableMap; +import org.apache.commons.lang3.StringUtils; +import org.slf4j.Logger; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +public interface QueryPassthroughSignature +{ + public static final String SCHEMA_FUNCTION_NAME = "schemaFunctionName"; + public static final String ENABLE_QUERY_PASSTHROUGH = "enable_query_passthrough"; + public static final String DEFAULT_ENABLE_QUERY_PASSTHROUGH_STATE = "true"; + /** + * + * @return QPT Function's Schema (also known as a domain or namepsace) + */ + String getFunctionSchema(); + + /** + * + * @return QPT Function's name + */ + String getFunctionName(); + + /** + * Returns the full function schema and name + * @return + */ + default String getFunctionSignature() + { + return getFunctionSchema().toUpperCase() + "." + getFunctionName().toUpperCase(); + } + + /** + * + * @return Query Passthrough Function's Arguments + */ + List getFunctionArguments(); + + /** + * note: due to needing to stay compatible with JDK8; we can't use JDK9 private method + * @return a logger + */ + Logger getLogger(); + + /** + * verify that the arguments returned by the engine are the same that the connector defined + * And call on any connector-custom verification + * @param engineQptArguments + * @throws IllegalArgumentException + */ + public default void verify(Map engineQptArguments) + throws IllegalArgumentException + { + //High-level Query Passthrough Function Argument verification + //First verifies that these arguments belong to this specific function + if (!verifyFunctionSignature(engineQptArguments)) { + throw new IllegalArgumentException("Function Signature doesn't match implementation's"); + } + //Ensuring the arguments received from the engine are the one defined by the connector + for (String argument : this.getFunctionArguments()) { + if (!engineQptArguments.containsKey(argument)) { + throw new IllegalArgumentException("Missing Query Passthrough Argument: " + argument); + } + if (StringUtils.isEmpty(engineQptArguments.get(argument))) { + throw new IllegalArgumentException("Missing Query Passthrough Value for Argument: " + argument); + } + } + //Finally, perform any connector-specific verification; + customConnectorVerifications(); + } + + /** + * Provides a mechanism to perform custom connector verification logic. + */ + default void customConnectorVerifications() + { + //No Op + } + + /** + * Verifying that the query passthrough function signature is the one expected by the connector + * @param argumentValues + * @return true if the signature matches, otherwise false + */ + default boolean verifyFunctionSignature(Map argumentValues) + { + if (argumentValues.containsKey(SCHEMA_FUNCTION_NAME)) { + String receivedSignature = argumentValues.get(SCHEMA_FUNCTION_NAME); + getLogger().info("Found signature: {}", receivedSignature); + return receivedSignature.equalsIgnoreCase(getFunctionSignature()); + } + getLogger().warn("No matching function signature found: {}", getFunctionSignature()); + return false; + } + + /** + * Creates a list of Optimization that includes the query passthrough definition + * @return list of capability describes of the signature of the current implementation + */ + default List getQueryPassthroughCapabilities() + { + List queryPassthroughDefinition = new ArrayList<>(3); + queryPassthroughDefinition.add(QueryPassthrough.QUERY_PASSTHROUGH_SCHEMA.withSchema(getFunctionSchema())); + queryPassthroughDefinition.add(QueryPassthrough.QUERY_PASSTHROUGH_NAME.withName(getFunctionName())); + queryPassthroughDefinition.add(QueryPassthrough.QUERY_PASSTHROUGH_ARGUMENTS.withArguments(getFunctionArguments())); + + return queryPassthroughDefinition; + } + + /** + * Adds the query passthrough implementation, if user has not disabled it, to the connector's capabilities + * @param capabilities + * @param configOptions + */ + default void addQueryPassthroughCapabilityIfEnabled(ImmutableMap.Builder> capabilities, Map configOptions) + { + if (allowQueryPassthrough(configOptions)) { + getLogger().info("Query Passthrough is enabled; adding implementation to connector's capabilities"); + capabilities.put(getFunctionSignature(), getQueryPassthroughCapabilities()); + } + else { + getLogger().info("Query Passthrough is disabled"); + } + } + + /** + * A method that checks the Lambda's environment variables to see if QPT is disabled/enabled + * @param configOptions + * @return true if enabled; otherwise false + */ + default boolean allowQueryPassthrough(Map configOptions) + { + String enableQueryPassthroughEnvVal = configOptions + .getOrDefault(ENABLE_QUERY_PASSTHROUGH, DEFAULT_ENABLE_QUERY_PASSTHROUGH_STATE) + .toLowerCase(); + return enableQueryPassthroughEnvVal.equals("true"); + } +} diff --git a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/serde/VersionedObjectMapperFactory.java b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/serde/VersionedObjectMapperFactory.java index 5f0a69a70c..e8f34865c8 100644 --- a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/serde/VersionedObjectMapperFactory.java +++ b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/serde/VersionedObjectMapperFactory.java @@ -7,9 +7,9 @@ * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at - * + * * http://www.apache.org/licenses/LICENSE-2.0 - * + * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. @@ -24,13 +24,17 @@ import com.amazonaws.athena.connector.lambda.serde.v2.ObjectMapperFactoryV2; import com.amazonaws.athena.connector.lambda.serde.v3.ObjectMapperFactoryV3; import com.amazonaws.athena.connector.lambda.serde.v4.ObjectMapperFactoryV4; +import com.amazonaws.athena.connector.lambda.serde.v5.ObjectMapperFactoryV5; import com.fasterxml.jackson.databind.ObjectMapper; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * Vends {@link ObjectMapper} instances that correspond to SerDe versions. */ public class VersionedObjectMapperFactory { + private static final Logger logger = LoggerFactory.getLogger(VersionedObjectMapperFactory.class); private VersionedObjectMapperFactory() {} /** @@ -53,6 +57,7 @@ public static ObjectMapper create(BlockAllocator allocator) */ public static ObjectMapper create(BlockAllocator allocator, int version) { + logger.debug("Athena SDK SerDe Version: " + version); switch (version) { case 1: return ObjectMapperFactory.create(allocator); @@ -62,6 +67,8 @@ public static ObjectMapper create(BlockAllocator allocator, int version) return ObjectMapperFactoryV3.create(allocator); case 4: return ObjectMapperFactoryV4.create(allocator); + case 5: + return ObjectMapperFactoryV5.create(allocator); default: throw new IllegalArgumentException("No serde version " + version); } diff --git a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/serde/v2/GetTableRequestSerDe.java b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/serde/v2/GetTableRequestSerDe.java index 02e4b232aa..6035cd77d5 100644 --- a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/serde/v2/GetTableRequestSerDe.java +++ b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/serde/v2/GetTableRequestSerDe.java @@ -31,6 +31,7 @@ import com.fasterxml.jackson.databind.SerializerProvider; import java.io.IOException; +import java.util.Collections; import static java.util.Objects.requireNonNull; @@ -82,7 +83,7 @@ protected MetadataRequest doRequestDeserialize(JsonParser jparser, Deserializati assertFieldName(jparser, TABLE_NAME_FIELD); TableName tableName = tableNameDeserializer.deserialize(jparser, ctxt); - return new GetTableRequest(identity, queryId, catalogName, tableName); + return new GetTableRequest(identity, queryId, catalogName, tableName, Collections.emptyMap()); } } } diff --git a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/serde/v5/ConstraintsSerDeV5.java b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/serde/v5/ConstraintsSerDeV5.java new file mode 100644 index 0000000000..2efd6655b5 --- /dev/null +++ b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/serde/v5/ConstraintsSerDeV5.java @@ -0,0 +1,157 @@ +/*- + * #%L + * Amazon Athena Query Federation SDK + * %% + * Copyright (C) 2019 - 2020 Amazon Web Services + * %% + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * #L% + */ +package com.amazonaws.athena.connector.lambda.serde.v5; + +import com.amazonaws.athena.connector.lambda.domain.predicate.Constraints; +import com.amazonaws.athena.connector.lambda.domain.predicate.OrderByField; +import com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet; +import com.amazonaws.athena.connector.lambda.domain.predicate.expression.FederationExpression; +import com.amazonaws.athena.connector.lambda.serde.BaseDeserializer; +import com.amazonaws.athena.connector.lambda.serde.BaseSerializer; +import com.amazonaws.athena.connector.lambda.serde.VersionedSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.ValueSetSerDe; +import com.fasterxml.jackson.core.JsonGenerator; +import com.fasterxml.jackson.core.JsonParser; +import com.fasterxml.jackson.core.JsonToken; +import com.fasterxml.jackson.databind.DeserializationContext; +import com.fasterxml.jackson.databind.SerializerProvider; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; + +import static java.util.Objects.requireNonNull; + +public final class ConstraintsSerDeV5 +{ + private static final String SUMMARY_FIELD = "summary"; + private static final String EXPRESSION_FIELD = "expression"; + private static final String ORDER_BY_CLAUSE = "orderByClause"; + private static final String LIMIT_FIELD = "limit"; + private static final String QUERY_PASSTHROUGH_ARGUMENTS = "queryPassthroughArguments"; + + private ConstraintsSerDeV5() {} + + public static final class Serializer extends BaseSerializer implements VersionedSerDe.Serializer + { + private final ValueSetSerDe.Serializer valueSetSerializer; + private final VersionedSerDe.Serializer federationExpressionSerializer; + private final VersionedSerDe.Serializer orderByFieldSerializer; + public Serializer(ValueSetSerDe.Serializer valueSetSerializer, + VersionedSerDe.Serializer federationExpressionSerializer, + VersionedSerDe.Serializer orderByFieldSerializer) + { + super(Constraints.class); + this.valueSetSerializer = requireNonNull(valueSetSerializer, "valueSetSerDe is null"); + this.federationExpressionSerializer = requireNonNull(federationExpressionSerializer, "federationExpressionSerDe is null"); + this.orderByFieldSerializer = requireNonNull(orderByFieldSerializer, "orderByFieldSerDe is null"); + } + + @Override + public void doSerialize(Constraints constraints, JsonGenerator jgen, SerializerProvider provider) + throws IOException + { + jgen.writeObjectFieldStart(SUMMARY_FIELD); + for (Map.Entry entry : constraints.getSummary().entrySet()) { + jgen.writeFieldName(entry.getKey()); + valueSetSerializer.serialize(entry.getValue(), jgen, provider); + } + jgen.writeEndObject(); + + jgen.writeArrayFieldStart(EXPRESSION_FIELD); + for (FederationExpression federationExpression : constraints.getExpression()) { + federationExpressionSerializer.serialize(federationExpression, jgen, provider); + } + jgen.writeEndArray(); + + jgen.writeArrayFieldStart(ORDER_BY_CLAUSE); + for (OrderByField orderByField : constraints.getOrderByClause()) { + orderByFieldSerializer.serialize(orderByField, jgen, provider); + } + jgen.writeEndArray(); + + jgen.writeNumberField(LIMIT_FIELD, constraints.getLimit()); + + writeStringMap(jgen, QUERY_PASSTHROUGH_ARGUMENTS, constraints.getQueryPassthroughArguments()); + } + } + + public static final class Deserializer extends BaseDeserializer implements VersionedSerDe.Deserializer + { + private final ValueSetSerDe.Deserializer valueSetDeserializer; + private final VersionedSerDe.Deserializer federationExpressionDeserializer; + private final VersionedSerDe.Deserializer orderByFieldDeserializer; + + public Deserializer(ValueSetSerDe.Deserializer valueSetDeserializer, + VersionedSerDe.Deserializer federationExpressionDeserializer, + VersionedSerDe.Deserializer orderByFieldDeserializer) + { + super(Constraints.class); + this.valueSetDeserializer = requireNonNull(valueSetDeserializer, "valueSetSerDe is null"); + this.federationExpressionDeserializer = requireNonNull(federationExpressionDeserializer, "federationExpressionSerDe is null"); + this.orderByFieldDeserializer = requireNonNull(orderByFieldDeserializer, "orderByFieldSerDe is null"); + } + + @Override + public Constraints doDeserialize(JsonParser jparser, DeserializationContext ctxt) + throws IOException + { + assertFieldName(jparser, SUMMARY_FIELD); + validateObjectStart(jparser.nextToken()); + ImmutableMap.Builder summaryMap = ImmutableMap.builder(); + while (jparser.nextToken() != JsonToken.END_OBJECT) { + String column = jparser.getCurrentName(); + summaryMap.put(column, valueSetDeserializer.deserialize(jparser, ctxt)); + } + + assertFieldName(jparser, EXPRESSION_FIELD); + validateArrayStart(jparser); + ImmutableList.Builder federationExpression = ImmutableList.builder(); + while (jparser.nextToken() != JsonToken.END_ARRAY) { + validateObjectStart(jparser.getCurrentToken()); + federationExpression.add(federationExpressionDeserializer.doDeserialize(jparser, ctxt)); + validateObjectEnd(jparser); + } + + assertFieldName(jparser, ORDER_BY_CLAUSE); + validateArrayStart(jparser); + ImmutableList.Builder orderByClauseBuilder = ImmutableList.builder(); + while (jparser.nextToken() != JsonToken.END_ARRAY) { + validateObjectStart(jparser.getCurrentToken()); + orderByClauseBuilder.add(orderByFieldDeserializer.doDeserialize(jparser, ctxt)); + validateObjectEnd(jparser); + } + + long limit = getNextLongField(jparser, LIMIT_FIELD); + + Map queryPassthroughArguments = new HashMap<>(); + assertFieldName(jparser, QUERY_PASSTHROUGH_ARGUMENTS); + validateObjectStart(jparser.nextToken()); + while (jparser.nextToken() != JsonToken.END_OBJECT) { + queryPassthroughArguments.put(jparser.getCurrentName(), jparser.getValueAsString()); + } + + Constraints constraints = new Constraints(summaryMap.build(), federationExpression.build(), orderByClauseBuilder.build(), limit, queryPassthroughArguments); + return constraints; + } + } +} diff --git a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/serde/v5/FederationRequestSerDeV5.java b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/serde/v5/FederationRequestSerDeV5.java new file mode 100644 index 0000000000..4bdba5698f --- /dev/null +++ b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/serde/v5/FederationRequestSerDeV5.java @@ -0,0 +1,95 @@ +/*- + * #%L + * Amazon Athena Query Federation SDK + * %% + * Copyright (C) 2019 - 2022 Amazon Web Services + * %% + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * #L% + */ +package com.amazonaws.athena.connector.lambda.serde.v5; + +import com.amazonaws.athena.connector.lambda.request.FederationRequest; +import com.amazonaws.athena.connector.lambda.serde.DelegatingDeserializer; +import com.amazonaws.athena.connector.lambda.serde.DelegatingSerializer; +import com.amazonaws.athena.connector.lambda.serde.PingRequestSerDe; +import com.amazonaws.athena.connector.lambda.serde.TypedDeserializer; +import com.amazonaws.athena.connector.lambda.serde.TypedSerializer; +import com.amazonaws.athena.connector.lambda.serde.VersionedSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.GetSplitsRequestSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.GetTableLayoutRequestSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.ListSchemasRequestSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.ListTablesRequestSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.ReadRecordsRequestSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.UserDefinedFunctionRequestSerDe; +import com.amazonaws.athena.connector.lambda.serde.v4.GetDataSourceCapabilitiesRequestSerDeV4; +import com.google.common.collect.ImmutableSet; + +public class FederationRequestSerDeV5 +{ + private FederationRequestSerDeV5() {} + + public static final class Serializer extends DelegatingSerializer implements VersionedSerDe.Serializer + { + public Serializer( + PingRequestSerDe.Serializer pingSerializer, + ListSchemasRequestSerDe.Serializer listSchemasSerializer, + ListTablesRequestSerDe.Serializer listTablesSerializer, + GetTableRequestSerDeV5.Serializer getTableSerializer, + GetTableLayoutRequestSerDe.Serializer getTableLayoutSerializer, + GetSplitsRequestSerDe.Serializer getSplitsSerializer, + ReadRecordsRequestSerDe.Serializer readRecordsSerializer, + UserDefinedFunctionRequestSerDe.Serializer userDefinedFunctionSerializer, + GetDataSourceCapabilitiesRequestSerDeV4.Serializer getDataSourceCapabilitiesSerializer) + { + super(FederationRequest.class, ImmutableSet.>builder() + .add(pingSerializer) + .add(listSchemasSerializer) + .add(listTablesSerializer) + .add(getTableSerializer) + .add(getTableLayoutSerializer) + .add(getSplitsSerializer) + .add(readRecordsSerializer) + .add(userDefinedFunctionSerializer) + .add(getDataSourceCapabilitiesSerializer) + .build()); + } + } + + public static final class Deserializer extends DelegatingDeserializer implements VersionedSerDe.Deserializer + { + public Deserializer( + PingRequestSerDe.Deserializer pingDeserializer, + ListSchemasRequestSerDe.Deserializer listSchemasDeserializer, + ListTablesRequestSerDe.Deserializer listTablesDeserializer, + GetTableRequestSerDeV5.Deserializer getTableDeserializer, + GetTableLayoutRequestSerDe.Deserializer getTableLayoutDeserializer, + GetSplitsRequestSerDe.Deserializer getSplitsDeserializer, + ReadRecordsRequestSerDe.Deserializer readRecordsDeserializer, + UserDefinedFunctionRequestSerDe.Deserializer userDefinedFunctionDeserializer, + GetDataSourceCapabilitiesRequestSerDeV4.Deserializer getDataSourceCapabilitiesDeserializer) + { + super(FederationRequest.class, ImmutableSet.>builder() + .add(pingDeserializer) + .add(listSchemasDeserializer) + .add(listTablesDeserializer) + .add(getTableDeserializer) + .add(getTableLayoutDeserializer) + .add(getSplitsDeserializer) + .add(readRecordsDeserializer) + .add(userDefinedFunctionDeserializer) + .add(getDataSourceCapabilitiesDeserializer) + .build()); + } + } +} diff --git a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/serde/v5/GetTableRequestSerDeV5.java b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/serde/v5/GetTableRequestSerDeV5.java new file mode 100644 index 0000000000..ef6014f53f --- /dev/null +++ b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/serde/v5/GetTableRequestSerDeV5.java @@ -0,0 +1,104 @@ +/*- + * #%L + * Amazon Athena Query Federation SDK + * %% + * Copyright (C) 2019 - 2020 Amazon Web Services + * %% + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * #L% + */ +package com.amazonaws.athena.connector.lambda.serde.v5; + +import com.amazonaws.athena.connector.lambda.domain.TableName; +import com.amazonaws.athena.connector.lambda.metadata.GetTableRequest; +import com.amazonaws.athena.connector.lambda.metadata.MetadataRequest; +import com.amazonaws.athena.connector.lambda.request.FederationRequest; +import com.amazonaws.athena.connector.lambda.security.FederatedIdentity; +import com.amazonaws.athena.connector.lambda.serde.FederatedIdentitySerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.MetadataRequestDeserializer; +import com.amazonaws.athena.connector.lambda.serde.v2.MetadataRequestSerializer; +import com.amazonaws.athena.connector.lambda.serde.v2.TableNameSerDe; +import com.fasterxml.jackson.core.JsonGenerator; +import com.fasterxml.jackson.core.JsonParser; +import com.fasterxml.jackson.core.JsonToken; +import com.fasterxml.jackson.databind.DeserializationContext; +import com.fasterxml.jackson.databind.SerializerProvider; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; + +import static java.util.Objects.requireNonNull; + +public final class GetTableRequestSerDeV5 +{ + private static final String TABLE_NAME_FIELD = "tableName"; + private static final String QUERY_PASSTHROUGH_ARGUMENTS = "queryPassthroughArguments"; + + private GetTableRequestSerDeV5() {} + + public static final class Serializer extends MetadataRequestSerializer + { + private final FederatedIdentitySerDe.Serializer identitySerializer; + private final TableNameSerDe.Serializer tableNameSerializer; + + public Serializer(FederatedIdentitySerDe.Serializer identitySerializer, TableNameSerDe.Serializer tableNameSerializer) + { + super(GetTableRequest.class, identitySerializer); + this.identitySerializer = requireNonNull(identitySerializer, "identitySerializer is null"); + this.tableNameSerializer = requireNonNull(tableNameSerializer, "tableNameSerializer is null"); + } + + @Override + protected void doRequestSerialize(FederationRequest federationRequest, JsonGenerator jgen, SerializerProvider provider) + throws IOException + { + GetTableRequest getTableRequest = (GetTableRequest) federationRequest; + + jgen.writeFieldName(TABLE_NAME_FIELD); + tableNameSerializer.serialize(getTableRequest.getTableName(), jgen, provider); + writeStringMap(jgen, QUERY_PASSTHROUGH_ARGUMENTS, getTableRequest.getQueryPassthroughArguments()); + } + } + + public static final class Deserializer extends MetadataRequestDeserializer + { + private final FederatedIdentitySerDe.Deserializer identityDeserializer; + private final TableNameSerDe.Deserializer tableNameDeserializer; + + public Deserializer(FederatedIdentitySerDe.Deserializer identityDeserializer, TableNameSerDe.Deserializer tableNameDeserializer) + { + super(GetTableRequest.class, identityDeserializer); + this.identityDeserializer = requireNonNull(identityDeserializer, "identityDeserializer is null"); + this.tableNameDeserializer = requireNonNull(tableNameDeserializer, "tableNameDeserializer is null"); + } + + @Override + protected MetadataRequest doRequestDeserialize(JsonParser jparser, DeserializationContext ctxt, FederatedIdentity identity, String queryId, String catalogName) + throws IOException + { + assertFieldName(jparser, TABLE_NAME_FIELD); + TableName tableName = tableNameDeserializer.deserialize(jparser, ctxt); + + Map queryPassthroughArguments = new HashMap<>(); + assertFieldName(jparser, QUERY_PASSTHROUGH_ARGUMENTS); + validateObjectStart(jparser.nextToken()); + while (jparser.nextToken() != JsonToken.END_OBJECT) { + queryPassthroughArguments.put(jparser.getCurrentName(), jparser.getValueAsString()); + } + + GetTableRequest getTableRequest = new GetTableRequest(identity, queryId, catalogName, tableName, queryPassthroughArguments); + return getTableRequest; + } + } +} diff --git a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/serde/v5/ObjectMapperFactoryV5.java b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/serde/v5/ObjectMapperFactoryV5.java new file mode 100644 index 0000000000..d47dc71c23 --- /dev/null +++ b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/serde/v5/ObjectMapperFactoryV5.java @@ -0,0 +1,378 @@ +/*- + * #%L + * Amazon Athena Query Federation SDK + * %% + * Copyright (C) 2019 - 2020 Amazon Web Services + * %% + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * #L% + */ +package com.amazonaws.athena.connector.lambda.serde.v5; + +import com.amazonaws.athena.connector.lambda.data.Block; +import com.amazonaws.athena.connector.lambda.data.BlockAllocator; +import com.amazonaws.athena.connector.lambda.domain.predicate.Constraints; +import com.amazonaws.athena.connector.lambda.domain.predicate.OrderByField; +import com.amazonaws.athena.connector.lambda.domain.predicate.expression.FederationExpression; +import com.amazonaws.athena.connector.lambda.domain.predicate.functions.FunctionName; +import com.amazonaws.athena.connector.lambda.metadata.optimizations.OptimizationSubType; +import com.amazonaws.athena.connector.lambda.request.FederationRequest; +import com.amazonaws.athena.connector.lambda.request.FederationResponse; +import com.amazonaws.athena.connector.lambda.serde.FederatedIdentitySerDe; +import com.amazonaws.athena.connector.lambda.serde.PingRequestSerDe; +import com.amazonaws.athena.connector.lambda.serde.PingResponseSerDe; +import com.amazonaws.athena.connector.lambda.serde.VersionedSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.AllOrNoneValueSetSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.ArrowTypeSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.EncryptionKeySerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.EquatableValueSetSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.GetSplitsRequestSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.GetSplitsResponseSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.GetTableLayoutRequestSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.GetTableLayoutResponseSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.GetTableResponseSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.LambdaFunctionExceptionSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.ListSchemasRequestSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.ListSchemasResponseSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.ListTablesRequestSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.ListTablesResponseSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.MarkerSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.RangeSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.ReadRecordsRequestSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.ReadRecordsResponseSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.RemoteReadRecordsResponseSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.S3SpillLocationSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.SortedRangeSetSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.SpillLocationSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.SplitSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.TableNameSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.UserDefinedFunctionRequestSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.UserDefinedFunctionResponseSerDe; +import com.amazonaws.athena.connector.lambda.serde.v2.ValueSetSerDe; +import com.amazonaws.athena.connector.lambda.serde.v4.BlockSerDeV4; +import com.amazonaws.athena.connector.lambda.serde.v4.ConstantExpressionSerDeV4; +import com.amazonaws.athena.connector.lambda.serde.v4.FederationExpressionSerDeV4; +import com.amazonaws.athena.connector.lambda.serde.v4.FederationResponseSerDeV4; +import com.amazonaws.athena.connector.lambda.serde.v4.FunctionCallExpressionSerDeV4; +import com.amazonaws.athena.connector.lambda.serde.v4.FunctionNameSerDeV4; +import com.amazonaws.athena.connector.lambda.serde.v4.GetDataSourceCapabilitiesRequestSerDeV4; +import com.amazonaws.athena.connector.lambda.serde.v4.GetDataSourceCapabilitiesResponseSerDeV4; +import com.amazonaws.athena.connector.lambda.serde.v4.OptimizationSubTypeSerDeV4; +import com.amazonaws.athena.connector.lambda.serde.v4.OrderByFieldSerDeV4; +import com.amazonaws.athena.connector.lambda.serde.v4.SchemaSerDeV4; +import com.amazonaws.athena.connector.lambda.serde.v4.VariableExpressionSerDeV4; +import com.amazonaws.services.lambda.invoke.LambdaFunctionException; +import com.fasterxml.jackson.core.JsonFactory; +import com.fasterxml.jackson.databind.BeanDescription; +import com.fasterxml.jackson.databind.DeserializationContext; +import com.fasterxml.jackson.databind.DeserializationFeature; +import com.fasterxml.jackson.databind.JavaType; +import com.fasterxml.jackson.databind.JsonDeserializer; +import com.fasterxml.jackson.databind.JsonMappingException; +import com.fasterxml.jackson.databind.JsonSerializer; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.SerializerProvider; +import com.fasterxml.jackson.databind.cfg.DeserializerFactoryConfig; +import com.fasterxml.jackson.databind.cfg.SerializerFactoryConfig; +import com.fasterxml.jackson.databind.deser.BeanDeserializerFactory; +import com.fasterxml.jackson.databind.deser.DefaultDeserializationContext; +import com.fasterxml.jackson.databind.deser.DeserializerFactory; +import com.fasterxml.jackson.databind.deser.Deserializers; +import com.fasterxml.jackson.databind.module.SimpleDeserializers; +import com.fasterxml.jackson.databind.module.SimpleSerializers; +import com.fasterxml.jackson.databind.ser.BeanSerializerFactory; +import com.fasterxml.jackson.databind.ser.SerializerFactory; +import com.fasterxml.jackson.databind.ser.Serializers; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import org.apache.arrow.vector.types.pojo.Schema; + +public class ObjectMapperFactoryV5 +{ + private static final JsonFactory JSON_FACTORY = new JsonFactory(); + private static final String LAMDA_EXCEPTION_CLASS_NAME = LambdaFunctionException.class.getName(); + + private static final SerializerFactory SERIALIZER_FACTORY; + + static { + // Serializers can be static since they don't need a BlockAllocator + ImmutableList> sers = ImmutableList.of(createRequestSerializer(), createResponseSerializer()); + SimpleSerializers serializers = new SimpleSerializers(sers); + SerializerFactoryConfig config = new SerializerFactoryConfig().withAdditionalSerializers(serializers); + SERIALIZER_FACTORY = new StrictSerializerFactory(config); + } + + private ObjectMapperFactoryV5() {} + + /** + * Custom SerializerFactory that *only* uses the custom serializers that we inject into the {@link ObjectMapper}. + */ + private static class StrictSerializerFactory extends BeanSerializerFactory + { + private StrictSerializerFactory(SerializerFactoryConfig config) + { + super(config); + } + + @Override + public StrictSerializerFactory withConfig(SerializerFactoryConfig config) + { + if (_factoryConfig == config) { + return this; + } + return new StrictSerializerFactory(config); + } + + @Override + @SuppressWarnings("unchecked") + public JsonSerializer createSerializer(SerializerProvider prov, JavaType origType) + throws JsonMappingException + { + for (Serializers serializers : customSerializers()) { + JsonSerializer ser = serializers.findSerializer(prov.getConfig(), origType, null); + if (ser != null) { + return (JsonSerializer) ser; + } + } + throw new IllegalArgumentException("No explicitly configured serializer for " + origType); + } + } + + /** + * Custom DeserializerFactory that *only* uses the custom deserializers that we inject into the {@link ObjectMapper}. + */ + private static class StrictDeserializerFactory extends BeanDeserializerFactory + { + private StrictDeserializerFactory(DeserializerFactoryConfig config) + { + super(config); + } + + @Override + public DeserializerFactory withConfig(DeserializerFactoryConfig config) + { + if (_factoryConfig == config) { + return this; + } + return new StrictDeserializerFactory(config); + } + + @Override + @SuppressWarnings("unchecked") + public JsonDeserializer createBeanDeserializer(DeserializationContext ctxt, JavaType type, BeanDescription beanDesc) + throws JsonMappingException + { + for (Deserializers d : _factoryConfig.deserializers()) { + JsonDeserializer deser = d.findBeanDeserializer(type, ctxt.getConfig(), beanDesc); + if (deser != null) { + return (JsonDeserializer) deser; + } + } + throw new IllegalArgumentException("No explicitly configured deserializer for " + type); + } + } + + /** + * Locked down ObjectMapper that only uses the serializers/deserializers provided and does not fall back to annotation or reflection + * based serialization. + */ + private static class StrictObjectMapper extends ObjectMapper + { + private StrictObjectMapper(BlockAllocator allocator) + { + super(JSON_FACTORY); + _serializerFactory = SERIALIZER_FACTORY; + + ImmutableMap, JsonDeserializer> desers = ImmutableMap.of( + FederationRequest.class, createRequestDeserializer(allocator), + FederationResponse.class, createResponseDeserializer(allocator), + LambdaFunctionException.class, new LambdaFunctionExceptionSerDe.Deserializer()); + SimpleDeserializers deserializers = new SimpleDeserializers(desers); + DeserializerFactoryConfig dConfig = new DeserializerFactoryConfig().withAdditionalDeserializers(deserializers); + _deserializationContext = new DefaultDeserializationContext.Impl(new StrictDeserializerFactory(dConfig)); + // required by LambdaInvokerFactory + disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES); + } + } + + public static ObjectMapper create(BlockAllocator allocator) + { + return new StrictObjectMapper(allocator); + } + + private static FederationRequestSerDeV5.Serializer createRequestSerializer() + { + FederatedIdentitySerDe.Serializer identity = new FederatedIdentitySerDe.Serializer(); + TableNameSerDe.Serializer tableName = new TableNameSerDe.Serializer(); + VersionedSerDe.Serializer schema = new SchemaSerDeV4.Serializer(); + VersionedSerDe.Serializer block = new BlockSerDeV4.Serializer(schema); + ArrowTypeSerDe.Serializer arrowType = new ArrowTypeSerDe.Serializer(); + MarkerSerDe.Serializer marker = new MarkerSerDe.Serializer(block); + RangeSerDe.Serializer range = new RangeSerDe.Serializer(marker); + EquatableValueSetSerDe.Serializer equatableValueSet = new EquatableValueSetSerDe.Serializer(block); + SortedRangeSetSerDe.Serializer sortedRangeSet = new SortedRangeSetSerDe.Serializer(arrowType, range); + AllOrNoneValueSetSerDe.Serializer allOrNoneValueSet = new AllOrNoneValueSetSerDe.Serializer(arrowType); + ValueSetSerDe.Serializer valueSet = new ValueSetSerDe.Serializer(equatableValueSet, sortedRangeSet, allOrNoneValueSet); + VersionedSerDe.Serializer functionName = new FunctionNameSerDeV4.Serializer(); + ConstantExpressionSerDeV4.Serializer constantExpression = new ConstantExpressionSerDeV4.Serializer(block, arrowType); + FunctionCallExpressionSerDeV4.Serializer functionCallExpression = new FunctionCallExpressionSerDeV4.Serializer(functionName, arrowType); + VariableExpressionSerDeV4.Serializer variableExpression = new VariableExpressionSerDeV4.Serializer(arrowType); + VersionedSerDe.Serializer federationExpression = new FederationExpressionSerDeV4.Serializer(constantExpression, functionCallExpression, variableExpression); + functionCallExpression.setFederationExpressionSerializer(federationExpression); + VersionedSerDe.Serializer orderByField = new OrderByFieldSerDeV4.Serializer(); + VersionedSerDe.Serializer constraints = new ConstraintsSerDeV5.Serializer(valueSet, federationExpression, orderByField); + S3SpillLocationSerDe.Serializer s3SpillLocation = new S3SpillLocationSerDe.Serializer(); + SpillLocationSerDe.Serializer spillLocation = new SpillLocationSerDe.Serializer(s3SpillLocation); + EncryptionKeySerDe.Serializer encryptionKey = new EncryptionKeySerDe.Serializer(); + SplitSerDe.Serializer split = new SplitSerDe.Serializer(spillLocation, encryptionKey); + PingRequestSerDe.Serializer ping = new PingRequestSerDe.Serializer(identity); + ListSchemasRequestSerDe.Serializer listSchemas = new ListSchemasRequestSerDe.Serializer(identity); + ListTablesRequestSerDe.Serializer listTables = new ListTablesRequestSerDe.Serializer(identity); + GetTableRequestSerDeV5.Serializer getTable = new GetTableRequestSerDeV5.Serializer(identity, tableName); + GetTableLayoutRequestSerDe.Serializer getTableLayout = new GetTableLayoutRequestSerDe.Serializer(identity, tableName, constraints, schema); + GetSplitsRequestSerDe.Serializer getSplits = new GetSplitsRequestSerDe.Serializer(identity, tableName, block, constraints); + ReadRecordsRequestSerDe.Serializer readRecords = new ReadRecordsRequestSerDe.Serializer(identity, tableName, constraints, schema, split); + UserDefinedFunctionRequestSerDe.Serializer userDefinedFunction = new UserDefinedFunctionRequestSerDe.Serializer(identity, block, schema); + GetDataSourceCapabilitiesRequestSerDeV4.Serializer getDataSourceCapabilities = new GetDataSourceCapabilitiesRequestSerDeV4.Serializer(identity); + return new FederationRequestSerDeV5.Serializer( + ping, + listSchemas, + listTables, + getTable, + getTableLayout, + getSplits, + readRecords, + userDefinedFunction, + getDataSourceCapabilities); + } + + private static FederationRequestSerDeV5.Deserializer createRequestDeserializer(BlockAllocator allocator) + { + FederatedIdentitySerDe.Deserializer identity = new FederatedIdentitySerDe.Deserializer(); + TableNameSerDe.Deserializer tableName = new TableNameSerDe.Deserializer(); + VersionedSerDe.Deserializer schema = new SchemaSerDeV4.Deserializer(); + VersionedSerDe.Deserializer block = new BlockSerDeV4.Deserializer(allocator, schema); + ArrowTypeSerDe.Deserializer arrowType = new ArrowTypeSerDe.Deserializer(); + MarkerSerDe.Deserializer marker = new MarkerSerDe.Deserializer(block); + RangeSerDe.Deserializer range = new RangeSerDe.Deserializer(marker); + EquatableValueSetSerDe.Deserializer equatableValueSet = new EquatableValueSetSerDe.Deserializer(block); + SortedRangeSetSerDe.Deserializer sortedRangeSet = new SortedRangeSetSerDe.Deserializer(arrowType, range); + AllOrNoneValueSetSerDe.Deserializer allOrNoneValueSet = new AllOrNoneValueSetSerDe.Deserializer(arrowType); + ValueSetSerDe.Deserializer valueSet = new ValueSetSerDe.Deserializer(equatableValueSet, sortedRangeSet, allOrNoneValueSet); + + VersionedSerDe.Deserializer functionName = new FunctionNameSerDeV4.Deserializer(); + ConstantExpressionSerDeV4.Deserializer constantExpression = new ConstantExpressionSerDeV4.Deserializer(block, arrowType); + FunctionCallExpressionSerDeV4.Deserializer functionCallExpression = new FunctionCallExpressionSerDeV4.Deserializer(functionName, arrowType); + VariableExpressionSerDeV4.Deserializer variableExpression = new VariableExpressionSerDeV4.Deserializer(arrowType); + VersionedSerDe.Deserializer federationExpression = new FederationExpressionSerDeV4.Deserializer(constantExpression, functionCallExpression, variableExpression); + functionCallExpression.setFederationExpressionSerializer(federationExpression); + VersionedSerDe.Deserializer orderByField = new OrderByFieldSerDeV4.Deserializer(); + VersionedSerDe.Deserializer constraints = new ConstraintsSerDeV5.Deserializer(valueSet, federationExpression, orderByField); + + S3SpillLocationSerDe.Deserializer s3SpillLocation = new S3SpillLocationSerDe.Deserializer(); + SpillLocationSerDe.Deserializer spillLocation = new SpillLocationSerDe.Deserializer(s3SpillLocation); + EncryptionKeySerDe.Deserializer encryptionKey = new EncryptionKeySerDe.Deserializer(); + SplitSerDe.Deserializer split = new SplitSerDe.Deserializer(spillLocation, encryptionKey); + + PingRequestSerDe.Deserializer ping = new PingRequestSerDe.Deserializer(identity); + ListSchemasRequestSerDe.Deserializer listSchemas = new ListSchemasRequestSerDe.Deserializer(identity); + ListTablesRequestSerDe.Deserializer listTables = new ListTablesRequestSerDe.Deserializer(identity); + GetTableRequestSerDeV5.Deserializer getTable = new GetTableRequestSerDeV5.Deserializer(identity, tableName); + GetTableLayoutRequestSerDe.Deserializer getTableLayout = new GetTableLayoutRequestSerDe.Deserializer(identity, tableName, constraints, schema); + GetSplitsRequestSerDe.Deserializer getSplits = new GetSplitsRequestSerDe.Deserializer(identity, tableName, block, constraints); + ReadRecordsRequestSerDe.Deserializer readRecords = new ReadRecordsRequestSerDe.Deserializer(identity, tableName, constraints, schema, split); + UserDefinedFunctionRequestSerDe.Deserializer userDefinedFunction = new UserDefinedFunctionRequestSerDe.Deserializer(identity, block, schema); + GetDataSourceCapabilitiesRequestSerDeV4.Deserializer getDataSourceCapabilities = new GetDataSourceCapabilitiesRequestSerDeV4.Deserializer(identity); + + return new FederationRequestSerDeV5.Deserializer( + ping, + listSchemas, + listTables, + getTable, + getTableLayout, + getSplits, + readRecords, + userDefinedFunction, + getDataSourceCapabilities); + } + + private static FederationResponseSerDeV4.Serializer createResponseSerializer() + { + TableNameSerDe.Serializer tableName = new TableNameSerDe.Serializer(); + VersionedSerDe.Serializer schema = new SchemaSerDeV4.Serializer(); + VersionedSerDe.Serializer block = new BlockSerDeV4.Serializer(schema); + S3SpillLocationSerDe.Serializer s3SpillLocation = new S3SpillLocationSerDe.Serializer(); + SpillLocationSerDe.Serializer spillLocation = new SpillLocationSerDe.Serializer(s3SpillLocation); + EncryptionKeySerDe.Serializer encryptionKey = new EncryptionKeySerDe.Serializer(); + SplitSerDe.Serializer split = new SplitSerDe.Serializer(spillLocation, encryptionKey); + + PingResponseSerDe.Serializer ping = new PingResponseSerDe.Serializer(); + ListSchemasResponseSerDe.Serializer listSchemas = new ListSchemasResponseSerDe.Serializer(); + ListTablesResponseSerDe.Serializer listTables = new ListTablesResponseSerDe.Serializer(tableName); + GetTableResponseSerDe.Serializer getTable = new GetTableResponseSerDe.Serializer(tableName, schema); + GetTableLayoutResponseSerDe.Serializer getTableLayout = new GetTableLayoutResponseSerDe.Serializer(tableName, block); + GetSplitsResponseSerDe.Serializer getSplits = new GetSplitsResponseSerDe.Serializer(split); + ReadRecordsResponseSerDe.Serializer readRecords = new ReadRecordsResponseSerDe.Serializer(block); + RemoteReadRecordsResponseSerDe.Serializer remoteReadRecords = new RemoteReadRecordsResponseSerDe.Serializer(schema, spillLocation, encryptionKey); + UserDefinedFunctionResponseSerDe.Serializer userDefinedFunction = new UserDefinedFunctionResponseSerDe.Serializer(block); + VersionedSerDe.Serializer optimizationSubtype = new OptimizationSubTypeSerDeV4.Serializer(); + GetDataSourceCapabilitiesResponseSerDeV4.Serializer getDataSourceCapabilities = new GetDataSourceCapabilitiesResponseSerDeV4.Serializer(optimizationSubtype); + + return new FederationResponseSerDeV4.Serializer( + ping, + listSchemas, + listTables, + getTable, + getTableLayout, + getSplits, + readRecords, + remoteReadRecords, + userDefinedFunction, + getDataSourceCapabilities); + } + + private static FederationResponseSerDeV4.Deserializer createResponseDeserializer(BlockAllocator allocator) + { + TableNameSerDe.Deserializer tableName = new TableNameSerDe.Deserializer(); + VersionedSerDe.Deserializer schema = new SchemaSerDeV4.Deserializer(); + VersionedSerDe.Deserializer block = new BlockSerDeV4.Deserializer(allocator, schema); + S3SpillLocationSerDe.Deserializer s3SpillLocation = new S3SpillLocationSerDe.Deserializer(); + SpillLocationSerDe.Deserializer spillLocation = new SpillLocationSerDe.Deserializer(s3SpillLocation); + EncryptionKeySerDe.Deserializer encryptionKey = new EncryptionKeySerDe.Deserializer(); + SplitSerDe.Deserializer split = new SplitSerDe.Deserializer(spillLocation, encryptionKey); + + PingResponseSerDe.Deserializer ping = new PingResponseSerDe.Deserializer(); + ListSchemasResponseSerDe.Deserializer listSchemas = new ListSchemasResponseSerDe.Deserializer(); + ListTablesResponseSerDe.Deserializer listTables = new ListTablesResponseSerDe.Deserializer(tableName); + GetTableResponseSerDe.Deserializer getTable = new GetTableResponseSerDe.Deserializer(tableName, schema); + GetTableLayoutResponseSerDe.Deserializer getTableLayout = new GetTableLayoutResponseSerDe.Deserializer(tableName, block); + GetSplitsResponseSerDe.Deserializer getSplits = new GetSplitsResponseSerDe.Deserializer(split); + ReadRecordsResponseSerDe.Deserializer readRecords = new ReadRecordsResponseSerDe.Deserializer(block); + RemoteReadRecordsResponseSerDe.Deserializer remoteReadRecords = new RemoteReadRecordsResponseSerDe.Deserializer(schema, spillLocation, encryptionKey); + UserDefinedFunctionResponseSerDe.Deserializer userDefinedFunction = new UserDefinedFunctionResponseSerDe.Deserializer(block); + VersionedSerDe.Deserializer optimizationSubtype = new OptimizationSubTypeSerDeV4.Deserializer(); + GetDataSourceCapabilitiesResponseSerDeV4.Deserializer getDataSourceCapabilities = new GetDataSourceCapabilitiesResponseSerDeV4.Deserializer(optimizationSubtype); + + return new FederationResponseSerDeV4.Deserializer( + ping, + listSchemas, + listTables, + getTable, + getTableLayout, + getSplits, + readRecords, + remoteReadRecords, + userDefinedFunction, + getDataSourceCapabilities); + } +} diff --git a/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/handlers/CompositeHandlerTest.java b/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/handlers/CompositeHandlerTest.java index 89475c0361..4df77e8052 100644 --- a/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/handlers/CompositeHandlerTest.java +++ b/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/handlers/CompositeHandlerTest.java @@ -65,7 +65,6 @@ import java.util.UUID; import static com.amazonaws.athena.connector.lambda.domain.predicate.Constraints.DEFAULT_NO_LIMIT; -import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.nullable; import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.mock; diff --git a/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/handlers/GlueMetadataHandlerTest.java b/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/handlers/GlueMetadataHandlerTest.java index 59e13204f3..4c8409877c 100644 --- a/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/handlers/GlueMetadataHandlerTest.java +++ b/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/handlers/GlueMetadataHandlerTest.java @@ -73,6 +73,7 @@ import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.Comparator; import java.util.HashMap; import java.util.List; @@ -363,7 +364,7 @@ public void doGetTable() return mockResult; }); - GetTableRequest req = new GetTableRequest(IdentityUtil.fakeIdentity(), queryId, catalog, new TableName(schema, table)); + GetTableRequest req = new GetTableRequest(IdentityUtil.fakeIdentity(), queryId, catalog, new TableName(schema, table), Collections.emptyMap()); GetTableResponse res = handler.doGetTable(allocator, req); logger.info("doGetTable - {}", res); @@ -446,7 +447,7 @@ public void doGetTableEmptyComment() return mockResult; }); - GetTableRequest req = new GetTableRequest(IdentityUtil.fakeIdentity(), queryId, catalog, new TableName(schema, table)); + GetTableRequest req = new GetTableRequest(IdentityUtil.fakeIdentity(), queryId, catalog, new TableName(schema, table), Collections.emptyMap()); GetTableResponse res = handler.doGetTable(allocator, req); logger.info("doGetTable - {}", res); @@ -461,7 +462,7 @@ public void doGetTableEmptyComment() @Test public void testGetCatalog() { // Catalog should be the account from the request - MetadataRequest req = new GetTableRequest(IdentityUtil.fakeIdentity(), queryId, catalog, new TableName(schema, table)); + MetadataRequest req = new GetTableRequest(IdentityUtil.fakeIdentity(), queryId, catalog, new TableName(schema, table), Collections.emptyMap()); String catalog = handler.getCatalog(req); assertEquals(IdentityUtil.fakeIdentity().getAccount(), catalog); diff --git a/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/serde/TypedSerDeTest.java b/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/serde/TypedSerDeTest.java index 55f1e42ddc..f83528d741 100644 --- a/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/serde/TypedSerDeTest.java +++ b/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/serde/TypedSerDeTest.java @@ -30,8 +30,8 @@ import java.io.IOException; import java.util.Collections; -import java.util.List; +import static com.amazonaws.athena.connector.lambda.utils.TestUtils.SERDE_VERSION_FIVE; import static com.amazonaws.athena.connector.lambda.utils.TestUtils.SERDE_VERSION_FOUR; import static com.amazonaws.athena.connector.lambda.utils.TestUtils.SERDE_VERSION_TWO; @@ -41,6 +41,7 @@ public abstract class TypedSerDeTest protected BlockAllocator allocator; protected ObjectMapper mapper; protected ObjectMapper mapperV4; + protected ObjectMapper mapperV5; protected FederatedIdentity federatedIdentity = new FederatedIdentity("testArn", "0123456789", Collections.emptyMap(), Collections.emptyList()); protected String expectedSerDeText; protected T expected; @@ -51,8 +52,10 @@ public void before() allocator = new BlockAllocatorImpl("test-allocator-id"); mapper = VersionedObjectMapperFactory.create(allocator, SERDE_VERSION_TWO); mapperV4 = VersionedObjectMapperFactory.create(allocator, SERDE_VERSION_FOUR); + mapperV5 = VersionedObjectMapperFactory.create(allocator, SERDE_VERSION_FIVE); mapper.enable(SerializationFeature.INDENT_OUTPUT); mapperV4.enable(SerializationFeature.INDENT_OUTPUT); + mapperV5.enable(SerializationFeature.INDENT_OUTPUT); } @After diff --git a/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/serde/v2/GetTableRequestSerDeTest.java b/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/serde/v2/GetTableRequestSerDeTest.java index 29d8a38589..71b6ae8e7b 100644 --- a/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/serde/v2/GetTableRequestSerDeTest.java +++ b/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/serde/v2/GetTableRequestSerDeTest.java @@ -33,6 +33,7 @@ import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.InputStream; +import java.util.Collections; import static org.junit.Assert.assertEquals; @@ -44,7 +45,7 @@ public class GetTableRequestSerDeTest extends TypedSerDeTest public void beforeTest() throws IOException { - expected = new GetTableRequest(federatedIdentity, "test-query-id", "test-catalog", new TableName("test-schema", "test-table")); + expected = new GetTableRequest(federatedIdentity, "test-query-id", "test-catalog", new TableName("test-schema", "test-table"), Collections.emptyMap()); String expectedSerDeFile = utils.getResourceOrFail("serde/v2", "GetTableRequest.json"); expectedSerDeText = utils.readAllAsString(expectedSerDeFile).trim(); diff --git a/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/serde/v5/GetTableRequestSerDeV5Test.java b/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/serde/v5/GetTableRequestSerDeV5Test.java new file mode 100644 index 0000000000..54b317f0e4 --- /dev/null +++ b/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/serde/v5/GetTableRequestSerDeV5Test.java @@ -0,0 +1,88 @@ +/*- + * #%L + * Amazon Athena Query Federation SDK + * %% + * Copyright (C) 2019 - 2020 Amazon Web Services + * %% + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * #L% + */ +package com.amazonaws.athena.connector.lambda.serde.v5; + +import com.amazonaws.athena.connector.lambda.domain.TableName; +import com.amazonaws.athena.connector.lambda.metadata.GetTableRequest; +import com.amazonaws.athena.connector.lambda.request.FederationRequest; +import com.amazonaws.athena.connector.lambda.serde.TypedSerDeTest; +import com.fasterxml.jackson.core.JsonEncoding; +import org.junit.Before; +import org.junit.Test; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.util.HashMap; + +import static org.junit.Assert.assertEquals; + +public class GetTableRequestSerDeV5Test extends TypedSerDeTest +{ + private static final Logger logger = LoggerFactory.getLogger(GetTableRequestSerDeV5Test.class); + + @Before + public void beforeTest() + throws IOException + { + HashMap queryPassthroughArguments = new HashMap<>(1); + queryPassthroughArguments.put("query", "SELECT * FROM DUMMY_TABLE;"); + expected = new GetTableRequest(federatedIdentity, "test-query-id", "test-catalog", new TableName("test-schema", "test-table"), queryPassthroughArguments); + + String expectedSerDeFile = utils.getResourceOrFail("serde/v5", "GetTableRequest.json"); + expectedSerDeText = utils.readAllAsString(expectedSerDeFile).trim(); + } + + @Test + public void serialize() + throws IOException + { + logger.info("serialize: enter"); + ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); + + mapperV5.writeValue(outputStream, expected); + + String actual = new String(outputStream.toByteArray(), JsonEncoding.UTF8.getJavaName()); + logger.info("serialize: serialized text[{}]", actual); + + assertEquals(expectedSerDeText, actual); + + logger.info("serialize: exit"); + } + + @Test + public void deserialize() + throws IOException + { + logger.info("deserialize: enter"); + InputStream input = new ByteArrayInputStream(expectedSerDeText.getBytes()); + + GetTableRequest actual = (GetTableRequest) mapperV5.readValue(input, FederationRequest.class); + + logger.info("deserialize: deserialized[{}]", actual); + + assertEquals(expected, actual); + + logger.info("deserialize: exit"); + } +} diff --git a/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/serde/v5/ObjectMapperFactoryV5Test.java b/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/serde/v5/ObjectMapperFactoryV5Test.java new file mode 100644 index 0000000000..b3df328828 --- /dev/null +++ b/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/serde/v5/ObjectMapperFactoryV5Test.java @@ -0,0 +1,56 @@ +/*- + * #%L + * Amazon Athena Query Federation SDK + * %% + * Copyright (C) 2019 - 2020 Amazon Web Services + * %% + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * #L% + */ +package com.amazonaws.athena.connector.lambda.serde.v5; + +import com.amazonaws.athena.connector.lambda.data.BlockAllocator; +import com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl; +import com.amazonaws.athena.connector.lambda.serde.VersionedObjectMapperFactory; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.JsonMappingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.junit.Test; + +import java.io.IOException; + +import static com.amazonaws.athena.connector.lambda.utils.TestUtils.SERDE_VERSION_FIVE; + +public class ObjectMapperFactoryV5Test +{ + @Test(expected = JsonMappingException.class) + public void testStrictSerializer() + throws JsonProcessingException + { + try (BlockAllocator allocator = new BlockAllocatorImpl()) { + ObjectMapper mapper = VersionedObjectMapperFactory.create(allocator, SERDE_VERSION_FIVE); + mapper.writeValueAsString(new ArrowType.Null()); + } + } + + @Test(expected = JsonMappingException.class) + public void testStrictDeserializer() + throws IOException + { + try (BlockAllocator allocator = new BlockAllocatorImpl()) { + ObjectMapper mapper = VersionedObjectMapperFactory.create(allocator, SERDE_VERSION_FIVE); + mapper.readValue("{\"@type\" : \"FloatingPoint\", \"precision\" : \"DOUBLE\"}", ArrowType.FloatingPoint.class); + } + } +} diff --git a/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/serde/v5/ReadRecordsRequestSerDeV5Test.java b/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/serde/v5/ReadRecordsRequestSerDeV5Test.java new file mode 100644 index 0000000000..ba813f9e94 --- /dev/null +++ b/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/serde/v5/ReadRecordsRequestSerDeV5Test.java @@ -0,0 +1,188 @@ +/*- + * #%L + * Amazon Athena Query Federation SDK + * %% + * Copyright (C) 2019 - 2020 Amazon Web Services + * %% + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * #L% + */ +package com.amazonaws.athena.connector.lambda.serde.v5; + +import com.amazonaws.athena.connector.lambda.data.Block; +import com.amazonaws.athena.connector.lambda.data.BlockUtils; +import com.amazonaws.athena.connector.lambda.data.SchemaBuilder; +import com.amazonaws.athena.connector.lambda.domain.Split; +import com.amazonaws.athena.connector.lambda.domain.TableName; +import com.amazonaws.athena.connector.lambda.domain.predicate.AllOrNoneValueSet; +import com.amazonaws.athena.connector.lambda.domain.predicate.Constraints; +import com.amazonaws.athena.connector.lambda.domain.predicate.EquatableValueSet; +import com.amazonaws.athena.connector.lambda.domain.predicate.OrderByField; +import com.amazonaws.athena.connector.lambda.domain.predicate.Range; +import com.amazonaws.athena.connector.lambda.domain.predicate.SortedRangeSet; +import com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet; +import com.amazonaws.athena.connector.lambda.domain.predicate.expression.ConstantExpression; +import com.amazonaws.athena.connector.lambda.domain.predicate.expression.FederationExpression; +import com.amazonaws.athena.connector.lambda.domain.predicate.expression.FunctionCallExpression; +import com.amazonaws.athena.connector.lambda.domain.predicate.expression.VariableExpression; +import com.amazonaws.athena.connector.lambda.domain.predicate.functions.StandardFunctions; +import com.amazonaws.athena.connector.lambda.domain.spill.S3SpillLocation; +import com.amazonaws.athena.connector.lambda.domain.spill.SpillLocation; +import com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest; +import com.amazonaws.athena.connector.lambda.request.FederationRequest; +import com.amazonaws.athena.connector.lambda.security.EncryptionKey; +import com.amazonaws.athena.connector.lambda.serde.TypedSerDeTest; +import com.fasterxml.jackson.core.JsonEncoding; +import com.google.common.collect.ImmutableList; +import org.apache.arrow.vector.types.Types; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.Before; +import org.junit.Test; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static com.amazonaws.athena.connector.lambda.domain.predicate.Constraints.DEFAULT_NO_LIMIT; +import static org.junit.Assert.assertEquals; + +public class ReadRecordsRequestSerDeV5Test extends TypedSerDeTest +{ + private static final Logger logger = LoggerFactory.getLogger(ReadRecordsRequestSerDeV5Test.class); + + @Before + public void beforeTest() + throws IOException + { + String yearCol = "year"; + String monthCol = "month"; + String dayCol = "day"; + + Schema schema = SchemaBuilder.newBuilder() + .addField(yearCol, new ArrowType.Int(32, true)) + .addField(monthCol, new ArrowType.Int(32, true)) + .addField(dayCol, new ArrowType.Int(32, true)) + .addField("col2", new ArrowType.Utf8()) + .addField("col3", Types.MinorType.FLOAT8.getType()) + .addField("col4", Types.MinorType.FLOAT8.getType()) + .addField("col5", Types.MinorType.FLOAT8.getType()) + .build(); + + Map constraintsMap = new HashMap<>(); + constraintsMap.put("col3", SortedRangeSet.copyOf(Types.MinorType.FLOAT8.getType(), + ImmutableList.of(Range.greaterThan(allocator, Types.MinorType.FLOAT8.getType(), -10000D)), false)); + constraintsMap.put("col4", EquatableValueSet.newBuilder(allocator, Types.MinorType.FLOAT8.getType(), false, true).add(1.1D).build()); + constraintsMap.put("col5", new AllOrNoneValueSet(Types.MinorType.FLOAT8.getType(), false, true)); + + HashMap queryPassthroughArguments = new HashMap<>(1); + queryPassthroughArguments.put("query", "SELECT * FROM DUMMY_TABLE;"); + + Block partitions = allocator.createBlock(schema); + + FederationExpression federationExpression = new FunctionCallExpression( + ArrowType.Bool.INSTANCE, + StandardFunctions.GREATER_THAN_OPERATOR_FUNCTION_NAME.getFunctionName(), + ImmutableList.of(new FunctionCallExpression( + Types.MinorType.FLOAT8.getType(), + StandardFunctions.ADD_FUNCTION_NAME.getFunctionName(), + ImmutableList.of(new VariableExpression("col3", Types.MinorType.FLOAT8.getType()), + new ConstantExpression( + BlockUtils.newBlock(allocator, "col1", new ArrowType.Int(32, true), ImmutableList.of(10)), + new ArrowType.Int(32, true)))), + new VariableExpression("col2", Types.MinorType.FLOAT8.getType()))); + + List orderByClause = ImmutableList.of( + new OrderByField("col3", OrderByField.Direction.ASC_NULLS_FIRST), + new OrderByField("col2", OrderByField.Direction.DESC_NULLS_FIRST) + ); + + Constraints constraints = new Constraints(constraintsMap, ImmutableList.of(federationExpression), orderByClause, DEFAULT_NO_LIMIT, queryPassthroughArguments); + + int num_partitions = 10; + for (int i = 0; i < num_partitions; i++) { + BlockUtils.setValue(partitions.getFieldVector(yearCol), i, 2016 + i); + BlockUtils.setValue(partitions.getFieldVector(monthCol), i, (i % 12) + 1); + BlockUtils.setValue(partitions.getFieldVector(dayCol), i, (i % 28) + 1); + } + partitions.setRowCount(num_partitions); + + SpillLocation spillLocation = S3SpillLocation.newBuilder() + .withBucket("athena-virtuoso-test") + .withPrefix("lambda-spill") + .withQueryId("test-query-id") + .withSplitId("test-split-id") + .withIsDirectory(true) + .build(); + EncryptionKey encryptionKey = new EncryptionKey("test-key".getBytes(), "test-nonce".getBytes()); + Split split = Split.newBuilder(spillLocation, encryptionKey) + .add("year", "2017") + .add("month", "11") + .add("day", "1") + .build(); + + expected = new ReadRecordsRequest(federatedIdentity, + "test-query-id", + "test-catalog", + new TableName("test-schema", "test-table"), + schema, + split, + constraints, + 100_000_000_000L, + 100_000_000_000L); + + + String expectedSerDeFile = utils.getResourceOrFail("serde/v5", "ReadRecordsRequest.json"); + expectedSerDeText = utils.readAllAsString(expectedSerDeFile).trim(); + } + + @Test + public void serialize() + throws Exception + { + logger.info("serialize: enter"); + ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); + + mapperV5.writeValue(outputStream, expected); + + String actual = new String(outputStream.toByteArray(), JsonEncoding.UTF8.getJavaName()); + logger.info("serialize: serialized text[{}]", actual); + + assertEquals(expectedSerDeText, actual); + expected.close(); + + logger.info("serialize: exit"); + } + + @Test + public void deserialize() + throws IOException + { + logger.info("deserialize: enter"); + InputStream input = new ByteArrayInputStream(expectedSerDeText.getBytes()); + + ReadRecordsRequest actual = (ReadRecordsRequest) mapperV5.readValue(input, FederationRequest.class); + + logger.info("deserialize: deserialized[{}]", actual); + + assertEquals(expected, actual); + + logger.info("deserialize: exit"); + } +} diff --git a/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/utils/TestUtils.java b/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/utils/TestUtils.java index 35777e3cf6..70fa45fb95 100644 --- a/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/utils/TestUtils.java +++ b/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/utils/TestUtils.java @@ -39,6 +39,7 @@ public class TestUtils public final static int SERDE_VERSION_TWO = 2; public final static int SERDE_VERSION_THREE = 3; public final static int SERDE_VERSION_FOUR = 4; + public final static int SERDE_VERSION_FIVE = 5; /** * Helper to retrieve resources from the class path and enforce they are found diff --git a/athena-federation-sdk/src/test/resources/serde/v5/GetTableRequest.json b/athena-federation-sdk/src/test/resources/serde/v5/GetTableRequest.json new file mode 100644 index 0000000000..3a34af1054 --- /dev/null +++ b/athena-federation-sdk/src/test/resources/serde/v5/GetTableRequest.json @@ -0,0 +1,20 @@ +{ + "@type" : "GetTableRequest", + "identity" : { + "id" : "UNKNOWN", + "principal" : "UNKNOWN", + "account" : "0123456789", + "arn" : "testArn", + "tags" : { }, + "groups" : [ ] + }, + "queryId" : "test-query-id", + "catalogName" : "test-catalog", + "tableName" : { + "schemaName" : "test-schema", + "tableName" : "test-table" + }, + "queryPassthroughArguments" : { + "query" : "SELECT * FROM DUMMY_TABLE;" + } +} \ No newline at end of file diff --git a/athena-federation-sdk/src/test/resources/serde/v5/ReadRecordsRequest.json b/athena-federation-sdk/src/test/resources/serde/v5/ReadRecordsRequest.json new file mode 100644 index 0000000000..e0455cc7ee --- /dev/null +++ b/athena-federation-sdk/src/test/resources/serde/v5/ReadRecordsRequest.json @@ -0,0 +1,145 @@ +{ + "@type" : "ReadRecordsRequest", + "identity" : { + "id" : "UNKNOWN", + "principal" : "UNKNOWN", + "account" : "0123456789", + "arn" : "testArn", + "tags" : { }, + "groups" : [ ] + }, + "queryId" : "test-catalog", + "catalogName" : "test-query-id", + "tableName" : { + "schemaName" : "test-schema", + "tableName" : "test-table" + }, + "schema" : "//////ABAAAQAAAAAAAKAA4ABgANAAgACgAAAAAABAAQAAAAAAEKAAwAAAAIAAQACgAAAAgAAAAIAAAAAAAAAAcAAABwAQAAJAEAAOwAAAC0AAAAdAAAADwAAAAEAAAAvv7//xQAAAAUAAAAFAAAAAAAAwEUAAAAAAAAAAAAAACW////AAACAAQAAABjb2w1AAAAAPL+//8UAAAAFAAAABQAAAAAAAMBFAAAAAAAAAAAAAAAyv///wAAAgAEAAAAY29sNAAAAAAm////FAAAABQAAAAcAAAAAAADARwAAAAAAAAAAAAAAAAABgAIAAYABgAAAAAAAgAEAAAAY29sMwAAAABi////FAAAABQAAAAYAAAAAAAFARQAAAAAAAAAAAAAAAQABAAEAAAABAAAAGNvbDIAAAAAlv///xQAAAAUAAAAFAAAAAAAAgEYAAAAAAAAAAAAAACE////AAAAASAAAAADAAAAZGF5AMr///8UAAAAFAAAABQAAAAAAAIBGAAAAAAAAAAAAAAAuP///wAAAAEgAAAABQAAAG1vbnRoABIAGAAUABMAEgAMAAAACAAEABIAAAAUAAAAFAAAABwAAAAAAAIBIAAAAAAAAAAAAAAACAAMAAgABwAIAAAAAAAAASAAAAAEAAAAeWVhcgAAAAAAAAAA", + "split" : { + "spillLocation" : { + "@type" : "S3SpillLocation", + "bucket" : "athena-virtuoso-test", + "key" : "lambda-spill/test-query-id/test-split-id", + "directory" : true + }, + "encryptionKey" : { + "key" : "dGVzdC1rZXk=", + "nonce" : "dGVzdC1ub25jZQ==" + }, + "properties" : { + "month" : "11", + "year" : "2017", + "day" : "1" + } + }, + "constraints" : { + "summary" : { + "col4" : { + "@type" : "EquatableValueSet", + "valueBlock" : { + "aId" : "test-allocator-id", + "schema" : "/////5AAAAAQAAAAAAAKAA4ABgANAAgACgAAAAAABAAQAAAAAAEKAAwAAAAIAAQACgAAAAgAAAAIAAAAAAAAAAEAAAAYAAAAAAASABgAFAATABIADAAAAAgABAASAAAAFAAAABQAAAAcAAAAAAADARwAAAAAAAAAAAAAAAAABgAIAAYABgAAAAAAAgAEAAAAY29sMQAAAAA=", + "records" : "/////4gAAAAUAAAAAAAAAAwAFgAOABUAEAAEAAwAAAAQAAAAAAAAAAAABAAQAAAAAAMKABgADAAIAAQACgAAABQAAAA4AAAAAQAAAAAAAAAAAAAAAgAAAAAAAAAAAAAAAQAAAAAAAAAIAAAAAAAAAAgAAAAAAAAAAAAAAAEAAAABAAAAAAAAAAAAAAAAAAAAAQAAAAAAAACamZmZmZnxPw==" + }, + "whiteList" : false, + "nullAllowed" : true + }, + "col5" : { + "@type" : "AllOrNoneValueSet", + "type" : { + "@type" : "FloatingPoint", + "precision" : "DOUBLE" + }, + "all" : false, + "nullAllowed" : true + }, + "col3" : { + "@type" : "SortedRangeSet", + "type" : { + "@type" : "FloatingPoint", + "precision" : "DOUBLE" + }, + "ranges" : [ { + "low" : { + "valueBlock" : { + "aId" : "test-allocator-id", + "schema" : "/////5AAAAAQAAAAAAAKAA4ABgANAAgACgAAAAAABAAQAAAAAAEKAAwAAAAIAAQACgAAAAgAAAAIAAAAAAAAAAEAAAAYAAAAAAASABgAFAATABIADAAAAAgABAASAAAAFAAAABQAAAAcAAAAAAADARwAAAAAAAAAAAAAAAAABgAIAAYABgAAAAAAAgAEAAAAY29sMQAAAAA=", + "records" : "/////4gAAAAUAAAAAAAAAAwAFgAOABUAEAAEAAwAAAAQAAAAAAAAAAAABAAQAAAAAAMKABgADAAIAAQACgAAABQAAAA4AAAAAQAAAAAAAAAAAAAAAgAAAAAAAAAAAAAAAQAAAAAAAAAIAAAAAAAAAAgAAAAAAAAAAAAAAAEAAAABAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAIjDwA==" + }, + "bound" : "ABOVE", + "nullValue" : false + }, + "high" : { + "valueBlock" : { + "aId" : "test-allocator-id", + "schema" : "/////5AAAAAQAAAAAAAKAA4ABgANAAgACgAAAAAABAAQAAAAAAEKAAwAAAAIAAQACgAAAAgAAAAIAAAAAAAAAAEAAAAYAAAAAAASABgAFAATABIADAAAAAgABAASAAAAFAAAABQAAAAcAAAAAAADARwAAAAAAAAAAAAAAAAABgAIAAYABgAAAAAAAgAEAAAAY29sMQAAAAA=", + "records" : "" + }, + "bound" : "BELOW", + "nullValue" : true + } + } ], + "nullAllowed" : false + } + }, + "expression" : [ { + "@type" : "FunctionCallExpression", + "type" : { + "@type" : "Bool" + }, + "functionName" : { + "functionName" : "$greater_than" + }, + "arguments" : [ { + "@type" : "FunctionCallExpression", + "type" : { + "@type" : "FloatingPoint", + "precision" : "DOUBLE" + }, + "functionName" : { + "functionName" : "$add" + }, + "arguments" : [ { + "@type" : "VariableExpression", + "columnName" : "col3", + "type" : { + "@type" : "FloatingPoint", + "precision" : "DOUBLE" + } + }, { + "@type" : "ConstantExpression", + "valueBlock" : { + "aId" : "test-allocator-id", + "schema" : "/////5gAAAAQAAAAAAAKAA4ABgANAAgACgAAAAAABAAQAAAAAAEKAAwAAAAIAAQACgAAAAgAAAAIAAAAAAAAAAEAAAAYAAAAAAASABgAFAATABIADAAAAAgABAASAAAAFAAAABQAAAAcAAAAAAACASAAAAAAAAAAAAAAAAgADAAIAAcACAAAAAAAAAEgAAAABAAAAGNvbDEAAAAAAAAAAA==", + "records" : "/////4gAAAAUAAAAAAAAAAwAFgAOABUAEAAEAAwAAAAQAAAAAAAAAAAABAAQAAAAAAMKABgADAAIAAQACgAAABQAAAA4AAAAAQAAAAAAAAAAAAAAAgAAAAAAAAAAAAAAAQAAAAAAAAAIAAAAAAAAAAQAAAAAAAAAAAAAAAEAAAABAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAKAAAAAAAAAA==" + }, + "type" : { + "@type" : "Int", + "bitWidth" : 32, + "isSigned" : true + } + } ] + }, { + "@type" : "VariableExpression", + "columnName" : "col2", + "type" : { + "@type" : "FloatingPoint", + "precision" : "DOUBLE" + } + } ] + } ], + "orderByClause" : [ { + "columnName" : "col3", + "direction" : "ASC_NULLS_FIRST" + }, { + "columnName" : "col2", + "direction" : "DESC_NULLS_FIRST" + } ], + "limit" : -1, + "queryPassthroughArguments" : { + "query" : "SELECT * FROM DUMMY_TABLE;" + } + }, + "maxBlockSize" : "100000000000", + "maxInlineBlockSize" : "100000000000" +} diff --git a/athena-gcs/src/test/java/com/amazonaws/athena/connectors/gcs/GcsMetadataHandlerTest.java b/athena-gcs/src/test/java/com/amazonaws/athena/connectors/gcs/GcsMetadataHandlerTest.java index 8855ccf8a8..84005f8d02 100644 --- a/athena-gcs/src/test/java/com/amazonaws/athena/connectors/gcs/GcsMetadataHandlerTest.java +++ b/athena-gcs/src/test/java/com/amazonaws/athena/connectors/gcs/GcsMetadataHandlerTest.java @@ -80,11 +80,13 @@ import java.util.ArrayList; import java.util.Collection; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.stream.Collectors; +import static com.amazonaws.athena.connector.lambda.domain.predicate.Constraints.DEFAULT_NO_LIMIT; import static com.amazonaws.athena.connectors.gcs.GcsConstants.CLASSIFICATION_GLUE_TABLE_PARAM; import static com.amazonaws.athena.connectors.gcs.GcsConstants.PARTITION_PATTERN_KEY; import static com.amazonaws.athena.connectors.gcs.GcsTestUtils.allocator; @@ -97,7 +99,6 @@ import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.mockStatic; -import static org.mockito.Mockito.nullable; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -250,7 +251,7 @@ public void doGetTable() Map metadataSchema = new HashMap<>(); metadataSchema.put("dataFormat", PARQUET); Schema schema = new Schema(asList(field), metadataSchema); - GetTableRequest getTableRequest = new GetTableRequest(federatedIdentity, QUERY_ID, "gcs", new TableName(SCHEMA_NAME, "testtable")); + GetTableRequest getTableRequest = new GetTableRequest(federatedIdentity, QUERY_ID, "gcs", new TableName(SCHEMA_NAME, "testtable"), Collections.emptyMap()); Table table = new Table(); table.setName(TABLE_1); table.setDatabaseName(DATABASE_NAME); @@ -301,7 +302,8 @@ public void testGetPartitions() throws Exception GetTableLayoutRequest getTableLayoutRequest = Mockito.mock(GetTableLayoutRequest.class); Mockito.when(getTableLayoutRequest.getTableName()).thenReturn(new TableName(DATABASE_NAME, TABLE_1)); Mockito.when(getTableLayoutRequest.getSchema()).thenReturn(schema); - Constraints constraints = new Constraints(createSummaryWithLValueRangeEqual("year", new ArrowType.Utf8(), 2000)); + Constraints constraints = new Constraints(createSummaryWithLValueRangeEqual("year", new ArrowType.Utf8(), 2000) , + Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()); Mockito.when(getTableLayoutRequest.getConstraints()).thenReturn(constraints); BlockWriter blockWriter = Mockito.mock(BlockWriter.class); gcsMetadataHandler.getPartitions(blockWriter, getTableLayoutRequest, null); @@ -314,7 +316,7 @@ public void testDoGetSplits() throws Exception Block partitions = BlockUtils.newBlock(blockAllocator, "year", Types.MinorType.VARCHAR.getType(), 2000, 2001); GetSplitsRequest request = new GetSplitsRequest(federatedIdentity, QUERY_ID, CATALOG, TABLE_NAME, - partitions, ImmutableList.of("year"), new Constraints(new HashMap<>()), null); + partitions, ImmutableList.of("year"), new Constraints(new HashMap<>(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), null); QueryStatusChecker queryStatusChecker = mock(QueryStatusChecker.class); GetTableResult getTableResult = mock(GetTableResult.class); StorageDescriptor storageDescriptor = mock(StorageDescriptor.class); @@ -351,7 +353,7 @@ public void testDoGetSplitsProperty() throws Exception partitions.setRowCount(num_partitions); GetSplitsRequest request = new GetSplitsRequest(federatedIdentity, QUERY_ID, CATALOG, TABLE_NAME, - partitions, ImmutableList.of("yearCol", "monthCol"), new Constraints(new HashMap<>()), null); + partitions, ImmutableList.of("yearCol", "monthCol"), new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), null); QueryStatusChecker queryStatusChecker = mock(QueryStatusChecker.class); GetTableResult getTableResult = mock(GetTableResult.class); StorageDescriptor storageDescriptor = mock(StorageDescriptor.class); @@ -377,7 +379,7 @@ public void testDoGetSplitsException() throws Exception Block partitions = BlockUtils.newBlock(blockAllocator, "gcs_file_format", Types.MinorType.VARCHAR.getType(), 2000, 2001); GetSplitsRequest request = new GetSplitsRequest(federatedIdentity, QUERY_ID, CATALOG, TABLE_NAME, - partitions, ImmutableList.of("gcs_file_format"), new Constraints(new HashMap<>()), null); + partitions, ImmutableList.of("gcs_file_format"), new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), null); QueryStatusChecker queryStatusChecker = mock(QueryStatusChecker.class); GetTableResult getTableResult = mock(GetTableResult.class); StorageDescriptor storageDescriptor = mock(StorageDescriptor.class); diff --git a/athena-gcs/src/test/java/com/amazonaws/athena/connectors/gcs/GcsRecordHandlerTest.java b/athena-gcs/src/test/java/com/amazonaws/athena/connectors/gcs/GcsRecordHandlerTest.java index ce712b1a06..60274ad4f7 100644 --- a/athena-gcs/src/test/java/com/amazonaws/athena/connectors/gcs/GcsRecordHandlerTest.java +++ b/athena-gcs/src/test/java/com/amazonaws/athena/connectors/gcs/GcsRecordHandlerTest.java @@ -57,6 +57,7 @@ import java.util.Collections; import java.util.UUID; +import static com.amazonaws.athena.connector.lambda.domain.predicate.Constraints.DEFAULT_NO_LIMIT; import static com.amazonaws.athena.connectors.gcs.GcsConstants.FILE_FORMAT; import static com.amazonaws.athena.connectors.gcs.GcsConstants.STORAGE_SPLIT_JSON; import static org.junit.jupiter.api.Assertions.assertEquals; @@ -165,7 +166,7 @@ public void testReadWithConstraint() new TableName("dataset1", "table1"), // dummy table GcsTestUtils.getDatatypeTestSchema(), split, - new Constraints(Collections.EMPTY_MAP), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), 0, //This is ignored when directly calling readWithConstraints. 0)) { //This is ignored when directly calling readWithConstraints. diff --git a/athena-gcs/src/test/java/com/amazonaws/athena/connectors/gcs/filter/FilterExpressionBuilderTest.java b/athena-gcs/src/test/java/com/amazonaws/athena/connectors/gcs/filter/FilterExpressionBuilderTest.java index e894048213..b114faa5b0 100644 --- a/athena-gcs/src/test/java/com/amazonaws/athena/connectors/gcs/filter/FilterExpressionBuilderTest.java +++ b/athena-gcs/src/test/java/com/amazonaws/athena/connectors/gcs/filter/FilterExpressionBuilderTest.java @@ -39,9 +39,11 @@ import org.mockito.junit.MockitoJUnitRunner; +import java.util.Collections; import java.util.List; import java.util.Map; +import static com.amazonaws.athena.connector.lambda.domain.predicate.Constraints.DEFAULT_NO_LIMIT; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.mockito.ArgumentMatchers.anyString; @@ -53,7 +55,8 @@ public void testGetExpressions() { Map>> result = FilterExpressionBuilder.getConstraintsForPartitionedColumns( com.google.common.collect.ImmutableList.of(new Column().withName("year")), - new Constraints(GcsTestUtils.createSummaryWithLValueRangeEqual("year", new ArrowType.Utf8(), "1"))); + new Constraints(GcsTestUtils.createSummaryWithLValueRangeEqual("year", new ArrowType.Utf8(), "1"), + Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap())); assertEquals(result.size(), 1); assertEquals(result.get("year").get(), com.google.common.collect.ImmutableSet.of("1")); assertEquals(result.get("yeAr").get(), com.google.common.collect.ImmutableSet.of("1")); diff --git a/athena-gcs/src/test/java/com/amazonaws/athena/connectors/gcs/storage/StorageMetadataTest.java b/athena-gcs/src/test/java/com/amazonaws/athena/connectors/gcs/storage/StorageMetadataTest.java index 46b512f156..55a8139b01 100644 --- a/athena-gcs/src/test/java/com/amazonaws/athena/connectors/gcs/storage/StorageMetadataTest.java +++ b/athena-gcs/src/test/java/com/amazonaws/athena/connectors/gcs/storage/StorageMetadataTest.java @@ -56,10 +56,12 @@ import java.util.ArrayList; import java.util.Collection; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; +import static com.amazonaws.athena.connector.lambda.domain.predicate.Constraints.DEFAULT_NO_LIMIT; import static com.amazonaws.athena.connectors.gcs.GcsConstants.CLASSIFICATION_GLUE_TABLE_PARAM; import static com.amazonaws.athena.connectors.gcs.GcsConstants.PARTITION_PATTERN_KEY; import static com.amazonaws.athena.connectors.gcs.GcsMetadataHandlerTest.LOCATION; @@ -168,7 +170,8 @@ public void testGetPartitionFolders() throws Exception List fieldList = ImmutableList.of(new Field("year", FieldType.nullable(new ArrowType.Int(64, true)), null)); List partKeys = ImmutableList.of(createColumn("year", "varchar")); Schema schema = getSchema(glue, fieldList, partKeys, "year=${year}/"); - List> partValue = storageMetadata.getPartitionFolders(schema, new TableName("testSchema", "testTable"), new Constraints(ImmutableMap.of()), glue); + List> partValue = storageMetadata.getPartitionFolders(schema, new TableName("testSchema", "testTable"), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), glue); assertEquals(1, partValue.size()); assertEquals(partValue, ImmutableList.of(ImmutableMap.of("year", "2000"))); @@ -188,7 +191,8 @@ public void testGetPartitionFolders() throws Exception new Field("month", FieldType.nullable(new ArrowType.Utf8()), null)); List partKeys1 = ImmutableList.of(createColumn("year", "varchar"), createColumn("month", "varchar")); Schema schema1 = getSchema(glue, fieldList1, partKeys1, "year=${year}/birth_month${month}/"); - List> partValue1 = storageMetadata.getPartitionFolders(schema1, new TableName("testSchema", "testTable"), new Constraints(ImmutableMap.of()), glue); + List> partValue1 = storageMetadata.getPartitionFolders(schema1, new TableName("testSchema", "testTable"), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), glue); assertEquals(4, partValue1.size()); assertEquals(partValue1, ImmutableList.of(ImmutableMap.of("year", "2000", "month", "1"), ImmutableMap.of("year", "2000", "month", "2"), ImmutableMap.of("year", "2001", "month", "1"), ImmutableMap.of("year", "2001", "month", "2"))); @@ -203,7 +207,8 @@ public void testGetPartitionFolders() throws Exception new Field("month", FieldType.nullable(new ArrowType.Utf8()), null)); List partKeys2 = ImmutableList.of(createColumn("year", "varchar"), createColumn("month", "varchar")); Schema schema2 = getSchema(glue, fieldList2, partKeys2, "year=${year}/birth_month${month}/"); - List> partValue2 = storageMetadata.getPartitionFolders(schema2, new TableName("testSchema", "testTable"), new Constraints(ImmutableMap.of()), glue); + List> partValue2 = storageMetadata.getPartitionFolders(schema2, new TableName("testSchema", "testTable"), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), glue); assertEquals(4, partValue2.size()); assertEquals(partValue2, ImmutableList.of(ImmutableMap.of("year", "2000", "month", "1"), ImmutableMap.of("year", "2000", "month", "2"), ImmutableMap.of("year", "2001", "month", "1"), ImmutableMap.of("year", "2001", "month", "2"))); @@ -213,7 +218,8 @@ public void testGetPartitionFolders() throws Exception List fieldList3 = ImmutableList.of(new Field("year", FieldType.nullable(new ArrowType.Int(64, true)), null)); List partKeys3 = ImmutableList.of(createColumn("year", "varchar")); Schema schema3 = getSchema(glue, fieldList3, partKeys3, "year=${year}/"); - List> partValue3 = storageMetadata.getPartitionFolders(schema3, new TableName("testSchema", "testTable"), new Constraints(ImmutableMap.of()), glue); + List> partValue3 = storageMetadata.getPartitionFolders(schema3, new TableName("testSchema", "testTable"), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), glue); assertEquals(0, partValue3.size()); assertEquals(partValue3, ImmutableList.of()); @@ -222,7 +228,8 @@ public void testGetPartitionFolders() throws Exception List fieldList4 = ImmutableList.of(new Field("year", FieldType.nullable(new ArrowType.Int(64, true)), null)); List partKeys4 = ImmutableList.of(createColumn("year", "varchar")); Schema schema4 = getSchema(glue, fieldList4, partKeys4, "year=${year}/"); - List> partValue4 = storageMetadata.getPartitionFolders(schema3, new TableName("testSchema", "testTable"), new Constraints(ImmutableMap.of()), glue); + List> partValue4 = storageMetadata.getPartitionFolders(schema3, new TableName("testSchema", "testTable"), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), glue); assertEquals(0, partValue4.size()); assertEquals(partValue4, ImmutableList.of()); } diff --git a/athena-google-bigquery/src/test/java/com/amazonaws/athena/connectors/google/bigquery/BigQueryMetadataHandlerTest.java b/athena-google-bigquery/src/test/java/com/amazonaws/athena/connectors/google/bigquery/BigQueryMetadataHandlerTest.java index 8cefbe3ddf..3b1dd1b044 100644 --- a/athena-google-bigquery/src/test/java/com/amazonaws/athena/connectors/google/bigquery/BigQueryMetadataHandlerTest.java +++ b/athena-google-bigquery/src/test/java/com/amazonaws/athena/connectors/google/bigquery/BigQueryMetadataHandlerTest.java @@ -198,7 +198,7 @@ public void testDoGetTable() throws IOException //Make the call GetTableRequest getTableRequest = new GetTableRequest(federatedIdentity, QUERY_ID, BigQueryTestUtils.PROJECT_1_NAME, - new TableName(datasetName, tableName)); + new TableName(datasetName, tableName), Collections.emptyMap()); GetTableResponse response = bigQueryMetadataHandler.doGetTable(blockAllocator, getTableRequest); @@ -216,7 +216,7 @@ public void testDoGetTable() throws IOException //Make the call GetTableRequest getTableRequest1 = new GetTableRequest(federatedIdentity, QUERY_ID, BigQueryTestUtils.PROJECT_1_NAME, - new TableName(datasetName, tableName)); + new TableName(datasetName, tableName), Collections.emptyMap()); GetTableResponse responseComplex = bigQueryMetadataHandler.doGetTable(blockAllocator, getTableRequest1); diff --git a/athena-google-bigquery/src/test/java/com/amazonaws/athena/connectors/google/bigquery/BigQueryRecordHandlerTest.java b/athena-google-bigquery/src/test/java/com/amazonaws/athena/connectors/google/bigquery/BigQueryRecordHandlerTest.java index d6bf9b6abc..74f97587bd 100644 --- a/athena-google-bigquery/src/test/java/com/amazonaws/athena/connectors/google/bigquery/BigQueryRecordHandlerTest.java +++ b/athena-google-bigquery/src/test/java/com/amazonaws/athena/connectors/google/bigquery/BigQueryRecordHandlerTest.java @@ -87,6 +87,7 @@ import java.util.Map; import java.util.UUID; +import static com.amazonaws.athena.connector.lambda.domain.predicate.Constraints.DEFAULT_NO_LIMIT; import static com.amazonaws.athena.connectors.google.bigquery.BigQueryTestUtils.getBlockTestSchema; import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; @@ -263,7 +264,7 @@ public void testReadWithConstraint() .withIsDirectory(true) .build(), keyFactory.create()).build(), - new Constraints(Collections.EMPTY_MAP), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), 0, //This is ignored when directly calling readWithConstraints. 0)) { diff --git a/athena-hbase/src/test/java/com/amazonaws/athena/connectors/hbase/HbaseMetadataHandlerTest.java b/athena-hbase/src/test/java/com/amazonaws/athena/connectors/hbase/HbaseMetadataHandlerTest.java index b42cb4f464..eb43f48841 100644 --- a/athena-hbase/src/test/java/com/amazonaws/athena/connectors/hbase/HbaseMetadataHandlerTest.java +++ b/athena-hbase/src/test/java/com/amazonaws/athena/connectors/hbase/HbaseMetadataHandlerTest.java @@ -67,11 +67,11 @@ import java.io.IOException; import java.util.ArrayList; import java.util.Collections; -import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Set; +import static com.amazonaws.athena.connector.lambda.domain.predicate.Constraints.DEFAULT_NO_LIMIT; import static com.amazonaws.athena.connector.lambda.metadata.ListTablesRequest.UNLIMITED_PAGE_SIZE_VALUE; import static org.junit.Assert.*; import static org.mockito.ArgumentMatchers.any; @@ -206,7 +206,7 @@ public void doGetTable() return processor.scan(mockScanner); }); - GetTableRequest req = new GetTableRequest(IDENTITY, QUERY_ID, DEFAULT_CATALOG, TABLE_NAME); + GetTableRequest req = new GetTableRequest(IDENTITY, QUERY_ID, DEFAULT_CATALOG, TABLE_NAME, Collections.emptyMap()); GetTableResponse res = handler.doGetTable(allocator, req); logger.info("doGetTable - {}", res); @@ -225,7 +225,7 @@ public void doGetTableLayout() QUERY_ID, DEFAULT_CATALOG, TABLE_NAME, - new Constraints(new HashMap<>()), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), SchemaBuilder.newBuilder().build(), Collections.EMPTY_SET); @@ -262,7 +262,7 @@ public void doGetSplits() TABLE_NAME, partitions, partitionCols, - new Constraints(new HashMap<>()), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), null); GetSplitsRequest req = new GetSplitsRequest(originalReq, continuationToken); diff --git a/athena-hortonworks-hive/src/test/java/com/amazonaws/athena/connectors/hortonworks/HiveMetadataHandlerTest.java b/athena-hortonworks-hive/src/test/java/com/amazonaws/athena/connectors/hortonworks/HiveMetadataHandlerTest.java index f1370ce736..ff9065debe 100644 --- a/athena-hortonworks-hive/src/test/java/com/amazonaws/athena/connectors/hortonworks/HiveMetadataHandlerTest.java +++ b/athena-hortonworks-hive/src/test/java/com/amazonaws/athena/connectors/hortonworks/HiveMetadataHandlerTest.java @@ -293,7 +293,7 @@ public void doGetTable() Mockito.when(this.connection.getMetaData().getColumns("testCatalog", inputTableName.getSchemaName(), inputTableName.getTableName(), null)).thenReturn(resultSet1); Mockito.when(this.connection.getCatalog()).thenReturn("testCatalog"); GetTableResponse getTableResponse = this.hiveMetadataHandler.doGetTable( - this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); Assert.assertEquals(inputTableName, getTableResponse.getTableName()); Assert.assertEquals("testCatalog", getTableResponse.getCatalogName()); } @@ -302,7 +302,7 @@ public void doGetTable() public void doGetTableNoColumns() throws Exception { TableName inputTableName = new TableName("testSchema", "testTable"); - this.hiveMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.hiveMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } @Test(expected = SQLException.class) @@ -312,6 +312,6 @@ public void doGetTableSQLException() TableName inputTableName = new TableName("testSchema", "testTable"); Mockito.when(this.connection.getMetaData().getColumns(nullable(String.class), nullable(String.class), nullable(String.class), nullable(String.class))) .thenThrow(new SQLException()); - this.hiveMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.hiveMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } } diff --git a/athena-jdbc/src/test/java/com/amazonaws/athena/connectors/jdbc/manager/JdbcMetadataHandlerTest.java b/athena-jdbc/src/test/java/com/amazonaws/athena/connectors/jdbc/manager/JdbcMetadataHandlerTest.java index 61faf3fedc..35bb083715 100644 --- a/athena-jdbc/src/test/java/com/amazonaws/athena/connectors/jdbc/manager/JdbcMetadataHandlerTest.java +++ b/athena-jdbc/src/test/java/com/amazonaws/athena/connectors/jdbc/manager/JdbcMetadataHandlerTest.java @@ -53,7 +53,7 @@ import java.sql.ResultSet; import java.sql.SQLException; import java.sql.Types; -import java.util.List; +import java.util.Collections; import java.util.concurrent.atomic.AtomicInteger; import static com.amazonaws.athena.connector.lambda.metadata.ListTablesRequest.UNLIMITED_PAGE_SIZE_VALUE; @@ -199,7 +199,7 @@ public void doGetTable() Mockito.when(connection.getMetaData().getColumns("testCatalog", inputTableName.getSchemaName(), inputTableName.getTableName(), null)).thenReturn(resultSet); GetTableResponse getTableResponse = this.jdbcMetadataHandler.doGetTable( - this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); Assert.assertEquals(expected, getTableResponse.getSchema()); Assert.assertEquals(inputTableName, getTableResponse.getTableName()); @@ -217,7 +217,7 @@ public void doGetTableCaseInsensitive() setupMocksDoGetTableCaseInsensitive(inputTableName, values1, "testTable"); GetTableResponse getTableResponse = this.jdbcMetadataHandler.doGetTable(this.blockAllocator, - new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); Assert.assertEquals("testTable", getTableResponse.getTableName().getTableName()); } @@ -228,7 +228,7 @@ public void doGetTableNoColumns() { TableName inputTableName = new TableName("testSchema", "testTable"); - this.jdbcMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.jdbcMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } @Test(expected = SQLException.class) @@ -238,7 +238,7 @@ public void doGetTableSQLException() TableName inputTableName = new TableName("testSchema", "testTable"); Mockito.when(this.connection.getMetaData().getColumns(nullable(String.class), nullable(String.class), nullable(String.class), Mockito.isNull())) .thenThrow(new SQLException()); - this.jdbcMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.jdbcMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } @Test(expected = SQLException.class) diff --git a/athena-kafka/src/test/java/com/amazonaws/athena/connectors/kafka/KafkaMetadataHandlerTest.java b/athena-kafka/src/test/java/com/amazonaws/athena/connectors/kafka/KafkaMetadataHandlerTest.java index 5064993863..27d8ad77d1 100644 --- a/athena-kafka/src/test/java/com/amazonaws/athena/connectors/kafka/KafkaMetadataHandlerTest.java +++ b/athena-kafka/src/test/java/com/amazonaws/athena/connectors/kafka/KafkaMetadataHandlerTest.java @@ -33,15 +33,12 @@ import com.amazonaws.services.glue.model.GetSchemaVersionResult; import com.amazonaws.services.glue.model.ListRegistriesResult; import com.amazonaws.services.glue.model.RegistryListItem; -import com.amazonaws.services.secretsmanager.AWSSecretsManagerClientBuilder; -import org.apache.kafka.clients.consumer.KafkaConsumer; import org.apache.kafka.clients.consumer.MockConsumer; import org.apache.kafka.clients.consumer.OffsetResetStrategy; import org.apache.kafka.common.PartitionInfo; import org.apache.kafka.common.TopicPartition; import org.junit.After; import org.junit.Before; -import org.junit.Rule; import org.junit.Test; import org.junit.runner.RunWith; import org.mockito.Mock; @@ -51,6 +48,7 @@ import org.mockito.junit.MockitoJUnitRunner; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -172,7 +170,7 @@ public void testDoGetTable() throws Exception { "}"); Mockito.when(awsGlue.getSchema(any())).thenReturn(getSchemaResult); Mockito.when(awsGlue.getSchemaVersion(any())).thenReturn(getSchemaVersionResult); - GetTableRequest getTableRequest = new GetTableRequest(federatedIdentity, QUERY_ID, "kafka", new TableName("default", "testtable")); + GetTableRequest getTableRequest = new GetTableRequest(federatedIdentity, QUERY_ID, "kafka", new TableName("default", "testtable"), Collections.emptyMap()); GetTableResponse getTableResponse = kafkaMetadataHandler.doGetTable(blockAllocator, getTableRequest); assertEquals(1, getTableResponse.getSchema().getFields().size()); } diff --git a/athena-kafka/src/test/java/com/amazonaws/athena/connectors/kafka/KafkaRecordHandlerTest.java b/athena-kafka/src/test/java/com/amazonaws/athena/connectors/kafka/KafkaRecordHandlerTest.java index 3ec2d65daa..149169158d 100644 --- a/athena-kafka/src/test/java/com/amazonaws/athena/connectors/kafka/KafkaRecordHandlerTest.java +++ b/athena-kafka/src/test/java/com/amazonaws/athena/connectors/kafka/KafkaRecordHandlerTest.java @@ -61,6 +61,7 @@ import java.util.UUID; +import static com.amazonaws.athena.connector.lambda.domain.predicate.Constraints.DEFAULT_NO_LIMIT; import static org.mockito.ArgumentMatchers.anyMap; import static org.mockito.Mockito.*; @@ -244,7 +245,7 @@ private ReadRecordsRequest createReadRecordsRequest(Schema schema) { .withIsDirectory(true) .build(), keyFactory.create()).build(), - new Constraints(Collections.EMPTY_MAP), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), 0, 0); } diff --git a/athena-msk/src/test/java/com/amazonaws/athena/connectors/msk/AmazonMskMetadataHandlerTest.java b/athena-msk/src/test/java/com/amazonaws/athena/connectors/msk/AmazonMskMetadataHandlerTest.java index 503e05ced5..06a2fae970 100644 --- a/athena-msk/src/test/java/com/amazonaws/athena/connectors/msk/AmazonMskMetadataHandlerTest.java +++ b/athena-msk/src/test/java/com/amazonaws/athena/connectors/msk/AmazonMskMetadataHandlerTest.java @@ -48,6 +48,7 @@ import org.mockito.junit.MockitoJUnitRunner; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -165,7 +166,7 @@ public void testDoGetTable() throws Exception { "}"); Mockito.when(awsGlue.getSchema(any())).thenReturn(getSchemaResult); Mockito.when(awsGlue.getSchemaVersion(any())).thenReturn(getSchemaVersionResult); - GetTableRequest getTableRequest = new GetTableRequest(federatedIdentity, QUERY_ID, "kafka", new TableName("default", "testtable")); + GetTableRequest getTableRequest = new GetTableRequest(federatedIdentity, QUERY_ID, "kafka", new TableName("default", "testtable"), Collections.emptyMap()); GetTableResponse getTableResponse = amazonMskMetadataHandler.doGetTable(blockAllocator, getTableRequest); assertEquals(1, getTableResponse.getSchema().getFields().size()); } diff --git a/athena-msk/src/test/java/com/amazonaws/athena/connectors/msk/AmazonMskRecordHandlerTest.java b/athena-msk/src/test/java/com/amazonaws/athena/connectors/msk/AmazonMskRecordHandlerTest.java index 66af336554..78d5f2da24 100644 --- a/athena-msk/src/test/java/com/amazonaws/athena/connectors/msk/AmazonMskRecordHandlerTest.java +++ b/athena-msk/src/test/java/com/amazonaws/athena/connectors/msk/AmazonMskRecordHandlerTest.java @@ -58,6 +58,7 @@ import java.util.HashMap; import java.util.UUID; +import static com.amazonaws.athena.connector.lambda.domain.predicate.Constraints.DEFAULT_NO_LIMIT; import static org.mockito.ArgumentMatchers.anyMap; import static org.mockito.Mockito.*; @@ -243,7 +244,7 @@ private ReadRecordsRequest createReadRecordsRequest(Schema schema) { .withIsDirectory(true) .build(), keyFactory.create()).build(), - new Constraints(Collections.EMPTY_MAP), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), 0, 0); } diff --git a/athena-mysql/src/test/java/com/amazonaws/athena/connectors/mysql/MySqlMetadataHandlerTest.java b/athena-mysql/src/test/java/com/amazonaws/athena/connectors/mysql/MySqlMetadataHandlerTest.java index 71642bf694..1f44feac4f 100644 --- a/athena-mysql/src/test/java/com/amazonaws/athena/connectors/mysql/MySqlMetadataHandlerTest.java +++ b/athena-mysql/src/test/java/com/amazonaws/athena/connectors/mysql/MySqlMetadataHandlerTest.java @@ -44,7 +44,6 @@ import com.amazonaws.services.secretsmanager.AWSSecretsManager; import com.amazonaws.services.secretsmanager.model.GetSecretValueRequest; import com.amazonaws.services.secretsmanager.model.GetSecretValueResult; -import net.jqwik.api.Table; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.junit.Assert; @@ -328,7 +327,7 @@ public void doGetTableCaseInsensitiveDuplicateTableNames() Mockito.when(this.connection.prepareStatement(sql).executeQuery()).thenReturn(resultSetName); GetTableResponse getTableResponse = this.mySqlMetadataHandler.doGetTable(this.blockAllocator, - new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } @org.testng.annotations.Test(expectedExceptions = {RuntimeException.class}, expectedExceptionsMessageRegExp = "During Case Insensitive look up could not find Table testtable in Database testSchema") @@ -342,6 +341,6 @@ public void doGetTableCaseInsensitiveNoTablesFound() Mockito.when(this.connection.prepareStatement(sql).executeQuery()).thenReturn(resultSetName); GetTableResponse getTableResponse = this.mySqlMetadataHandler.doGetTable(this.blockAllocator, - new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } } diff --git a/athena-neptune/src/test/java/com/amazonaws/athena/connectors/neptune/NeptuneMetadataHandlerTest.java b/athena-neptune/src/test/java/com/amazonaws/athena/connectors/neptune/NeptuneMetadataHandlerTest.java index 9a4eb31a77..a6c31c88ae 100644 --- a/athena-neptune/src/test/java/com/amazonaws/athena/connectors/neptune/NeptuneMetadataHandlerTest.java +++ b/athena-neptune/src/test/java/com/amazonaws/athena/connectors/neptune/NeptuneMetadataHandlerTest.java @@ -47,6 +47,7 @@ import org.slf4j.LoggerFactory; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -165,7 +166,7 @@ public void doGetTable() throws Exception { storageDescriptor.setColumns(columns); table.setStorageDescriptor(storageDescriptor); - GetTableRequest req = new GetTableRequest(IDENTITY, "queryId", "default", new TableName("schema1", "table1")); + GetTableRequest req = new GetTableRequest(IDENTITY, "queryId", "default", new TableName("schema1", "table1"), Collections.emptyMap()); GetTableResult getTableResult = new GetTableResult(); getTableResult.setTable(table); diff --git a/athena-neptune/src/test/java/com/amazonaws/athena/connectors/neptune/NeptuneRecordHandlerTest.java b/athena-neptune/src/test/java/com/amazonaws/athena/connectors/neptune/NeptuneRecordHandlerTest.java index 2aee628c78..5f4ed5ea24 100644 --- a/athena-neptune/src/test/java/com/amazonaws/athena/connectors/neptune/NeptuneRecordHandlerTest.java +++ b/athena-neptune/src/test/java/com/amazonaws/athena/connectors/neptune/NeptuneRecordHandlerTest.java @@ -19,6 +19,7 @@ */ package com.amazonaws.athena.connectors.neptune; +import static com.amazonaws.athena.connector.lambda.domain.predicate.Constraints.DEFAULT_NO_LIMIT; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; @@ -275,7 +276,7 @@ public void doReadRecordsSpill() throws Exception { ReadRecordsRequest request = new ReadRecordsRequest(IDENTITY, DEFAULT_CATALOG, QUERY_ID, TABLE_NAME, schemaPGVertexForRead, Split.newBuilder(splitLoc, keyFactory.create()).build(), - new Constraints(constraintsMap), 1_500_000L, // ~1.5MB so we should see some spill + new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), 1_500_000L, // ~1.5MB so we should see some spill 0L); RecordResponse rawResponse = handler.doReadRecords(allocator, request); diff --git a/athena-oracle/src/test/java/com/amazonaws/athena/connectors/oracle/OracleMetadataHandlerTest.java b/athena-oracle/src/test/java/com/amazonaws/athena/connectors/oracle/OracleMetadataHandlerTest.java index 47d977c6a2..63a4e6b7bd 100644 --- a/athena-oracle/src/test/java/com/amazonaws/athena/connectors/oracle/OracleMetadataHandlerTest.java +++ b/athena-oracle/src/test/java/com/amazonaws/athena/connectors/oracle/OracleMetadataHandlerTest.java @@ -33,7 +33,6 @@ import com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig; import com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory; import com.amazonaws.athena.connectors.jdbc.connection.JdbcCredentialProvider; -import com.amazonaws.athena.connectors.oracle.OracleMetadataHandler; import com.amazonaws.services.athena.AmazonAthena; import com.amazonaws.services.secretsmanager.AWSSecretsManager; import com.amazonaws.services.secretsmanager.model.GetSecretValueRequest; @@ -330,7 +329,7 @@ public void doGetTable() Mockito.when(connection.getCatalog()).thenReturn("testCatalog"); GetTableResponse getTableResponse = this.oracleMetadataHandler.doGetTable( - blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); Assert.assertEquals(expected, getTableResponse.getSchema()); Assert.assertEquals(inputTableName, getTableResponse.getTableName()); diff --git a/athena-postgresql/src/test/java/com/amazonaws/athena/connectors/postgresql/PostGreSqlMetadataHandlerTest.java b/athena-postgresql/src/test/java/com/amazonaws/athena/connectors/postgresql/PostGreSqlMetadataHandlerTest.java index 474f59afbb..557619ad7f 100644 --- a/athena-postgresql/src/test/java/com/amazonaws/athena/connectors/postgresql/PostGreSqlMetadataHandlerTest.java +++ b/athena-postgresql/src/test/java/com/amazonaws/athena/connectors/postgresql/PostGreSqlMetadataHandlerTest.java @@ -392,7 +392,7 @@ public void doGetTableWithArrayColumns() Mockito.when(connection.getCatalog()).thenReturn("testCatalog"); GetTableResponse getTableResponse = this.postGreSqlMetadataHandler.doGetTable(new BlockAllocatorImpl(), - new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); logger.info("Schema: {}", getTableResponse.getSchema()); @@ -474,7 +474,7 @@ public void doGetTableMaterializedView() Mockito.when(connection.getCatalog()).thenReturn("testCatalog"); GetTableResponse getTableResponse = this.postGreSqlMetadataHandler.doGetTable(new BlockAllocatorImpl(), - new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); logger.info("Schema: {}", getTableResponse.getSchema()); diff --git a/athena-redshift/src/test/java/com/amazonaws/athena/connectors/redshift/RedshiftMetadataHandlerTest.java b/athena-redshift/src/test/java/com/amazonaws/athena/connectors/redshift/RedshiftMetadataHandlerTest.java index d7638757cd..027c716876 100644 --- a/athena-redshift/src/test/java/com/amazonaws/athena/connectors/redshift/RedshiftMetadataHandlerTest.java +++ b/athena-redshift/src/test/java/com/amazonaws/athena/connectors/redshift/RedshiftMetadataHandlerTest.java @@ -394,7 +394,7 @@ public void doGetTableWithArrayColumns() Mockito.when(connection.getCatalog()).thenReturn("testCatalog"); GetTableResponse getTableResponse = this.redshiftMetadataHandler.doGetTable(new BlockAllocatorImpl(), - new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); logger.info("Schema: {}", getTableResponse.getSchema()); diff --git a/athena-saphana/src/test/java/com/amazonaws/athena/connectors/saphana/SaphanaMetadataHandlerTest.java b/athena-saphana/src/test/java/com/amazonaws/athena/connectors/saphana/SaphanaMetadataHandlerTest.java index 1571277c16..bf8b86c007 100644 --- a/athena-saphana/src/test/java/com/amazonaws/athena/connectors/saphana/SaphanaMetadataHandlerTest.java +++ b/athena-saphana/src/test/java/com/amazonaws/athena/connectors/saphana/SaphanaMetadataHandlerTest.java @@ -23,7 +23,6 @@ import java.sql.ResultSet; import java.sql.SQLException; import java.sql.Types; -import java.sql.Statement; import java.util.*; import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; @@ -288,7 +287,7 @@ public void doGetTable() Mockito.when(connection.getMetaData().getColumns("testCatalog", inputTableName.getSchemaName(), inputTableName.getTableName(), null)).thenReturn(resultSet); Mockito.when(connection.getCatalog()).thenReturn("testCatalog"); GetTableResponse getTableResponse = this.saphanaMetadataHandler.doGetTable( - this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); Assert.assertEquals(expected, getTableResponse.getSchema()); Assert.assertEquals(inputTableName, getTableResponse.getTableName()); Assert.assertEquals("testCatalog", getTableResponse.getCatalogName()); @@ -355,13 +354,13 @@ public void doGetTableSQLException() TableName inputTableName = new TableName("testSchema", "testTable"); Mockito.when(this.connection.getMetaData().getColumns(nullable(String.class), nullable(String.class), nullable(String.class), nullable(String.class))) .thenThrow(new SQLException()); - this.saphanaMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.saphanaMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } @Test (expected = RuntimeException.class) public void doGetTableNoColumns() throws Exception { TableName inputTableName = new TableName("testSchema", "testTable"); - this.saphanaMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.saphanaMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } } diff --git a/athena-snowflake/src/test/java/com/amazonaws/athena/connectors/snowflake/SnowflakeMetadataHandlerTest.java b/athena-snowflake/src/test/java/com/amazonaws/athena/connectors/snowflake/SnowflakeMetadataHandlerTest.java index e3b5bd5431..b03df0f097 100644 --- a/athena-snowflake/src/test/java/com/amazonaws/athena/connectors/snowflake/SnowflakeMetadataHandlerTest.java +++ b/athena-snowflake/src/test/java/com/amazonaws/athena/connectors/snowflake/SnowflakeMetadataHandlerTest.java @@ -29,8 +29,6 @@ import com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig; import com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory; import com.amazonaws.athena.connectors.jdbc.connection.JdbcCredentialProvider; -import com.amazonaws.athena.connectors.snowflake.SnowflakeConstants; -import com.amazonaws.athena.connectors.snowflake.SnowflakeMetadataHandler; import com.amazonaws.services.athena.AmazonAthena; import com.amazonaws.services.secretsmanager.AWSSecretsManager; import com.amazonaws.services.secretsmanager.model.GetSecretValueRequest; @@ -39,9 +37,7 @@ import org.apache.arrow.vector.types.pojo.Schema; import org.junit.Assert; import org.junit.Before; -import org.junit.Rule; import org.junit.Test; -import org.mockito.Mock; import org.mockito.Mockito; import java.sql.*; @@ -324,7 +320,7 @@ public void doGetTableNoColumns() throws Exception { TableName inputTableName = new TableName("testSchema", "testTable"); - this.snowflakeMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.snowflakeMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } @Test(expected = SQLException.class) @@ -334,21 +330,21 @@ public void doGetTableSQLException() TableName inputTableName = new TableName("testSchema", "testTable"); Mockito.when(this.connection.getMetaData().getColumns(nullable(String.class), nullable(String.class), nullable(String.class), nullable(String.class))) .thenThrow(new SQLException()); - this.snowflakeMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.snowflakeMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } @Test(expected = RuntimeException.class) public void doGetTableException() throws Exception { TableName inputTableName = new TableName("testSchema", "test@schema"); - this.snowflakeMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.snowflakeMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } @Test(expected = RuntimeException.class) public void doGetTableNoColumnsException() throws Exception { TableName inputTableName = new TableName("testSchema", "test@table"); - this.snowflakeMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.snowflakeMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } @Test @@ -375,7 +371,7 @@ public void doGetTable() Mockito.when(connection.getCatalog()).thenReturn("testCatalog"); GetTableResponse getTableResponse = this.snowflakeMetadataHandler.doGetTable( - this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); Assert.assertEquals(expected, getTableResponse.getSchema()); Assert.assertEquals(inputTableName, getTableResponse.getTableName()); diff --git a/athena-sqlserver/src/test/java/com/amazonaws/athena/connectors/sqlserver/SqlServerMetadataHandlerTest.java b/athena-sqlserver/src/test/java/com/amazonaws/athena/connectors/sqlserver/SqlServerMetadataHandlerTest.java index 0151f14870..27704cfdb1 100644 --- a/athena-sqlserver/src/test/java/com/amazonaws/athena/connectors/sqlserver/SqlServerMetadataHandlerTest.java +++ b/athena-sqlserver/src/test/java/com/amazonaws/athena/connectors/sqlserver/SqlServerMetadataHandlerTest.java @@ -439,7 +439,7 @@ public void doGetTable() Mockito.when(connection.getMetaData().getColumns("testCatalog", inputTableName.getSchemaName(), inputTableName.getTableName(), null)).thenReturn(resultSet); Mockito.when(connection.getCatalog()).thenReturn("testCatalog"); GetTableResponse getTableResponse = this.sqlServerMetadataHandler.doGetTable( - blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); Assert.assertEquals(expected, getTableResponse.getSchema()); Assert.assertEquals(inputTableName, getTableResponse.getTableName()); Assert.assertEquals("testCatalog", getTableResponse.getCatalogName()); diff --git a/athena-synapse/src/test/java/com/amazonaws/athena/connectors/synapse/SynapseMetadataHandlerTest.java b/athena-synapse/src/test/java/com/amazonaws/athena/connectors/synapse/SynapseMetadataHandlerTest.java index 3ead850c51..037b0540e5 100644 --- a/athena-synapse/src/test/java/com/amazonaws/athena/connectors/synapse/SynapseMetadataHandlerTest.java +++ b/athena-synapse/src/test/java/com/amazonaws/athena/connectors/synapse/SynapseMetadataHandlerTest.java @@ -368,7 +368,7 @@ public void doGetTable() Mockito.when(connection.getMetaData().getColumns("testCatalog", inputTableName.getSchemaName(), inputTableName.getTableName(), null)).thenReturn(resultSet2); GetTableResponse getTableResponse = this.synapseMetadataHandler.doGetTable( - blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); Assert.assertEquals(expected, getTableResponse.getSchema()); Assert.assertEquals(inputTableName, getTableResponse.getTableName()); Assert.assertEquals("testCatalog", getTableResponse.getCatalogName()); @@ -402,7 +402,7 @@ public void doDataTypeConversion() Mockito.when(connection.getMetaData().getColumns("testCatalog", inputTableName.getSchemaName(), inputTableName.getTableName(), null)).thenReturn(resultSet2); GetTableResponse getTableResponse = this.synapseMetadataHandler.doGetTable( - blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); Assert.assertEquals(inputTableName, getTableResponse.getTableName()); Assert.assertEquals("testCatalog", getTableResponse.getCatalogName()); } diff --git a/athena-teradata/src/test/java/com/amazonaws/athena/connectors/teradata/TeradataMetadataHandlerTest.java b/athena-teradata/src/test/java/com/amazonaws/athena/connectors/teradata/TeradataMetadataHandlerTest.java index 8f78aa61aa..d7ad931f00 100644 --- a/athena-teradata/src/test/java/com/amazonaws/athena/connectors/teradata/TeradataMetadataHandlerTest.java +++ b/athena-teradata/src/test/java/com/amazonaws/athena/connectors/teradata/TeradataMetadataHandlerTest.java @@ -37,7 +37,6 @@ import org.apache.arrow.vector.types.pojo.Schema; import org.junit.Assert; import org.junit.Before; -import org.junit.Rule; import org.junit.Test; import org.mockito.Mockito; @@ -308,7 +307,7 @@ public void doGetTable() Mockito.when(connection.getCatalog()).thenReturn("testCatalog"); GetTableResponse getTableResponse = this.teradataMetadataHandler.doGetTable( - this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); Assert.assertEquals(expected, getTableResponse.getSchema()); Assert.assertEquals(inputTableName, getTableResponse.getTableName()); @@ -322,13 +321,13 @@ public void doGetTableSQLException() TableName inputTableName = new TableName("testSchema", "testTable"); Mockito.when(this.connection.getMetaData().getColumns(nullable(String.class), nullable(String.class), nullable(String.class), nullable(String.class))) .thenThrow(new SQLException()); - this.teradataMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.teradataMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } @Test (expected = RuntimeException.class) public void doGetTableNoColumns() throws Exception { TableName inputTableName = new TableName("testSchema", "testTable"); - this.teradataMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName)); + this.teradataMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap())); } } diff --git a/athena-timestream/src/test/java/com/amazonaws/athena/connectors/timestream/TimestreamMetadataHandlerTest.java b/athena-timestream/src/test/java/com/amazonaws/athena/connectors/timestream/TimestreamMetadataHandlerTest.java index ac54c7e3cd..45478f3b84 100644 --- a/athena-timestream/src/test/java/com/amazonaws/athena/connectors/timestream/TimestreamMetadataHandlerTest.java +++ b/athena-timestream/src/test/java/com/amazonaws/athena/connectors/timestream/TimestreamMetadataHandlerTest.java @@ -72,10 +72,10 @@ import java.util.ArrayList; import java.util.Collections; -import java.util.HashMap; import java.util.Iterator; import java.util.List; +import static com.amazonaws.athena.connector.lambda.domain.predicate.Constraints.DEFAULT_NO_LIMIT; import static com.amazonaws.athena.connector.lambda.handlers.GlueMetadataHandler.VIEW_METADATA_FIELD; import static com.amazonaws.athena.connector.lambda.metadata.ListTablesRequest.UNLIMITED_PAGE_SIZE_VALUE; import static org.junit.Assert.*; @@ -266,7 +266,7 @@ public void doGetTable() GetTableRequest req = new GetTableRequest(identity, "query-id", "default", - new TableName(defaultSchema, "table1")); + new TableName(defaultSchema, "table1"), Collections.emptyMap()); GetTableResponse res = handler.doGetTable(allocator, req); logger.info("doGetTable - {}", res); @@ -316,7 +316,7 @@ public void doGetTableGlue() GetTableRequest req = new GetTableRequest(identity, "query-id", "default", - new TableName(defaultSchema, "table1")); + new TableName(defaultSchema, "table1"), Collections.emptyMap()); GetTableResponse res = handler.doGetTable(allocator, req); logger.info("doGetTable - {}", res); @@ -364,7 +364,7 @@ public void doGetTimeSeriesTableGlue() GetTableRequest req = new GetTableRequest(identity, "query-id", "default", - new TableName(defaultSchema, "table1")); + new TableName(defaultSchema, "table1"), Collections.emptyMap()); GetTableResponse res = handler.doGetTable(allocator, req); logger.info("doGetTable - {}", res); @@ -409,7 +409,7 @@ public void doGetTableLayout() "query-id", defaultSchema, new TableName("database1", "table1"), - new Constraints(new HashMap<>()), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), schema, Collections.EMPTY_SET); @@ -442,7 +442,7 @@ public void doGetSplits() new TableName("database1", "table1"), partitions, partitionCols, - new Constraints(new HashMap<>()), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), null); GetSplitsRequest req = new GetSplitsRequest(originalReq, continuationToken); diff --git a/athena-timestream/src/test/java/com/amazonaws/athena/connectors/timestream/TimestreamRecordHandlerTest.java b/athena-timestream/src/test/java/com/amazonaws/athena/connectors/timestream/TimestreamRecordHandlerTest.java index e82505ad20..1682b4ef53 100644 --- a/athena-timestream/src/test/java/com/amazonaws/athena/connectors/timestream/TimestreamRecordHandlerTest.java +++ b/athena-timestream/src/test/java/com/amazonaws/athena/connectors/timestream/TimestreamRecordHandlerTest.java @@ -227,7 +227,7 @@ public void doReadRecordsNoSpill() new TableName(DEFAULT_SCHEMA, TEST_TABLE), schemaForRead, splitBuilder.build(), - new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT), + new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), 100_000_000_000L, //100GB don't expect this to spill 100_000_000_000L ); @@ -282,7 +282,7 @@ public void doReadRecordsSpill() new TableName(DEFAULT_SCHEMA, TEST_TABLE), schemaForRead, splitBuilder.build(), - new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT), + new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), 1_500_000L, //~1.5MB so we should see some spill 0L ); @@ -356,7 +356,7 @@ public void readRecordsView() new TableName(DEFAULT_SCHEMA, TEST_VIEW), schemaForReadView, split, - new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT), + new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), 100_000_000_000L, //100GB don't expect this to spill 100_000_000_000L ); @@ -424,7 +424,7 @@ public void readRecordsTimeSeriesView() new TableName(DEFAULT_SCHEMA, TEST_TABLE), schemaForReadView, split, - new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT), + new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), 100_000_000_000L, //100GB don't expect this to spill 100_000_000_000L ); @@ -477,7 +477,7 @@ public void doReadRecordsNoSpillValidateTimeStamp() new TableName(DEFAULT_SCHEMA, TEST_TABLE), schemaForRead, splitBuilder.build(), - new Constraints(constraintsMap), + new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), 100_000_000_000L, //100GB don't expect this to spill 100_000_000_000L ); diff --git a/athena-tpcds/src/test/java/com/amazonaws/athena/connectors/tpcds/TPCDSMetadataHandlerTest.java b/athena-tpcds/src/test/java/com/amazonaws/athena/connectors/tpcds/TPCDSMetadataHandlerTest.java index a9ae6274ee..7e39835921 100644 --- a/athena-tpcds/src/test/java/com/amazonaws/athena/connectors/tpcds/TPCDSMetadataHandlerTest.java +++ b/athena-tpcds/src/test/java/com/amazonaws/athena/connectors/tpcds/TPCDSMetadataHandlerTest.java @@ -135,7 +135,7 @@ public void doGetTable() GetTableRequest req = new GetTableRequest(identity, "queryId", "default", - new TableName(expectedSchema, "customer")); + new TableName(expectedSchema, "customer"), Collections.emptyMap()); GetTableResponse res = handler.doGetTable(allocator, req); logger.info("doGetTable - {} {}", res.getTableName(), res.getSchema()); @@ -192,7 +192,7 @@ public void doGetSplits() new TableName("tpcds1", "customer"), partitions, Collections.EMPTY_LIST, - new Constraints(new HashMap<>()), + new Constraints(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), continuationToken); int numContinuations = 0; diff --git a/athena-tpcds/src/test/java/com/amazonaws/athena/connectors/tpcds/TPCDSRecordHandlerTest.java b/athena-tpcds/src/test/java/com/amazonaws/athena/connectors/tpcds/TPCDSRecordHandlerTest.java index 8051afa4bc..2fc214cee8 100644 --- a/athena-tpcds/src/test/java/com/amazonaws/athena/connectors/tpcds/TPCDSRecordHandlerTest.java +++ b/athena-tpcds/src/test/java/com/amazonaws/athena/connectors/tpcds/TPCDSRecordHandlerTest.java @@ -294,7 +294,7 @@ public void doReadRecordForTPCDSTIMETypeColumn() .add(SPLIT_TOTAL_NUMBER_FIELD, "1000") .add(SPLIT_SCALE_FACTOR_FIELD, "1") .build(), - new Constraints(ImmutableMap.of()), + new Constraints(ImmutableMap.of(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap()), 100_000_000_000L, 100_000_000_000L ); diff --git a/pom.xml b/pom.xml index 77ecb8d54d..4acb3f4cfd 100644 --- a/pom.xml +++ b/pom.xml @@ -13,18 +13,21 @@ 11 3.12.1 - 1.12.661 + + 1.12.533 1.2.2 1.6.0 1.204.0 1.94.0 - 2.0.12 + + 2.0.7 4.11.0 4.13.2 1.8.2 3.25.3 7.9.0 - 2.16.1 + + 2.12.6 3.2.5 2.22.1 13.0.0