Repository: mabel-dev/opteryx

Issue:
✨ GCS improvements
Create the client object once and reuse

List blobs should only return the name of the blob and not any other details 


Discussion:
On a test dataset of 49 files on GCS - before the change, planning took 1.45s (fastest of four) and reading 40.3s (fastest of four). following the changes, planning took 1.2s (fastest of four) and reading took 20.5s (fastest of four)

Resolving patch:
diff --git a/opteryx/__version__.py b/opteryx/__version__.py
--- a/opteryx/__version__.py
+++ b/opteryx/__version__.py
@@ -1,4 +1,4 @@
-__build__ = 203
+__build__ = 206
 
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/opteryx/connectors/gcp_cloudstorage_connector.py b/opteryx/connectors/gcp_cloudstorage_connector.py
--- a/opteryx/connectors/gcp_cloudstorage_connector.py
+++ b/opteryx/connectors/gcp_cloudstorage_connector.py
@@ -11,10 +11,13 @@
 # limitations under the License.
 
 import os
+import urllib.request
 from typing import Dict
 from typing import List
 
 import pyarrow
+import requests
+from google.auth.transport.requests import Request
 from orso.schema import FlatColumn
 from orso.schema import RelationSchema
 from orso.tools import single_item_cache
@@ -65,6 +68,17 @@ def __init__(self, credentials=None, **kwargs):
         # sometimes both start here
         self.cached_first_blob = None
         self.client = self._get_storage_client()
+        self.client_credentials = self.client._credentials
+
+        # Cache access tokens for accessing GCS
+        if not self.client_credentials.valid:
+            request = Request()
+            self.client_credentials.refresh(request)
+            self.access_token = self.client_credentials.token
+
+        # Create a HTTP connection session to reduce effort for
+        # each fetch
+        self.session = requests.Session()
 
     def _get_storage_client(self):
         from google.cloud import storage
@@ -76,22 +90,25 @@ def _get_storage_client(self):
         else:  # pragma: no cover
             return storage.Client()
 
-    def _get_blob(self, bucket: str, blob_name: str):
-        gcs_bucket = self.client.get_bucket(bucket)
-        blob = gcs_bucket.get_blob(blob_name)
-        return blob
-
     def read_blob(self, *, blob_name, **kwargs):
-        bucket, object_path, name, extension = paths.get_parts(blob_name)
+        # For performance we use the GCS API directly, this is roughly 10%
+        # faster than using the SDK.
+        bucket, _, _, _ = paths.get_parts(blob_name)
+
+        # Ensure the credentials are valid, refreshing them if necessary
+        if not self.client_credentials.valid:
+            request = Request()
+            self.client_credentials.refresh(request)
+            self.access_token = self.client_credentials.token
 
         bucket = bucket.replace("va_data", "va-data")
         bucket = bucket.replace("data_", "data-")
+        object_full_path = urllib.parse.quote(blob_name[(len(bucket) + 1) :], safe="")
 
-        blob = self._get_blob(
-            bucket=bucket,
-            blob_name=object_path + "/" + name + extension,
-        )
-        return blob.download_as_bytes()
+        url = f"https://storage.googleapis.com/storage/v1/b/{bucket}/o/{object_full_path}?alt=media"
+
+        response = self.session.get(url, headers={"Authorization": f"Bearer {self.access_token}"})
+        return response.content
 
     @single_item_cache
     def get_list_of_blob_names(self, *, prefix: str) -> List[str]:


Test patch:
diff --git a/tests/sql_battery/test_shapes_and_errors_battery.py b/tests/sql_battery/test_shapes_and_errors_battery.py
--- a/tests/sql_battery/test_shapes_and_errors_battery.py
+++ b/tests/sql_battery/test_shapes_and_errors_battery.py
@@ -1537,12 +1537,14 @@ def test_sql_battery(statement, rows, columns, exception):
             failures.append((statement, err))
 
     print("--- ✅ \033[0;32mdone\033[0m")
+
+    if failed > 0:
+        print("\n\033[38;2;139;233;253m\033[3mFAILURES\033[0m")
+        for statement, err in failures:
+            print(err)
+
     print(
         f"\n\033[38;2;139;233;253m\033[3mCOMPLETE\033[0m ({((time.monotonic_ns() - start_suite) / 1e9):.2f} seconds)\n"
         f"  \033[38;2;26;185;67m{passed} passed ({(passed * 100) // (passed + failed)}%)\033[0m\n"
         f"  \033[38;2;255;121;198m{failed} failed\033[0m"
     )
-    if failed > 0:
-        print("\n\033[38;2;139;233;253m\033[3mFAILURES\033[0m")
-        for statement, err in failures:
-            print(err)