From 1b4f2eaba759fa0666d6f9e81e87812798fbb990 Mon Sep 17 00:00:00 2001 From: data-engineer Date: Sat, 6 Jun 2026 23:25:20 -0700 Subject: [PATCH] fix: add friendly Gradata auth error --- Gradata/src/gradata/__init__.py | 2 ++ Gradata/src/gradata/brain.py | 24 ++++++++++++++++---- Gradata/src/gradata/cloud/client.py | 10 ++++++--- Gradata/src/gradata/exceptions.py | 4 ++++ Gradata/src/gradata/onboard.py | 2 +- Gradata/tests/test_auth_error.py | 35 +++++++++++++++++++++++++++++ 6 files changed, 69 insertions(+), 8 deletions(-) create mode 100644 Gradata/tests/test_auth_error.py diff --git a/Gradata/src/gradata/__init__.py b/Gradata/src/gradata/__init__.py index ecd13e31..cc26ff4b 100644 --- a/Gradata/src/gradata/__init__.py +++ b/Gradata/src/gradata/__init__.py @@ -62,6 +62,7 @@ EmbeddingError, EventPersistenceError, ExportError, + GradataAuthError, GradataError, TaxonomyError, ValidationError, @@ -80,6 +81,7 @@ "EmbeddingError", "EventPersistenceError", "ExportError", + "GradataAuthError", "GradataError", "Lesson", "LessonState", diff --git a/Gradata/src/gradata/brain.py b/Gradata/src/gradata/brain.py index 367b4307..d47d2e0f 100644 --- a/Gradata/src/gradata/brain.py +++ b/Gradata/src/gradata/brain.py @@ -54,17 +54,25 @@ from gradata._env import env_str from gradata.brain_inspection import BrainInspectionMixin +AUTH_ERROR_MESSAGE = ( + "No API key found. Set GRADATA_API_KEY or pass api_key=... to Brain(). " + "Get one at https://gradata.ai/keys" +) -def _resolve_sync_api_key() -> str | None: + +def _resolve_sync_api_key(api_key: str | None = None) -> str | None: """Resolve cloud API key for the write-through sync worker. Order matches ``gradata.daemon._resolve_api_key``: ``GRADATA_API_KEY`` env var, then ``~/.gradata/key`` file. - Returns ``None`` when neither is set — write-through is silently - disabled (the local correct() path still works normally). + An explicit ``api_key=`` wins, then ``GRADATA_API_KEY``, then + ``~/.gradata/key``. Returns ``None`` when neither is set so callers + can decide whether local-only operation is allowed. """ import os + if api_key: + return api_key env_key = os.environ.get("GRADATA_API_KEY") if env_key: return env_key @@ -90,9 +98,17 @@ def __init__( brain_dir: str | Path | None = None, working_dir: str | Path | None = None, encryption_key: str | None = None, + api_key: str | None = None, + _skip_auth_check: bool = False, ): from gradata._paths import resolve_brain_dir + self._api_key = _resolve_sync_api_key(api_key) + if not _skip_auth_check and not self._api_key: + from gradata.exceptions import GradataAuthError + + raise GradataAuthError(AUTH_ERROR_MESSAGE) + self.dir = resolve_brain_dir(brain_dir) if not self.dir.exists(): from gradata.exceptions import BrainNotFoundError @@ -279,7 +295,7 @@ def _maybe_start_sync_worker(self) -> None: if os.environ.get("GRADATA_DISABLE_WRITE_THROUGH") == "1": return - api_key = _resolve_sync_api_key() + api_key = self._api_key or _resolve_sync_api_key() if not api_key: return diff --git a/Gradata/src/gradata/cloud/client.py b/Gradata/src/gradata/cloud/client.py index 3c1b836c..39c8b5ad 100644 --- a/Gradata/src/gradata/cloud/client.py +++ b/Gradata/src/gradata/cloud/client.py @@ -24,6 +24,8 @@ from urllib.request import Request, urlopen from gradata._http import require_https +from gradata.brain import AUTH_ERROR_MESSAGE +from gradata.exceptions import GradataAuthError logger = logging.getLogger("gradata.cloud") @@ -334,9 +336,11 @@ def _post(self, path: str, data: dict) -> dict[str, Any]: try: with urlopen(req, timeout=30) as resp: return json.loads(resp.read().decode("utf-8")) - except HTTPError as e: - if e.code == 413: - raise _TooLargeError() from e + except HTTPError as e: + if e.code in (401, 403): + raise GradataAuthError(AUTH_ERROR_MESSAGE) from e + if e.code == 413: + raise _TooLargeError() from e try: body = e.read().decode("utf-8", errors="replace")[:500] except Exception: diff --git a/Gradata/src/gradata/exceptions.py b/Gradata/src/gradata/exceptions.py index f3d09095..c656d9b4 100644 --- a/Gradata/src/gradata/exceptions.py +++ b/Gradata/src/gradata/exceptions.py @@ -22,6 +22,10 @@ class BrainNotConfiguredError(BrainError): """Brain directory is required but could not be resolved.""" +class GradataAuthError(BrainError): + """Authentication is required for a Gradata Cloud operation.""" + + class BrainLockedError(BrainError): """Brain directory is already held by another daemon/server process.""" diff --git a/Gradata/src/gradata/onboard.py b/Gradata/src/gradata/onboard.py index 0ec2abd8..55a892aa 100644 --- a/Gradata/src/gradata/onboard.py +++ b/Gradata/src/gradata/onboard.py @@ -487,7 +487,7 @@ def onboard( # ── Success output ───────────────────────────────────────────────── - brain = Brain(brain_dir) + brain = Brain(brain_dir, _skip_auth_check=True) if interactive: stats = brain.stats() diff --git a/Gradata/tests/test_auth_error.py b/Gradata/tests/test_auth_error.py new file mode 100644 index 00000000..9d862509 --- /dev/null +++ b/Gradata/tests/test_auth_error.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +import pytest + +from gradata import Brain, GradataAuthError +from gradata.brain import AUTH_ERROR_MESSAGE + + +def test_brain_constructor_without_api_key_raises_auth_error(tmp_path, monkeypatch): + brain_dir = tmp_path / "brain" + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.delenv("GRADATA_API_KEY", raising=False) + + Brain.init(brain_dir, name="TestBrain", domain="Testing", embedding="local", interactive=False) + + with pytest.raises(GradataAuthError) as exc: + Brain(brain_dir) + + assert "GRADATA_API_KEY" in str(exc.value) + assert str(exc.value) == AUTH_ERROR_MESSAGE + + +def test_brain_constructor_accepts_explicit_api_key(tmp_path, monkeypatch): + brain_dir = tmp_path / "brain" + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.delenv("GRADATA_API_KEY", raising=False) + monkeypatch.setenv("GRADATA_DISABLE_WRITE_THROUGH", "1") + + Brain.init(brain_dir, name="TestBrain", domain="Testing", embedding="local", interactive=False) + + brain = Brain(brain_dir, api_key="gd_test") + try: + assert brain._api_key == "gd_test" + finally: + brain.close()