Merge pull request #1084 from data-integrations/feature/add-sampling-options

ellliao · web-flow · commit cf173b4bba83 · 2022-09-01T10:48:48.000-07:00
Add sampling options to BigQuery connector
diff --git a/src/main/java/io/cdap/plugin/gcp/bigquery/connector/BigQueryConnector.java b/src/main/java/io/cdap/plugin/gcp/bigquery/connector/BigQueryConnector.java
@@ -18,6 +18,7 @@
 
 import com.google.api.gax.paging.Page;
 import com.google.auth.oauth2.GoogleCredentials;
+import com.google.cloud.RetryOption;
 import com.google.cloud.bigquery.BigQuery;
 import com.google.cloud.bigquery.BigQueryException;
 import com.google.cloud.bigquery.Dataset;
@@ -49,6 +50,7 @@
 import io.cdap.cdap.etl.api.connector.DirectConnector;
 import io.cdap.cdap.etl.api.connector.PluginSpec;
 import io.cdap.cdap.etl.api.connector.SampleRequest;
+import io.cdap.cdap.etl.api.connector.SampleType;
 import io.cdap.cdap.etl.api.engine.sql.BatchSQLEngine;
 import io.cdap.cdap.etl.api.validation.ValidationException;
 import io.cdap.plugin.common.ConfigUtil;
@@ -62,6 +64,7 @@
 import io.cdap.plugin.gcp.bigquery.util.BigQueryDataParser;
 import io.cdap.plugin.gcp.bigquery.util.BigQueryUtil;
 import io.cdap.plugin.gcp.common.GCPUtils;
+import org.threeten.bp.Duration;
 
 import java.io.IOException;
 import java.util.HashMap;
@@ -95,8 +98,13 @@ public List<StructuredRecord> sample(ConnectorContext context, SampleRequest sam
       throw new IllegalArgumentException("Path should contain both dataset and table name.");
     }
     String dataset = path.getDataset();
-    return getTableData(getBigQuery(config.getProject()), config.getDatasetProject(), dataset, table,
-      sampleRequest.getLimit());
+    String query = getTableQuery(String.format("`%s.%s.%s`", config.getDatasetProject(), dataset, table),
+                                 sampleRequest.getLimit(),
+                                 SampleType.fromString(sampleRequest.getProperties().get("sampleType")),
+                                 sampleRequest.getProperties().get("strata"),
+                                 UUID.randomUUID().toString().replace("-", "_"));
+    String id = UUID.randomUUID().toString();
+    return getQueryResult(waitForJob(getBigQuery(config.getProject()), query, sampleRequest.getTimeoutMs(), id), id);
   }
 
   @Override
@@ -117,7 +125,7 @@ public void test(ConnectorContext context) throws ValidationException {
           GCPUtils.loadServiceAccountCredentials(config.getServiceAccount(), config.isServiceAccountFilePath());
       } catch (Exception e) {
         failureCollector.addFailure(String.format("Service account key provided is not valid: %s", e.getMessage()),
-          "Please provide a valid service account key.");
+                                    "Please provide a valid service account key.");
       }
     }
     // if either project or credentials cannot be loaded , no need to continue
@@ -130,7 +138,7 @@ public void test(ConnectorContext context) throws ValidationException {
       bigQuery.listDatasets(BigQuery.DatasetListOption.pageSize(1));
     } catch (Exception e) {
       failureCollector.addFailure(String.format("Could not connect to BigQuery: %s", e.getMessage()),
-        "Please specify correct connection properties.");
+                                  "Please specify correct connection properties.");
     }
   }
 
@@ -143,15 +151,15 @@ public BrowseDetail browse(ConnectorContext context, BrowseRequest browseRequest
     if (dataset == null) {
       // browse project to list all datasets
       return config.rootDataset == null ?
-               listDatasets(getBigQuery(config.getDatasetProject()), browseRequest.getLimit()) :
-               BrowseDetail.builder().setTotalCount(1).addEntity(
-                 BrowseEntity.builder(config.rootDataset, "/" + config.rootDataset, ENTITY_TYPE_DATASET)
-                   .canBrowse(true).build()).build();
+        listDatasets(getBigQuery(config.getDatasetProject()), browseRequest.getLimit()) :
+        BrowseDetail.builder().setTotalCount(1).addEntity(
+          BrowseEntity.builder(config.rootDataset, "/" + config.rootDataset, ENTITY_TYPE_DATASET)
+            .canBrowse(true).build()).build();
     }
     String table = path.getTable();
     if (table == null) {
       return listTables(getBigQuery(config.getProject()), config.getDatasetProject(), dataset,
-        browseRequest.getLimit());
+                        browseRequest.getLimit());
     }
     return getTableDetail(getBigQuery(config.getProject()), config.getDatasetProject(), dataset, table);
   }
@@ -202,7 +210,7 @@ private BrowseDetail listTables(BigQuery bigQuery, String datasetProject, String
 
   private BrowseDetail listDatasets(BigQuery bigQuery, Integer limit) {
     Page<Dataset> datasetPage = config.showHiddenDatasets() ?
-                                  bigQuery.listDatasets(BigQuery.DatasetListOption.all()) : bigQuery.listDatasets();
+      bigQuery.listDatasets(BigQuery.DatasetListOption.all()) : bigQuery.listDatasets();
     int countLimit = limit == null || limit <= 0 ? Integer.MAX_VALUE : limit;
     int count = 0;
     BrowseDetail.Builder browseDetailBuilder = BrowseDetail.builder();
@@ -233,31 +241,99 @@ private BigQuery getBigQuery(String project) throws IOException {
     return GCPUtils.getBigQuery(project, credentials);
   }
 
-  private List<StructuredRecord> getTableData(BigQuery bigQuery, String datasetProject, String dataset, String table,
-    int limit)
-    throws IOException {
-    String query =
-      String.format("SELECT * FROM `%s.%s.%s` LIMIT %d", datasetProject, dataset, table, limit);
+  /**
+   * Get the SQL query used to sample the table
+   * @param tableName name of the table
+   * @param limit limit on rows returned
+   * @param sampleType sampling method
+   * @param strata strata column (if given)
+   * @param sessionID UUID
+   * @return String
+   * @throws IllegalArgumentException if no strata column is given for a stratified query
+   */
+  protected String getTableQuery(String tableName, int limit, SampleType sampleType, @Nullable String strata,
+                               String sessionID) {
+    switch (sampleType) {
+      case RANDOM:
+        return String.format("WITH table AS (\n" +
+                                "  SELECT *, RAND() AS r_%s\n" +
+                                "  FROM %s\n" +
+                                "  WHERE RAND() < 2*%d/(SELECT COUNT(*) FROM %s)\n" +
+                                ")\n" +
+                                "SELECT * EXCEPT (r_%s)\n" +
+                                "FROM table\n" +
+                                "ORDER BY r_%s\n" +
+                                "LIMIT %d",
+                              sessionID, tableName, limit, tableName, sessionID, sessionID, limit);
+      case STRATIFIED:
+        if (strata == null) {
+          throw new IllegalArgumentException("No strata column given.");
+        }
+        return String.format("SELECT * EXCEPT (`sqn_%s`, `c_%s`)\n" +
+                                "FROM (\n" +
+                                "SELECT *, row_number() OVER (ORDER BY %s, RAND()) AS sqn_%s,\n" +
+                                "COUNT(*) OVER () as c_%s,\n" +
+                                "FROM %s\n" +
+                                ") %s\n" +
+                                "WHERE MOD(sqn_%s, CAST(c_%s / %d AS INT64)) = 1\n" +
+                                "ORDER BY %s\n" +
+                                "LIMIT %d",
+                              sessionID, sessionID, strata, sessionID, sessionID, tableName, tableName, sessionID,
+                              sessionID, limit, strata, limit);
+      default:
+        return String.format("SELECT * FROM %s LIMIT %d", tableName, limit);
+    }
+  }
+
+  /**
+   * Wait for job to complete or time out (if timeout is given)
+   * @param bigQuery BigQuery client
+   * @param query SQL query
+   * @param timeoutMs timeout (if given)
+   * @param id job ID
+   * @return Job
+   * @throws IOException if the job is interrupted
+   */
+  private Job waitForJob(BigQuery bigQuery, String query, @Nullable Long timeoutMs, String id) throws IOException {
+
+    // set up job
     QueryJobConfiguration queryConfig = QueryJobConfiguration.newBuilder(query).build();
-    String id = UUID.randomUUID().toString();
     JobId jobId = JobId.of(id);
     Job queryJob = bigQuery.create(JobInfo.newBuilder(queryConfig).setJobId(jobId).build());
-    // Wait for the job to finish
+
+    // wait for job
     try {
-      queryJob = queryJob.waitFor();
+      if (timeoutMs == null) {
+        return queryJob.waitFor();
+      } else {
+        return queryJob.waitFor(RetryOption.totalTimeout(Duration.ofMillis(timeoutMs)));
+      }
     } catch (InterruptedException e) {
       throw new IOException(String.format("Query job %s interrupted.", id), e);
     }
+  }
 
-    // check for errors
+  /**
+   * Retrieve the results of a SQL query
+   * @param queryJob query job after completion or timeout
+   * @param id job ID
+   * @return List of structured records
+   * @throws IOException if query encounters an error or times out
+   */
+  protected List<StructuredRecord> getQueryResult(@Nullable Job queryJob, String id) throws IOException {
+
+    // Check for errors
     if (queryJob == null) {
       throw new IOException(String.format("Job %s no longer exists.", id));
+    } else if (!queryJob.isDone()) {
+      queryJob.cancel();
+      throw new IOException(String.format("Job %s timed out.", id));
     } else if (queryJob.getStatus().getError() != null) {
       throw new IOException(String.format("Failed to query table : %s", queryJob.getStatus().getError().toString()));
     }
 
     // Get the results
-    TableResult result = null;
+    TableResult result;
     try {
       result = queryJob.getQueryResults();
     } catch (InterruptedException e) {
@@ -266,7 +342,6 @@ private List<StructuredRecord> getTableData(BigQuery bigQuery, String datasetPro
     return BigQueryDataParser.parse(result);
   }
 
-
   @Override
   public ConnectorSpec generateSpec(ConnectorContext context,
                                     ConnectorSpecRequest connectorSpecRequest) throws IOException {
@@ -297,6 +372,9 @@ public ConnectorSpec generateSpec(ConnectorContext context,
       .addRelatedPlugin(new PluginSpec(BigQuerySink.NAME, BatchSink.PLUGIN_TYPE, properties))
       .addRelatedPlugin(new PluginSpec(BigQueryMultiSink.NAME, BatchSink.PLUGIN_TYPE, properties))
       .addRelatedPlugin(new PluginSpec(BigQuerySQLEngine.NAME, BatchSQLEngine.PLUGIN_TYPE, properties))
+      .addSupportedSampleType(SampleType.RANDOM)
+      .addSupportedSampleType(SampleType.STRATIFIED)
       .build();
   }
 }
+
diff --git a/src/test/java/io/cdap/plugin/gcp/bigquery/connector/BigQueryConnectorUnitTest.java b/src/test/java/io/cdap/plugin/gcp/bigquery/connector/BigQueryConnectorUnitTest.java
@@ -0,0 +1,163 @@
+/*
+ * Copyright © 2022 Cask Data, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+
+package io.cdap.plugin.gcp.bigquery.connector;
+
+import com.google.cloud.bigquery.BigQueryError;
+import com.google.cloud.bigquery.EmptyTableResult;
+import com.google.cloud.bigquery.Job;
+import com.google.cloud.bigquery.JobStatus;
+import com.google.cloud.bigquery.TableResult;
+import io.cdap.cdap.api.data.format.StructuredRecord;
+import io.cdap.cdap.etl.api.connector.SampleType;
+import io.cdap.plugin.gcp.bigquery.util.BigQueryDataParser;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+import org.junit.runner.RunWith;
+import org.mockito.junit.MockitoJUnitRunner;
+
+import static org.powermock.api.mockito.PowerMockito.doAnswer;
+import static org.powermock.api.mockito.PowerMockito.doThrow;
+import static org.powermock.api.mockito.PowerMockito.mock;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.UUID;
+
+@RunWith(MockitoJUnitRunner.class)
+public class BigQueryConnectorUnitTest {
+
+  @Rule
+  public ExpectedException expectedEx = ExpectedException.none();
+
+  private static final BigQueryConnector CONNECTOR = new BigQueryConnector(null);
+  private String tableName;
+  private String sessionID;
+  private int limit;
+  private String strata;
+  private Job queryJob;
+
+  @Before
+  public void setUp() {
+    tableName = "`project.dataset.table`";
+    sessionID = UUID.randomUUID().toString().replace('-', '_');
+    limit = 100;
+    strata = "strata";
+  }
+
+  /**
+   * Unit tests for getTableQuery()
+   */
+  @Test
+  public void getTableQueryTest() {
+    // random query
+    Assert.assertEquals(String.format("WITH table AS (\n" +
+                                        "  SELECT *, RAND() AS r_%s\n" +
+                                        "  FROM %s\n" +
+                                        "  WHERE RAND() < 2*%d/(SELECT COUNT(*) FROM %s)\n" +
+                                        ")\n" +
+                                        "SELECT * EXCEPT (r_%s)\n" +
+                                        "FROM table\n" +
+                                        "ORDER BY r_%s\n" +
+                                        "LIMIT %d",
+                                      sessionID, tableName, limit, tableName, sessionID, sessionID, limit),
+                        CONNECTOR.getTableQuery(tableName, limit, SampleType.RANDOM, null, sessionID));
+
+    // stratified query
+    Assert.assertEquals(String.format("SELECT * EXCEPT (`sqn_%s`, `c_%s`)\n" +
+                                        "FROM (\n" +
+                                        "SELECT *, row_number() OVER (ORDER BY %s, RAND()) AS sqn_%s,\n" +
+                                        "COUNT(*) OVER () as c_%s,\n" +
+                                        "FROM %s\n" +
+                                        ") %s\n" +
+                                        "WHERE MOD(sqn_%s, CAST(c_%s / %d AS INT64)) = 1\n" +
+                                        "ORDER BY %s\n" +
+                                        "LIMIT %d",
+                                      sessionID, sessionID, strata, sessionID, sessionID, tableName, tableName,
+                                      sessionID, sessionID, limit, strata, limit),
+                        CONNECTOR.getTableQuery(tableName, limit, SampleType.STRATIFIED, strata, sessionID));
+
+    // default query
+    Assert.assertEquals(String.format("SELECT * FROM %s LIMIT %d", tableName, limit),
+                        CONNECTOR.getTableQuery(tableName, limit, SampleType.DEFAULT, null, sessionID));
+  }
+
+  /**
+   * Test for {@link IllegalArgumentException} from getTableQuery when attempting stratified query with null strata
+   * @throws IllegalArgumentException expected
+   */
+  @Test
+  public void getTableQueryNullStrataTest() throws IllegalArgumentException {
+    expectedEx.expect(IllegalArgumentException.class);
+    CONNECTOR.getTableQuery(tableName, limit, SampleType.STRATIFIED, null, sessionID);
+  }
+
+  /**
+   * Test for {@link IOException} from getQueryResult() when attempting on null query job
+   * @throws IOException expected
+   */
+  @Test
+  public void getQueryResultNullJobTest() throws IOException {
+    expectedEx.expect(IOException.class);
+    CONNECTOR.getQueryResult(null, sessionID);
+  }
+
+  /**
+   * Test for {@link IOException} from getQueryResult() if job timed out
+   * @throws IOException expected
+   */
+  @Test
+  public void getQueryResultTimedOutTest() throws IOException {
+    expectedEx.expect(IOException.class);
+    queryJob = mock(Job.class);
+    doAnswer(invocation -> false).when(queryJob).isDone();
+    CONNECTOR.getQueryResult(queryJob, sessionID);
+  }
+
+  /**
+   * Test for {@link IOException} from getQueryResult() if query has error
+   * @throws IOException expected
+   */
+  @Test
+  public void getQueryResultErrorTest() throws IOException {
+    expectedEx.expect(IOException.class);
+    queryJob = mock(Job.class);
+    doAnswer(invocation -> true).when(queryJob).isDone();
+    JobStatus status = mock(JobStatus.class);
+    doAnswer(invocation -> mock(BigQueryError.class)).when(status).getError();
+    doAnswer(invocation -> status).when(queryJob).getStatus();
+    CONNECTOR.getQueryResult(queryJob, sessionID);
+  }
+
+  /**
+   * Test for {@link IOException} from getQueryResult() if query is interrupted
+   * @throws InterruptedException to get IOException
+   * @throws IOException expected
+   */
+  @Test
+  public void getQueryResultInterruptedTest() throws InterruptedException, IOException {
+    expectedEx.expect(IOException.class);
+    queryJob = mock(Job.class);
+    doAnswer(invocation -> true).when(queryJob).isDone();
+    doAnswer(invocation -> mock(JobStatus.class)).when(queryJob).getStatus();
+    doThrow(mock(InterruptedException.class)).when(queryJob).getQueryResults();
+    CONNECTOR.getQueryResult(queryJob, sessionID);
+  }
+
+}
