From 7d048f9115f9fec36dc9aea0316863583cde43be Mon Sep 17 00:00:00 2001
From: Lain <hhichri60@gmail.com>
Date: Fri, 3 Jan 2025 19:33:59 +0100
Subject: [PATCH] feat(auth): implement Google Colab token retrieval mechanism

---
 projects/fal_client/src/fal_client/auth.py | 46 ++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/projects/fal_client/src/fal_client/auth.py b/projects/fal_client/src/fal_client/auth.py
index d83a23d0..6887ae01 100644
--- a/projects/fal_client/src/fal_client/auth.py
+++ b/projects/fal_client/src/fal_client/auth.py
@@ -1,4 +1,45 @@
 import os
+from threading import Lock
+from typing import Optional
+
+
+class GoogleColabState:
+    def __init__(self):
+        self.is_checked = False
+        self.lock = Lock()
+        self.secret: Optional[str] = None
+
+
+_colab_state = GoogleColabState()
+
+
+def is_google_colab() -> bool:
+    try:
+        return "google.colab" in str(get_ipython())  # noqa: F821
+    except NameError:
+        return False
+
+
+def get_colab_token() -> Optional[str]:
+    if not is_google_colab():
+        return None
+    with _colab_state.lock:
+        if _colab_state.is_checked:  # request access only once
+            return _colab_state.secret
+
+        try:
+            from google.colab import userdata  # noqa: I001
+        except ImportError:
+            return None
+
+        try:
+            token = userdata.get("FAL_KEY")
+            _colab_state.secret = token.strip()
+        except Exception:
+            _colab_state.secret = None
+
+        _colab_state.is_checked = True
+        return _colab_state.secret
 
 
 class MissingCredentialsError(Exception):
@@ -9,6 +50,11 @@ class MissingCredentialsError(Exception):
 
 
 def fetch_credentials() -> str:
+    # First check for Colab token
+    if colab_token := get_colab_token():
+        return colab_token
+
+    # Then check environment variables
     if key := os.getenv("FAL_KEY"):
         return key
     elif (key_id := os.getenv("FAL_KEY_ID")) and (