Search in sources :

Example 1 with Table

use of com.amazonaws.services.timestreamwrite.model.Table 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");
}
Also used : Table(com.amazonaws.services.timestreamwrite.model.Table) ArrayList(java.util.ArrayList) StorageDescriptor(com.amazonaws.services.glue.model.StorageDescriptor) GetTableRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableRequest) TableName(com.amazonaws.athena.connector.lambda.domain.TableName) Field(org.apache.arrow.vector.types.pojo.Field) Column(com.amazonaws.services.glue.model.Column) GetTableResponse(com.amazonaws.athena.connector.lambda.metadata.GetTableResponse) InvocationOnMock(org.mockito.invocation.InvocationOnMock) GetTableResult(com.amazonaws.services.glue.model.GetTableResult) Test(org.junit.Test)

Example 2 with Table

use of com.amazonaws.services.timestreamwrite.model.Table 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");
}
Also used : Table(com.amazonaws.services.timestreamwrite.model.Table) ArrayList(java.util.ArrayList) StorageDescriptor(com.amazonaws.services.glue.model.StorageDescriptor) GetTableRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableRequest) TableName(com.amazonaws.athena.connector.lambda.domain.TableName) Field(org.apache.arrow.vector.types.pojo.Field) Column(com.amazonaws.services.glue.model.Column) GetTableResponse(com.amazonaws.athena.connector.lambda.metadata.GetTableResponse) InvocationOnMock(org.mockito.invocation.InvocationOnMock) GetTableResult(com.amazonaws.services.glue.model.GetTableResult) Test(org.junit.Test)

Example 3 with Table

use of com.amazonaws.services.timestreamwrite.model.Table in project aws-athena-query-federation by awslabs.

the class TimestreamMetadataHandlerTest method doListTables.

@Test
public void doListTables() throws Exception {
    logger.info("doListTables - enter");
    when(mockTsMeta.listTables(any(com.amazonaws.services.timestreamwrite.model.ListTablesRequest.class))).thenAnswer((InvocationOnMock invocation) -> {
        com.amazonaws.services.timestreamwrite.model.ListTablesRequest request = invocation.getArgumentAt(0, com.amazonaws.services.timestreamwrite.model.ListTablesRequest.class);
        String newNextToken = null;
        List<Table> tables = new ArrayList<>();
        if (request.getNextToken() == null) {
            for (int i = 0; i < 10; i++) {
                tables.add(new Table().withDatabaseName(request.getDatabaseName()).withTableName("table_" + i));
            }
            newNextToken = "1";
        } else if (request.getNextToken().equals("1")) {
            for (int i = 10; i < 100; i++) {
                tables.add(new Table().withDatabaseName(request.getDatabaseName()).withTableName("table_" + i));
            }
            newNextToken = "2";
        } else if (request.getNextToken().equals("2")) {
            for (int i = 100; i < 1000; i++) {
                tables.add(new Table().withDatabaseName(request.getDatabaseName()).withTableName("table_" + i));
            }
            newNextToken = null;
        }
        return new ListTablesResult().withTables(tables).withNextToken(newNextToken);
    });
    ListTablesRequest req = new ListTablesRequest(identity, "queryId", "default", defaultSchema, null, UNLIMITED_PAGE_SIZE_VALUE);
    ListTablesResponse res = handler.doListTables(allocator, req);
    logger.info("doListTables - {}", res.getTables());
    assertEquals(1000, res.getTables().size());
    verify(mockTsMeta, times(3)).listTables(any(com.amazonaws.services.timestreamwrite.model.ListTablesRequest.class));
    Iterator<TableName> schemaItr = res.getTables().iterator();
    for (int i = 0; i < 1000; i++) {
        TableName tableName = schemaItr.next();
        assertEquals(defaultSchema, tableName.getSchemaName());
        assertEquals("table_" + i, tableName.getTableName());
    }
    logger.info("doListTables - exit");
}
Also used : Table(com.amazonaws.services.timestreamwrite.model.Table) ArrayList(java.util.ArrayList) ListTablesResult(com.amazonaws.services.timestreamwrite.model.ListTablesResult) TableName(com.amazonaws.athena.connector.lambda.domain.TableName) InvocationOnMock(org.mockito.invocation.InvocationOnMock) ListTablesRequest(com.amazonaws.athena.connector.lambda.metadata.ListTablesRequest) ListTablesResponse(com.amazonaws.athena.connector.lambda.metadata.ListTablesResponse) Test(org.junit.Test)

Aggregations

TableName (com.amazonaws.athena.connector.lambda.domain.TableName)3 Table (com.amazonaws.services.timestreamwrite.model.Table)3 ArrayList (java.util.ArrayList)3 Test (org.junit.Test)3 InvocationOnMock (org.mockito.invocation.InvocationOnMock)3 GetTableRequest (com.amazonaws.athena.connector.lambda.metadata.GetTableRequest)2 GetTableResponse (com.amazonaws.athena.connector.lambda.metadata.GetTableResponse)2 Column (com.amazonaws.services.glue.model.Column)2 GetTableResult (com.amazonaws.services.glue.model.GetTableResult)2 StorageDescriptor (com.amazonaws.services.glue.model.StorageDescriptor)2 Field (org.apache.arrow.vector.types.pojo.Field)2 ListTablesRequest (com.amazonaws.athena.connector.lambda.metadata.ListTablesRequest)1 ListTablesResponse (com.amazonaws.athena.connector.lambda.metadata.ListTablesResponse)1 ListTablesResult (com.amazonaws.services.timestreamwrite.model.ListTablesResult)1