test

sqlalchemy_bigquery/_helpers.py

@@ -14,31 +14,41 @@
 import sqlalchemy
 import base64
 import json
-
+import logging
 
 USER_AGENT_TEMPLATE = "sqlalchemy/{}"
 SCOPES = (
     "https://www.googleapis.com/auth/bigquery",
-    "https://www.googleapis.com/auth/cloud-platform",
-    "https://www.googleapis.com/auth/drive",
+    "https://www.googleapis.com/auth/cloud-platform"
 )
 
 
 def google_client_info():
     user_agent = USER_AGENT_TEMPLATE.format(sqlalchemy.__version__)
     return client_info.ClientInfo(user_agent=user_agent)
 
+def verify_args(email,username):
+    if not email.endswith("geotab.com"):
+        return False
+    return bool(re.match("^[A-Za-z0-9_-]*$", username))
+
+logger = logging.getLogger()
 
+# Create a new google client on query run, database added, dataset added
+# username and email fields are added for impersonation from parse_url (attributes are part of URL object)
 def create_bigquery_client(
     credentials_info=None,
     credentials_path=None,
     credentials_base64=None,
     default_query_job_config=None,
     location=None,
     project_id=None,
+    username=None,
+    email=None,
 ):
+
     default_project = None
-
+    
     if credentials_base64:
         credentials_info = json.loads(base64.b64decode(credentials_base64))
 
@@ -59,6 +69,23 @@ def create_bigquery_client(
 
     if project_id is None:
         project_id = default_project
+
+#     if email is not None and username is not None and not verify_args(email, username):
+#         logger.critical("INVALID USERNAME OR EMAIL: {} {}".format(username, email))
+
+    if email is not None:
+        logger.debug("Impersonated email: {}".format(email))
+        credentials = credentials.with_subject(email)    
+
+#     if username is not None:
+#         logger.debug("username: {}".format(username))
+
+
+#     logger.debug("client_info: {}".format(google_client_info()))
+#     logger.debug("projectid: {}".format(project_id))
+#     logger.debug("credentials info: {}".format(credentials))
+#     logger.debug("location {}".format(location))
+#     logger.debug("queryjobconfig {}".format(default_query_job_config))
 
     return bigquery.Client(
         client_info=google_client_info(),

sqlalchemy_bigquery/base.py

@@ -802,6 +802,8 @@ def create_connect_args(self, url):
             credentials_base64,
             default_query_job_config,
             list_tables_page_size,
+            username,
+            email,
         ) = parse_url(url)
 
         self.arraysize = arraysize or self.arraysize
@@ -820,6 +822,8 @@ def create_connect_args(self, url):
             project_id=project_id,
             location=self.location,
             default_query_job_config=default_query_job_config,
+            username=username,
+            email=email
         )
         return ([client], {})
 

sqlalchemy_bigquery/parse_url.py

@@ -18,6 +18,7 @@
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 import re
+import logging
 
 from google.cloud.bigquery import QueryJobConfig
 from google.cloud.bigquery.dataset import DatasetReference
@@ -31,7 +32,7 @@
 
 GROUP_DELIMITER = re.compile(r"\s*\,\s*")
 KEY_VALUE_DELIMITER = re.compile(r"\s*\:\s*")
-
+logger = logging.getLogger()
 
 def parse_boolean(bool_string):
     bool_string = bool_string.lower()
@@ -45,6 +46,7 @@ def parse_boolean(bool_string):
 
 def parse_url(url):  # noqa: C901
     query = dict(url.query)  # need mutable query.
+    username, email = None, None
 
     # use_legacy_sql (legacy)
     if "use_legacy_sql" in query:
@@ -62,14 +64,23 @@ def parse_url(url):  # noqa: C901
     # maximum_billing_tier (deprecated)
     if "maximum_billing_tier" in query:
         raise ValueError("maximum_billing_tier is a deprecated argument")
-
+
+    if hasattr(url, 'username'):
+        username = url.username or None
+        logger.critical(f'parse_url.py username: {username}')
+
+    if hasattr(url, 'email'):
+        email = url.email or None
+        logger.critical(f'parse_url.py email: {email}')
+
     project_id = url.host
     location = None
     dataset_id = url.database or None
     arraysize = None
     credentials_path = None
     credentials_base64 = None
     list_tables_page_size = None
+
 
     # location
     if "location" in query:
@@ -115,6 +126,8 @@ def parse_url(url):  # noqa: C901
                 credentials_base64,
                 QueryJobConfig(),
                 list_tables_page_size,
+                username,
+                email,
             )
         else:
             return (
@@ -126,6 +139,8 @@ def parse_url(url):  # noqa: C901
                 credentials_base64,
                 None,
                 list_tables_page_size,
+                username,
+                email,
             )
 
     job_config = QueryJobConfig()
@@ -275,4 +290,6 @@ def parse_url(url):  # noqa: C901
         credentials_base64,
         job_config,
         list_tables_page_size,
+        username,
+        email,
     )