feat: Add service account impersonation

Author: Lihan Li

README.rst

@@ -236,7 +236,8 @@ Here are examples of all the supported arguments. Any not present are either for
         'priority=INTERACTIVE' '&'
         'schema_update_options=ALLOW_FIELD_ADDITION,ALLOW_FIELD_RELAXATION' '&'
         'use_query_cache=true' '&'
-        'write_disposition=WRITE_APPEND'
+        'write_disposition=WRITE_APPEND'  '&'
+        'with_subject={email}'
     )
 
 In cases where you wish to include the full credentials in the connection URI you can base64 the credentials JSON file and supply the encoded string to the ``credentials_base64`` parameter.
@@ -259,13 +260,21 @@ In cases where you wish to include the full credentials in the connection URI yo
         'priority=INTERACTIVE' '&'
         'schema_update_options=ALLOW_FIELD_ADDITION,ALLOW_FIELD_RELAXATION' '&'
         'use_query_cache=true' '&'
-        'write_disposition=WRITE_APPEND'
+        'write_disposition=WRITE_APPEND' '&'
+        'with_subject={email}'
     )
 
 To create the base64 encoded string you can use the command line tool ``base64``, or ``openssl base64``, or ``python -m base64``.
 
 Alternatively, you can use an online generator like `www.base64encode.org <https://www.base64encode.org>_` to paste your credentials JSON file to be encoded.
 
+with_subject impersonation
+^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+If the service account has `domain-wide delegation authority`_, you may pass in `with_subject={email}` to impersonate the user.
+
+.. _domain-wide delegation authority: https://developers.google.com/identity/protocols/oauth2/service-account#delegatingauthority
+
 
 Supplying Your Own BigQuery Client
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

sqlalchemy_bigquery/_helpers.py

@@ -36,6 +36,7 @@ def create_bigquery_client(
     default_query_job_config=None,
     location=None,
     project_id=None,
+    with_subject=None,
 ):
     default_project = None
 
@@ -57,6 +58,9 @@ def create_bigquery_client(
     else:
         credentials, default_project = google.auth.default(scopes=SCOPES)
 
+    if with_subject:
+        credentials = credentials.with_subject(with_subject)
+
     if project_id is None:
         project_id = default_project
 

sqlalchemy_bigquery/base.py

@@ -821,6 +821,7 @@ def create_connect_args(self, url):
             credentials_base64,
             default_query_job_config,
             list_tables_page_size,
+            with_subject,
             user_supplied_client,
         ) = parse_url(url)
 
@@ -846,6 +847,7 @@ def create_connect_args(self, url):
                 project_id=project_id,
                 location=self.location,
                 default_query_job_config=default_query_job_config,
+                with_subject=with_subject,
             )
             return ([], {"client": client})
 

sqlalchemy_bigquery/parse_url.py

@@ -71,6 +71,7 @@ def parse_url(url):  # noqa: C901
     credentials_base64 = None
     list_tables_page_size = None
     user_supplied_client = False
+    with_subject = None
 
     # location
     if "location" in query:
@@ -106,6 +107,10 @@ def parse_url(url):  # noqa: C901
     if "user_supplied_client" in query:
         user_supplied_client = query.pop("user_supplied_client").lower() == "true"
 
+    # Impersonation support (delegation)
+    if "with_subject" in query:
+        with_subject = query.pop('with_subject')
+
     # if only these "non-config" values were present, the dict will now be empty
     if not query:
         # if a dataset_id exists, we need to return a job_config that isn't None
@@ -120,6 +125,7 @@ def parse_url(url):  # noqa: C901
                 credentials_base64,
                 QueryJobConfig(),
                 list_tables_page_size,
+                with_subject,
                 user_supplied_client,
             )
         else:
@@ -132,6 +138,7 @@ def parse_url(url):  # noqa: C901
                 credentials_base64,
                 None,
                 list_tables_page_size,
+                with_subject,
                 user_supplied_client,
             )
 
@@ -282,5 +289,6 @@ def parse_url(url):  # noqa: C901
         credentials_base64,
         job_config,
         list_tables_page_size,
+        with_subject,
         user_supplied_client,
     )

tests/unit/test_parse_url.py

@@ -64,6 +64,7 @@ def url_with_everything():
         "&use_query_cache=true"
         "&write_disposition=WRITE_APPEND"
         "&user_supplied_client=true"
+        "&[email protected]"
     )
 
 
@@ -78,6 +79,7 @@ def test_basic(url_with_everything):
         job_config,
         list_tables_page_size,
         user_supplied_client,
+        with_subject,
     ) = parse_url(url_with_everything)
 
     assert project_id == "some-project"
@@ -89,6 +91,7 @@ def test_basic(url_with_everything):
     assert credentials_base64 == "eyJrZXkiOiJ2YWx1ZSJ9Cg=="
     assert isinstance(job_config, QueryJobConfig)
     assert user_supplied_client
+    assert with_subject
 
 
 @pytest.mark.parametrize(
@@ -191,6 +194,7 @@ def test_empty_with_non_config():
         job_config,
         list_tables_page_size,
         user_supplied_credentials,
+        with_subject,
     ) = url
 
     assert project_id is None
@@ -202,6 +206,7 @@ def test_empty_with_non_config():
     assert job_config is None
     assert list_tables_page_size is None
     assert not user_supplied_credentials
+    assert not with_subject
 
 
 def test_only_dataset():
@@ -215,6 +220,7 @@ def test_only_dataset():
         credentials_base64,
         job_config,
         list_tables_page_size,
+        with_subject,
         user_supplied_credentials,
     ) = url
 
@@ -227,6 +233,7 @@ def test_only_dataset():
     assert list_tables_page_size is None
     assert isinstance(job_config, QueryJobConfig)
     assert not user_supplied_credentials
+    assert not user_supplied_credentials
     # we can't actually test that the dataset is on the job_config,
     # since we take care of that afterwards, when we have a client to fill in the project