use of com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet in project aws-athena-query-federation by awslabs.
the class OracleRecordHandlerTest method buildSplitSql.
@Test
public void buildSplitSql() 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_name", Types.MinorType.VARCHAR.getType()).build());
Schema schema = schemaBuilder.build();
Split split = Mockito.mock(Split.class);
Mockito.when(split.getProperties()).thenReturn(Collections.singletonMap("partition_name", "p0"));
Mockito.when(split.getProperty(Mockito.eq("partition_name"))).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());
String expectedSql = "SELECT \"testCol1\", \"testCol2\", \"testCol3\", \"testCol4\", \"testCol5\", \"testCol6\", \"testCol7\", \"testCol8\" FROM \"testSchema\".\"testTable\" PARTITION (p0) WHERE (\"testCol1\" IN (?,?)) AND ((\"testCol2\" >= ? AND \"testCol2\" < ?)) AND ((\"testCol3\" > ? AND \"testCol3\" <= ?)) AND (\"testCol4\" = ?) AND (\"testCol5\" = ?) AND (\"testCol6\" = ?) AND (\"testCol7\" = ?) AND (\"testCol8\" = ?)";
PreparedStatement expectedPreparedStatement = Mockito.mock(PreparedStatement.class);
Mockito.when(this.connection.prepareStatement(Mockito.eq(expectedSql))).thenReturn(expectedPreparedStatement);
PreparedStatement preparedStatement = this.oracleRecordHandler.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);
}
use of com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet in project aws-athena-query-federation by awslabs.
the class OracleRecordHandlerTest method getSingleValueSet.
private ValueSet getSingleValueSet(Object value) {
Range range = Mockito.mock(Range.class, Mockito.RETURNS_DEEP_STUBS);
Mockito.when(range.isSingleValue()).thenReturn(true);
Mockito.when(range.getLow().getValue()).thenReturn(value);
ValueSet valueSet = Mockito.mock(SortedRangeSet.class, Mockito.RETURNS_DEEP_STUBS);
Mockito.when(valueSet.getRanges().getOrderedRanges()).thenReturn(Collections.singletonList(range));
return valueSet;
}
use of com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet in project aws-athena-query-federation by awslabs.
the class RedisRecordHandlerTest method doReadRecordsHash.
@Test
public void doReadRecordsHash() throws Exception {
// 4 keys per prefix
when(mockSyncCommands.scan(any(ScanCursor.class), any(ScanArgs.class))).then((InvocationOnMock invocationOnMock) -> {
ScanCursor cursor = (ScanCursor) invocationOnMock.getArguments()[0];
if (cursor == null || cursor.getCursor().equals("0")) {
List<String> result = new ArrayList<>();
result.add(UUID.randomUUID().toString());
result.add(UUID.randomUUID().toString());
result.add(UUID.randomUUID().toString());
result.add(UUID.randomUUID().toString());
result.add(UUID.randomUUID().toString());
MockKeyScanCursor<String> scanCursor = new MockKeyScanCursor<>();
scanCursor.setCursor("1");
scanCursor.setKeys(result);
return scanCursor;
} else {
List<String> result = new ArrayList<>();
result.add(UUID.randomUUID().toString());
result.add(UUID.randomUUID().toString());
MockKeyScanCursor<String> scanCursor = new MockKeyScanCursor<>();
scanCursor.setCursor("0");
scanCursor.setKeys(result);
scanCursor.setFinished(true);
return scanCursor;
}
});
// 4 columns per key
AtomicLong intColVal = new AtomicLong(0);
when(mockSyncCommands.hgetall(anyString())).then((InvocationOnMock invocationOnMock) -> {
Map<String, String> result = new HashMap<>();
result.put("intcol", String.valueOf(intColVal.getAndIncrement()));
result.put("stringcol", UUID.randomUUID().toString());
result.put("extracol", UUID.randomUUID().toString());
return result;
});
AtomicLong value = new AtomicLong(0);
when(mockSyncCommands.get(anyString())).thenAnswer((InvocationOnMock invocationOnMock) -> String.valueOf(value.getAndIncrement()));
S3SpillLocation splitLoc = S3SpillLocation.newBuilder().withBucket(UUID.randomUUID().toString()).withSplitId(UUID.randomUUID().toString()).withQueryId(UUID.randomUUID().toString()).withIsDirectory(true).build();
Split split = Split.newBuilder(splitLoc, keyFactory.create()).add(REDIS_ENDPOINT_PROP, endpoint).add(KEY_TYPE, KeyType.PREFIX.getId()).add(KEY_PREFIX_TABLE_PROP, "key-*").add(VALUE_TYPE_TABLE_PROP, ValueType.HASH.getId()).build();
Schema schemaForRead = SchemaBuilder.newBuilder().addField("_key_", Types.MinorType.VARCHAR.getType()).addField("intcol", Types.MinorType.INT.getType()).addField("stringcol", Types.MinorType.VARCHAR.getType()).build();
Map<String, ValueSet> constraintsMap = new HashMap<>();
constraintsMap.put("intcol", SortedRangeSet.copyOf(Types.MinorType.INT.getType(), ImmutableList.of(Range.greaterThan(allocator, Types.MinorType.INT.getType(), 1)), false));
ReadRecordsRequest request = new ReadRecordsRequest(IDENTITY, DEFAULT_CATALOG, "queryId-" + System.currentTimeMillis(), TABLE_NAME, schemaForRead, split, new Constraints(constraintsMap), // 100GB don't expect this to spill
100_000_000_000L, 100_000_000_000L);
RecordResponse rawResponse = handler.doReadRecords(allocator, request);
assertTrue(rawResponse instanceof ReadRecordsResponse);
ReadRecordsResponse response = (ReadRecordsResponse) rawResponse;
logger.info("doReadRecordsHash: rows[{}]", response.getRecordCount());
logger.info("doReadRecordsHash: {}", BlockUtils.rowToString(response.getRecords(), 0));
assertTrue(response.getRecords().getRowCount() == 5);
assertTrue(response.getRecords().getFields().size() == schemaForRead.getFields().size());
FieldReader keyReader = response.getRecords().getFieldReader(KEY_COLUMN_NAME);
keyReader.setPosition(0);
assertNotNull(keyReader.readText());
FieldReader intCol = response.getRecords().getFieldReader("intcol");
intCol.setPosition(0);
assertNotNull(intCol.readInteger());
FieldReader stringCol = response.getRecords().getFieldReader("stringcol");
stringCol.setPosition(0);
assertNotNull(stringCol.readText());
}
use of com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet in project aws-athena-query-federation by awslabs.
the class RedshiftRecordHandlerTest method getRangeSet.
private ValueSet getRangeSet(Marker.Bound lowerBound, Object lowerValue, Marker.Bound upperBound, Object upperValue) {
Range range = Mockito.mock(Range.class, Mockito.RETURNS_DEEP_STUBS);
Mockito.when(range.isSingleValue()).thenReturn(false);
Mockito.when(range.getLow().getBound()).thenReturn(lowerBound);
Mockito.when(range.getLow().getValue()).thenReturn(lowerValue);
Mockito.when(range.getHigh().getBound()).thenReturn(upperBound);
Mockito.when(range.getHigh().getValue()).thenReturn(upperValue);
ValueSet valueSet = Mockito.mock(SortedRangeSet.class, Mockito.RETURNS_DEEP_STUBS);
Mockito.when(valueSet.getRanges().getOrderedRanges()).thenReturn(Collections.singletonList(range));
return valueSet;
}
use of com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet 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");
}
Aggregations