Examples with SchemaBuilder com.amazonaws.athena.connector.lambda.data.SchemaBuilder used on opensource projects

Example 51 with SchemaBuilder

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

the class RedshiftMetadataHandlerTest 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());
    redshiftMetadataHandler.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.redshiftMetadataHandler.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");
}

Example 52 with SchemaBuilder

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

the class RedshiftRecordHandlerTest method buildSplitSqlForDateTest.

@Test
public void buildSplitSqlForDateTest() throws SQLException {
    logger.info("buildSplitSqlForDateTest - enter");
    TableName tableName = new TableName("testSchema", "testTable");
    SchemaBuilder schemaBuilder = SchemaBuilder.newBuilder();
    schemaBuilder.addField(FieldBuilder.newBuilder("testDate", Types.MinorType.DATEDAY.getType()).build());
    schemaBuilder.addField(FieldBuilder.newBuilder("partition_schema_name", Types.MinorType.VARCHAR.getType()).build());
    schemaBuilder.addField(FieldBuilder.newBuilder("partition_name", Types.MinorType.VARCHAR.getType()).build());
    Schema schema = schemaBuilder.build();
    Split split = Mockito.mock(Split.class);
    Mockito.when(split.getProperties()).thenReturn(ImmutableMap.of("partition_schema_name", "s0", "partition_name", "p0"));
    Mockito.when(split.getProperty(Mockito.eq(PostGreSqlMetadataHandler.BLOCK_PARTITION_SCHEMA_COLUMN_NAME))).thenReturn("s0");
    Mockito.when(split.getProperty(Mockito.eq(PostGreSqlMetadataHandler.BLOCK_PARTITION_COLUMN_NAME))).thenReturn("p0");
    final long dateDays = TimeUnit.MILLISECONDS.toDays(Date.valueOf("2020-01-05").getTime());
    ValueSet valueSet = getSingleValueSet(dateDays);
    Constraints constraints = Mockito.mock(Constraints.class);
    Mockito.when(constraints.getSummary()).thenReturn(Collections.singletonMap("testDate", valueSet));
    String expectedSql = "SELECT \"testDate\" FROM \"s0\".\"p0\"  WHERE (\"testDate\" = ?)";
    PreparedStatement expectedPreparedStatement = Mockito.mock(PreparedStatement.class);
    Mockito.when(this.connection.prepareStatement(Mockito.eq(expectedSql))).thenReturn(expectedPreparedStatement);
    PreparedStatement preparedStatement = this.redshiftRecordHandler.buildSplitSql(this.connection, "testCatalogName", tableName, schema, constraints, split);
    Assert.assertEquals(expectedPreparedStatement, preparedStatement);
    Mockito.verify(preparedStatement, Mockito.times(1)).setDate(1, new Date(TimeUnit.DAYS.toMillis(dateDays)));
    logger.info("buildSplitSqlForDateTest - exit");
}

Example 53 with SchemaBuilder

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

the class SaphanaRecordHandlerTest method buildSplitSqlNew.

@Test
public void buildSplitSqlNew() throws SQLException {
    TableName tableName = new TableName("testSchema", "testTable");
    SchemaBuilder schemaBuilder = SchemaBuilder.newBuilder();
    schemaBuilder.addField(FieldBuilder.newBuilder("testCol1", Types.MinorType.INT.getType()).build());
    schemaBuilder.addField(FieldBuilder.newBuilder("testCol2", Types.MinorType.VARCHAR.getType()).build());
    schemaBuilder.addField(FieldBuilder.newBuilder("testCol3", Types.MinorType.BIGINT.getType()).build());
    schemaBuilder.addField(FieldBuilder.newBuilder("testCol4", Types.MinorType.FLOAT4.getType()).build());
    schemaBuilder.addField(FieldBuilder.newBuilder("testCol5", Types.MinorType.SMALLINT.getType()).build());
    schemaBuilder.addField(FieldBuilder.newBuilder("testCol6", Types.MinorType.TINYINT.getType()).build());
    schemaBuilder.addField(FieldBuilder.newBuilder("testCol7", Types.MinorType.FLOAT8.getType()).build());
    schemaBuilder.addField(FieldBuilder.newBuilder("testCol8", Types.MinorType.BIT.getType()).build());
    schemaBuilder.addField(FieldBuilder.newBuilder("partition", Types.MinorType.VARCHAR.getType()).build());
    Schema schema = schemaBuilder.build();
    Split split = Mockito.mock(Split.class);
    Mockito.when(split.getProperties()).thenReturn(Collections.singletonMap("partition", "p0"));
    Mockito.when(split.getProperty(Mockito.eq("partition"))).thenReturn("p0");
    Range range1a = Mockito.mock(Range.class, Mockito.RETURNS_DEEP_STUBS);
    Mockito.when(range1a.isSingleValue()).thenReturn(true);
    Mockito.when(range1a.getLow().getValue()).thenReturn(1);
    Range range1b = Mockito.mock(Range.class, Mockito.RETURNS_DEEP_STUBS);
    Mockito.when(range1b.isSingleValue()).thenReturn(true);
    Mockito.when(range1b.getLow().getValue()).thenReturn(2);
    ValueSet valueSet1 = Mockito.mock(SortedRangeSet.class, Mockito.RETURNS_DEEP_STUBS);
    Mockito.when(valueSet1.getRanges().getOrderedRanges()).thenReturn(ImmutableList.of(range1a, range1b));
    ValueSet valueSet2 = getRangeSet(Marker.Bound.EXACTLY, "1", Marker.Bound.BELOW, "10");
    ValueSet valueSet3 = getRangeSet(Marker.Bound.ABOVE, 2L, Marker.Bound.EXACTLY, 20L);
    ValueSet valueSet4 = getSingleValueSet(1.1F);
    ValueSet valueSet5 = getSingleValueSet(1);
    ValueSet valueSet6 = getSingleValueSet(0);
    ValueSet valueSet7 = getSingleValueSet(1.2d);
    ValueSet valueSet8 = getSingleValueSet(true);
    Constraints constraints = Mockito.mock(Constraints.class);
    Mockito.when(constraints.getSummary()).thenReturn(new ImmutableMap.Builder<String, ValueSet>().put("testCol1", valueSet1).put("testCol2", valueSet2).put("testCol3", valueSet3).put("testCol4", valueSet4).put("testCol5", valueSet5).put("testCol6", valueSet6).put("testCol7", valueSet7).put("testCol8", valueSet8).build());
    PreparedStatement expectedPreparedStatement = Mockito.mock(PreparedStatement.class);
    Mockito.when(this.connection.prepareStatement(Mockito.anyString())).thenReturn(expectedPreparedStatement);
    PreparedStatement preparedStatement = this.saphanaRecordHandler.buildSplitSql(this.connection, "testCatalogName", tableName, schema, constraints, split);
    Assert.assertEquals(expectedPreparedStatement, preparedStatement);
    Mockito.verify(preparedStatement, Mockito.times(1)).setInt(1, 1);
    Mockito.verify(preparedStatement, Mockito.times(1)).setInt(2, 2);
    Mockito.verify(preparedStatement, Mockito.times(1)).setString(3, "1");
    Mockito.verify(preparedStatement, Mockito.times(1)).setString(4, "10");
    Mockito.verify(preparedStatement, Mockito.times(1)).setLong(5, 2L);
    Mockito.verify(preparedStatement, Mockito.times(1)).setLong(6, 20L);
    Mockito.verify(preparedStatement, Mockito.times(1)).setFloat(7, 1.1F);
    Mockito.verify(preparedStatement, Mockito.times(1)).setShort(8, (short) 1);
    Mockito.verify(preparedStatement, Mockito.times(1)).setByte(9, (byte) 0);
    Mockito.verify(preparedStatement, Mockito.times(1)).setDouble(10, 1.2d);
    Mockito.verify(preparedStatement, Mockito.times(1)).setBoolean(11, true);
}

Example 54 with SchemaBuilder

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

the class SynapseMetadataHandlerTest method doGetTableLayoutWithNoPartitions.

@Test
public void doGetTableLayoutWithNoPartitions() throws Exception {
    BlockAllocator blockAllocator = new BlockAllocatorImpl();
    Constraints constraints = Mockito.mock(Constraints.class);
    TableName tableName = new TableName("testSchema", "testTable");
    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);
    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);
    GetTableLayoutResponse getTableLayoutResponse = this.synapseMetadataHandler.doGetTableLayout(blockAllocator, getTableLayoutRequest);
    Assert.assertEquals(values.length, getTableLayoutResponse.getPartitions().getRowCount());
    List<String> actualValues = new ArrayList<>();
    for (int i = 0; i < getTableLayoutResponse.getPartitions().getRowCount(); i++) {
        actualValues.add(BlockUtils.rowToString(getTableLayoutResponse.getPartitions(), i));
    }
    Assert.assertEquals(Collections.singletonList("[PARTITION_NUMBER : 0]"), actualValues);
    SchemaBuilder expectedSchemaBuilder = SchemaBuilder.newBuilder();
    expectedSchemaBuilder.addField(FieldBuilder.newBuilder(SynapseMetadataHandler.PARTITION_NUMBER, org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build());
    Schema expectedSchema = expectedSchemaBuilder.build();
    Assert.assertEquals(expectedSchema, getTableLayoutResponse.getPartitions().getSchema());
    Assert.assertEquals(tableName, getTableLayoutResponse.getTableName());
}

Example 55 with SchemaBuilder

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

the class SynapseMetadataHandlerTest method doGetTableLayout.

@Test
public void doGetTableLayout() throws Exception {
    BlockAllocator blockAllocator = new BlockAllocatorImpl();
    Constraints constraints = Mockito.mock(Constraints.class);
    TableName tableName = new TableName("testSchema", "testTable");
    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);
    String[] columns = { "ROW_COUNT", SynapseMetadataHandler.PARTITION_NUMBER, SynapseMetadataHandler.PARTITION_COLUMN, "PARTITION_BOUNDARY_VALUE" };
    int[] types = { Types.INTEGER, Types.VARCHAR, Types.VARCHAR, Types.VARCHAR };
    Object[][] values = { { 2, null, null, null }, { 0, "1", "id", "100000" }, { 0, "2", "id", "300000" } };
    ResultSet resultSet = mockResultSet(columns, types, 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);
    GetTableLayoutResponse getTableLayoutResponse = this.synapseMetadataHandler.doGetTableLayout(blockAllocator, getTableLayoutRequest);
    List<String> actualValues = new ArrayList<>();
    for (int i = 0; i < getTableLayoutResponse.getPartitions().getRowCount(); i++) {
        actualValues.add(BlockUtils.rowToString(getTableLayoutResponse.getPartitions(), i));
    }
    Assert.assertEquals(Arrays.asList("[PARTITION_NUMBER : 1::: :::100000:::id]", "[PARTITION_NUMBER : 2:::100000:::300000:::id]"), actualValues);
    SchemaBuilder expectedSchemaBuilder = SchemaBuilder.newBuilder();
    expectedSchemaBuilder.addField(FieldBuilder.newBuilder(SynapseMetadataHandler.PARTITION_NUMBER, org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build());
    Schema expectedSchema = expectedSchemaBuilder.build();
    Assert.assertEquals(expectedSchema, getTableLayoutResponse.getPartitions().getSchema());
    Assert.assertEquals(tableName, getTableLayoutResponse.getTableName());
}