Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 9 additions & 2 deletions src/py/mat3ra/notebooks_utils/auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,12 @@

from mat3ra.api_client import ACCESS_TOKEN_ENV_VAR

from .core.api.auth import authenticate_oidc
from .core.api.auth import authenticate_oidc, get_oidc_base_url, store_token_data_in_environment
from .io import get_data
from .ipython.ui import show_device_flow_popup
from .primitive.environment import ENVIRONMENT, EnvironmentsEnum
from .pyodide.api.auth import authenticate_jupyterlite
from .token_store import load_token, save_token

REFRESH_TOKEN_ENV_VAR = "OIDC_REFRESH_TOKEN"

Expand All @@ -25,4 +26,10 @@ async def authenticate(force=False, globals_dict=None):
if data_from_host:
await authenticate_jupyterlite(data_from_host)
elif ACCESS_TOKEN_ENV_VAR not in os.environ or force:
await authenticate_oidc(show_popup=show_device_flow_popup)
oidc_url = get_oidc_base_url()
cached = None if force else await load_token(oidc_url)
if cached:
store_token_data_in_environment(cached)
else:
token_data = await authenticate_oidc(show_popup=show_device_flow_popup)
await save_token(oidc_url, token_data)
51 changes: 51 additions & 0 deletions src/py/mat3ra/notebooks_utils/core/api/token_store.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,51 @@
"""OIDC token cache — business logic + file-based storage (default)."""

import json
import os
import sys
import time
from typing import Optional

_EXPIRY_BUFFER = 60 # seconds before actual expiry to consider token stale
_FILE_PATH = os.path.join(os.path.expanduser("~"), ".mat3ra", "oidc_token_cache.json")

# Storage backend: defaults to this module (file-based).
# Overridden by top-level token_store.py for Pyodide (IndexedDB).
_storage = sys.modules[__name__]


async def read() -> dict:
try:
with open(_FILE_PATH) as f:
return json.load(f)
except Exception:
return {}


async def write(cache: dict) -> None:
os.makedirs(os.path.dirname(_FILE_PATH), mode=0o700, exist_ok=True)
with open(_FILE_PATH, "w") as f:
json.dump(cache, f)
os.chmod(_FILE_PATH, 0o600)


async def save_token(oidc_url: str, token_data: dict) -> None:
token_cache_entry = dict(token_data)
token_cache_entry["expires_at"] = time.time() + token_cache_entry.get("expires_in", 3600)

token_cache = await _storage.read()
token_cache[oidc_url] = token_cache_entry
await _storage.write(token_cache)


async def load_token(oidc_url: str) -> Optional[dict]:
token_cache_entry = (await _storage.read()).get(oidc_url)

if not token_cache_entry:
return None

expires_at = token_cache_entry.get("expires_at", 0)
if expires_at <= time.time() + _EXPIRY_BUFFER:
return None

return token_cache_entry
45 changes: 45 additions & 0 deletions src/py/mat3ra/notebooks_utils/pyodide/api/token_store.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
"""IndexedDB-based OIDC token cache for Pyodide Web Workers."""

import json

from pyodide.code import run_js # type: ignore

_idb = run_js(
"""
(function() {
const open = () => new Promise(resolve => {
const req = indexedDB.open("mat3ra", 1);
req.onupgradeneeded = () => req.result.createObjectStore("tokens");
req.onsuccess = () => resolve(req.result);
req.onerror = () => resolve(null);
});
return {
get: async () => {
const db = await open();
if (!db) return null;
return new Promise(resolve => {
const g = db.transaction("tokens").objectStore("tokens").get("oidc_cache");
g.onsuccess = () => { db.close(); resolve(g.result || null); };
g.onerror = () => { db.close(); resolve(null); };
});
},
set: async (data) => {
const db = await open();
if (!db) return;
const tx = db.transaction("tokens", "readwrite");
tx.objectStore("tokens").put(data, "oidc_cache");
return new Promise(resolve => { tx.oncomplete = () => { db.close(); resolve(); }; });
}
};
})()
"""
)


async def read() -> dict:
result = await _idb.get()
return json.loads(str(result)) if result else {}


async def write(cache: dict) -> None:
await _idb.set(json.dumps(cache))
10 changes: 10 additions & 0 deletions src/py/mat3ra/notebooks_utils/token_store.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
"""Persist OIDC tokens across notebook kernels — thin routing adapter."""

from .core.api.token_store import load_token, save_token # noqa: F401 — re-export
from .primitive.environment import is_pyodide_environment

if is_pyodide_environment():
from .core.api import token_store as _core_ts
from .pyodide.api import token_store as _pyodide_storage

_core_ts._storage = _pyodide_storage
Loading