Copilot commented on code in PR #56324:
URL: https://github.com/apache/airflow/pull/56324#discussion_r2662090249
##########
providers/google/src/airflow/providers/google/common/hooks/base_google.py:
##########
@@ -265,17 +266,34 @@ def get_ui_field_behaviour(cls) -> dict[str, Any]:
return {
"hidden_fields": ["host", "schema", "login", "password", "port",
"extra"],
"relabeling": {},
+ "placeholders": {
+ "quota_project_id": "Project ID to use for API quota and
billing (optional)",
+ },
+ "tooltips": {
+ "quota_project_id": "Specify a different project for
quota/billing purposes. Useful when using shared service accounts.",
+ },
}
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, the project ID from the GCP connection is used.
Review Comment:
The docstring says "If None, the project ID from the GCP connection is
used." However, this is not entirely accurate. When quota_project_id is None,
no quota project is configured at all - it doesn't default to using the
connection's project ID for quota purposes. The default behavior is to not set
a quota project, which means Google Cloud will use the default project
associated with the credentials.
Consider clarifying the docstring to: "If None, no separate quota project is
configured and the default behavior of the credentials is used."
```suggestion
If None, no separate quota project is configured and the default
behavior of the
credentials is used.
```
##########
providers/google/tests/unit/google/common/hooks/test_base_google.py:
##########
@@ -541,6 +541,78 @@ 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("google.auth._default._load_credentials_from_file")
+ @mock.patch("os.environ", {CREDENTIALS: "/test/key-path"})
+ def test_quota_project_id_from_connection(self, mock_load_creds):
+ """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_load_creds.return_value = (mock_creds, "test-project")
+
+ # Mock connection with quota_project_id in extras
+ uri =
f"google-cloud-platform://?extras={json.dumps({'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("google.auth._default._load_credentials_from_file")
+ @mock.patch("os.environ", {CREDENTIALS: "/test/key-path"})
+ def test_quota_project_id_param_overrides_connection(self,
mock_load_creds):
+ """Test that quota project ID from param overrides connection value."""
+ mock_creds = mock.MagicMock()
+ mock_creds.with_quota_project.return_value = mock_creds
+ mock_load_creds.return_value = (mock_creds, "test-project")
+
+ # Connection with quota_project_id in extras
+ conn_quota = "connection-quota-project"
+ uri =
f"google-cloud-platform://?extras={json.dumps({'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")
+
+ def test_quota_project_invalid_format(self):
+ """Test validation of quota project ID format."""
+ invalid_ids = [
+ "UPPERCASE", # Must be lowercase
+ "special@chars", # Invalid characters
+ "a", # Too short
Review Comment:
The test includes "a" (single character) as an invalid ID. While this
correctly tests that short IDs are rejected, the regex pattern requires at
minimum 6 characters total (1 start + 4 middle + 1 end). A single "a" would
fail the pattern, but not because it's "too short" in the sense being tested -
it fails because it doesn't match the structural requirements.
Consider adding test cases that better represent the boundary conditions:
- A 5-character valid structure like "ab-cd" (too short, should fail)
- A 6-character valid structure like "ab-cde" (minimum valid, should pass -
though this might not be a valid test case if testing invalid IDs)
This would make the test more explicit about what aspects of validation are
being tested.
```suggestion
"ab-cd", # Structurally valid but too short
```
##########
providers/google/src/airflow/providers/google/common/hooks/base_google.py:
##########
@@ -347,9 +377,51 @@ def get_credentials_and_project_id(self) ->
tuple[Credentials, str | None]:
return credentials, project_id
+ def _validate_quota_project(self, quota_project: str | None) -> None:
+ """Validate the quota project ID format and existence.
Review Comment:
The docstring says "Validate the quota project ID format and existence" but
the method only validates the format through a regex check. It does not
validate the existence of the project (i.e., it doesn't check if the project
actually exists in Google Cloud). This could be misleading for developers.
Change "and existence" to just "format" in the docstring, or implement
actual existence checking if that's the intended behavior.
```suggestion
"""Validate the quota project ID format.
```
##########
providers/google/src/airflow/providers/google/common/hooks/base_google.py:
##########
@@ -347,9 +377,51 @@ def get_credentials_and_project_id(self) ->
tuple[Credentials, str | None]:
return credentials, project_id
+ def _validate_quota_project(self, quota_project: str | None) -> None:
+ """Validate the quota project ID format and existence.
+
+ :param quota_project: The quota project ID to validate
+ :raises AirflowException: If the quota project ID is invalid
+ """
+ if quota_project is not None:
+ if not isinstance(quota_project, str):
+ raise AirflowException(f"quota_project_id must be a string,
got {type(quota_project)}")
+ if not quota_project.strip():
+ raise AirflowException("quota_project_id cannot be empty")
+ # Check for valid GCP project ID format
+ if not re.match(r"^[a-z][-a-z0-9]{4,28}[a-z0-9]$", quota_project):
Review Comment:
The regex pattern "^[a-z][-a-z0-9]{4,28}[a-z0-9]$" enforces project IDs to
be 6-30 characters (1 starting letter + 4-28 middle characters + 1 ending
character). However, according to Google Cloud documentation, project IDs must
be between 6 and 30 characters, which means the minimum middle section should
allow for 4 characters (when you have 1 starting + 4 middle + 1 ending = 6
total).
But the pattern allows {4,28} for the middle section, which would mean:
- Minimum: 1 + 4 + 1 = 6 characters ✓
- Maximum: 1 + 28 + 1 = 30 characters ✓
This appears correct, but consider adding a comment explaining this
calculation to make the regex more maintainable and easier to understand for
future developers.
##########
providers/google/tests/unit/google/common/hooks/test_base_google.py:
##########
@@ -541,6 +541,78 @@ 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("google.auth._default._load_credentials_from_file")
+ @mock.patch("os.environ", {CREDENTIALS: "/test/key-path"})
+ def test_quota_project_id_from_connection(self, mock_load_creds):
+ """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_load_creds.return_value = (mock_creds, "test-project")
+
+ # Mock connection with quota_project_id in extras
+ uri =
f"google-cloud-platform://?extras={json.dumps({'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("google.auth._default._load_credentials_from_file")
+ @mock.patch("os.environ", {CREDENTIALS: "/test/key-path"})
+ def test_quota_project_id_param_overrides_connection(self,
mock_load_creds):
+ """Test that quota project ID from param overrides connection value."""
+ mock_creds = mock.MagicMock()
+ mock_creds.with_quota_project.return_value = mock_creds
+ mock_load_creds.return_value = (mock_creds, "test-project")
+
+ # Connection with quota_project_id in extras
+ conn_quota = "connection-quota-project"
+ uri =
f"google-cloud-platform://?extras={json.dumps({'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")
+
+ def test_quota_project_invalid_format(self):
+ """Test validation of quota project ID format."""
+ invalid_ids = [
+ "UPPERCASE", # Must be lowercase
+ "special@chars", # Invalid characters
+ "a", # Too short
+ "a" * 31, # Too long
Review Comment:
The test includes "a" * 31 (31 characters) as an invalid ID in the test
list. However, this value would not match the regex pattern
"^[a-z][-a-z0-9]{4,28}[a-z0-9]$" because it's a string of all 'a' characters
with no hyphens or digits, meaning it doesn't satisfy the pattern structure
which requires specific character types in different positions.
The test should use a more realistic invalid ID that would be 31 characters
but still matches the pattern structure, such as "a" + "-" * 28 + "a" (which
would be 30 characters with valid structure) or better yet "a" + "b" * 29 (31
lowercase letters). This would properly test the length validation aspect of
the regex.
```suggestion
"a" + "b" * 30, # Too long (31 characters, structurally valid
but exceeds max length)
```
##########
providers/google/docs/connections/gcp.rst:
##########
@@ -309,3 +328,95 @@ 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 operators or hooks (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"
+ }
+ }'
+
+Operator / Hook parameter
+^^^^^^^^^^^^^^^^^^^^^^^^^
+
+You can also pass the quota project directly when creating an operator or
hook. This takes
+precedence over the connection extras:
+
+.. code-block:: python
+
+ from airflow.providers.google.cloud.operators.bigquery import
BigQueryExecuteQueryOperator
+
+ task = BigQueryExecuteQueryOperator(
+ task_id='execute_query',
+ sql='SELECT * FROM `my_project.dataset.table`',
+ quota_project_id='your-billing-project-id'
+ )
+
+ # Or when creating a hook
+ 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 an
operator/hook
+parameter, the operator/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
+- 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.
+
+Examples
+~~~~~~~~
+
+**Examples**
+
+See the example DAG:
``airflow/providers/google/cloud/example_dags/example_quota_project.py``
Review Comment:
The file path reference
"airflow/providers/google/cloud/example_dags/example_quota_project.py" is
incorrect. Based on the actual file structure, the example is located at
"tests/system/providers/google/common/example_quota_project_system.py". This
will cause confusion for users trying to find the example.
Update the path to the correct location where the example file exists.
```suggestion
See the example DAG:
``tests/system/providers/google/common/example_quota_project_system.py``
```
##########
providers/google/src/airflow/providers/google/common/hooks/base_google.py:
##########
@@ -338,6 +356,18 @@ def get_credentials_and_project_id(self) ->
tuple[Credentials, str | None]:
idp_extra_params_dict=idp_extra_params_dict,
)
+ # Support for quota project ID
+ quota_project = self.quota_project_id or
self._get_field("quota_project_id")
+ if quota_project:
+ self._validate_quota_project(quota_project)
+ try:
+ credentials = credentials.with_quota_project(quota_project)
+ except Exception as e:
+ raise AirflowException(
+ f"Failed to configure quota project '{quota_project}':
{str(e)}. "
+ "Ensure the service account has permission to use the
quota project."
+ ) from e
Review Comment:
The get_credentials_and_project_id() method does not check for anonymous
credentials before applying the quota project (lines 359-369), while
get_credentials() does have this check (line 408). This inconsistency means
that if anonymous credentials are used, get_credentials_and_project_id() might
attempt to call with_quota_project() on anonymous credentials, which could fail
or behave unexpectedly.
For consistency and correctness, add a check for is_anonymous before
applying the quota project in get_credentials_and_project_id(), similar to the
check in get_credentials(). Alternatively, if the duplicate logic in
get_credentials() is removed as suggested in another comment, ensure the check
for anonymous credentials is retained in get_credentials_and_project_id().
##########
providers/google/src/airflow/providers/google/common/hooks/base_google.py:
##########
@@ -347,9 +377,51 @@ def get_credentials_and_project_id(self) ->
tuple[Credentials, str | None]:
return credentials, project_id
+ def _validate_quota_project(self, quota_project: str | None) -> None:
+ """Validate the quota project ID format and existence.
+
+ :param quota_project: The quota project ID to validate
+ :raises AirflowException: If the quota project ID is invalid
+ """
+ if quota_project is not None:
+ if not isinstance(quota_project, str):
+ raise AirflowException(f"quota_project_id must be a string,
got {type(quota_project)}")
+ if not quota_project.strip():
+ raise AirflowException("quota_project_id cannot be empty")
+ # Check for valid GCP project ID format
+ if not re.match(r"^[a-z][-a-z0-9]{4,28}[a-z0-9]$", quota_project):
+ raise AirflowException(
+ f"Invalid quota_project_id '{quota_project}'. "
+ "Project IDs must start with a lowercase letter and can
contain "
+ "only lowercase letters, digits, and hyphens."
Review Comment:
The validation error message states "Project IDs must start with a lowercase
letter and can contain only lowercase letters, digits, and hyphens." However,
this doesn't mention the length requirements (6-30 characters) which are also
enforced by the regex. This could leave users confused about why their project
ID is rejected if it has valid characters but incorrect length.
Update the error message to include the length requirement: "Project IDs
must be 6-30 characters long, start with a lowercase letter, and can contain
only lowercase letters, digits, and hyphens."
```suggestion
"Project IDs must be 6-30 characters long, start with a
lowercase letter, and can "
"contain only lowercase letters, digits, and hyphens."
```
##########
providers/google/docs/connections/gcp.rst:
##########
@@ -121,6 +121,25 @@ 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, the default project from the connection is used.
Must be a
+valid GCP project ID (lowercase letters, digits, hyphens, 6–30 characters,
starting
+with a letter).
Review Comment:
The documentation states "If not specified, the default project from the
connection is used." This is misleading. When quota_project_id is not
specified, no quota project is set on the credentials, which means Google Cloud
uses the default behavior of the credentials (typically the project that owns
the service account). It does not automatically use the "project" field from
the connection configuration as the quota project.
Clarify the documentation to explain that when not specified, no separate
quota project is configured and Google Cloud's default behavior applies.
```suggestion
project. If not specified, no separate quota project is configured on the
credentials and
Google Cloud's default behavior applies (typically the project that owns the
service
account/credentials); it is not automatically derived from the connection's
``Project Id`` field.
```
##########
providers/google/src/airflow/providers/google/common/hooks/base_google.py:
##########
@@ -347,9 +377,51 @@ def get_credentials_and_project_id(self) ->
tuple[Credentials, str | None]:
return credentials, project_id
+ def _validate_quota_project(self, quota_project: str | None) -> None:
+ """Validate the quota project ID format and existence.
+
+ :param quota_project: The quota project ID to validate
+ :raises AirflowException: If the quota project ID is invalid
+ """
+ if quota_project is not None:
+ if not isinstance(quota_project, str):
+ raise AirflowException(f"quota_project_id must be a string,
got {type(quota_project)}")
+ if not quota_project.strip():
+ raise AirflowException("quota_project_id cannot be empty")
+ # Check for valid GCP project ID format
+ if not re.match(r"^[a-z][-a-z0-9]{4,28}[a-z0-9]$", quota_project):
+ raise AirflowException(
+ f"Invalid quota_project_id '{quota_project}'. "
+ "Project IDs must 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
+ :raises AirflowException: If credentials cannot be configured with
quota project
+ """
credentials, _ = self.get_credentials_and_project_id()
+
+ # If credentials are anonymous, skip quota project logic
+ is_anonymous = self._get_field("is_anonymous")
+ 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 AirflowException(
+ 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."
+ )
+ try:
+ credentials = credentials.with_quota_project(quota_project)
+ except Exception as e:
+ raise AirflowException(
+ f"Failed to configure quota project '{quota_project}':
{str(e)}"
+ ) from e
+
Review Comment:
This method duplicates the quota project logic that already exists in
get_credentials_and_project_id() (lines 359-369). The get_credentials() method
calls get_credentials_and_project_id() which already applies the quota project
to credentials. This means the quota project configuration is attempted twice,
which is redundant and could lead to confusion during debugging.
Consider removing the duplicate quota project logic from get_credentials()
since get_credentials_and_project_id() already handles this configuration. The
get_credentials() method can simply return the credentials obtained from
get_credentials_and_project_id().
```suggestion
:return: Google Cloud credentials object
"""
credentials, _ = self.get_credentials_and_project_id()
```
##########
tests/system/providers/google/common/example_quota_project_system.py:
##########
@@ -0,0 +1,112 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you 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.
+"""
+Example Airflow DAG that demonstrates Google Cloud quota project functionality.
+This DAG shows how to use quota project settings with Google Cloud connections
+and operators.
+"""
+from __future__ import annotations
+
+import os
+from datetime import datetime
+
+from airflow import DAG
+from airflow.operators.python import PythonOperator
+from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
+from airflow.utils.email import send_email
+
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID", "default")
+DAG_ID = "example_quota_project"
+
+def test_quota_project(gcp_conn_id: str = "google_cloud_default",
quota_project_id: str | None = None, **context):
+ """Test quota project configuration with Google Cloud credentials."""
+ hook = BigQueryHook(gcp_conn_id=gcp_conn_id,
quota_project_id=quota_project_id)
+ credentials, project_id = hook.get_credentials_and_project_id()
+ quota_project = getattr(credentials, "_quota_project_id", None)
+
+ hook.log.info("Testing Google Cloud quota project configuration:")
+ hook.log.info("Credentials project ID: %s", project_id)
+ hook.log.info("Quota project ID: %s", quota_project)
+
+ return {
+ "project_id": project_id,
+ "quota_project_id": quota_project,
+ }
+
+with DAG(
+ DAG_ID,
+ schedule="@once",
+ start_date=datetime(2021, 1, 1),
+ catchup=False,
+ tags=["example", "quota-project"],
+) as dag:
+ # [START howto_quota_project_from_connection]
+ # This task uses quota project from connection if configured
+ test_conn_quota = PythonOperator(
+ task_id="test_connection_quota",
+ python_callable=test_quota_project,
+ op_kwargs={"gcp_conn_id": "google_cloud_default"}, # Assumes
quota_project_id in connection extras
+ )
+ # [END howto_quota_project_from_connection]
+
+ # [START howto_quota_project_from_param]
+ # This task explicitly sets quota project
+ test_param_quota = PythonOperator(
+ task_id="test_parameter_quota",
+ python_callable=test_quota_project,
+ op_kwargs={
+ "gcp_conn_id": "google_cloud_default",
+ "quota_project_id": "explicit-quota-project",
+ },
+ )
+ # [END howto_quota_project_from_param]
+
+ # Task dependencies
+ test_conn_quota >> test_param_quota
+
+# [START howto_quota_project_email_on_failure]
+def _handle_quota_failure(context):
+ """Send an email if quota project test fails."""
+ dag_id = context["task_instance"].dag_id
+ task_id = context["task_instance"].task_id
+ logs_url = context["task_instance"].log_url
+
+ send_email(
+ to="[email protected]",
+ subject=f"[Airflow] Quota Project Test Failed - {dag_id}.{task_id}",
+ html_content=f"""
+ Quota project test failed.<br/>
+ DAG: {dag_id}<br/>
+ Task: {task_id}<br/>
+ Logs: <a href="{logs_url}">Link</a>
+ """,
+ )
+# [END howto_quota_project_email_on_failure]
+
+for task in dag.tasks:
+ task.on_failure_callback = _handle_quota_failure
Review Comment:
The send_email import and the _handle_quota_failure function with email
functionality are included in this system test example. However, this
introduces unnecessary complexity and dependencies (email configuration) for
what should be a simple demonstration of quota project functionality. System
test examples should focus on demonstrating the feature being tested, not
generic failure handling patterns.
Consider removing the send_email import and the email failure callback
section (lines 81-101) to keep the example focused on the quota project feature
itself. If failure handling examples are needed, they should be in separate,
dedicated example files.
##########
providers/google/tests/unit/google/common/hooks/test_base_google.py:
##########
@@ -541,6 +541,78 @@ 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("google.auth._default._load_credentials_from_file")
+ @mock.patch("os.environ", {CREDENTIALS: "/test/key-path"})
+ def test_quota_project_id_from_connection(self, mock_load_creds):
+ """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_load_creds.return_value = (mock_creds, "test-project")
+
+ # Mock connection with quota_project_id in extras
+ uri =
f"google-cloud-platform://?extras={json.dumps({'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("google.auth._default._load_credentials_from_file")
+ @mock.patch("os.environ", {CREDENTIALS: "/test/key-path"})
+ def test_quota_project_id_param_overrides_connection(self,
mock_load_creds):
+ """Test that quota project ID from param overrides connection value."""
+ mock_creds = mock.MagicMock()
+ mock_creds.with_quota_project.return_value = mock_creds
+ mock_load_creds.return_value = (mock_creds, "test-project")
+
+ # Connection with quota_project_id in extras
+ conn_quota = "connection-quota-project"
+ uri =
f"google-cloud-platform://?extras={json.dumps({'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")
+
+ def test_quota_project_invalid_format(self):
+ """Test validation of quota project ID format."""
+ invalid_ids = [
+ "UPPERCASE", # Must be lowercase
+ "special@chars", # Invalid characters
+ "a", # Too short
+ "a" * 31, # Too long
+ "1starts-with-number", # Must start with letter
+ "", # Empty string
+ None, # None value
Review Comment:
The validation logic checks if quota_project is not None at line 386, but
then the test list includes None as an invalid value (line 592). This creates
an inconsistency: the validation method will not raise an exception for None
(it returns early), but the test expects an exception.
If None should be treated as invalid, the validation method should raise an
exception for None values. If None is acceptable (which makes more sense since
it means "no quota project"), then None should be removed from the invalid_ids
test list.
```suggestion
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
