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

Example 26 with GetTableLayoutRequest

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

the class SqlServerMuxMetadataHandlerTest method doGetTableLayout.

@Test
public void doGetTableLayout() throws Exception {
    GetTableLayoutRequest getTableLayoutRequest = Mockito.mock(GetTableLayoutRequest.class);
    Mockito.when(getTableLayoutRequest.getTableName()).thenReturn(new TableName("testSchema", "testTable"));
    Mockito.when(getTableLayoutRequest.getCatalogName()).thenReturn("fakedatabase");
    this.jdbcMetadataHandler.doGetTableLayout(this.allocator, getTableLayoutRequest);
    Mockito.verify(this.sqlServerMetadataHandler, Mockito.times(1)).doGetTableLayout(Mockito.eq(this.allocator), Mockito.eq(getTableLayoutRequest));
}

Example 27 with GetTableLayoutRequest

use of com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest 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");
}

Example 28 with GetTableLayoutRequest

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

the class AwsCmdbMetadataHandlerTest method doGetTableLayout.

@Test
public void doGetTableLayout() throws Exception {
    GetTableLayoutRequest request = new GetTableLayoutRequest(identity, queryId, catalog, new TableName("schema1", "table1"), mockConstraints, SchemaBuilder.newBuilder().build(), Collections.EMPTY_SET);
    GetTableLayoutResponse response = handler.doGetTableLayout(blockAllocator, request);
    assertNotNull(response);
    assertEquals(1, response.getPartitions().getRowCount());
}

Example 29 with GetTableLayoutRequest

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

the class CloudwatchMetadataHandlerTest method doGetTableLayout.

@Test
public void doGetTableLayout() throws Exception {
    logger.info("doGetTableLayout - enter");
    when(mockAwsLogs.describeLogStreams(any(DescribeLogStreamsRequest.class))).thenAnswer((InvocationOnMock invocationOnMock) -> {
        DescribeLogStreamsRequest request = (DescribeLogStreamsRequest) invocationOnMock.getArguments()[0];
        DescribeLogStreamsResult result = new DescribeLogStreamsResult();
        Integer nextToken;
        if (request.getNextToken() == null) {
            nextToken = 1;
        } else if (Integer.valueOf(request.getNextToken()) < 3) {
            nextToken = Integer.valueOf(request.getNextToken()) + 1;
        } else {
            nextToken = null;
        }
        List<LogStream> logStreams = new ArrayList<>();
        if (request.getNextToken() == null || Integer.valueOf(request.getNextToken()) < 3) {
            int continuation = request.getNextToken() == null ? 0 : Integer.valueOf(request.getNextToken());
            for (int i = 0 + continuation * 100; i < 300; i++) {
                LogStream nextLogStream = new LogStream();
                nextLogStream.setLogStreamName("table-" + String.valueOf(i));
                nextLogStream.setStoredBytes(i * 1000L);
                logStreams.add(nextLogStream);
            }
        }
        result.withLogStreams(logStreams);
        if (nextToken != null) {
            result.setNextToken(String.valueOf(nextToken));
        }
        return result;
    });
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put("log_stream", EquatableValueSet.newBuilder(allocator, Types.MinorType.VARCHAR.getType(), true, false).add("table-10").build());
    Schema schema = SchemaBuilder.newBuilder().addStringField("log_stream").build();
    GetTableLayoutRequest req = new GetTableLayoutRequest(identity, "queryId", "default", new TableName("schema-1", "all_log_streams"), new Constraints(constraintsMap), schema, Collections.singleton("log_stream"));
    GetTableLayoutResponse res = handler.doGetTableLayout(allocator, req);
    logger.info("doGetTableLayout - {}", res.getPartitions().getSchema());
    logger.info("doGetTableLayout - {}", res.getPartitions());
    assertTrue(res.getPartitions().getSchema().findField("log_stream") != null);
    assertTrue(res.getPartitions().getRowCount() == 1);
    verify(mockAwsLogs, times(4)).describeLogStreams(any(DescribeLogStreamsRequest.class));
    logger.info("doGetTableLayout - exit");
}

Example 30 with GetTableLayoutRequest

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

the class MetricsMetadataHandlerTest method doGetTableLayout.

@Test
public void doGetTableLayout() throws Exception {
    logger.info("doGetTableLayout - enter");
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put(METRIC_NAME_FIELD, EquatableValueSet.newBuilder(allocator, Types.MinorType.VARCHAR.getType(), true, false).add("MyMetric").build());
    GetTableLayoutRequest req = new GetTableLayoutRequest(identity, "queryId", "default", new TableName(defaultSchema, "metrics"), new Constraints(constraintsMap), SchemaBuilder.newBuilder().build(), Collections.EMPTY_SET);
    GetTableLayoutResponse res = handler.doGetTableLayout(allocator, req);
    logger.info("doGetTableLayout - {}", res.getPartitions().getSchema());
    logger.info("doGetTableLayout - {}", res.getPartitions());
    assertTrue(res.getPartitions().getRowCount() == 1);
    logger.info("doGetTableLayout - exit");
}