local-commit

Author: vikasrathee-cs

src/main/java/io/cdap/plugin/gcp/bigquery/source/BigQuerySource.java

@@ -261,6 +261,10 @@ private void configureBigQuerySource() {
     if (config.getFilter() != null) {
       configuration.set(BigQueryConstants.CONFIG_FILTER, config.getFilter());
     }
+    // TODO Add code of parameter map
+    if (config.getParameterMap() != null) {
+      configuration.set(BigQueryConstants.CONFIG_FILTER_PARAMETER_MAP, config.getParameterMap());
+    }
     if (config.getViewMaterializationProject() != null) {
       configuration.set(BigQueryConstants.CONFIG_VIEW_MATERIALIZATION_PROJECT, config.getViewMaterializationProject());
     }

src/main/java/io/cdap/plugin/gcp/bigquery/source/BigQuerySourceConfig.java

@@ -68,6 +68,7 @@ public final class BigQuerySourceConfig extends BigQueryBaseConfig {
   public static final String NAME_PARTITION_FROM = "partitionFrom";
   public static final String NAME_PARTITION_TO = "partitionTo";
   public static final String NAME_FILTER = "filter";
+  public static final String NAME_FILTER_PARAMETER_MAP = "parameterMap";
   public static final String NAME_ENABLE_QUERYING_VIEWS = "enableQueryingViews";
   public static final String NAME_VIEW_MATERIALIZATION_PROJECT = "viewMaterializationProject";
   public static final String NAME_VIEW_MATERIALIZATION_DATASET = "viewMaterializationDataset";
@@ -111,6 +112,12 @@ public final class BigQuerySourceConfig extends BigQueryBaseConfig {
     "and discards all rows that do not return TRUE (that is, rows that return FALSE or NULL).")
   private String filter;
 
+  @Name(NAME_FILTER_PARAMETER_MAP)
+  @Macro
+  @Nullable
+  @Description("While using parameterized query, provide values in key-value pair ")
+  private String parameterMap;
+
   @Name(NAME_ENABLE_QUERYING_VIEWS)
   @Macro
   @Nullable
@@ -295,6 +302,11 @@ public String getFilter() {
     return filter;
   }
 
+  @Nullable
+  public String getParameterMap() {
+    return parameterMap;
+  }
+
   public boolean isEnableQueryingViews() {
     return "true".equalsIgnoreCase(enableQueryingViews);
   }

src/main/java/io/cdap/plugin/gcp/bigquery/source/PartitionedBigQueryInputFormat.java

@@ -20,8 +20,12 @@
 import com.google.api.services.bigquery.model.JobConfiguration;
 import com.google.api.services.bigquery.model.JobConfigurationQuery;
 import com.google.api.services.bigquery.model.JobReference;
+import com.google.api.services.bigquery.model.QueryParameter;
+import com.google.api.services.bigquery.model.QueryParameterType;
+import com.google.api.services.bigquery.model.QueryParameterValue;
 import com.google.api.services.bigquery.model.Table;
 import com.google.api.services.bigquery.model.TableReference;
+import com.google.cloud.bigquery.FieldList;
 import com.google.cloud.bigquery.StandardTableDefinition;
 import com.google.cloud.bigquery.TableDefinition.Type;
 import com.google.cloud.bigquery.TimePartitioning;
@@ -40,6 +44,7 @@
 import io.cdap.cdap.api.exception.ErrorCategory;
 import io.cdap.cdap.api.exception.ErrorType;
 import io.cdap.cdap.api.exception.ErrorUtils;
+import io.cdap.plugin.common.ConfigUtil;
 import io.cdap.plugin.gcp.bigquery.util.BigQueryConstants;
 import io.cdap.plugin.gcp.bigquery.util.BigQueryUtil;
 import io.cdap.plugin.gcp.common.GCPUtils;
@@ -134,6 +139,7 @@ private void processQuery(JobContext context) throws IOException, InterruptedExc
     String filter = configuration.get(BigQueryConstants.CONFIG_FILTER, null);
     Integer readTimeout = configuration.getInt(BigQueryConstants.CONFIG_BQ_HTTP_READ_TIMEOUT,
         GCPUtils.BQ_DEFAULT_READ_TIMEOUT_SECONDS);
+    String parameterMap = configuration.get(BigQueryConstants.CONFIG_FILTER_PARAMETER_MAP, null);
 
     com.google.cloud.bigquery.Table bigQueryTable = BigQueryUtil.getBigQueryTable(
       datasetProjectId, datasetId, tableName, serviceAccount, isServiceAccountFilePath, null, readTimeout);
@@ -150,11 +156,16 @@ private void processQuery(JobContext context) throws IOException, InterruptedExc
     if (query != null) {
       TableReference sourceTable = new TableReference().setDatasetId(datasetId).setProjectId(datasetProjectId)
         .setTableId(tableName);
+      com.google.cloud.bigquery.Table bqTable = BigQueryUtil.getBigQueryTable(datasetProjectId, datasetId, tableName,
+                                                                              serviceAccount, isServiceAccountFilePath,
+                                                                              null,
+                                                                              null);
       String location = bigQueryHelper.getTable(sourceTable).getLocation();
       String temporaryTableName = configuration.get(BigQueryConstants.CONFIG_TEMPORARY_TABLE_NAME);
       TableReference exportTableReference = createExportTableReference(type, datasetProjectId, datasetId,
                                                                        temporaryTableName, configuration);
-      runQuery(configuration, bigQueryHelper, projectId, exportTableReference, query, location);
+      runQuery(configuration, bigQueryHelper, projectId, exportTableReference, query, location, bqTable,
+               parameterMap);
 
       // Default values come from BigquerySource config, and can be overridden by config.
       configuration.set(BigQueryConfiguration.INPUT_PROJECT_ID.getKey(),
@@ -244,14 +255,30 @@ private static void runQuery(Configuration configuration,
                                String projectId,
                                TableReference tableRef,
                                String query,
-                               String location)
+                               String location,
+                               com.google.cloud.bigquery.Table bqTable,
+                               @Nullable String parameterMapString)
     throws IOException, InterruptedException {
 
     // Create a query statement and query request object.
     JobConfigurationQuery queryConfig = new JobConfigurationQuery();
     queryConfig.setAllowLargeResults(true);
     queryConfig.setQuery(query);
     queryConfig.setUseLegacySql(false);
+    if (!Strings.isNullOrEmpty(parameterMapString)) {
+      Map<String, String> parameterMap = ConfigUtil.parseKeyValueConfig(parameterMapString, ",", "=");
+      List<QueryParameter> queryParameters = new ArrayList<>();
+      FieldList fieldList = bqTable.getDefinition().getSchema().getFields();
+      for (String columnName : parameterMap.keySet()) {
+        String parameterType = fieldList.get(columnName).getType().name();
+        QueryParameter queryParameter = new QueryParameter().setName(columnName)
+          .setParameterType(new QueryParameterType().setType(parameterType));
+        QueryParameterValue value = new QueryParameterValue().setValue(parameterMap.get(columnName));
+        queryParameters.add(queryParameter.setParameterValue(value));
+      }
+      queryConfig.setParameterMode("NAMED");
+      queryConfig.setQueryParameters(queryParameters);
+    }
 
     // Set the table to put results into.
     queryConfig.setDestinationTable(tableRef);

src/main/java/io/cdap/plugin/gcp/bigquery/util/BigQueryConstants.java

@@ -36,6 +36,8 @@ public interface BigQueryConstants {
   String CONFIG_TABLE_FIELDS = "cdap.bq.sink.table.fields";
   String CONFIG_JSON_STRING_FIELDS = "cdap.bq.sink.json.string.fields";
   String CONFIG_FILTER = "cdap.bq.source.filter";
+
+  String CONFIG_FILTER_PARAMETER_MAP = "cdap.bq.source.filter.parameter.map";
   String CONFIG_PARTITION_FILTER = "cdap.bq.sink.partition.filter";
   String CONFIG_JOB_ID = "cdap.bq.sink.job.id";
   String CONFIG_VIEW_MATERIALIZATION_PROJECT = "cdap.bq.source.view.materialization.project";

widgets/BigQueryTable-batchsource.json

@@ -151,6 +151,21 @@
             "placeholder": ""
           }
         },
+        {
+          "widget-type": "ds-multiplevalues",
+          "label": "Parameterized Query Map",
+          "name": "parameterMap",
+          "widget-attributes": {
+            "delimiter": ",",
+            "values-delimiter": ":",
+            "numValues": "3",
+            "placeholders": [
+              "Field Alias",
+              "Field Name",
+              "Field Value"
+            ]
+          }
+        },
         {
           "widget-type": "textbox",
           "label": "Temporary Bucket Name",