feat: introduce new configuration properties for bigquery source

v2/googlecloud-to-neo4j/src/main/java/com/google/cloud/teleport/v2/neo4j/model/helpers/SqlQuerySpec.java → v2/googlecloud-to-neo4j/src/main/java/com/google/cloud/teleport/v2/neo4j/model/helpers/BigQuerySpec.java

@@ -19,16 +19,25 @@
  * Convenience object for invoking SQL query as well as providing descriptions for read and cast
  * phase of transform.
  */
-public class SqlQuerySpec {
+public class BigQuerySpec {
 
   private final String readDescription;
   private final String castDescription;
   private final String sql;
+  private final String queryTempProject;
+  private final String queryTempDataset;
 
-  public SqlQuerySpec(String readDescription, String castDescription, String sql) {
+  public BigQuerySpec(
+      String readDescription,
+      String castDescription,
+      String sql,
+      String queryTempProject,
+      String queryTempDataset) {
     this.readDescription = readDescription;
     this.castDescription = castDescription;
     this.sql = sql;
+    this.queryTempProject = queryTempProject;
+    this.queryTempDataset = queryTempDataset;
   }
 
   public String getReadDescription() {
@@ -43,29 +52,50 @@
     return sql;
   }
 
-  public static class SqlQuerySpecBuilder {
+  public String getQueryTempProject() {
+    return queryTempProject;
+  }
+
+  public String getQueryTempDataset() {
+    return queryTempDataset;
+  }
+
+  public static class BigQuerySpecBuilder {
 
     private String readDescription;
     private String castDescription;
     private String sql;
+    private String queryTempProject;
+    private String queryTempDataset;
 
-    public SqlQuerySpecBuilder readDescription(String readDescription) {
+    public BigQuerySpecBuilder readDescription(String readDescription) {
       this.readDescription = readDescription;
       return this;
     }
 
-    public SqlQuerySpecBuilder castDescription(String castDescription) {
+    public BigQuerySpecBuilder castDescription(String castDescription) {
       this.castDescription = castDescription;
       return this;
     }
 
-    public SqlQuerySpecBuilder sql(String sql) {
+    public BigQuerySpecBuilder sql(String sql) {
       this.sql = sql;
       return this;
     }
 
-    public SqlQuerySpec build() {
-      return new SqlQuerySpec(readDescription, castDescription, sql);
+    public BigQuerySpecBuilder queryTempProject(String queryTempProject) {
+      this.queryTempProject = queryTempProject;
+      return this;
+    }
+
+    public BigQuerySpecBuilder queryTempDataset(String queryTempDataset) {
+      this.queryTempDataset = queryTempDataset;
+      return this;
+    }
+
+    public BigQuerySpec build() {
+      return new BigQuerySpec(
+          readDescription, castDescription, sql, queryTempProject, queryTempDataset);
     }
   }
 }

v2/googlecloud-to-neo4j/src/main/java/com/google/cloud/teleport/v2/neo4j/model/sources/BigQuerySource.java

@@ -22,11 +22,19 @@
 
   private final String name;
   private final String query;
+  private final String queryTempProject;
+  private final String queryTempDataset;
 
   public BigQuerySource(String name, String query) {
+    this(name, query, null, null);
+  }
 
+  public BigQuerySource(
+      String name, String query, String queryTempProject, String queryTempDataset) {
     this.name = name;
     this.query = query;
+    this.queryTempProject = queryTempProject;
+    this.queryTempDataset = queryTempDataset;
   }
 
   @Override
@@ -43,6 +51,14 @@
     return query;
   }
 
+  public String getQueryTempProject() {
+    return queryTempProject;
+  }
+
+  public String getQueryTempDataset() {
+    return queryTempDataset;
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) {
@@ -62,6 +78,19 @@
 
   @Override
   public String toString() {
-    return "BigQuerySource{" + "name='" + name + '\'' + ", query='" + query + '\'' + '}';
+    return "BigQuerySource{"
+        + "name='"
+        + name
+        + '\''
+        + ", query='"
+        + query
+        + '\''
+        + ", queryTempProject='"
+        + queryTempProject
+        + '\''
+        + ", queryTempDataset='"
+        + queryTempDataset
+        + '\''
+        + '}';
   }
 }

v2/googlecloud-to-neo4j/src/main/java/com/google/cloud/teleport/v2/neo4j/model/sources/BigQuerySourceProvider.java

@@ -15,18 +15,23 @@
  */
 package com.google.cloud.teleport.v2.neo4j.model.sources;
 
+import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.ObjectNode;
+import java.util.Optional;
 import org.neo4j.importer.v1.sources.SourceProvider;
 
 public class BigQuerySourceProvider implements SourceProvider<BigQuerySource> {
-
   @Override
   public String supportedType() {
     return "bigquery";
   }
 
   @Override
   public BigQuerySource provide(ObjectNode node) {
-    return new BigQuerySource(node.get("name").textValue(), node.get("query").textValue());
+    return new BigQuerySource(
+        node.get("name").textValue(),
+        node.get("query").textValue(),
+        Optional.ofNullable(node.get("query_temp_project")).map(JsonNode::textValue).orElse(null),
+        Optional.ofNullable(node.get("query_temp_dataset")).map(JsonNode::textValue).orElse(null));
   }
 }

v2/googlecloud-to-neo4j/src/main/java/com/google/cloud/teleport/v2/neo4j/model/validation/BigQuerySourceProjectDatasetValidator.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright (C) 2024 Google LLC
+ *
+ * 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 com.google.cloud.teleport.v2.neo4j.model.validation;
+
+import com.google.cloud.teleport.v2.neo4j.model.sources.BigQuerySource;
+import java.util.LinkedHashSet;
+import java.util.Set;
+import org.neo4j.importer.v1.sources.Source;
+import org.neo4j.importer.v1.validation.SpecificationValidationResult;
+import org.neo4j.importer.v1.validation.SpecificationValidator;
+
+public class BigQuerySourceProjectDatasetValidator implements SpecificationValidator {
+
+  private static final String ERROR_CODE = "DFBQ-001";
+  private final Set<String> paths = new LinkedHashSet<>();
+
+  @Override
+  public void visitSource(int index, Source source) {
+    if (!(source instanceof BigQuerySource)) {
+      return;
+    }
+
+    var queryTempProject = ((BigQuerySource) source).getQueryTempProject();
+    var queryTempDataset = ((BigQuerySource) source).getQueryTempDataset();
+
+    if (queryTempProject != null && queryTempDataset == null) {
+      paths.add(String.format("$.sources[%d]", index));
+    }
+  }
+
+  @Override
+  public boolean report(SpecificationValidationResult.Builder builder) {
+    paths.forEach(
+        path ->
+            builder.addError(
+                path,
+                ERROR_CODE,
+                String.format(
+                    "%s query_temp_project is provided, but query_temp_dataset is missing", path)));
+    return paths.isEmpty();
+  }
+}

v2/googlecloud-to-neo4j/src/main/java/com/google/cloud/teleport/v2/neo4j/providers/bigquery/BigQueryImpl.java

@@ -15,8 +15,8 @@
  */
 package com.google.cloud.teleport.v2.neo4j.providers.bigquery;
 
-import com.google.cloud.teleport.v2.neo4j.model.helpers.SqlQuerySpec;
-import com.google.cloud.teleport.v2.neo4j.model.helpers.SqlQuerySpec.SqlQuerySpecBuilder;
+import com.google.cloud.teleport.v2.neo4j.model.helpers.BigQuerySpec;
+import com.google.cloud.teleport.v2.neo4j.model.helpers.BigQuerySpec.BigQuerySpecBuilder;
 import com.google.cloud.teleport.v2.neo4j.model.helpers.TargetQuerySpec;
 import com.google.cloud.teleport.v2.neo4j.model.helpers.TargetSequence;
 import com.google.cloud.teleport.v2.neo4j.model.job.OptionsParams;
@@ -33,11 +33,9 @@
 
 /** Provider implementation for reading and writing BigQuery. */
 public class BigQueryImpl implements Provider {
-
   private static final Logger LOG = LoggerFactory.getLogger(BigQueryImpl.class);
   private final BigQuerySource source;
   private final TargetSequence targetSequence;
-
   private OptionsParams optionsParams;
 
   public BigQueryImpl(BigQuerySource source, TargetSequence targetSequence) {
@@ -67,15 +65,15 @@
 
   @Override
   public PTransform<PBegin, PCollection<Row>> queryMetadata() {
-    return new BqQueryToRow(getMetadataQueryBeamSpec(source));
+    return new BqQueryToRow(getMetadataQueryBeamSpec());
   }
 
   /**
    * Returns zero rows metadata query based on original query.
    *
    * @return helper object includes metadata and SQL
    */
-  public SqlQuerySpec getMetadataQueryBeamSpec(BigQuerySource source) {
+  public BigQuerySpec getMetadataQueryBeamSpec() {
 
     String baseQuery = source.getQuery();
 
@@ -85,10 +83,12 @@
     String zeroRowSql = "SELECT * FROM (" + baseQuery + ") LIMIT 0";
     LOG.info("Reading BQ metadata with query: {}", zeroRowSql);
 
-    return new SqlQuerySpecBuilder()
+    return new BigQuerySpecBuilder()
         .readDescription("Read from BQ " + source.getName())
         .castDescription("Cast to BeamRow " + source.getName())
         .sql(zeroRowSql)
+        .queryTempProject(source.getQueryTempProject())
+        .queryTempDataset(source.getQueryTempDataset())
         .build();
   }
 
@@ -97,11 +97,13 @@
    *
    * @return helper object includes metadata and SQL
    */
-  private SqlQuerySpec getSourceQueryBeamSpec() {
-    return new SqlQuerySpecBuilder()
+  private BigQuerySpec getSourceQueryBeamSpec() {
+    return new BigQuerySpecBuilder()
         .castDescription("Cast to BeamRow " + source.getName())
         .readDescription("Read from BQ " + source.getName())
         .sql(source.getQuery())
+        .queryTempProject(source.getQueryTempProject())
+        .queryTempDataset(source.getQueryTempDataset())
         .build();
   }
 
@@ -110,20 +112,22 @@
    *
    * @return helper object includes metadata and SQL
    */
-  private SqlQuerySpec getTargetQueryBeamSpec(TargetQuerySpec spec) {
+  private BigQuerySpec getTargetQueryBeamSpec(TargetQuerySpec spec) {
     var sourceFields = ModelUtils.getBeamFieldSet(spec.getSourceBeamSchema());
     var target = spec.getTarget();
     var startNodeTarget = spec.getStartNodeTarget();
     var endNodeTarget = spec.getEndNodeTarget();
     String sql =
         ModelUtils.getTargetSql(
             target, startNodeTarget, endNodeTarget, sourceFields, true, source.getQuery());
-    return new SqlQuerySpecBuilder()
+    return new BigQuerySpecBuilder()
         .readDescription(
             targetSequence.getSequenceNumber(target) + ": Read from BQ " + target.getName())
         .castDescription(
             targetSequence.getSequenceNumber(target) + ": Cast to BeamRow " + target.getName())
         .sql(sql)
+        .queryTempProject(source.getQueryTempProject())
+        .queryTempDataset(source.getQueryTempDataset())
         .build();
   }
 }

v2/googlecloud-to-neo4j/src/main/java/com/google/cloud/teleport/v2/neo4j/providers/bigquery/BqQueryToRow.java

@@ -16,7 +16,7 @@
 package com.google.cloud.teleport.v2.neo4j.providers.bigquery;
 
 import com.google.api.services.bigquery.model.TableRow;
-import com.google.cloud.teleport.v2.neo4j.model.helpers.SqlQuerySpec;
+import com.google.cloud.teleport.v2.neo4j.model.helpers.BigQuerySpec;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO;
 import org.apache.beam.sdk.schemas.Schema;
@@ -34,25 +34,33 @@
 public class BqQueryToRow extends PTransform<PBegin, PCollection<Row>> {
 
   private static final Logger LOG = LoggerFactory.getLogger(BqQueryToRow.class);
-  private final SqlQuerySpec bqQuerySpec;
+  private final BigQuerySpec bqQuerySpec;
 
-  public BqQueryToRow(SqlQuerySpec bqQuerySpec) {
+  public BqQueryToRow(BigQuerySpec bqQuerySpec) {
     this.bqQuerySpec = bqQuerySpec;
   }
 
   @Override
   public PCollection<Row> expand(PBegin input) {
-
     String rewrittenSql = this.bqQuerySpec.getSql();
     LOG.info("Reading BQ with query: {}", rewrittenSql);
 
-    PCollection<TableRow> sourceRows =
-        input.apply(
-            bqQuerySpec.getReadDescription(),
-            BigQueryIO.readTableRowsWithSchema()
-                .fromQuery(rewrittenSql)
-                .usingStandardSql()
-                .withTemplateCompatibility());
+    var read =
+        BigQueryIO.readTableRowsWithSchema()
+            .fromQuery(rewrittenSql)
+            .usingStandardSql()
+            .withTemplateCompatibility();
+
+    var queryTempProject = this.bqQuerySpec.getQueryTempProject();
+    var queryTempDataset = this.bqQuerySpec.getQueryTempDataset();
+
+    if (queryTempProject != null && queryTempDataset != null) {
+      read = read.withQueryTempProjectAndDataset(queryTempProject, queryTempDataset);
+    } else if (queryTempDataset != null) {
+      read = read.withQueryTempDataset(queryTempDataset);
+    }
+
+    PCollection<TableRow> sourceRows = input.apply(bqQuerySpec.getReadDescription(), read);
 
     Schema beamSchema = sourceRows.getSchema();
     Coder<Row> rowCoder = SchemaCoder.of(beamSchema);

v2/googlecloud-to-neo4j/src/main/resources/META-INF/services/org.neo4j.importer.v1.validation.SpecificationValidator

@@ -1,4 +1,5 @@
 # keep this sorted
+com.google.cloud.teleport.v2.neo4j.model.validation.BigQuerySourceProjectDatasetValidator
 com.google.cloud.teleport.v2.neo4j.model.validation.DuplicateAggregateFieldNameValidator
 com.google.cloud.teleport.v2.neo4j.model.validation.DuplicateTextHeaderValidator
 com.google.cloud.teleport.v2.neo4j.model.validation.InlineSourceDataValidator