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

Example 16 with GetTableRequest

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

the class SqlServerMetadataHandlerTest method doGetTable.

@Test
public void doGetTable() throws Exception {
    BlockAllocator blockAllocator = new BlockAllocatorImpl();
    String[] schema = { "DATA_TYPE", "COLUMN_SIZE", "COLUMN_NAME", "DECIMAL_DIGITS", "NUM_PREC_RADIX" };
    Object[][] values = { { Types.INTEGER, 12, "testCol1", 0, 0 }, { Types.VARCHAR, 25, "testCol2", 0, 0 }, { Types.TIMESTAMP, 93, "testCol3", 0, 0 }, { Types.TIMESTAMP_WITH_TIMEZONE, 93, "testCol4", 0, 0 } };
    AtomicInteger rowNumber = new AtomicInteger(-1);
    ResultSet resultSet = mockResultSet(schema, values, rowNumber);
    SchemaBuilder expectedSchemaBuilder = SchemaBuilder.newBuilder();
    expectedSchemaBuilder.addField(FieldBuilder.newBuilder("testCol1", org.apache.arrow.vector.types.Types.MinorType.INT.getType()).build());
    expectedSchemaBuilder.addField(FieldBuilder.newBuilder("testCol2", org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build());
    expectedSchemaBuilder.addField(FieldBuilder.newBuilder("testCol3", org.apache.arrow.vector.types.Types.MinorType.DATEMILLI.getType()).build());
    expectedSchemaBuilder.addField(FieldBuilder.newBuilder("testCol4", org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build());
    PARTITION_SCHEMA.getFields().forEach(expectedSchemaBuilder::addField);
    Schema expected = expectedSchemaBuilder.build();
    TableName inputTableName = new TableName("TESTSCHEMA", "TESTTABLE");
    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));
    Assert.assertEquals(expected, getTableResponse.getSchema());
    Assert.assertEquals(inputTableName, getTableResponse.getTableName());
    Assert.assertEquals("testCatalog", getTableResponse.getCatalogName());
}

Example 17 with GetTableRequest

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

the class TimestreamMetadataHandlerTest method doGetTableGlue.

@Test
public void doGetTableGlue() throws Exception {
    logger.info("doGetTable - enter");
    when(mockGlue.getTable(any(com.amazonaws.services.glue.model.GetTableRequest.class))).thenAnswer((InvocationOnMock invocation) -> {
        com.amazonaws.services.glue.model.GetTableRequest request = invocation.getArgumentAt(0, com.amazonaws.services.glue.model.GetTableRequest.class);
        List<Column> columns = new ArrayList<>();
        columns.add(new Column().withName("col1").withType("varchar"));
        columns.add(new Column().withName("col2").withType("double"));
        com.amazonaws.services.glue.model.Table table = new com.amazonaws.services.glue.model.Table();
        table.setName(request.getName());
        table.setDatabaseName(request.getDatabaseName());
        StorageDescriptor storageDescriptor = new StorageDescriptor();
        storageDescriptor.setColumns(columns);
        table.setStorageDescriptor(storageDescriptor);
        table.setViewOriginalText("view text");
        table.setParameters(Collections.singletonMap("timestream-metadata-flag", "timestream-metadata-flag"));
        return new GetTableResult().withTable(table);
    });
    GetTableRequest req = new GetTableRequest(identity, "query-id", "default", new TableName(defaultSchema, "table1"));
    GetTableResponse res = handler.doGetTable(allocator, req);
    logger.info("doGetTable - {}", res);
    assertEquals(2, res.getSchema().getFields().size());
    Field measureName = res.getSchema().findField("col1");
    assertEquals(Types.MinorType.VARCHAR, Types.getMinorTypeForArrowType(measureName.getType()));
    Field measureValue = res.getSchema().findField("col2");
    assertEquals(Types.MinorType.FLOAT8, Types.getMinorTypeForArrowType(measureValue.getType()));
    assertEquals("view text", res.getSchema().getCustomMetadata().get(VIEW_METADATA_FIELD));
    logger.info("doGetTable - exit");
}

Example 18 with GetTableRequest

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

the class TimestreamMetadataHandlerTest method doGetTimeSeriesTableGlue.

@Test
public void doGetTimeSeriesTableGlue() throws Exception {
    logger.info("doGetTimeSeriesTableGlue - enter");
    when(mockGlue.getTable(any(com.amazonaws.services.glue.model.GetTableRequest.class))).thenAnswer((InvocationOnMock invocation) -> {
        com.amazonaws.services.glue.model.GetTableRequest request = invocation.getArgumentAt(0, com.amazonaws.services.glue.model.GetTableRequest.class);
        List<Column> columns = new ArrayList<>();
        columns.add(new Column().withName("az").withType("varchar"));
        columns.add(new Column().withName("hostname").withType("varchar"));
        columns.add(new Column().withName("region").withType("varchar"));
        columns.add(new Column().withName("cpu_utilization").withType("ARRAY<STRUCT<time: timestamp, measure_value\\:\\:double: double>>"));
        com.amazonaws.services.glue.model.Table table = new com.amazonaws.services.glue.model.Table();
        table.setName(request.getName());
        table.setDatabaseName(request.getDatabaseName());
        StorageDescriptor storageDescriptor = new StorageDescriptor();
        storageDescriptor.setColumns(columns);
        table.setStorageDescriptor(storageDescriptor);
        table.setViewOriginalText("SELECT az, hostname, region, cpu_utilization FROM TIMESERIES(metrics_table,'cpu_utilization')");
        table.setParameters(Collections.singletonMap("timestream-metadata-flag", "timestream-metadata-flag"));
        return new GetTableResult().withTable(table);
    });
    GetTableRequest req = new GetTableRequest(identity, "query-id", "default", new TableName(defaultSchema, "table1"));
    GetTableResponse res = handler.doGetTable(allocator, req);
    logger.info("doGetTable - {}", res);
    assertEquals(4, res.getSchema().getFields().size());
    Field measureName = res.getSchema().findField("az");
    assertEquals(Types.MinorType.VARCHAR, Types.getMinorTypeForArrowType(measureName.getType()));
    Field hostname = res.getSchema().findField("hostname");
    assertEquals(Types.MinorType.VARCHAR, Types.getMinorTypeForArrowType(hostname.getType()));
    Field region = res.getSchema().findField("region");
    assertEquals(Types.MinorType.VARCHAR, Types.getMinorTypeForArrowType(region.getType()));
    Field cpuUtilization = res.getSchema().findField("cpu_utilization");
    assertEquals(Types.MinorType.LIST, Types.getMinorTypeForArrowType(cpuUtilization.getType()));
    Field timeseries = cpuUtilization.getChildren().get(0);
    assertEquals(Types.MinorType.STRUCT, Types.getMinorTypeForArrowType(timeseries.getType()));
    Field time = timeseries.getChildren().get(0);
    assertEquals(Types.MinorType.DATEMILLI, Types.getMinorTypeForArrowType(time.getType()));
    Field value = timeseries.getChildren().get(1);
    assertEquals(Types.MinorType.FLOAT8, Types.getMinorTypeForArrowType(value.getType()));
    assertEquals("SELECT az, hostname, region, cpu_utilization FROM TIMESERIES(metrics_table,'cpu_utilization')", res.getSchema().getCustomMetadata().get(VIEW_METADATA_FIELD));
    logger.info("doGetTimeSeriesTableGlue - exit");
}

Example 19 with GetTableRequest

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

the class TimestreamMetadataHandlerTest method doGetTable.

@Test
public void doGetTable() throws Exception {
    logger.info("doGetTable - enter");
    when(mockGlue.getTable(any(com.amazonaws.services.glue.model.GetTableRequest.class))).thenReturn(mock(GetTableResult.class));
    when(mockTsQuery.query(any(QueryRequest.class))).thenAnswer((InvocationOnMock invocation) -> {
        QueryRequest request = invocation.getArgumentAt(0, QueryRequest.class);
        assertEquals("DESCRIBE \"default\".\"table1\"", request.getQueryString());
        List<Row> rows = new ArrayList<>();
        // TODO: Add types here
        rows.add(new Row().withData(new Datum().withScalarValue("availability_zone"), new Datum().withScalarValue("varchar"), new Datum().withScalarValue("dimension")));
        rows.add(new Row().withData(new Datum().withScalarValue("measure_value"), new Datum().withScalarValue("double"), new Datum().withScalarValue("measure_value")));
        rows.add(new Row().withData(new Datum().withScalarValue("measure_name"), new Datum().withScalarValue("varchar"), new Datum().withScalarValue("measure_name")));
        rows.add(new Row().withData(new Datum().withScalarValue("time"), new Datum().withScalarValue("timestamp"), new Datum().withScalarValue("timestamp")));
        return new QueryResult().withRows(rows);
    });
    GetTableRequest req = new GetTableRequest(identity, "query-id", "default", new TableName(defaultSchema, "table1"));
    GetTableResponse res = handler.doGetTable(allocator, req);
    logger.info("doGetTable - {}", res);
    assertEquals(4, res.getSchema().getFields().size());
    Field measureName = res.getSchema().findField("measure_name");
    assertEquals(Types.MinorType.VARCHAR, Types.getMinorTypeForArrowType(measureName.getType()));
    Field measureValue = res.getSchema().findField("measure_value");
    assertEquals(Types.MinorType.FLOAT8, Types.getMinorTypeForArrowType(measureValue.getType()));
    Field availabilityZone = res.getSchema().findField("availability_zone");
    assertEquals(Types.MinorType.VARCHAR, Types.getMinorTypeForArrowType(availabilityZone.getType()));
    Field time = res.getSchema().findField("time");
    assertEquals(Types.MinorType.DATEMILLI, Types.getMinorTypeForArrowType(time.getType()));
    logger.info("doGetTable - exit");
}

Example 20 with GetTableRequest

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

the class AbstractTableProviderTest method readTableTest.

@Test
public void readTableTest() {
    GetTableRequest request = new GetTableRequest(identity, expectedQuery, expectedCatalog, expectedTableName);
    GetTableResponse response = provider.getTable(allocator, request);
    assertTrue(response.getSchema().getFields().size() > 1);
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put(idField, EquatableValueSet.newBuilder(allocator, Types.MinorType.VARCHAR.getType(), true, false).add(idValue).build());
    Constraints constraints = new Constraints(constraintsMap);
    ConstraintEvaluator evaluator = new ConstraintEvaluator(allocator, response.getSchema(), constraints);
    S3SpillLocation spillLocation = S3SpillLocation.newBuilder().withBucket("bucket").withPrefix("prefix").withSplitId(UUID.randomUUID().toString()).withQueryId(UUID.randomUUID().toString()).withIsDirectory(true).build();
    ReadRecordsRequest readRequest = new ReadRecordsRequest(identity, expectedCatalog, "queryId", expectedTableName, response.getSchema(), Split.newBuilder(spillLocation, keyFactory.create()).build(), constraints, 100_000_000, 100_000_000);
    SpillConfig spillConfig = SpillConfig.newBuilder().withSpillLocation(spillLocation).withMaxBlockBytes(3_000_000).withMaxInlineBlockBytes(0).withRequestId("queryid").withEncryptionKey(keyFactory.create()).build();
    setUpRead();
    BlockSpiller spiller = new S3BlockSpiller(amazonS3, spillConfig, allocator, response.getSchema(), evaluator);
    provider.readWithConstraint(spiller, readRequest, queryStatusChecker);
    validateRead(response.getSchema(), blockSpillReader, spiller.getSpillLocations(), spillConfig.getEncryptionKey());
}