feat: Add support for quota project

Author: nickcaballero

README.md

@@ -197,14 +197,17 @@ Follow [this document](https://cloud.google.com/dataproc/docs/concepts/component
 * This includes the table's schema (column names and data types), partitioning information, constraints etc.
   It doesn't contain the actual table data.
 * SQL Command for Catalog Table Creation
-  ```java
+  ```sql
     CREATE TABLE sample_catalog_table
-    (name STRING) // Schema Details
+    (name STRING) -- Schema Details
     WITH
-    ('connector' = 'bigquery',
-    'project' = '<bigquery_project_name>',
-    'dataset' = '<bigquery_dataset_name>',
-    'table' = '<bigquery_table_name>');
+    (
+        'connector' = 'bigquery',
+        'project' = '<bigquery_project_name>',
+        'dataset' = '<bigquery_dataset_name>',
+        'table' = '<bigquery_table_name>',
+        'quota-project-id' = '<gcp_project_name>' -- Optional, defaults to credentials or project
+    );
   ```
 
 
@@ -296,7 +299,7 @@ wait for the job to complete.
 
 ### Sink Configurations
 
-The connector supports a number of options to configure the source.
+The connector supports a number of options to configure the sink.
 
 | Property                                     | Data Type              | Description                                                                                                                                                            |
 |----------------------------------------------|------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
@@ -437,18 +440,18 @@ sourceTable = sourceTable.select($("*"));
 
 The connector supports a number of options to configure the source.
 
-| Property                                     | Data Type          | Description                                                                                                                                                                                                                                                                                                                                               |
-|----------------------------------------------|--------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| `projectId`                                  | String             | Google Cloud Project ID of the table. This config is required, and assumes no default value.                                                                                                                                                                                                                                                              |
-| `dataset`                                    | String             | Dataset containing the table. This config is required, and assumes no default value.                                                                                                                                                                                                                                                                      |
-| `table`                                      | String             | BigQuery table name (not the full ID). This config is required, and assumes no default value.                                                                                                                                                                                                                                                             |
-| `credentialsOptions`                         | CredentialsOptions | Google credentials for connecting to BigQuery. This config is optional, and default behavior is to use the `GOOGLE_APPLICATION_CREDENTIALS` environment variable.<br/>**Note**: The query bounded source only uses default application credentials.                                                                                                       |
-| `columnNames`                                | List&lt;String&gt; | Columns to project from the table. If unspecified, all columns are fetched.                                                                                                                                                                                                                                                                               |
-| `limit`                                      | Integer            | Maximum number of rows to read from source table **per task slot**. If unspecified, all rows are fetched.                                                                                                                                                                                                                                                                          |
-| `maxRecordsPerSplitFetch`                    | Integer            | Maximum number of records to read from a split once Flink requests fetch. If unspecified, the default value used is 10000. <br/>**Note**: Configuring this number too high may cause memory pressure in the task manager, depending on the BigQuery record's size and total rows on the stream.                                                           |
-| `maxStreamCount`                             | Integer            | Maximum read streams to open during a read session. BigQuery can return a lower number of streams than specified based on internal optimizations. If unspecified, this config is not set and BigQuery has complete control over the number of read streams created.                                                                                       |
-| `rowRestriction`                             | String             | BigQuery SQL query for row filter pushdown. If unspecified, all rows are fetched.                                                                                                                                                                                                                                                                         |
-| `snapshotTimeInMillis`                       | Long               | Time (in milliseconds since epoch) for the BigQuery table snapshot to read. If unspecified, the latest snapshot is read.                                                                                                                                                                                                                                  |
+| Property                  | Data Type          | Description                                                                                                                                                                                                                                                                                     |
+|---------------------------|--------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| `projectId`               | String             | Google Cloud Project ID of the table. This config is required, and assumes no default value.                                                                                                                                                                                                    |
+| `dataset`                 | String             | Dataset containing the table. This config is required, and assumes no default value.                                                                                                                                                                                                            |
+| `table`                   | String             | BigQuery table name (not the full ID). This config is required, and assumes no default value.                                                                                                                                                                                                   |
+| `credentialsOptions`      | CredentialsOptions | Google credentials for connecting to BigQuery. This config is optional, and default behavior is to use the `GOOGLE_APPLICATION_CREDENTIALS` environment variable.<br/>**Note**: The query bounded source only uses default application credentials.                                             |
+| `columnNames`             | List&lt;String&gt; | Columns to project from the table. If unspecified, all columns are fetched.                                                                                                                                                                                                                     |
+| `limit`                   | Integer            | Maximum number of rows to read from source table **per task slot**. If unspecified, all rows are fetched.                                                                                                                                                                                       |
+| `maxRecordsPerSplitFetch` | Integer            | Maximum number of records to read from a split once Flink requests fetch. If unspecified, the default value used is 10000. <br/>**Note**: Configuring this number too high may cause memory pressure in the task manager, depending on the BigQuery record's size and total rows on the stream. |
+| `maxStreamCount`          | Integer            | Maximum read streams to open during a read session. BigQuery can return a lower number of streams than specified based on internal optimizations. If unspecified, this config is not set and BigQuery has complete control over the number of read streams created.                             |
+| `rowRestriction`          | String             | BigQuery SQL query for row filter pushdown. If unspecified, all rows are fetched.                                                                                                                                                                                                               |
+| `snapshotTimeInMillis`    | Long               | Time (in milliseconds since epoch) for the BigQuery table snapshot to read. If unspecified, the latest snapshot is read.                                                                                                                                                                        |
 
 
 ### Datatypes
@@ -568,6 +571,7 @@ to provide it:
 [here](https://cloud.google.com/docs/authentication/client-libraries).
 - In case the environment variable cannot be changed, the credentials file can be configured as a connector option. The 
 file should reside on the same path on all the nodes of the cluster.
+- The quota project can be set in the connector options or through the credentials.
 
 ### How to fix classloader error in Flink application?
 

flink-1.17-connector-bigquery/flink-connector-bigquery/src/main/java/com/google/cloud/flink/bigquery/table/BigQueryDynamicTableFactory.java

@@ -72,6 +72,7 @@ public Set<ConfigOption<?>> optionalOptions() {
         additionalOptions.add(BigQueryConnectorOptions.CREDENTIALS_ACCESS_TOKEN);
         additionalOptions.add(BigQueryConnectorOptions.CREDENTIALS_FILE);
         additionalOptions.add(BigQueryConnectorOptions.CREDENTIALS_KEY);
+        additionalOptions.add(BigQueryConnectorOptions.QUOTA_PROJECT_ID);
         additionalOptions.add(BigQueryConnectorOptions.TEST_MODE);
         additionalOptions.add(BigQueryConnectorOptions.DELIVERY_GUARANTEE);
         additionalOptions.add(BigQueryConnectorOptions.SINK_PARALLELISM);
@@ -100,6 +101,7 @@ public Set<ConfigOption<?>> forwardOptions() {
         forwardOptions.add(BigQueryConnectorOptions.CREDENTIALS_ACCESS_TOKEN);
         forwardOptions.add(BigQueryConnectorOptions.CREDENTIALS_FILE);
         forwardOptions.add(BigQueryConnectorOptions.CREDENTIALS_KEY);
+        forwardOptions.add(BigQueryConnectorOptions.QUOTA_PROJECT_ID);
         forwardOptions.add(BigQueryConnectorOptions.DELIVERY_GUARANTEE);
         forwardOptions.add(BigQueryConnectorOptions.SINK_PARALLELISM);
         forwardOptions.add(BigQueryConnectorOptions.ENABLE_TABLE_CREATION);

flink-1.17-connector-bigquery/flink-connector-bigquery/src/main/java/com/google/cloud/flink/bigquery/table/config/BigQueryConnectorOptions.java

@@ -59,6 +59,15 @@ private BigQueryConnectorOptions() {}
                     .noDefaultValue()
                     .withDescription("Specifies the BigQuery table name.");
 
+    /**
+     * [OPTIONAL] The quota project ID to use when connecting.
+     */
+    public static final ConfigOption<String> QUOTA_PROJECT_ID =
+            ConfigOptions.key("quota-project-id")
+                    .stringType()
+                    .noDefaultValue()
+                    .withDescription("Specifies the BigQuery quota project to use.");
+
     /**
      * [OPTIONAL, Read Configuration] Integer value indicating the maximum number of rows/records to
      * be read from source. <br>

flink-1.17-connector-bigquery/flink-connector-bigquery/src/main/java/com/google/cloud/flink/bigquery/table/config/BigQueryTableConfigurationProvider.java

@@ -126,6 +126,8 @@ public BigQueryConnectOptions translateBigQueryConnectOptions() {
                                         config.get(BigQueryConnectorOptions.CREDENTIALS_FILE))
                                 .setCredentialsKey(
                                         config.get(BigQueryConnectorOptions.CREDENTIALS_KEY))
+                                .setQuotaProjectId(
+                                        config.get(BigQueryConnectorOptions.QUOTA_PROJECT_ID))
                                 .build())
                 .build();
     }

flink-connector-bigquery-common/src/main/java/com/google/cloud/flink/bigquery/common/config/CredentialsOptions.java

@@ -42,6 +42,9 @@ public abstract class CredentialsOptions implements Serializable {
     @Nullable
     public abstract String getAccessToken();
 
+    @Nullable
+    public abstract String getQuotaProjectId();
+
     /**
      * Returns the Google Credentials created given the provided configuration.
      *
@@ -60,6 +63,7 @@ public final int hashCode() {
         hash = 61 * hash + Objects.hashCode(getCredentialsFile());
         hash = 61 * hash + Objects.hashCode(getCredentialsKey());
         hash = 61 * hash + Objects.hashCode(getAccessToken());
+        hash = 61 * hash + Objects.hashCode(getQuotaProjectId());
         return hash;
     }
 
@@ -77,6 +81,7 @@ public final boolean equals(Object obj) {
         final CredentialsOptions other = (CredentialsOptions) obj;
         return Objects.equals(this.getCredentialsFile(), other.getCredentialsFile())
                 && Objects.equals(this.getCredentialsKey(), other.getCredentialsKey())
+                && Objects.equals(this.getQuotaProjectId(), other.getQuotaProjectId())
                 && Objects.equals(this.getAccessToken(), other.getAccessToken());
     }
 
@@ -117,6 +122,14 @@ public abstract static class Builder {
          */
         public abstract Builder setAccessToken(String credentialsToken);
 
+        /**
+         * Sets the BigQuery quota project ID to use.
+         *
+         * @param quotaProjectId The BigQuery project ID
+         * @return this builder's instance
+         */
+        public abstract Builder setQuotaProjectId(String quotaProjectId);
+
         /**
          * Builds a fully initialized {@link CredentialsOptions} instance.
          *

flink-connector-bigquery-common/src/main/java/com/google/cloud/flink/bigquery/services/BigQueryServicesImpl.java

@@ -29,6 +29,7 @@
 import com.google.api.services.bigquery.Bigquery;
 import com.google.api.services.bigquery.model.Table;
 import com.google.api.services.bigquery.model.TableSchema;
+import com.google.auth.Credentials;
 import com.google.cloud.bigquery.BigQuery;
 import com.google.cloud.bigquery.BigQueryOptions;
 import com.google.cloud.bigquery.Dataset;
@@ -135,12 +136,16 @@ public void cancel() {
     /** Implementation of a BigQuery read client wrapper. */
     public static class StorageReadClientImpl implements StorageReadClient {
         private final BigQueryReadClient client;
+        private final String quotaProjectId;
 
         private StorageReadClientImpl(CredentialsOptions options) throws IOException {
+            Credentials credentials = options.getCredentials();
+            quotaProjectId = BigQueryUtils.getQuotaProjectId(options, credentials);
+
             BigQueryReadSettings.Builder settingsBuilder =
                     BigQueryReadSettings.newBuilder()
-                            .setCredentialsProvider(
-                                    FixedCredentialsProvider.create(options.getCredentials()))
+                            .setCredentialsProvider(FixedCredentialsProvider.create(credentials))
+                            .setQuotaProjectId(quotaProjectId)
                             .setHeaderProvider(USER_AGENT_HEADER_PROVIDER)
                             .setTransportChannelProvider(
                                     BigQueryReadSettings.defaultGrpcTransportProviderBuilder()
@@ -177,7 +182,9 @@ private StorageReadClientImpl(CredentialsOptions options) throws IOException {
 
         @Override
         public ReadSession createReadSession(CreateReadSessionRequest request) {
-            return client.createReadSession(request);
+            CreateReadSessionRequest updatedRequest =
+                    BigQueryUtils.updateWithQuotaProject(request, quotaProjectId);
+            return client.createReadSession(updatedRequest);
         }
 
         @Override
@@ -196,10 +203,12 @@ public static class StorageWriteClientImpl implements StorageWriteClient {
         private final BigQueryWriteClient client;
 
         private StorageWriteClientImpl(CredentialsOptions options) throws IOException {
+            Credentials credentials = options.getCredentials();
+            String quotaProjectId = BigQueryUtils.getQuotaProjectId(options, credentials);
             BigQueryWriteSettings.Builder settingsBuilder =
                     BigQueryWriteSettings.newBuilder()
-                            .setCredentialsProvider(
-                                    FixedCredentialsProvider.create(options.getCredentials()))
+                            .setCredentialsProvider(FixedCredentialsProvider.create(credentials))
+                            .setQuotaProjectId(quotaProjectId)
                             .setHeaderProvider(USER_AGENT_HEADER_PROVIDER)
                             .setTransportChannelProvider(
                                     BigQueryReadSettings.defaultGrpcTransportProviderBuilder()
@@ -320,9 +329,12 @@ public static class QueryDataClientImpl implements QueryDataClient {
         private final Bigquery bigquery;
 
         public QueryDataClientImpl(CredentialsOptions options) {
+            Credentials credentials = options.getCredentials();
+            String quotaProjectId = BigQueryUtils.getQuotaProjectId(options, credentials);
             bigQuery =
                     BigQueryOptions.newBuilder()
-                            .setCredentials(options.getCredentials())
+                            .setCredentials(credentials)
+                            .setQuotaProjectId(quotaProjectId)
                             .setHeaderProvider(USER_AGENT_HEADER_PROVIDER)
                             .build()
                             .getService();