Add quota project id support to Google cloud base hook

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.

providers/google/src/airflow/providers/google/common/hooks/base_google.py

@@ -25,6 +25,7 @@
 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 @@ def get_connection_form_widgets(cls) -> dict[str, Any]:
             "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 @@ def __init__(
         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 @@ def get_credentials_and_project_id(self) -> tuple[Credentials, str | None]:
             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 @@ def get_credentials_and_project_id(self) -> tuple[Credentials, str | None]:
 
         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
 

providers/google/tests/unit/google/common/hooks/test_base_google.py

@@ -541,6 +541,59 @@ def test_get_credentials_and_project_id_with_default_auth_and_overridden_project
         )
         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",