Fix BigQuery configuration with parent-project-id

.github/workflows/ci.yml

@@ -639,11 +639,15 @@ jobs:
         env:
           BIGQUERY_CREDENTIALS_KEY: ${{ secrets.BIGQUERY_CREDENTIALS_KEY }}
           GCP_STORAGE_BUCKET: ${{ vars.GCP_STORAGE_BUCKET }}
+          BIGQUERY_TESTING_PROJECT_ID: ${{ vars.BIGQUERY_TESTING_PROJECT_ID }}
+          BIGQUERY_TESTING_PARENT_PROJECT_ID: ${{ vars.BIGQUERY_TESTING_PARENT_PROJECT_ID }}
         if: matrix.modules == 'plugin/trino-bigquery' && !contains(matrix.profile, 'cloud-tests-2') && (env.CI_SKIP_SECRETS_PRESENCE_CHECKS != '' || env.BIGQUERY_CREDENTIALS_KEY != '')
         run: |
           $MAVEN test ${MAVEN_TEST} -pl :trino-bigquery -Pcloud-tests-1 \
             -Dbigquery.credentials-key="${BIGQUERY_CREDENTIALS_KEY}" \
-            -Dtesting.gcp-storage-bucket="${GCP_STORAGE_BUCKET}"
+            -Dtesting.gcp-storage-bucket="${GCP_STORAGE_BUCKET}" \
+            -Dtesting.bigquery-project-id="${BIGQUERY_TESTING_PROJECT_ID}" \
+            -Dtesting.bigquery-parent-project-id="${BIGQUERY_TESTING_PARENT_PROJECT_ID}"
       - name: Cloud BigQuery Smoke Tests
         id: tests-bq-smoke
         env:

plugin/trino-bigquery/pom.xml

@@ -537,6 +537,7 @@
                                 <exclude>**/TestBigQueryCaseInsensitiveMapping.java</exclude>
                                 <exclude>**/TestBigQuery*FailureRecoveryTest.java</exclude>
                                 <exclude>**/TestBigQueryWithProxyTest.java</exclude>
+                                <exclude>**/TestBigQueryParentProjectId.java</exclude>
                             </excludes>
                         </configuration>
                     </plugin>
@@ -557,6 +558,7 @@
                         <configuration>
                             <includes>
                                 <include>**/TestBigQueryAvroConnectorTest.java</include>
+                                <include>**/TestBigQueryParentProjectId.java</include>
                             </includes>
                         </configuration>
                     </plugin>

plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryClient.java

@@ -259,7 +259,7 @@ public TableInfo getCachedTable(Duration viewExpiration, TableInfo remoteTableId
      */
     public String getParentProjectId()
     {
-        return bigQuery.getOptions().getProjectId();
+        return Optional.ofNullable(bigQuery.getOptions().getQuotaProjectId()).orElse(bigQuery.getOptions().getProjectId());
     }
 
     /**

plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/CredentialsOptionsConfigurer.java

@@ -31,22 +31,26 @@ public class CredentialsOptionsConfigurer
         implements BigQueryOptionsConfigurer
 {
     private final BigQueryCredentialsSupplier credentialsSupplier;
-    private final Optional<String> parentProjectId;
+    private final Optional<String> configProjectId;
+    private final Optional<String> configParentProjectId;
 
     @Inject
     public CredentialsOptionsConfigurer(BigQueryConfig bigQueryConfig, BigQueryCredentialsSupplier credentialsSupplier)
     {
-        this.parentProjectId = requireNonNull(bigQueryConfig, "bigQueryConfig is null").getParentProjectId();
+        this.configProjectId = bigQueryConfig.getProjectId();
+        this.configParentProjectId = bigQueryConfig.getParentProjectId();
         this.credentialsSupplier = requireNonNull(credentialsSupplier, "credentialsSupplier is null");
     }
 
     @Override
     public BigQueryOptions.Builder configure(BigQueryOptions.Builder builder, ConnectorSession session)
     {
         Optional<Credentials> credentials = credentialsSupplier.getCredentials(session);
-        String billingProjectId = calculateBillingProjectId(parentProjectId, credentials);
+        String projectId = resolveProjectId(configProjectId, credentials);
         credentials.ifPresent(builder::setCredentials);
-        builder.setProjectId(billingProjectId);
+        builder.setProjectId(projectId);
+        // Quota project id is different name for parent project id, both indicates project used for quota and billing purposes.
+        configParentProjectId.ifPresent(builder::setQuotaProjectId);
         return builder;
     }
 
@@ -69,10 +73,10 @@ public BigQueryWriteSettings.Builder configure(BigQueryWriteSettings.Builder bui
 
     // Note that at this point the config has been validated, which means that option 2 or option 3 will always be valid
     @VisibleForTesting
-    static String calculateBillingProjectId(Optional<String> configParentProjectId, Optional<Credentials> credentials)
+    static String resolveProjectId(Optional<String> configProjectId, Optional<Credentials> credentials)
     {
         // 1. Get from configuration
-        return configParentProjectId
+        return configProjectId
                 // 2. Get from the provided credentials, but only ServiceAccountCredentials contains the project id.
                 // All other credentials types (User, AppEngine, GCE, CloudShell, etc.) take it from the environment
                 .orElseGet(() -> credentials

plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryParentProjectId.java

@@ -0,0 +1,60 @@
+/*
+ * 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.trino.plugin.bigquery;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import io.trino.testing.AbstractTestQueryFramework;
+import io.trino.testing.QueryRunner;
+import org.junit.jupiter.api.Test;
+
+import static io.trino.testing.TestingProperties.requiredNonEmptySystemProperty;
+import static io.trino.tpch.TpchTable.NATION;
+import static java.lang.String.format;
+import static org.assertj.core.api.Assertions.assertThat;
+
+class TestBigQueryParentProjectId
+        extends AbstractTestQueryFramework
+{
+    private final String testingProjectId;
+    private final String testingParentProjectId;
+
+    TestBigQueryParentProjectId()
+    {
+        testingProjectId = requiredNonEmptySystemProperty("testing.bigquery-project-id");
+        testingParentProjectId = requiredNonEmptySystemProperty("testing.bigquery-parent-project-id");
+    }
+
+    @Override
+    protected QueryRunner createQueryRunner()
+            throws Exception
+    {
+        return BigQueryQueryRunner.builder()
+                .setConnectorProperties(ImmutableMap.<String, String>builder()
+                        .put("bigquery.project-id", testingProjectId)
+                        .put("bigquery.parent-project-id", testingParentProjectId)
+                        .buildOrThrow())
+                .setInitialTables(ImmutableList.of(NATION))
+                .build();
+    }
+
+    @Test
+    void testQueriesWithParentProjectId()
+    {
+        assertThat(computeScalar("SELECT name FROM bigquery.tpch.nation WHERE nationkey = 0")).isEqualTo("ALGERIA");
+        assertThat(computeScalar("SELECT * FROM TABLE(bigquery.system.query(query => 'SELECT name FROM tpch.nation WHERE nationkey = 0'))")).isEqualTo("ALGERIA");
+        assertThat(computeScalar(format("SELECT * FROM TABLE(bigquery.system.query(query => 'SELECT name FROM %s.tpch.nation WHERE nationkey = 0'))", testingProjectId)))
+                .isEqualTo("ALGERIA");
+    }
+}

plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestCredentialsOptionsConfigurer.java

@@ -20,31 +20,31 @@
 import java.io.IOException;
 import java.util.Optional;
 
-import static io.trino.plugin.bigquery.CredentialsOptionsConfigurer.calculateBillingProjectId;
+import static io.trino.plugin.bigquery.CredentialsOptionsConfigurer.resolveProjectId;
 import static org.assertj.core.api.Assertions.assertThat;
 
 public class TestCredentialsOptionsConfigurer
 {
     @Test
     public void testConfigurationOnly()
     {
-        String projectId = calculateBillingProjectId(Optional.of("pid"), Optional.empty());
+        String projectId = resolveProjectId(Optional.of("pid"), Optional.empty());
         assertThat(projectId).isEqualTo("pid");
     }
 
     @Test
     public void testCredentialsOnly()
             throws Exception
     {
-        String projectId = calculateBillingProjectId(Optional.empty(), credentials());
+        String projectId = resolveProjectId(Optional.empty(), credentials());
         assertThat(projectId).isEqualTo("presto-bq-credentials-test");
     }
 
     @Test
     public void testBothConfigurationAndCredentials()
             throws Exception
     {
-        String projectId = calculateBillingProjectId(Optional.of("pid"), credentials());
+        String projectId = resolveProjectId(Optional.of("pid"), credentials());
         assertThat(projectId).isEqualTo("pid");
     }