Skip to content

Commit

Permalink
feat(auth): implement Google Colab token retrieval mechanism
Browse files Browse the repository at this point in the history
  • Loading branch information
not-lain committed Jan 3, 2025
1 parent 003a4a7 commit 7d048f9
Showing 1 changed file with 46 additions and 0 deletions.
46 changes: 46 additions & 0 deletions projects/fal_client/src/fal_client/auth.py
Original file line number Diff line number Diff line change
@@ -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):
Expand All @@ -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 (
Expand Down

0 comments on commit 7d048f9

Please sign in to comment.