Bigquery sqlalchemy metastore (#1352)

* [bigquery] bump bigquery engine dependencies * [bigquery] Allow use of Application defaults credentials if no credentials specified * [metastore] Split table name on '.' when using sqlalchemy-bigquery metasore
pinterest · Nov 29, 2023 · 91b2daf · 91b2daf
1 parent f8d8dd4
commit 91b2daf
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 14 deletions.
diff --git a/querybook/server/clients/google_client.py b/querybook/server/clients/google_client.py
@@ -21,11 +21,11 @@ def get_google_credentials(creds_info=None):
         )
 
     cred_to_use = creds_info or QuerybookSettings.GOOGLE_CREDS
-    assert cred_to_use is not None, "Invalid Google credentials"
 
-    credentials = service_account.Credentials.from_service_account_info(cred_to_use)
-
-    return credentials
+    if cred_to_use is not None:
+        return service_account.Credentials.from_service_account_info(cred_to_use)
+    else:
+        return None
 
 
 GOOGLE_AUTH_CONFIG = "https://accounts.google.com/.well-known/openid-configuration"

diff --git a/querybook/server/lib/metastore/loaders/sqlalchemy_metastore_loader.py b/querybook/server/lib/metastore/loaders/sqlalchemy_metastore_loader.py
@@ -24,7 +24,10 @@ def get_all_schema_names(self) -> List[str]:
         return self._inspect.get_schema_names()
 
     def get_all_table_names_in_schema(self, schema_name: str) -> List[str]:
-        return self._inspect.get_table_names(schema=schema_name)
+        if self._engine.dialect.name == "bigquery":
+            return [table.split(".")[1] for table in self._inspect.get_table_names(schema=schema_name)]
+        else:
+            return self._inspect.get_table_names(schema=schema_name)
 
     def get_table_and_columns(
         self, schema_name, table_name

diff --git a/querybook/server/lib/query_executor/clients/bigquery.py b/querybook/server/lib/query_executor/clients/bigquery.py
@@ -12,9 +12,11 @@ def __init__(self, google_credentials_json=None, *args, **kwargs):
             if google_credentials_json is not None
             else None
         )
-        cred = get_google_credentials(parsed_google_json)
-
-        client = Client(project=cred.project_id, credentials=cred)
+        if parsed_google_json is not None:
+            cred = get_google_credentials(parsed_google_json)
+            client = Client(project=cred.project_id, credentials=cred)
+        else:
+            client = Client()
 
         self._conn = dbapi.connect(client=client)
         super(BigQueryClient, self).__init__()

diff --git a/querybook/server/lib/query_executor/executor_template/templates.py b/querybook/server/lib/query_executor/executor_template/templates.py
@@ -120,7 +120,11 @@
     (
         "google_credentials_json",
         FormField(
-            helper="The JSON string used to log in as service account. If not provided then **GOOGLE_CREDS** from settings will be used.",
+            helper="""
+<p>The JSON string used to log in as service account.</p>
+<p>If not provided then **GOOGLE_CREDS** from settings will be used</p>
+<p>If both are empty, Application default credentials are used</p>
+""",
         ),
     )
 )
diff --git a/requirements/engine/bigquery.txt b/requirements/engine/bigquery.txt
@@ -1,5 +1,4 @@
-google-cloud-bigquery==1.28.0
-
-# Downgrade Protobuf to fix error with bigquery:
-# TypeError: Descriptors cannot not be created directly.
-protobuf==3.20.1
+google-cloud-bigquery==3.12.0
+google-cloud-bigquery-storage==2.22.0
+pyarrow==13.0.0
+sqlalchemy-bigquery==1.8.0