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

Example 6 with GetTableRequest

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

the class JdbcMetadataHandlerTest method doGetTable.

@Test
public void doGetTable() throws SQLException {
    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.jdbcMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName));
    Assert.assertEquals(expected, getTableResponse.getSchema());
    Assert.assertEquals(inputTableName, getTableResponse.getTableName());
    Assert.assertEquals("testCatalog", getTableResponse.getCatalogName());
}

Example 7 with GetTableRequest

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

the class JdbcMetadataHandlerTest method doGetTableSQLException.

@Test(expected = RuntimeException.class)
public void doGetTableSQLException() throws SQLException {
    TableName inputTableName = new TableName("testSchema", "testTable");
    Mockito.when(this.connection.getMetaData().getColumns(Mockito.anyString(), Mockito.anyString(), Mockito.anyString(), Mockito.anyString())).thenThrow(new SQLException());
    this.jdbcMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName));
}

Example 8 with GetTableRequest

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

the class JdbcMetadataHandlerTest method doGetTableNoColumns.

@Test(expected = RuntimeException.class)
public void doGetTableNoColumns() {
    TableName inputTableName = new TableName("testSchema", "testTable");
    this.jdbcMetadataHandler.doGetTable(this.blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName));
}

Example 9 with GetTableRequest

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

the class MultiplexingJdbcMetadataHandlerTest method doGetTable.

@Test
public void doGetTable() {
    GetTableRequest getTableRequest = Mockito.mock(GetTableRequest.class);
    Mockito.when(getTableRequest.getCatalogName()).thenReturn("fakedatabase");
    this.jdbcMetadataHandler.doGetTable(this.allocator, getTableRequest);
    Mockito.verify(this.fakeDatabaseHandler, Mockito.times(1)).doGetTable(Mockito.eq(this.allocator), Mockito.eq(getTableRequest));
}

Example 10 with GetTableRequest

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

the class PostGreSqlMetadataHandlerTest method doGetTableWithArrayColumns.

@Test
public void doGetTableWithArrayColumns() throws Exception {
    logger.info("doGetTableWithArrayColumns - enter");
    String[] schema = { "DATA_TYPE", "COLUMN_NAME", "COLUMN_SIZE", "DECIMAL_DIGITS", "TYPE_NAME" };
    Object[][] values = { { Types.ARRAY, "bool_array", 0, 0, "_bool" }, { Types.ARRAY, "smallint_array", 0, 0, "_int2" }, { Types.ARRAY, "int_array", 0, 0, "_int4" }, { Types.ARRAY, "bigint_array", 0, 0, "_int8" }, { Types.ARRAY, "float_array", 0, 0, "_float4" }, { Types.ARRAY, "double_array", 0, 0, "_float8" }, { Types.ARRAY, "date_array", 0, 0, "_date" }, { Types.ARRAY, "timestamp_array", 0, 0, "_timestamp" }, { Types.ARRAY, "binary_array", 0, 0, "_bytea" }, { Types.ARRAY, "decimal_array", 38, 2, "_numeric" }, { Types.ARRAY, "string_array", 0, 0, "_text" }, { Types.ARRAY, "uuid_array", 0, 0, "_uuid" } };
    AtomicInteger rowNumber = new AtomicInteger(-1);
    ResultSet resultSet = mockResultSet(schema, values, rowNumber);
    SchemaBuilder expectedSchemaBuilder = SchemaBuilder.newBuilder();
    expectedSchemaBuilder.addListField("bool_array", new ArrowType.Bool()).addListField("smallint_array", new ArrowType.Int(16, true)).addListField("int_array", new ArrowType.Int(32, true)).addListField("bigint_array", new ArrowType.Int(64, true)).addListField("float_array", new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)).addListField("double_array", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)).addListField("date_array", new ArrowType.Date(DateUnit.DAY)).addListField("timestamp_array", new ArrowType.Date(DateUnit.MILLISECOND)).addListField("binary_array", new ArrowType.Utf8()).addListField("decimal_array", new ArrowType.Decimal(38, 2)).addListField("string_array", new ArrowType.Utf8()).addListField("uuid_array", new ArrowType.Utf8());
    postGreSqlMetadataHandler.getPartitionSchema("testCatalog").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.postGreSqlMetadataHandler.doGetTable(new BlockAllocatorImpl(), new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName));
    logger.info("Schema: {}", getTableResponse.getSchema());
    Assert.assertEquals(expected, getTableResponse.getSchema());
    Assert.assertEquals(inputTableName, getTableResponse.getTableName());
    Assert.assertEquals("testCatalog", getTableResponse.getCatalogName());
    logger.info("doGetTableWithArrayColumns - exit");
}