Lazy load GCS client. (#227)

datacommonsorg · Oct 10, 2023 · f61070d · f61070d
1 parent 2aa387a
commit f61070d
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/simple/stats/main.py b/simple/stats/main.py
@@ -24,7 +24,6 @@
 _CODEDIR = os.path.dirname(os.path.realpath(__file__))
 sys.path.insert(1, os.path.join(_CODEDIR, "../"))
 
-
 FLAGS = flags.FLAGS
 
 flags.DEFINE_string(

diff --git a/simple/util/filehandler.py b/simple/util/filehandler.py
@@ -94,15 +94,24 @@ def list_files(self, extension: str = None) -> list[str]:
                       all_files)
 
 
-class GcsFileHandler(FileHandler):
-    # Using print instead of logging since the class is loaded before logging is initialized.
+class GcsMeta(type):
+
+    @property
+    def gcs_client(cls) -> storage.Client:
+        if getattr(cls, "_GCS_CLIENT", None) is None:
+            gcs_client = storage.Client()
+            logging.info("Using GCS project: %s", gcs_client.project)
+            cls._GCS_CLIENT = gcs_client
+        return cls._GCS_CLIENT
+
+
+class GcsFileHandler(FileHandler, metaclass=GcsMeta):
 
     def __init__(self, path: str) -> None:
         if not path.startswith(_GCS_PATH_PREFIX):
             raise ValueError(f"Expected {_GCS_PATH_PREFIX} prefix, got {path}")
         bucket_name, blob_name = path[len(_GCS_PATH_PREFIX):].split('/', 1)
-        gcs_client = storage.Client()
-        self.bucket = gcs_client.bucket(bucket_name)
+        self.bucket = GcsFileHandler.gcs_client.bucket(bucket_name)
         self.blob = self.bucket.blob(blob_name)
         isdir = path.endswith("/")
         super().__init__(path, isdir)