Lee-W commented on code in PR #36849:
URL: https://github.com/apache/airflow/pull/36849#discussion_r1458183618
##########
airflow/providers/google/cloud/hooks/bigquery.py:
##########
@@ -3513,11 +3518,11 @@ async def get_table_client(
access to the specified project.
:param session: aiohttp ClientSession
"""
- with await self.service_file_as_context() as file:
- return Table_async(
- dataset_name=dataset,
- table_name=table_id,
- project=project_id,
- service_file=file,
Review Comment:
Will this affected how `service_files` was called?
##########
tests/providers/google/common/hooks/test_base_google.py:
##########
@@ -874,3 +875,95 @@ def
test_should_fallback_when_empty_string_in_env_var(self):
instance = hook.GoogleBaseHook(gcp_conn_id="google_cloud_default")
assert isinstance(instance.num_retries, int)
assert 5 == instance.num_retries
+
+
+class TestCredentialsToken:
+ @pytest.mark.asyncio
+ async def test_get_project(self):
+ mock_credentials =
mock.MagicMock(spec=google.auth.compute_engine.Credentials)
+ token = hook.CredentialsToken(mock_credentials, project=PROJECT_ID)
+ assert await token.get_project() == PROJECT_ID
+
+ @pytest.mark.asyncio
+ async def test_get(self):
+ mock_credentials =
mock.MagicMock(spec=google.auth.compute_engine.Credentials)
+ mock_credentials.token = "ACCESS_TOKEN"
+ token = hook.CredentialsToken(mock_credentials, project=PROJECT_ID)
+ assert await token.get() == "ACCESS_TOKEN"
+ mock_credentials.refresh.assert_called_once()
+
+ @pytest.mark.asyncio
+ @mock.patch(MODULE_NAME + ".get_credentials_and_project_id",
return_value=("CREDENTIALS", "PROJECT_ID"))
Review Comment:
```suggestion
@mock.patch(f"{MODULE_NAME}.get_credentials_and_project_id",
return_value=("CREDENTIALS", "PROJECT_ID"))
```
##########
airflow/providers/google/common/hooks/base_google.py:
##########
@@ -623,6 +628,51 @@ def test_connection(self):
return status, message
+class CredentialsToken(Token):
+ """A token implementation which makes Google credentials objects
accessible to [gcloud-aio](https://talkiq.github.io/gcloud-aio/) clients.
+
+ This class allows us to create token instances from credentials objects
and thus supports a variety of use cases for Google
+ credentials in Airflow (i.e. impersonation chain). By relying on a
existing credentials object we leverage functionality provided by the
GoogleBaseHook
+ for generating credentials objects.
+ """
+
+ def __init__(
+ self,
+ credentials: Credentials,
+ *,
+ project: str | None = None,
+ session: ClientSession | None = None,
+ ) -> None:
+ super().__init__(session=cast(Session, session))
+ self.credentials = credentials
+ self.project = project
+
+ @classmethod
+ async def from_hook(
+ cls,
+ hook: GoogleBaseHook,
+ *,
+ session: ClientSession | None = None,
+ ) -> CredentialsToken:
+ credentials, project = hook.get_credentials_and_project_id()
+ return cls(
+ credentials=credentials,
+ project=project,
+ session=session,
+ )
+
+ async def get_project(self) -> str | None:
+ return self.project
Review Comment:
May I know what is this method for? I do not see it's being used elsewhere
--
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: commits-unsubscr...@airflow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
