Skip to content

Commit

Permalink
SDK-Side: Query Passthrough Implementation (awslabs#1769)
Browse files Browse the repository at this point in the history 
Co-authored-by: AbdulRehman Faraj <arfaraj@amazon.com>
  • Loading branch information
@AbdulR3hman
AbdulR3hman and AbdulRehman Faraj authored Feb 23, 2024
1 parent 6a6d738 commit f364862
Show file tree
Hide file tree
Showing 74 changed files with 1,778 additions and 188 deletions.
3 changes: 2 additions & 1 deletion .gitignore
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
*/target
*/dependency-reduced-pom.xml
dependency-reduced-pom.xml
*/*.xml.unformatted
.idea/
/target/
*/*.iml
**/*.iml
.classpath
.factorypath
.project
Expand Down
4 changes: 1 addition & 3 deletions ...db/src/test/java/com/amazonaws/athena/connectors/aws/cmdb/AwsCmdbMetadataHandlerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down
2 changes: 1 addition & 1 deletion .../test/java/com/amazonaws/athena/connectors/aws/cmdb/tables/AbstractTableProviderTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down
7 changes: 3 additions & 4 deletions ...-hive/src/test/java/com/amazonaws/athena/connectors/cloudera/HiveMetadataHandlerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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());
}
Expand All @@ -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)
Expand All @@ -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()));
}
}
6 changes: 3 additions & 3 deletions ...ala/src/test/java/com/amazonaws/athena/connectors/cloudera/ImpalaMetadataHandlerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -278,15 +278,15 @@ 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());
}
@Test
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)
Expand All @@ -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()));
}
}
4 changes: 2 additions & 2 deletions ...t/java/com/amazonaws/athena/connectors/cloudwatch/metrics/MetricsMetadataHandlerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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());

Expand All @@ -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());
Expand Down
2 changes: 1 addition & 1 deletion ...c/test/java/com/amazonaws/athena/connectors/cloudwatch/CloudwatchMetadataHandlerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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());

Expand Down
2 changes: 1 addition & 1 deletion ...st/java/com/amazonaws/athena/connectors/datalakegen2/DataLakeGen2MetadataHandlerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down
10 changes: 5 additions & 5 deletions ...0/src/test/java/com/amazonaws/athena/connectors/db2as400/Db2As400MetadataHandlerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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
Expand Down
10 changes: 5 additions & 5 deletions athena-db2/src/test/java/com/amazonaws/athena/connectors/db2/Db2MetadataHandlerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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
Expand Down
8 changes: 4 additions & 4 deletions ...a-docdb/src/test/java/com/amazonaws/athena/connectors/docdb/DocDBMetadataHandlerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down Expand Up @@ -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());
Expand Down Expand Up @@ -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");
Expand Down Expand Up @@ -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());
Expand Down
3 changes: 1 addition & 2 deletions athena-docdb/src/test/java/com/amazonaws/athena/connectors/docdb/DocDBRecordHandlerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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);

Expand Down
Loading

0 comments on commit f364862

Please sign in to comment.