This is an automated email from the ASF dual-hosted git repository.
Lee-W pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 1e14a3d6bcd Add quota project id support to Google cloud base hook
(#56324)
1e14a3d6bcd is described below
commit 1e14a3d6bcd10fa1c735615f4d2c961051ec7c03
Author: Ankurdeewan <[email protected]>
AuthorDate: Tue Apr 21 12:03:03 2026 +0530
Add quota project id support to Google cloud base hook (#56324)
Co-authored-by: Wei Lee <[email protected]>
Co-authored-by: Shahar Epstein <[email protected]>
---
providers/google/docs/connections/gcp.rst | 84 +++++++++++++++++++++-
.../providers/google/common/hooks/base_google.py | 81 ++++++++++++++++++++-
.../unit/google/common/hooks/test_base_google.py | 53 ++++++++++++++
3 files changed, 216 insertions(+), 2 deletions(-)
diff --git a/providers/google/docs/connections/gcp.rst
b/providers/google/docs/connections/gcp.rst
index 776f53a50a0..aedf9af1e23 100644
--- a/providers/google/docs/connections/gcp.rst
+++ b/providers/google/docs/connections/gcp.rst
@@ -16,7 +16,6 @@
under the License.
-
.. _howto/connection:google_cloud_platform:
Google Cloud Connection
@@ -121,6 +120,18 @@ Scopes (comma separated)
<https://developers.google.com/identity/protocols/googlescopes>`_ to
authenticate with.
+
+Quota Project ID (optional)
+ The Google Cloud project ID to use for API quota and billing purposes.
This is useful
+ when using a shared service account but want to attribute quota/billing to
a different
+ project. If not specified, no separate quota project is configured on the
credentials
+ and Google Cloud's default behavior applies. Must be a valid GCP project
ID (lowercase
+ letters, digits, hyphens, 6-30 characters, starting with a letter).
+
+ .. note:: If using anonymous credentials, quota project logic is ignored.
+
+ .. warning:: Ensure the service account has permission to use the
specified quota project; an invalid or unauthorized project ID will result in
an API error.
+
Number of Retries
Integer, number of times to retry with randomized
exponential backoff. If all retries fail, the
:class:`googleapiclient.errors.HttpError`
@@ -309,3 +320,74 @@ Note that as domain-wide delegation is currently supported
by most of the Google
* All of Google Cloud operators and hooks.
* Firebase hooks.
* All transfer operators that involve Google cloud in different providers, for
example:
:class:`airflow.providers.amazon.aws.transfers.gcs_to_s3.GCSToS3Operator`.
+
+
+Quota Project Support
+---------------------
+
+Airflow's Google Cloud providers support specifying a "quota project" (a
billing project) for
+API calls. That lets API usage be billed to a different Google Cloud project
than the one that
+owns the service account. This is useful for organizations that share service
accounts but
+centralize billing in specific projects.
+
+Usage
+~~~~~
+
+There are two ways to set a quota project in Airflow:
+
+- Via connection extras (recommended for environment-wide defaults).
+- Directly on supported hook instances (recommended when a single task must
bill to a different project).
+
+Connection extras
+^^^^^^^^^^^^^^^^^
+
+Add the quota project ID to the Google Cloud connection extras. For example:
+
+.. code-block:: json
+
+ {
+ "quota_project_id": "your-billing-project-id"
+ }
+
+You can set this via the Airflow UI, the Connections REST API, or an
environment variable, for
+example:
+
+.. code-block:: bash
+
+ export AIRFLOW_CONN_GOOGLE_CLOUD_DEFAULT='{
+ "conn_type": "google_cloud_platform",
+ "extra": {
+ "quota_project_id": "your-billing-project-id"
+ }
+ }'
+
+Hook parameter
+^^^^^^^^^^^^^^
+
+You can also pass the quota project directly when creating a hook. This takes
+precedence over the connection extras:
+
+.. code-block:: python
+
+ from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
+
+ hook = BigQueryHook(quota_project_id="your-billing-project-id")
+
+Priority
+^^^^^^^^
+
+If a quota project is provided both in the connection extras and as a hook
+parameter, the hook parameter wins.
+
+Compatibility
+^^^^^^^^^^^^^
+
+This setting works with Google Cloud services that support the quota project
mechanism (the
+``x-goog-user-project`` header), for example BigQuery, Cloud Storage, Dataflow,
+and other Google Cloud APIs that accept quota project headers.
+
+Impact
+^^^^^^
+
+Using a quota project affects where API usage is billed, which quotas are
applied, and how
+usage is reported for monitoring and auditing.
diff --git
a/providers/google/src/airflow/providers/google/common/hooks/base_google.py
b/providers/google/src/airflow/providers/google/common/hooks/base_google.py
index 0a7a0805a4f..579300d18bb 100644
--- a/providers/google/src/airflow/providers/google/common/hooks/base_google.py
+++ b/providers/google/src/airflow/providers/google/common/hooks/base_google.py
@@ -25,6 +25,7 @@ import functools
import json
import logging
import os
+import re
import tempfile
from collections.abc import Callable, Generator, Sequence
from contextlib import ExitStack, contextmanager
@@ -123,6 +124,26 @@ def is_refresh_credentials_exception(exception: Exception)
-> bool:
return False
+_GCP_PROJECT_ID_PATTERN = re.compile(r"^[a-z][a-z0-9\-]{4,28}[a-z0-9]$")
+
+
+def is_valid_gcp_project_id(project_id: str) -> bool:
+ """
+ Validate a Google Cloud Project ID format.
+
+ A valid project ID must:
+
+ - Be 6 to 30 characters long
+ - Start with a lowercase letter
+ - Contain only lowercase letters, digits, and hyphens
+ - Not end with a hyphen
+
+ :param project_id: The project ID string to validate.
+ :return: True if the project ID is valid, False otherwise.
+ """
+ return bool(_GCP_PROJECT_ID_PATTERN.match(project_id))
+
+
class retry_if_temporary_quota(tenacity.retry_if_exception):
"""Retries if there was an exception for exceeding the temporary quote
limit."""
@@ -206,6 +227,10 @@ class GoogleBaseHook(BaseHook):
If set as a sequence, the identities from the list must grant
Service Account Token Creator IAM role to the directly preceding
identity, with first
account from the list granting this role to the originating account.
+ :param quota_project_id: Optional Google Cloud project ID to use for
billing and
+ quota purposes. If set, API usage will be charged to this project
instead of
+ the project associated with the credentials. If not set, the
connection's
+ ``quota_project_id`` field is used as a fallback.
"""
conn_name_attr = "gcp_conn_id"
@@ -261,6 +286,7 @@ class GoogleBaseHook(BaseHook):
"is_anonymous": BooleanField(
lazy_gettext("Anonymous credentials (ignores all other
settings)"), default=False
),
+ "quota_project_id": StringField(lazy_gettext("Quota Project ID"),
widget=BS3TextFieldWidget()),
}
@classmethod
@@ -275,11 +301,27 @@ class GoogleBaseHook(BaseHook):
self,
gcp_conn_id: str = "google_cloud_default",
impersonation_chain: str | Sequence[str] | None = None,
+ *,
+ quota_project_id: str | None = None,
**kwargs,
) -> None:
+ """
+ Initialize the Google Cloud Base Hook.
+
+ :param gcp_conn_id: The connection ID to use when fetching connection
info.
+ :param impersonation_chain: Optional service account to impersonate
using short-term
+ credentials.
+ :param quota_project_id: Optional Project ID to use for quota/billing
purposes.
+ If None, no separate quota project is configured and the default
behavior of the
+ credentials is used.
+ :param kwargs: Additional arguments to pass to parent constructor.
+ """
super().__init__(**kwargs)
self.gcp_conn_id = gcp_conn_id
self.impersonation_chain = impersonation_chain
+ if quota_project_id is not None:
+ self._validate_quota_project(quota_project_id)
+ self.quota_project_id = quota_project_id
self.extras: dict = self.get_connection(self.gcp_conn_id).extra_dejson
self._cached_credentials: Credentials | None = None
self._cached_project_id: str | None = None
@@ -342,6 +384,19 @@ class GoogleBaseHook(BaseHook):
idp_extra_params_dict=idp_extra_params_dict,
)
+ # Apply quota project before caching credentials
+ quota_project = self.quota_project_id or
self._get_field("quota_project_id")
+ if quota_project and not is_anonymous:
+ self._validate_quota_project(quota_project)
+ if not hasattr(credentials, "with_quota_project"):
+ raise ValueError(
+ f"Credentials of type {type(credentials).__name__} do not
support "
+ "quota project configuration. Please use a different
authentication method "
+ "or remove the quota_project_id setting."
+ )
+ credentials = credentials.with_quota_project(quota_project)
+
+ # Override project_id if set in extras
overridden_project_id = self._get_field("project")
if overridden_project_id:
project_id = overridden_project_id
@@ -351,8 +406,32 @@ class GoogleBaseHook(BaseHook):
return credentials, project_id
+ def _validate_quota_project(self, quota_project: str) -> None:
+ """
+ Validate the quota Project ID format.
+
+ :param quota_project: The quota Project ID to validate
+ :raises TypeError: If the quota Project ID is not a string
+ :raises ValueError: If the quota Project ID is empty or does not match
the expected format
+ """
+ if not isinstance(quota_project, str):
+ raise TypeError(f"quota_project_id must be a string, got
{type(quota_project)}")
+ if not quota_project.strip():
+ raise ValueError("quota_project_id cannot be empty")
+ # Check for valid GCP Project ID format
+ if not is_valid_gcp_project_id(quota_project):
+ raise ValueError(
+ f"Invalid quota_project_id '{quota_project}'. "
+ "Project IDs must be 6-30 characters long, start with a
lowercase letter, "
+ "and can contain only lowercase letters, digits, and hyphens."
+ )
+
def get_credentials(self) -> Credentials:
- """Return the Credentials object for Google API."""
+ """
+ Return the Credentials object for Google API.
+
+ :return: Google Cloud credentials object
+ """
credentials, _ = self.get_credentials_and_project_id()
return credentials
diff --git
a/providers/google/tests/unit/google/common/hooks/test_base_google.py
b/providers/google/tests/unit/google/common/hooks/test_base_google.py
index 46adde9b8f9..4a662143333 100644
--- a/providers/google/tests/unit/google/common/hooks/test_base_google.py
+++ b/providers/google/tests/unit/google/common/hooks/test_base_google.py
@@ -541,6 +541,59 @@ class TestGoogleBaseHook:
)
assert result == ("CREDENTIALS", "SECOND_PROJECT_ID")
+ def test_quota_project_id_init(self):
+ """Test that quota project ID is properly initialized."""
+ hook = GoogleBaseHook(gcp_conn_id="google_cloud_default",
quota_project_id="test-quota-project")
+ assert hook.quota_project_id == "test-quota-project"
+
+ @mock.patch(MODULE_NAME + ".get_credentials_and_project_id")
+ def test_quota_project_id_from_connection(self,
mock_get_creds_and_proj_id):
+ """Test that quota project ID from connection is applied to
credentials."""
+ mock_creds = mock.MagicMock()
+ mock_creds.with_quota_project.return_value = mock_creds
+ mock_get_creds_and_proj_id.return_value = (mock_creds, "test-project")
+
+ # Mock connection with quota_project_id in extras
+ uri = "google-cloud-platform://?quota_project_id=test-quota-project"
+ with mock.patch.dict("os.environ",
{"AIRFLOW_CONN_GOOGLE_CLOUD_DEFAULT": uri}):
+ hook = GoogleBaseHook(gcp_conn_id="google_cloud_default")
+ creds, _ = hook.get_credentials_and_project_id()
+
mock_creds.with_quota_project.assert_called_once_with("test-quota-project")
+
+ @mock.patch(MODULE_NAME + ".get_credentials_and_project_id")
+ def test_quota_project_id_param_overrides_connection(self,
mock_get_creds_and_proj_id):
+ """Test that quota project ID from param overrides connection value."""
+ mock_creds = mock.MagicMock()
+ mock_creds.with_quota_project.return_value = mock_creds
+ mock_get_creds_and_proj_id.return_value = (mock_creds, "test-project")
+
+ # Connection with quota_project_id in extras
+ conn_quota = "connection-quota-project"
+ uri = f"google-cloud-platform://?quota_project_id={conn_quota}"
+
+ with mock.patch.dict("os.environ",
{"AIRFLOW_CONN_GOOGLE_CLOUD_DEFAULT": uri}):
+ hook = GoogleBaseHook(gcp_conn_id="google_cloud_default",
quota_project_id="test-quota-project")
+ creds, _ = hook.get_credentials_and_project_id()
+
+ # Should use param quota project, not connection quota project
+
mock_creds.with_quota_project.assert_called_once_with("test-quota-project")
+
+ @pytest.mark.parametrize(
+ "invalid_id",
+ [
+ "UPPERCASE", # Must be lowercase
+ "special@chars", # Invalid characters
+ "ab-cd", # Too short
+ "a" + "b" * 30, # Too long
+ "1starts-with-number", # Must start with letter
+ "", # Empty string
+ ],
+ )
+ def test_quota_project_invalid_format(self, invalid_id):
+ """Test validation of quota project ID format."""
+ with pytest.raises((TypeError, ValueError)):
+ GoogleBaseHook(quota_project_id=invalid_id)
+
def
test_get_credentials_and_project_id_with_mutually_exclusive_configuration(self):
self.instance.extras = {
"project": "PROJECT_ID",
