make sure project is where job is run and datasetProject is where dataset is (#719) (#721)

seanzhougoogle · web-flow · commit 13439aadc376 · 2021-07-23T10:30:21.000-07:00
diff --git a/src/main/java/io/cdap/plugin/gcp/bigquery/sink/BigQuerySink.java b/src/main/java/io/cdap/plugin/gcp/bigquery/sink/BigQuerySink.java
@@ -40,7 +40,6 @@
 
 import java.io.IOException;
 import java.util.List;
-import java.util.Map;
 import java.util.Objects;
 import java.util.UUID;
 import java.util.stream.Collectors;
@@ -218,7 +217,7 @@ private void configureBigQuerySink() {
    */
   private void configureTable(Schema schema) {
     AbstractBigQuerySinkConfig config = getConfig();
-    Table table = BigQueryUtil.getBigQueryTable(config.getDatasetProject(), config.getDataset(),
+    Table table = BigQueryUtil.getBigQueryTable(config.getProject(), config.getDatasetProject(), config.getDataset(),
                                                 config.getTable(),
                                                 config.getServiceAccount(),
                                                 config.isServiceAccountFilePath());
diff --git a/src/main/java/io/cdap/plugin/gcp/bigquery/source/BigQuerySource.java b/src/main/java/io/cdap/plugin/gcp/bigquery/source/BigQuerySource.java
@@ -123,7 +123,7 @@ public void prepareRun(BatchSourceContext context) throws Exception {
     // Create BigQuery client
     String serviceAccount = config.getServiceAccount();
     Credentials credentials = BigQuerySourceUtils.getCredentials(config.getConnection());
-    BigQuery bigQuery = GCPUtils.getBigQuery(config.getDatasetProject(), credentials);
+    BigQuery bigQuery = GCPUtils.getBigQuery(config.getProject(), credentials);
 
     // Get Configuration for this run
     bucketPath = UUID.randomUUID().toString();
diff --git a/src/main/java/io/cdap/plugin/gcp/bigquery/source/BigQuerySourceConfig.java b/src/main/java/io/cdap/plugin/gcp/bigquery/source/BigQuerySourceConfig.java
@@ -242,7 +242,7 @@ private void validateTable(FailureCollector collector) {
    */
   public Type getSourceTableType() {
     Table sourceTable =
-      BigQueryUtil.getBigQueryTable(
+      BigQueryUtil.getBigQueryTable(getProject(),
         getDatasetProject(), getDataset(), table, getServiceAccount(), isServiceAccountFilePath());
     return sourceTable != null ? sourceTable.getDefinition().getType() : null;
   }
diff --git a/src/main/java/io/cdap/plugin/gcp/bigquery/source/BigQuerySourceUtils.java b/src/main/java/io/cdap/plugin/gcp/bigquery/source/BigQuerySourceUtils.java
@@ -19,6 +19,7 @@
 import com.google.auth.Credentials;
 import com.google.cloud.bigquery.BigQuery;
 import com.google.cloud.bigquery.Dataset;
+import com.google.cloud.bigquery.DatasetId;
 import com.google.cloud.bigquery.TableId;
 import com.google.cloud.hadoop.io.bigquery.BigQueryConfiguration;
 import io.cdap.plugin.gcp.bigquery.connector.BigQueryConnectorConfig;
@@ -78,7 +79,7 @@ public static String getOrCreateBucket(Configuration configuration,
       configuration.setBoolean("fs.gs.bucket.delete.enable", true);
 
       // the dataset existence is validated before, so this cannot be null
-      Dataset dataset = bigQuery.getDataset(config.getDataset());
+      Dataset dataset = bigQuery.getDataset(DatasetId.of(config.getDatasetProject(), config.getDataset()));
       GCPUtils.createBucket(GCPUtils.getStorage(config.getProject(), credentials),
                             bucket,
                             dataset.getLocation(),
@@ -158,8 +159,8 @@ public static void deleteBigQueryTemporaryTable(Configuration configuration, Big
     String temporaryTable = configuration.get(BigQueryConstants.CONFIG_TEMPORARY_TABLE_NAME);
     try {
       Credentials credentials = getCredentials(config.getConnection());
-      BigQuery bigQuery = GCPUtils.getBigQuery(config.getDatasetProject(), credentials);
-      bigQuery.delete(TableId.of(config.getProject(), config.getDataset(), temporaryTable));
+      BigQuery bigQuery = GCPUtils.getBigQuery(config.getProject(), credentials);
+      bigQuery.delete(TableId.of(config.getDatasetProject(), config.getDataset(), temporaryTable));
       LOG.debug("Deleted temporary table '{}'", temporaryTable);
     } catch (IOException e) {
       LOG.error("Failed to load service account credentials: {}", e.getMessage(), e);
diff --git a/src/main/java/io/cdap/plugin/gcp/bigquery/source/PartitionedBigQueryInputFormat.java b/src/main/java/io/cdap/plugin/gcp/bigquery/source/PartitionedBigQueryInputFormat.java
@@ -82,7 +82,8 @@ private void processQuery(JobContext context) throws IOException, InterruptedExc
     }
     Map<String, String> mandatoryConfig = ConfigurationUtil.getMandatoryConfig(
       configuration, BigQueryConfiguration.MANDATORY_CONFIG_PROPERTIES_INPUT);
-    String inputProjectId = mandatoryConfig.get(BigQueryConfiguration.INPUT_PROJECT_ID_KEY);
+    String projectId = mandatoryConfig.get(BigQueryConfiguration.PROJECT_ID_KEY);
+    String datasetProjectId = mandatoryConfig.get(BigQueryConfiguration.INPUT_PROJECT_ID_KEY);
     String datasetId = mandatoryConfig.get(BigQueryConfiguration.INPUT_DATASET_ID_KEY);
     String tableName = mandatoryConfig.get(BigQueryConfiguration.INPUT_TABLE_ID_KEY);
     String serviceAccount = configuration.get(BigQueryConstants.CONFIG_SERVICE_ACCOUNT, null);
@@ -93,26 +94,26 @@ private void processQuery(JobContext context) throws IOException, InterruptedExc
     String partitionToDate = configuration.get(BigQueryConstants.CONFIG_PARTITION_TO_DATE, null);
     String filter = configuration.get(BigQueryConstants.CONFIG_FILTER, null);
 
-    com.google.cloud.bigquery.Table bigQueryTable =
-      BigQueryUtil.getBigQueryTable(inputProjectId, datasetId, tableName, serviceAccount, isServiceAccountFilePath);
+    com.google.cloud.bigquery.Table bigQueryTable = BigQueryUtil.getBigQueryTable(
+      projectId, datasetProjectId, datasetId, tableName, serviceAccount, isServiceAccountFilePath);
     Type type = Objects.requireNonNull(bigQueryTable).getDefinition().getType();
 
     String query;
     if (type == Type.VIEW || type == Type.MATERIALIZED_VIEW) {
-      query = generateQueryForMaterializingView(datasetId, tableName, filter);
+      query = generateQueryForMaterializingView(datasetProjectId, datasetId, tableName, filter);
     } else {
-      query = generateQuery(partitionFromDate, partitionToDate, filter, inputProjectId, datasetId, tableName,
-                            serviceAccount, isServiceAccountFilePath);
+      query = generateQuery(partitionFromDate, partitionToDate, filter, projectId, datasetProjectId, datasetId,
+                            tableName, serviceAccount, isServiceAccountFilePath);
     }
 
     if (query != null) {
-      TableReference sourceTable = new TableReference().setDatasetId(datasetId).setProjectId(inputProjectId)
+      TableReference sourceTable = new TableReference().setDatasetId(datasetId).setProjectId(datasetProjectId)
         .setTableId(tableName);
       String location = bigQueryHelper.getTable(sourceTable).getLocation();
       String temporaryTableName = configuration.get(BigQueryConstants.CONFIG_TEMPORARY_TABLE_NAME);
-      TableReference exportTableReference = createExportTableReference(type, inputProjectId, datasetId,
+      TableReference exportTableReference = createExportTableReference(type, datasetProjectId, datasetId,
                                                                        temporaryTableName, configuration);
-      runQuery(configuration, bigQueryHelper, inputProjectId, exportTableReference, query, location);
+      runQuery(configuration, bigQueryHelper, projectId, exportTableReference, query, location);
       if (type == Type.VIEW || type == Type.MATERIALIZED_VIEW) {
         configuration.set(BigQueryConfiguration.INPUT_PROJECT_ID_KEY,
                           configuration.get(BigQueryConstants.CONFIG_VIEW_MATERIALIZATION_PROJECT));
@@ -123,14 +124,14 @@ private void processQuery(JobContext context) throws IOException, InterruptedExc
     }
   }
 
-  private String generateQuery(String partitionFromDate, String partitionToDate, String filter,
-                               String project, String dataset, String table, @Nullable String serviceAccount,
+  private String generateQuery(String partitionFromDate, String partitionToDate, String filter, String project,
+                               String datasetProject, String dataset, String table, @Nullable String serviceAccount,
                                @Nullable Boolean isServiceAccountFilePath) {
     if (partitionFromDate == null && partitionToDate == null && filter == null) {
       return null;
     }
     String queryTemplate = "select * from %s where %s";
-    com.google.cloud.bigquery.Table sourceTable = BigQueryUtil.getBigQueryTable(project, dataset, table,
+    com.google.cloud.bigquery.Table sourceTable = BigQueryUtil.getBigQueryTable(project, datasetProject, dataset, table,
                                                                                 serviceAccount,
                                                                                 isServiceAccountFilePath);
     StandardTableDefinition tableDefinition = Objects.requireNonNull(sourceTable).getDefinition();
@@ -154,35 +155,35 @@ private String generateQuery(String partitionFromDate, String partitionToDate, S
       }
     }
 
-    String tableName = dataset + "." + table;
+    String tableName = datasetProject + "." + dataset + "." + table;
     return String.format(queryTemplate, tableName, condition.toString());
   }
 
-  private String generateQueryForMaterializingView(String dataset, String table, String filter) {
+  private String generateQueryForMaterializingView(String datasetProject, String dataset, String table, String filter) {
     String queryTemplate = "select * from %s %s";
     StringBuilder condition = new StringBuilder();
 
     if (!Strings.isNullOrEmpty(filter)) {
       condition.append(String.format(" where %s", filter));
     }
 
-    String tableName = dataset + "." + table;
+    String tableName = datasetProject + "." + dataset + "." + table;
     return String.format(queryTemplate, tableName, condition.toString());
   }
 
   /**
    * Create {@link TableReference} to export Table or View
    *
    * @param type           BigQuery table type
-   * @param inputProjectId project id of source table
+   * @param datasetProjectId project id of source table
    * @param datasetId      dataset id of source table
    * @param tableId        The ID of the table
    * @param configuration  Configuration that contains View Materialization ProjectId and View
    *                       Materialization Dataset Id
    * @return {@link TableReference}
    */
   private TableReference createExportTableReference(
-    Type type, String inputProjectId,
+    Type type, String datasetProjectId,
     String datasetId,
     String tableId,
     Configuration configuration) {
@@ -193,7 +194,7 @@ private TableReference createExportTableReference(
       tableReference.setProjectId(configuration.get(BigQueryConstants.CONFIG_VIEW_MATERIALIZATION_PROJECT));
       tableReference.setDatasetId(configuration.get(BigQueryConstants.CONFIG_VIEW_MATERIALIZATION_DATASET));
     } else {
-      tableReference.setProjectId(inputProjectId);
+      tableReference.setProjectId(datasetProjectId);
       tableReference.setDatasetId(datasetId);
     }
 
diff --git a/src/main/java/io/cdap/plugin/gcp/bigquery/util/BigQueryUtil.java b/src/main/java/io/cdap/plugin/gcp/bigquery/util/BigQueryUtil.java
@@ -476,33 +476,18 @@ public static ValidationFailure validateArraySchema(Schema arraySchema, String n
 
   /**
    * Get BigQuery table.
-   *
-   * @param projectId BigQuery project ID
-   * @param datasetId BigQuery dataset ID
-   * @param tableName BigQuery table name
-   * @param serviceAccountPath service account file path
-   * @return BigQuery table
-   */
-  @Nullable
-  public static Table getBigQueryTable(String projectId, String datasetId, String tableName,
-                                       @Nullable String serviceAccountPath) {
-    return getBigQueryTable(projectId, datasetId, tableName, serviceAccountPath, true);
-  }
-
-  /**
-   * Get BigQuery table.
-   *
-   * @param projectId BigQuery project ID
+   * @param project project where the BQ job will be run
+   * @param datasetProject project where dataset is in
    * @param datasetId BigQuery dataset ID
    * @param tableName BigQuery table name
    * @param serviceAccount service account file path or JSON content
    * @param isServiceAccountFilePath indicator for whether service account is file or json
    * @return BigQuery table
    */
   @Nullable
-  public static Table getBigQueryTable(String projectId, String datasetId, String tableName,
+  public static Table getBigQueryTable(String project, String datasetProject, String datasetId, String tableName,
                                        @Nullable String serviceAccount, boolean isServiceAccountFilePath) {
-    TableId tableId = TableId.of(projectId, datasetId, tableName);
+    TableId tableId = TableId.of(datasetProject, datasetId, tableName);
 
     com.google.auth.Credentials credentials = null;
     if (serviceAccount != null) {
@@ -514,7 +499,7 @@ public static Table getBigQueryTable(String projectId, String datasetId, String
           "serviceFilePath");
       }
     }
-    BigQuery bigQuery = GCPUtils.getBigQuery(projectId, credentials);
+    BigQuery bigQuery = GCPUtils.getBigQuery(project, credentials);
 
     Table table;
     try {

Original file line number	Diff line number	Diff line change
`@@ -242,7 +242,7 @@ private void validateTable(FailureCollector collector) {`
`242`	`242`	`*/`
`243`	`243`	`public Type getSourceTableType() {`
`244`	`244`	`Table sourceTable =`
`245`		`- BigQueryUtil.getBigQueryTable(`
	`245`	`+ BigQueryUtil.getBigQueryTable(getProject(),`
`246`	`246`	`getDatasetProject(), getDataset(), table, getServiceAccount(), isServiceAccountFilePath());`
`247`	`247`	`return sourceTable != null ? sourceTable.getDefinition().getType() : null;`
`248`	`248`	`}`