Added service account impersonation feature

Author: Lihan Li

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