Examples with GetSplitsRequest com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest used on opensource projects

Example 11 with GetSplitsRequest

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

the class OracleMuxJdbcMetadataHandlerTest method doGetSplits.

@Test
public void doGetSplits() {
    GetSplitsRequest getSplitsRequest = Mockito.mock(GetSplitsRequest.class);
    Mockito.when(getSplitsRequest.getCatalogName()).thenReturn("fakedatabase");
    this.jdbcMetadataHandler.doGetSplits(this.allocator, getSplitsRequest);
    Mockito.verify(this.oracleMetadataHandler, Mockito.times(1)).doGetSplits(Mockito.eq(this.allocator), Mockito.eq(getSplitsRequest));
}

Example 12 with GetSplitsRequest

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

the class RedisMetadataHandlerTest method doGetSplitsZset.

@Test
public void doGetSplitsZset() {
    // 3 prefixes for this table
    String prefixes = "prefix1-*,prefix2-*, prefix3-*";
    // 4 zsets per prefix
    when(mockSyncCommands.scan(any(ScanCursor.class), any(ScanArgs.class))).then((InvocationOnMock invocationOnMock) -> {
        ScanCursor cursor = (ScanCursor) invocationOnMock.getArguments()[0];
        if (cursor == null || cursor.getCursor().equals("0")) {
            List<String> result = new ArrayList<>();
            result.add(UUID.randomUUID().toString());
            result.add(UUID.randomUUID().toString());
            result.add(UUID.randomUUID().toString());
            MockKeyScanCursor<String> scanCursor = new MockKeyScanCursor<>();
            scanCursor.setCursor("1");
            scanCursor.setKeys(result);
            return scanCursor;
        } else {
            List<String> result = new ArrayList<>();
            result.add(UUID.randomUUID().toString());
            MockKeyScanCursor<String> scanCursor = new MockKeyScanCursor<>();
            scanCursor.setCursor("0");
            scanCursor.setKeys(result);
            scanCursor.setFinished(true);
            return scanCursor;
        }
    });
    // 100 keys per zset
    when(mockSyncCommands.zcount(anyString(), any(Range.class))).thenReturn(200L);
    List<String> partitionCols = new ArrayList<>();
    Schema schema = SchemaBuilder.newBuilder().addField("partitionId", Types.MinorType.INT.getType()).addStringField(REDIS_ENDPOINT_PROP).addStringField(VALUE_TYPE_TABLE_PROP).addStringField(KEY_PREFIX_TABLE_PROP).addStringField(ZSET_KEYS_TABLE_PROP).addStringField(REDIS_SSL_FLAG).addStringField(REDIS_CLUSTER_FLAG).addStringField(REDIS_DB_NUMBER).build();
    Block partitions = allocator.createBlock(schema);
    partitions.setValue(REDIS_ENDPOINT_PROP, 0, endpoint);
    partitions.setValue(VALUE_TYPE_TABLE_PROP, 0, "literal");
    partitions.setValue(KEY_PREFIX_TABLE_PROP, 0, null);
    partitions.setValue(ZSET_KEYS_TABLE_PROP, 0, prefixes);
    partitions.setValue(REDIS_SSL_FLAG, 0, null);
    partitions.setValue(REDIS_CLUSTER_FLAG, 0, null);
    partitions.setValue(REDIS_DB_NUMBER, 0, null);
    partitions.setRowCount(1);
    String continuationToken = null;
    GetSplitsRequest originalReq = new GetSplitsRequest(IDENTITY, QUERY_ID, DEFAULT_CATALOG, TABLE_NAME, partitions, partitionCols, new Constraints(new HashMap<>()), null);
    GetSplitsRequest req = new GetSplitsRequest(originalReq, continuationToken);
    logger.info("doGetSplitsPrefix: req[{}]", req);
    MetadataResponse rawResponse = handler.doGetSplits(allocator, req);
    assertEquals(MetadataRequestType.GET_SPLITS, rawResponse.getRequestType());
    GetSplitsResponse response = (GetSplitsResponse) rawResponse;
    continuationToken = response.getContinuationToken();
    logger.info("doGetSplitsPrefix: continuationToken[{}] - numSplits[{}]", new Object[] { continuationToken, response.getSplits().size() });
    assertEquals("Continuation criteria violated", 120, response.getSplits().size());
    assertTrue("Continuation criteria violated", response.getContinuationToken() == null);
    verify(mockSyncCommands, times(6)).scan(any(ScanCursor.class), any(ScanArgs.class));
}

Example 13 with GetSplitsRequest

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

the class RedisMetadataHandlerTest method doGetSplitsPrefix.

@Test
public void doGetSplitsPrefix() {
    Schema schema = SchemaBuilder.newBuilder().addField("partitionId", Types.MinorType.INT.getType()).addStringField(REDIS_ENDPOINT_PROP).addStringField(VALUE_TYPE_TABLE_PROP).addStringField(KEY_PREFIX_TABLE_PROP).addStringField(ZSET_KEYS_TABLE_PROP).addStringField(REDIS_SSL_FLAG).addStringField(REDIS_CLUSTER_FLAG).addStringField(REDIS_DB_NUMBER).build();
    Block partitions = allocator.createBlock(schema);
    partitions.setValue(REDIS_ENDPOINT_PROP, 0, endpoint);
    partitions.setValue(VALUE_TYPE_TABLE_PROP, 0, "literal");
    partitions.setValue(KEY_PREFIX_TABLE_PROP, 0, "prefix1-*,prefix2-*, prefix3-*");
    partitions.setValue(ZSET_KEYS_TABLE_PROP, 0, null);
    partitions.setValue(REDIS_SSL_FLAG, 0, null);
    partitions.setValue(REDIS_CLUSTER_FLAG, 0, null);
    partitions.setValue(REDIS_DB_NUMBER, 0, null);
    partitions.setRowCount(1);
    String continuationToken = null;
    GetSplitsRequest originalReq = new GetSplitsRequest(IDENTITY, QUERY_ID, DEFAULT_CATALOG, TABLE_NAME, partitions, new ArrayList<>(), new Constraints(new HashMap<>()), null);
    GetSplitsRequest req = new GetSplitsRequest(originalReq, continuationToken);
    logger.info("doGetSplitsPrefix: req[{}]", req);
    MetadataResponse rawResponse = handler.doGetSplits(allocator, req);
    assertEquals(MetadataRequestType.GET_SPLITS, rawResponse.getRequestType());
    GetSplitsResponse response = (GetSplitsResponse) rawResponse;
    continuationToken = response.getContinuationToken();
    logger.info("doGetSplitsPrefix: continuationToken[{}] - numSplits[{}]", new Object[] { continuationToken, response.getSplits().size() });
    assertTrue("Continuation criteria violated", response.getSplits().size() == 3);
    assertTrue("Continuation criteria violated", response.getContinuationToken() == null);
}

Example 14 with GetSplitsRequest

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

the class RedshiftMetadataHandlerTest 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.redshiftMetadataHandler.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 preparedStatement = Mockito.mock(PreparedStatement.class);
    Mockito.when(this.connection.prepareStatement(PostGreSqlMetadataHandler.GET_PARTITIONS_QUERY)).thenReturn(preparedStatement);
    String[] columns = { "child_schema", "child" };
    int[] types = { Types.VARCHAR, Types.VARCHAR };
    Object[][] values = { { "s0", "p0" }, { "s1", "p1" } };
    ResultSet resultSet = mockResultSet(columns, types, values, new AtomicInteger(-1));
    final String expectedQuery = String.format(PostGreSqlMetadataHandler.GET_PARTITIONS_QUERY, tableName.getTableName(), tableName.getSchemaName());
    Mockito.when(preparedStatement.executeQuery()).thenReturn(resultSet);
    Mockito.when(this.connection.getMetaData().getSearchStringEscape()).thenReturn(null);
    GetTableLayoutResponse getTableLayoutResponse = this.redshiftMetadataHandler.doGetTableLayout(blockAllocator, getTableLayoutRequest);
    BlockAllocator splitBlockAllocator = new BlockAllocatorImpl();
    GetSplitsRequest getSplitsRequest = new GetSplitsRequest(this.federatedIdentity, "testQueryId", "testCatalogName", tableName, getTableLayoutResponse.getPartitions(), new ArrayList<>(partitionCols), constraints, "1");
    GetSplitsResponse getSplitsResponse = this.redshiftMetadataHandler.doGetSplits(splitBlockAllocator, getSplitsRequest);
    Set<Map<String, String>> expectedSplits = new HashSet<>();
    expectedSplits.add(ImmutableMap.of("partition_schema_name", "s1", "partition_name", "p1"));
    Assert.assertEquals(expectedSplits.size(), getSplitsResponse.getSplits().size());
    Set<Map<String, String>> actualSplits = getSplitsResponse.getSplits().stream().map(Split::getProperties).collect(Collectors.toSet());
    Assert.assertEquals(expectedSplits, actualSplits);
}

Example 15 with GetSplitsRequest

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

the class SynapseMetadataHandlerTest method doGetSplitsWithNoPartition.

@Test
public void doGetSplitsWithNoPartition() throws Exception {
    BlockAllocator blockAllocator = new BlockAllocatorImpl();
    Constraints constraints = Mockito.mock(Constraints.class);
    TableName tableName = new TableName("testSchema", "testTable");
    Object[][] values = { {} };
    ResultSet resultSet = mockResultSet(new String[] { "ROW_COUNT" }, new int[] { Types.INTEGER }, values, new AtomicInteger(-1));
    Statement st = Mockito.mock(Statement.class);
    Mockito.when(this.connection.createStatement()).thenReturn(st);
    Mockito.when(st.executeQuery(Mockito.anyString())).thenReturn(resultSet);
    Schema partitionSchema = this.synapseMetadataHandler.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);
    GetTableLayoutResponse getTableLayoutResponse = this.synapseMetadataHandler.doGetTableLayout(blockAllocator, getTableLayoutRequest);
    BlockAllocator splitBlockAllocator = new BlockAllocatorImpl();
    GetSplitsRequest getSplitsRequest = new GetSplitsRequest(this.federatedIdentity, "testQueryId", "testCatalogName", tableName, getTableLayoutResponse.getPartitions(), new ArrayList<>(partitionCols), constraints, null);
    GetSplitsResponse getSplitsResponse = this.synapseMetadataHandler.doGetSplits(splitBlockAllocator, getSplitsRequest);
    Set<Map<String, String>> expectedSplits = new HashSet<>();
    expectedSplits.add(Collections.singletonMap(SynapseMetadataHandler.PARTITION_NUMBER, "0"));
    Assert.assertEquals(expectedSplits.size(), getSplitsResponse.getSplits().size());
    Set<Map<String, String>> actualSplits = getSplitsResponse.getSplits().stream().map(Split::getProperties).collect(Collectors.toSet());
    Assert.assertEquals(expectedSplits, actualSplits);
}