add impersonation_chain support when calling Bigquery Operators in deferrable mode by Lee-W · Pull Request #36341 · apache/airflow

Lee-W · 2023-12-21T07:11:11Z

Add missing impersonation_chain argument when calling trigger in Bigquery Operators add impersonation_chain support when calling Bigquery Operators in deferrable mode #36341
Simulate BaseToken in gcloud-aio with ImpersonationToken and generate desired token

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

nathadfield · 2023-12-21T08:24:11Z

@Lee-W Thanks for taking a look at this but I don't think this will fully solve the problem. If you follow the trigger call through, it leads to GoogleBaseHook.provide_gcp_credential_file_as_context() but there isn't any logic to deal with the impersonation chain.

https://github.com/apache/airflow/blob/main/airflow/providers/google/common/hooks/base_google.py#L499-L532

Lee-W · 2023-12-21T13:49:55Z

I dig it a bit more. It seems the lib we're using does not support Impersonated credentials. talkiq/gcloud-aio#421 Might need to dig a bit t see how we add it

m1racoli · 2023-12-21T14:14:06Z

So we might first need to look into impersonation support in gcloud-aio?

nathadfield · 2023-12-21T14:39:07Z

It's not the most active library to rely upon for a key part of GCP authentication and deferrable task capabilities.

Lee-W · 2023-12-22T00:10:34Z

@Lee-W Thanks for taking a look at this but I don't think this will fully solve the problem. If you follow the trigger call through, it leads to GoogleBaseHook.provide_gcp_credential_file_as_context() but there isn't any logic to deal with the impersonation chain.

https://github.com/apache/airflow/blob/main/airflow/providers/google/common/hooks/base_google.py#L499-L532

Thanks for reminding me! Yep, I've done some test afterward and find it did not work.

So we might first need to look into impersonation support in gcloud-aio?

Yep, I think we could take a look at how gcloud implement this feature and see what we can do

It's not the most active library to rely upon for a key part of GCP authentication and deferrable task capabilities.

I tried to did the investigation but did not find a more active one 😞 aiogoogle does not seem to fit our need either

Lee-W · 2023-12-22T11:28:30Z

after digging a bit deeper, I doubt we can use a similar way to use this impersonation feature https://github.com/googleapis/google-auth-library-python/blob/776d634ac6d989b224f8dbfb11d166cb3025a342/google/auth/_default_async.py#L29

melugoyal · 2023-12-22T22:24:22Z

airflow/providers/google/cloud/hooks/bigquery.py

impersonation_chain can have multiple service accounts in the chain. you'll have to handle that by doing something similar to _get_target_principal_and_delegates:

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

Lines 273 to 285 in e2393ee

target_principal, delegates = _get_target_principal_and_delegates(self.impersonation_chain)

credentials, project_id = get_credentials_and_project_id(

key_path=key_path,

keyfile_dict=keyfile_dict_json,

credential_config_file=credential_config_file,

key_secret_name=key_secret_name,

key_secret_project_id=key_secret_project_id,

scopes=self.scopes,

delegate_to=self.delegate_to,

target_principal=target_principal,

delegates=delegates,

)

Thanks for reminding me! just addressed

m1racoli · 2023-12-23T13:04:57Z

FYI We might get better impersonation support in gcloud-aio soon. :)
talkiq/gcloud-aio#665

Lee-W · 2023-12-24T08:54:58Z

FYI We might get better impersonation support in gcloud-aio soon. :) talkiq/gcloud-aio#665

This looks great! Should we use the current PR as a quick fix? Or is it ok to wait for the release of the next gcloud-aio?

eladkal · 2023-12-25T05:55:13Z

Lets wait for a fix in upstream first

Lee-W · 2023-12-25T09:29:09Z

Lets wait for a fix in upstream first

Sure 🙂

…perators

…perator async

phanikumv

Blocking it till the upstream changes are ready

…operators

m1racoli · 2024-01-17T08:34:11Z

airflow/providers/google/cloud/hooks/bigquery.py

+        return project
+
+    async def get(self) -> str | None:
+        creds, _ = google.auth.default()


Instead of getting default credentials, we should retrieve the credentials for a given connection ID here, no?

Furthermore, I was looking at other Google hooks which work fine asyncronously with impersonation chain and noticed that for example the DataprocAsyncHook (subclass of GoogleBaseHook) just calls self.get_credentials().

I suppose this itself (creating a credentials object) is non-blocking (unless you consider file IO blocking) until we want to actually generate a token for the given credentials, or? So I wonder if we maybe just have to implement credentials -> token asynchronously? That's probably what Google's async clients do under the hood, but unfortunately we're here also because there is no async client for BigQuery to begin with.

I would like to add the GoogleBaseHook.get_credentials() probably not only respects impersonation chain set on hook level, but also on connection level. If we can rely on that, then we would need to cover those cases individually.

m1racoli · 2024-01-17T09:13:01Z

airflow/providers/google/cloud/hooks/bigquery.py

+        return project
+
+    async def get(self) -> str | None:
+        creds, _ = google.auth.default()


I would like to add the GoogleBaseHook.get_credentials() probably not only respects impersonation chain set on hook level, but also on connection level. If we can rely on that, then we would need to cover those cases individually.

m1racoli · 2024-01-17T09:14:15Z

airflow/providers/google/cloud/hooks/bigquery.py

+            target_scopes=["https://www.googleapis.com/auth/cloud-platform"],
+        )
+
+        impersonated_creds.refresh(google_auth_requests.Request())


This need's to be done asynchronously. Otherwise it will block the entire triggerer process.

m1racoli · 2024-01-17T16:55:13Z

Based on my review comments I've taken another attempt on this in #36849. That one should make it easier to use credentials from hooks in gcloud-aio clients without much extra work.

Lee-W · 2024-01-19T02:09:04Z

Hi @m1racoli, thanks for your feedback! I just took a look at your new PRs. Both of them look good. I think we might no longer need this PR, but I'll mark it as draft for now before we merge yours

Lee-W · 2024-01-24T00:39:03Z

as #36849 has been merged, close this one

boring-cyborg bot added area:providers provider:google Google (including GCP) related issues labels Dec 21, 2023

Lee-W force-pushed the add-missing-impersonation_chain-to-bigquery-operators branch from 1b64ce8 to 7d2dfc7 Compare December 22, 2023 15:49

melugoyal reviewed Dec 22, 2023

View reviewed changes

Lee-W force-pushed the add-missing-impersonation_chain-to-bigquery-operators branch 2 times, most recently from 9172959 to aeed629 Compare December 23, 2023 03:34

Lee-W marked this pull request as ready for review December 23, 2023 03:37

Lee-W changed the title ~~add missing impersonation_chain argument when calling trigger in Bigquery Operators~~ add impersonation_chain support when calling Bigquery Operators in deferrable mode Dec 23, 2023

Lee-W force-pushed the add-missing-impersonation_chain-to-bigquery-operators branch from aeed629 to 74525d4 Compare December 23, 2023 04:13

Lee-W force-pushed the add-missing-impersonation_chain-to-bigquery-operators branch 2 times, most recently from 067bccf to 6f81639 Compare December 25, 2023 03:32

Lee-W added 3 commits December 29, 2023 17:49

feat(providers/google): add missing impersonation_chain to bigquery o…

b7c2fe1

…perators

feat(providers/google): add impersonation chain support to bigquery o…

3685186

…perator async

feat(providers/google): get target principal from impersonation_chain

796c59c

Lee-W force-pushed the add-missing-impersonation_chain-to-bigquery-operators branch from 6f81639 to 796c59c Compare December 29, 2023 09:49

phanikumv requested changes Dec 29, 2023

View reviewed changes

Merge branch 'main' into add-missing-impersonation_chain-to-bigquery-…

f54c97f

…operators

m1racoli reviewed Jan 17, 2024

View reviewed changes

m1racoli suggested changes Jan 17, 2024

View reviewed changes

m1racoli mentioned this pull request Jan 17, 2024

feat: full support for google credentials in gcloud-aio clients #36849

Merged

Lee-W marked this pull request as draft January 19, 2024 02:07

Lee-W closed this Jan 24, 2024

Lee-W deleted the add-missing-impersonation_chain-to-bigquery-operators branch June 5, 2024 06:20

	target_principal, delegates = _get_target_principal_and_delegates(self.impersonation_chain)

	credentials, project_id = get_credentials_and_project_id(
	key_path=key_path,
	keyfile_dict=keyfile_dict_json,
	credential_config_file=credential_config_file,
	key_secret_name=key_secret_name,
	key_secret_project_id=key_secret_project_id,
	scopes=self.scopes,
	delegate_to=self.delegate_to,
	target_principal=target_principal,
	delegates=delegates,
	)

Conversation

Lee-W commented Dec 21, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

nathadfield commented Dec 21, 2023

Uh oh!

Lee-W commented Dec 21, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

m1racoli commented Dec 21, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

nathadfield commented Dec 21, 2023

Uh oh!

Lee-W commented Dec 22, 2023

Uh oh!

Lee-W commented Dec 22, 2023

Uh oh!

melugoyal Dec 22, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Lee-W Dec 23, 2023

Choose a reason for hiding this comment

Uh oh!

m1racoli commented Dec 23, 2023

Uh oh!

Lee-W commented Dec 24, 2023

Uh oh!

eladkal commented Dec 25, 2023

Uh oh!

Lee-W commented Dec 25, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

phanikumv left a comment

Choose a reason for hiding this comment

Uh oh!

m1racoli Jan 17, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

m1racoli Jan 17, 2024

Choose a reason for hiding this comment

Uh oh!

m1racoli Jan 17, 2024

Choose a reason for hiding this comment

Uh oh!

m1racoli Jan 17, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

m1racoli commented Jan 17, 2024

Uh oh!

Lee-W commented Jan 19, 2024

Uh oh!

Lee-W commented Jan 24, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Lee-W commented Dec 21, 2023 •

edited

Loading

Lee-W commented Dec 21, 2023 •

edited

Loading

m1racoli commented Dec 21, 2023 •

edited

Loading

melugoyal Dec 22, 2023 •

edited

Loading

Lee-W commented Dec 25, 2023 •

edited

Loading

m1racoli Jan 17, 2024 •

edited

Loading

m1racoli Jan 17, 2024 •

edited

Loading