Examples with Constraints com.amazonaws.athena.connector.lambda.domain.predicate.Constraints used on opensource projects

Example 46 with Constraints

use of com.amazonaws.athena.connector.lambda.domain.predicate.Constraints in project aws-athena-query-federation by awslabs.

the class SqlServerMetadataHandlerTest method doGetSplitsContinuation.

@Test
public void doGetSplitsContinuation() throws Exception {
    BlockAllocator blockAllocator = new BlockAllocatorImpl();
    Constraints constraints = Mockito.mock(Constraints.class);
    TableName tableName = new TableName("testSchema", "testTable");
    Schema partitionSchema = this.sqlServerMetadataHandler.getPartitionSchema("testCatalogName");
    Set<String> partitionCols = partitionSchema.getFields().stream().map(Field::getName).collect(Collectors.toSet());
    GetTableLayoutRequest getTableLayoutRequest = new GetTableLayoutRequest(this.federatedIdentity, "testQueryId", "testCatalogName", tableName, constraints, partitionSchema, partitionCols);
    PreparedStatement rowCountPreparedStatement = Mockito.mock(PreparedStatement.class);
    Mockito.when(this.connection.prepareStatement(sqlServerMetadataHandler.ROW_COUNT_QUERY)).thenReturn(rowCountPreparedStatement);
    ResultSet rowCountResultSet = mockResultSet(new String[] { "ROW_COUNT" }, new int[] { Types.INTEGER }, new Object[][] { { 2 } }, new AtomicInteger(-1));
    Mockito.when(rowCountPreparedStatement.executeQuery()).thenReturn(rowCountResultSet);
    PreparedStatement preparedStatement = Mockito.mock(PreparedStatement.class);
    Mockito.when(this.connection.prepareStatement(sqlServerMetadataHandler.GET_PARTITIONS_QUERY)).thenReturn(preparedStatement);
    String[] columns = { sqlServerMetadataHandler.PARTITION_NUMBER };
    int[] types = { Types.INTEGER };
    Object[][] values = { { 2 }, { 3 } };
    ResultSet resultSet = mockResultSet(columns, types, values, new AtomicInteger(-1));
    Mockito.when(preparedStatement.executeQuery()).thenReturn(resultSet);
    PreparedStatement partFuncPreparedStatement = Mockito.mock(PreparedStatement.class);
    Mockito.when(this.connection.prepareStatement(sqlServerMetadataHandler.GET_PARTITION_FUNCTION_QUERY)).thenReturn(partFuncPreparedStatement);
    ResultSet partFuncResultSet = mockResultSet(new String[] { "PARTITION FUNCTION", "PARTITIONING COLUMN" }, new int[] { Types.VARCHAR, Types.VARCHAR }, new Object[][] { { "pf", "pc" } }, new AtomicInteger(-1));
    Mockito.when(partFuncPreparedStatement.executeQuery()).thenReturn(partFuncResultSet);
    Mockito.when(this.connection.getMetaData().getSearchStringEscape()).thenReturn(null);
    GetTableLayoutResponse getTableLayoutResponse = this.sqlServerMetadataHandler.doGetTableLayout(blockAllocator, getTableLayoutRequest);
    BlockAllocator splitBlockAllocator = new BlockAllocatorImpl();
    GetSplitsRequest getSplitsRequest = new GetSplitsRequest(this.federatedIdentity, "testQueryId", "testCatalogName", tableName, getTableLayoutResponse.getPartitions(), new ArrayList<>(partitionCols), constraints, "2");
    GetSplitsResponse getSplitsResponse = this.sqlServerMetadataHandler.doGetSplits(splitBlockAllocator, getSplitsRequest);
    Set<Map<String, String>> expectedSplits = new HashSet<>();
    expectedSplits.add(Map.ofEntries(Map.entry(sqlServerMetadataHandler.PARTITION_NUMBER, "3"), Map.entry("PARTITIONING_COLUMN", "pc"), Map.entry("PARTITION_FUNCTION", "pf")));
    Set<Map<String, String>> actualSplits = getSplitsResponse.getSplits().stream().map(Split::getProperties).collect(Collectors.toSet());
    Assert.assertEquals(expectedSplits, actualSplits);
}

Example 47 with Constraints

use of com.amazonaws.athena.connector.lambda.domain.predicate.Constraints in project aws-athena-query-federation by awslabs.

the class SynapseMuxRecordHandlerTest method buildSplitSql.

@Test
public void buildSplitSql() throws SQLException {
    ReadRecordsRequest readRecordsRequest = Mockito.mock(ReadRecordsRequest.class);
    Mockito.when(readRecordsRequest.getCatalogName()).thenReturn(SynapseConstants.NAME);
    Connection jdbcConnection = Mockito.mock(Connection.class);
    TableName tableName = new TableName("testSchema", "tableName");
    Schema schema = Mockito.mock(Schema.class);
    Constraints constraints = Mockito.mock(Constraints.class);
    Split split = Mockito.mock(Split.class);
    this.jdbcRecordHandler.buildSplitSql(jdbcConnection, SynapseConstants.NAME, tableName, schema, constraints, split);
    Mockito.verify(this.synapseRecordHandler, Mockito.times(1)).buildSplitSql(Mockito.eq(jdbcConnection), Mockito.eq(SynapseConstants.NAME), Mockito.eq(tableName), Mockito.eq(schema), Mockito.eq(constraints), Mockito.eq(split));
}

Example 48 with Constraints

use of com.amazonaws.athena.connector.lambda.domain.predicate.Constraints in project aws-athena-query-federation by awslabs.

the class SqlServerMuxRecordHandlerTest method buildSplitSql.

@Test
public void buildSplitSql() throws SQLException {
    ReadRecordsRequest readRecordsRequest = Mockito.mock(ReadRecordsRequest.class);
    Mockito.when(readRecordsRequest.getCatalogName()).thenReturn(SqlServerConstants.NAME);
    Connection jdbcConnection = Mockito.mock(Connection.class);
    TableName tableName = new TableName("testSchema", "tableName");
    Schema schema = Mockito.mock(Schema.class);
    Constraints constraints = Mockito.mock(Constraints.class);
    Split split = Mockito.mock(Split.class);
    this.jdbcRecordHandler.buildSplitSql(jdbcConnection, SqlServerConstants.NAME, tableName, schema, constraints, split);
    Mockito.verify(this.sqlServerRecordHandler, Mockito.times(1)).buildSplitSql(Mockito.eq(jdbcConnection), Mockito.eq(SqlServerConstants.NAME), Mockito.eq(tableName), Mockito.eq(schema), Mockito.eq(constraints), Mockito.eq(split));
}

Example 49 with Constraints

use of com.amazonaws.athena.connector.lambda.domain.predicate.Constraints in project aws-athena-query-federation by awslabs.

the class TimestreamMetadataHandlerTest method doGetSplits.

@Test
public void doGetSplits() throws Exception {
    logger.info("doGetSplits - enter");
    List<String> partitionCols = new ArrayList<>();
    Block partitions = BlockUtils.newBlock(allocator, "partition_id", Types.MinorType.INT.getType(), 0);
    String continuationToken = null;
    GetSplitsRequest originalReq = new GetSplitsRequest(identity, "query-id", defaultSchema, new TableName("database1", "table1"), partitions, partitionCols, new Constraints(new HashMap<>()), null);
    GetSplitsRequest req = new GetSplitsRequest(originalReq, continuationToken);
    logger.info("doGetSplits: req[{}]", req);
    MetadataResponse rawResponse = handler.doGetSplits(allocator, req);
    assertEquals(MetadataRequestType.GET_SPLITS, rawResponse.getRequestType());
    GetSplitsResponse response = (GetSplitsResponse) rawResponse;
    continuationToken = response.getContinuationToken();
    logger.info("doGetSplits: continuationToken[{}] - numSplits[{}]", new Object[] { continuationToken, response.getSplits().size() });
    assertTrue("Continuation criteria violated", response.getSplits().size() == 1);
    assertTrue("Continuation criteria violated", response.getContinuationToken() == null);
    logger.info("doGetSplits - exit");
}

Example 50 with Constraints

use of com.amazonaws.athena.connector.lambda.domain.predicate.Constraints in project aws-athena-query-federation by awslabs.

the class TimestreamMetadataHandlerTest method doGetTableLayout.

@Test
public void doGetTableLayout() throws Exception {
    logger.info("doGetTableLayout - enter");
    Schema schema = SchemaBuilder.newBuilder().build();
    GetTableLayoutRequest req = new GetTableLayoutRequest(identity, "query-id", defaultSchema, new TableName("database1", "table1"), new Constraints(new HashMap<>()), schema, Collections.EMPTY_SET);
    GetTableLayoutResponse res = handler.doGetTableLayout(allocator, req);
    logger.info("doGetTableLayout - {}", res);
    Block partitions = res.getPartitions();
    for (int row = 0; row < partitions.getRowCount() && row < 10; row++) {
        logger.info("doGetTableLayout:{} {}", row, BlockUtils.rowToString(partitions, row));
    }
    assertTrue(partitions.getRowCount() == 1);
    logger.info("doGetTableLayout - exit");
}