shahar1 commented on code in PR #56324:
URL: https://github.com/apache/airflow/pull/56324#discussion_r2724021857
##########
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:
Here I'd avoid the usage of the `try..except` block at all - it obscures the
original exception for no good reason.
##########
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:
Same comment as before regarding `try..except`
##########
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:
Nice! It seems like a useful regex expression which we could utilize to any
input which is a a GCP project ID across the project. Maybe you could extract
it to a common utility function? (you don't have to implement the validation
for the other inputs)
##########
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:
```suggestion
"""Validate the quota project ID format
```
It doesn't validate existence (not sure that it should though)
##########
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()
Review Comment:
The logic related to quota project should be relocated to the inner method
`self.get_credentials_and_project_id()`, so we could also handle the caching
logic correctly by applying quota project configuration before caching
credentials.
--
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]
