Examples with ReadRecordsRequest com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest used on opensource projects

Example 46 with ReadRecordsRequest

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

the class BigQueryRecordHandlerTest method testReadWithConstraint.

@Test
public void testReadWithConstraint() throws Exception {
    try (ReadRecordsRequest request = new ReadRecordsRequest(federatedIdentity, BigQueryTestUtils.PROJECT_1_NAME, "queryId", new TableName("dataset1", "table1"), BigQueryTestUtils.getBlockTestSchema(), Split.newBuilder(S3SpillLocation.newBuilder().withBucket(bucket).withPrefix(prefix).withSplitId(UUID.randomUUID().toString()).withQueryId(UUID.randomUUID().toString()).withIsDirectory(true).build(), keyFactory.create()).build(), new Constraints(Collections.EMPTY_MAP), // This is ignored when directly calling readWithConstraints.
    0, 0)) {
        // This is ignored when directly calling readWithConstraints.
        // Always return try for the evaluator to keep all rows.
        ConstraintEvaluator evaluator = mock(ConstraintEvaluator.class);
        when(evaluator.apply(any(String.class), any(Object.class))).thenAnswer((InvocationOnMock invocationOnMock) -> {
            return true;
        });
        // Populate the schema and data that the mocked Google BigQuery client will return.
        com.google.cloud.bigquery.Schema tableSchema = BigQueryTestUtils.getTestSchema();
        List<FieldValueList> tableRows = Arrays.asList(BigQueryTestUtils.getBigQueryFieldValueList(false, 1000, "test1", 123123.12312), BigQueryTestUtils.getBigQueryFieldValueList(true, 500, "test2", 5345234.22111), BigQueryTestUtils.getBigQueryFieldValueList(false, 700, "test3", 324324.23423), BigQueryTestUtils.getBigQueryFieldValueList(true, 900, null, null), BigQueryTestUtils.getBigQueryFieldValueList(null, null, "test5", 2342.234234), BigQueryTestUtils.getBigQueryFieldValueList(true, 1200, "test6", 1123.12312), BigQueryTestUtils.getBigQueryFieldValueList(false, 100, "test7", 1313.12312), BigQueryTestUtils.getBigQueryFieldValueList(true, 120, "test8", 12313.1312), BigQueryTestUtils.getBigQueryFieldValueList(false, 300, "test9", 12323.1312));
        Page<FieldValueList> fieldValueList = new BigQueryPage<>(tableRows);
        TableResult result = new TableResult(tableSchema, tableRows.size(), fieldValueList);
        // Mock out the Google BigQuery Job.
        Job mockBigQueryJob = mock(Job.class);
        when(mockBigQueryJob.isDone()).thenReturn(false).thenReturn(true);
        when(mockBigQueryJob.getQueryResults()).thenReturn(result);
        when(bigQuery.create(any(JobInfo.class))).thenReturn(mockBigQueryJob);
        QueryStatusChecker queryStatusChecker = mock(QueryStatusChecker.class);
        when(queryStatusChecker.isQueryRunning()).thenReturn(true);
        // Execute the test
        bigQueryRecordHandler.readWithConstraint(spillWriter, request, queryStatusChecker);
        PowerMockito.mockStatic(System.class);
        PowerMockito.when(System.getenv(anyString())).thenReturn("test");
        logger.info("Project Name: " + BigQueryUtils.getProjectName(request.getCatalogName()));
        // Ensure that there was a spill so that we can read the spilled block.
        assertTrue(spillWriter.spilled());
    }
}

Example 47 with ReadRecordsRequest

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

the class HbaseRecordHandlerTest method doReadRecordsSpill.

@Test
public void doReadRecordsSpill() throws Exception {
    List<Result> results = TestUtils.makeResults(10_000);
    ResultScanner mockScanner = mock(ResultScanner.class);
    when(mockScanner.iterator()).thenReturn(results.iterator());
    when(mockClient.scanTable(anyObject(), any(Scan.class), anyObject())).thenAnswer((InvocationOnMock invocationOnMock) -> {
        ResultProcessor processor = (ResultProcessor) invocationOnMock.getArguments()[2];
        return processor.scan(mockScanner);
    });
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put("family1:col3", SortedRangeSet.copyOf(Types.MinorType.BIGINT.getType(), ImmutableList.of(Range.greaterThan(allocator, Types.MinorType.BIGINT.getType(), 0L)), true));
    S3SpillLocation splitLoc = S3SpillLocation.newBuilder().withBucket(UUID.randomUUID().toString()).withSplitId(UUID.randomUUID().toString()).withQueryId(UUID.randomUUID().toString()).withIsDirectory(true).build();
    Split.Builder splitBuilder = Split.newBuilder(splitLoc, keyFactory.create()).add(HBASE_CONN_STR, "fake_con_str").add(START_KEY_FIELD, "fake_start_key").add(END_KEY_FIELD, "fake_end_key").add(REGION_ID_FIELD, "fake_region_id").add(REGION_NAME_FIELD, "fake_region_name");
    ReadRecordsRequest request = new ReadRecordsRequest(IDENTITY, DEFAULT_CATALOG, "queryId-" + System.currentTimeMillis(), new TableName(DEFAULT_SCHEMA, TEST_TABLE), schemaForRead, splitBuilder.build(), new Constraints(constraintsMap), // ~1.5MB so we should see some spill
    1_500_000L, 0L);
    RecordResponse rawResponse = handler.doReadRecords(allocator, request);
    assertTrue(rawResponse instanceof RemoteReadRecordsResponse);
    try (RemoteReadRecordsResponse response = (RemoteReadRecordsResponse) rawResponse) {
        logger.info("doReadRecordsSpill: remoteBlocks[{}]", response.getRemoteBlocks().size());
        assertTrue(response.getNumberBlocks() > 1);
        int blockNum = 0;
        for (SpillLocation next : response.getRemoteBlocks()) {
            S3SpillLocation spillLocation = (S3SpillLocation) next;
            try (Block block = spillReader.read(spillLocation, response.getEncryptionKey(), response.getSchema())) {
                logger.info("doReadRecordsSpill: blockNum[{}] and recordCount[{}]", blockNum++, block.getRowCount());
                // assertTrue(++blockNum < response.getRemoteBlocks().size() && block.getRowCount() > 10_000);
                logger.info("doReadRecordsSpill: {}", BlockUtils.rowToString(block, 0));
                assertNotNull(BlockUtils.rowToString(block, 0));
            }
        }
    }
}

Example 48 with ReadRecordsRequest

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

the class HiveMuxRecordHandlerTest method readWithConstraint.

@Test
public void readWithConstraint() throws SQLException {
    BlockSpiller blockSpiller = Mockito.mock(BlockSpiller.class);
    ReadRecordsRequest readRecordsRequest = Mockito.mock(ReadRecordsRequest.class);
    Mockito.when(readRecordsRequest.getCatalogName()).thenReturn("recordHive");
    this.jdbcRecordHandler.readWithConstraint(blockSpiller, readRecordsRequest, queryStatusChecker);
    Mockito.verify(this.hiveRecordHandler, Mockito.times(1)).readWithConstraint(Mockito.eq(blockSpiller), Mockito.eq(readRecordsRequest), Mockito.eq(queryStatusChecker));
}

Example 49 with ReadRecordsRequest

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

the class HiveMuxRecordHandlerTest method readWithConstraintWithUnsupportedCatalog.

@Test(expected = RuntimeException.class)
public void readWithConstraintWithUnsupportedCatalog() throws SQLException {
    BlockSpiller blockSpiller = Mockito.mock(BlockSpiller.class);
    ReadRecordsRequest readRecordsRequest = Mockito.mock(ReadRecordsRequest.class);
    Mockito.when(readRecordsRequest.getCatalogName()).thenReturn("unsupportedCatalog");
    this.jdbcRecordHandler.readWithConstraint(blockSpiller, readRecordsRequest, queryStatusChecker);
}

Example 50 with ReadRecordsRequest

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

the class MultiplexingJdbcRecordHandlerTest method buildSplitSql.

@Test
public void buildSplitSql() throws SQLException {
    ReadRecordsRequest readRecordsRequest = Mockito.mock(ReadRecordsRequest.class);
    Mockito.when(readRecordsRequest.getCatalogName()).thenReturn("fakedatabase");
    Connection jdbcConnection = Mockito.mock(Connection.class);
    TableName tableName = new TableName("testSchema", "tableName");
    Schema schema = Mockito.mock(Schema.class);
    Constraints constraints = Mockito.mock(Constraints.class);
    Split split = Mockito.mock(Split.class);
    this.jdbcRecordHandler.buildSplitSql(jdbcConnection, "fakedatabase", tableName, schema, constraints, split);
    Mockito.verify(this.fakeJdbcRecordHandler, Mockito.times(1)).buildSplitSql(Mockito.eq(jdbcConnection), Mockito.eq("fakedatabase"), Mockito.eq(tableName), Mockito.eq(schema), Mockito.eq(constraints), Mockito.eq(split));
}