Skip to content

Commit

Permalink
fix(auth): implement Google Colab token retrieval with state management
Browse files Browse the repository at this point in the history
  • Loading branch information
not-lain committed Jan 4, 2025
1 parent 9ceafc9 commit 2fe103d
Showing 1 changed file with 47 additions and 0 deletions.
47 changes: 47 additions & 0 deletions projects/fal_client/src/fal_client/auth.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,49 @@
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:
from IPython import get_ipython

return "google.colab" in str(get_ipython())
except ModuleNotFoundError:
return False
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 @@ -15,6 +60,8 @@ def fetch_credentials() -> str:
fal_key_secret := os.getenv("FAL_KEY_SECRET")
):
return f"{key_id}:{fal_key_secret}"
elif colab_token := get_colab_token():
return colab_token
else:
raise MissingCredentialsError(
"Please set the FAL_KEY environment variable to your API key, or set the FAL_KEY_ID and FAL_KEY_SECRET environment variables."
Expand Down

0 comments on commit 2fe103d

Please sign in to comment.