Merge pull request #997 from data-integrations/pushdown-metrics

Added support for context metrics in the BigQuery SQL Engine

Author: fernst

src/main/java/io/cdap/plugin/gcp/bigquery/sqlengine/BigQueryJoinDataset.java

@@ -27,13 +27,15 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Strings;
 import io.cdap.cdap.api.RuntimeContext;
+import io.cdap.cdap.api.SQLEngineContext;
 import io.cdap.cdap.api.annotation.Description;
 import io.cdap.cdap.api.annotation.Metadata;
 import io.cdap.cdap.api.annotation.MetadataProperty;
 import io.cdap.cdap.api.annotation.Name;
 import io.cdap.cdap.api.annotation.Plugin;
 import io.cdap.cdap.api.data.format.StructuredRecord;
 import io.cdap.cdap.api.data.schema.Schema;
+import io.cdap.cdap.api.metrics.Metrics;
 import io.cdap.cdap.etl.api.PipelineConfigurer;
 import io.cdap.cdap.etl.api.connector.Connector;
 import io.cdap.cdap.etl.api.engine.sql.BatchSQLEngine;
@@ -115,8 +117,8 @@ public class BigQuerySQLEngine
   private String dataset;
   private String bucket;
   private String runId;
-  private Map<String, String> tableNames;
   private Map<String, BigQuerySQLDataset> datasets;
+  private Metrics metrics;
 
   @SuppressWarnings("unused")
   public BigQuerySQLEngine(BigQuerySQLEngineConfig sqlEngineConfig) {
@@ -132,14 +134,13 @@ public void configurePipeline(PipelineConfigurer pipelineConfigurer) {
   }
 
   @Override
-  public void prepareRun(RuntimeContext context) throws Exception {
+  public void prepareRun(SQLEngineContext context) throws Exception {
     super.prepareRun(context);
 
     // Validate configuration and throw exception if the supplied configuration is invalid.
     sqlEngineConfig.validate();
 
     runId = BigQuerySQLEngineUtils.newIdentifier();
-    tableNames = new HashMap<>();
     datasets = new HashMap<>();
 
     String serviceAccount = sqlEngineConfig.getServiceAccount();
@@ -171,10 +172,13 @@ public void prepareRun(RuntimeContext context) throws Exception {
 
     // Configure credentials for the source
     BigQuerySourceUtils.configureServiceAccount(configuration, sqlEngineConfig.connection);
+
+    // Get metrics instance
+    metrics = context.getMetrics();
   }
 
   @Override
-  public void onRunFinish(boolean succeeded, RuntimeContext context) {
+  public void onRunFinish(boolean succeeded, SQLEngineContext context) {
     super.onRunFinish(succeeded, context);
 
     String gcsPath;
@@ -358,7 +362,8 @@ public SQLWriteResult write(SQLWriteRequest writeRequest) {
                                                             sqlEngineConfig,
                                                             bigQuery,
                                                             writeRequest,
-                                                            sourceTableId);
+                                                            sourceTableId,
+                                                            metrics);
     return bigQueryWrite.write();
   }
 
@@ -529,7 +534,8 @@ private BigQuerySelectDataset executeSelect(String datasetName,
       table,
       jobId,
       jobType,
-      query
+      query,
+      metrics
     ).execute();
 
     datasets.put(datasetName, selectDataset);

src/main/java/io/cdap/plugin/gcp/bigquery/sqlengine/BigQuerySQLEngine.java

@@ -11,6 +11,7 @@
 import com.google.cloud.bigquery.TableDefinition;
 import com.google.cloud.bigquery.TableId;
 import io.cdap.cdap.api.data.schema.Schema;
+import io.cdap.cdap.api.metrics.Metrics;
 import io.cdap.cdap.etl.api.engine.sql.SQLEngineException;
 import io.cdap.cdap.etl.api.engine.sql.dataset.SQLDataset;
 import io.cdap.plugin.gcp.bigquery.sink.BigQuerySinkUtils;
@@ -19,6 +20,8 @@
 import org.slf4j.LoggerFactory;
 
 import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
 import javax.annotation.Nullable;
 
 /**
@@ -38,6 +41,7 @@ public class BigQuerySelectDataset implements SQLDataset, BigQuerySQLDataset {
   private final String jobId;
   private final BigQueryJobType operation;
   private final String selectQuery;
+  private final Metrics metrics;
   private Long numRows;
 
   public static BigQuerySelectDataset getInstance(String datasetName,
@@ -49,7 +53,8 @@ public static BigQuerySelectDataset getInstance(String datasetName,
                                                   String bqTable,
                                                   String jobId,
                                                   BigQueryJobType jobType,
-                                                  String selectQuery) {
+                                                  String selectQuery,
+                                                  Metrics metrics) {
 
     return new BigQuerySelectDataset(datasetName,
                                      outputSchema,
@@ -60,7 +65,8 @@ public static BigQuerySelectDataset getInstance(String datasetName,
                                      bqTable,
                                      jobId,
                                      jobType,
-                                     selectQuery);
+                                     selectQuery,
+                                     metrics);
   }
 
   private BigQuerySelectDataset(String datasetName,
@@ -72,7 +78,8 @@ private BigQuerySelectDataset(String datasetName,
                                 String bqTable,
                                 String jobId,
                                 BigQueryJobType operation,
-                                String selectQuery) {
+                                String selectQuery,
+                                Metrics metrics) {
     this.datasetName = datasetName;
     this.outputSchema = outputSchema;
     this.sqlEngineConfig = sqlEngineConfig;
@@ -83,6 +90,7 @@ private BigQuerySelectDataset(String datasetName,
     this.jobId = jobId;
     this.operation = operation;
     this.selectQuery = selectQuery;
+    this.metrics = metrics;
   }
 
   public BigQuerySelectDataset execute() {
@@ -124,14 +132,14 @@ public BigQuerySelectDataset execute() {
     if (queryJob == null) {
       throw new SQLEngineException("BigQuery job not found: " + jobId);
     } else if (queryJob.getStatus().getError() != null) {
-      BigQuerySQLEngineUtils.logJobMetrics(queryJob);
+      BigQuerySQLEngineUtils.logJobMetrics(queryJob, metrics);
       throw new SQLEngineException(String.format(
         "Error executing BigQuery Job: '%s' in Project '%s', Dataset '%s', Location'%s' : %s",
         jobId, project, bqDataset, location, queryJob.getStatus().getError().toString()));
     }
 
     LOG.info("Created BigQuery table `{}` using Job: {}", bqTable, jobId);
-    BigQuerySQLEngineUtils.logJobMetrics(queryJob);
+    BigQuerySQLEngineUtils.logJobMetrics(queryJob, metrics);
     return this;
   }