wherobots-python-sdk 0.6.0__tar.gz → 0.7.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (48) hide show
  1. {wherobots_python_sdk-0.6.0/wherobots_python_sdk.egg-info → wherobots_python_sdk-0.7.0}/PKG-INFO +14 -5
  2. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/README.md +12 -3
  3. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/pyproject.toml +6 -6
  4. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/tests/test_base_client.py +57 -4
  5. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/tests/test_config.py +12 -0
  6. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/tests/test_files_api.py +14 -1
  7. wherobots_python_sdk-0.7.0/tests/test_workload_auth.py +151 -0
  8. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/wherobots/__version__.py +1 -1
  9. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/wherobots/api/base.py +15 -7
  10. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/wherobots/api/drive.py +1 -1
  11. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/wherobots/api/files.py +8 -1
  12. wherobots_python_sdk-0.7.0/wherobots/api/workload_auth.py +116 -0
  13. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/wherobots/config.py +12 -2
  14. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/wherobots/exceptions.py +1 -1
  15. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0/wherobots_python_sdk.egg-info}/PKG-INFO +14 -5
  16. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/wherobots_python_sdk.egg-info/SOURCES.txt +2 -0
  17. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/wherobots_python_sdk.egg-info/requires.txt +1 -1
  18. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/LICENSE +0 -0
  19. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/setup.cfg +0 -0
  20. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/tests/test_api.py +0 -0
  21. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/tests/test_client.py +0 -0
  22. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/tests/test_client_header.py +0 -0
  23. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/tests/test_enums.py +0 -0
  24. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/tests/test_exceptions.py +0 -0
  25. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/tests/test_files_drive.py +0 -0
  26. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/tests/test_init_exports.py +0 -0
  27. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/tests/test_integration.py +0 -0
  28. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/tests/test_logger.py +0 -0
  29. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/tests/test_models.py +0 -0
  30. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/tests/test_regressions.py +0 -0
  31. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/tests/test_security.py +0 -0
  32. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/tests/test_tasks.py +0 -0
  33. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/tests/test_utils.py +0 -0
  34. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/wherobots/__init__.py +0 -0
  35. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/wherobots/api/__init__.py +0 -0
  36. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/wherobots/api/runs.py +0 -0
  37. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/wherobots/api/tasks.py +0 -0
  38. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/wherobots/client.py +0 -0
  39. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/wherobots/enums.py +0 -0
  40. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/wherobots/models.py +0 -0
  41. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/wherobots/py.typed +0 -0
  42. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/wherobots/sdk.py +0 -0
  43. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/wherobots/utils/__init__.py +0 -0
  44. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/wherobots/utils/client_header.py +0 -0
  45. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/wherobots/utils/logger.py +0 -0
  46. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/wherobots/utils/validation.py +0 -0
  47. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/wherobots_python_sdk.egg-info/dependency_links.txt +0 -0
  48. {wherobots_python_sdk-0.6.0 → wherobots_python_sdk-0.7.0}/wherobots_python_sdk.egg-info/top_level.txt +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: wherobots-python-sdk
3
- Version: 0.6.0
3
+ Version: 0.7.0
4
4
  Summary: Python SDK for Wherobots (currently covers the Jobs REST API)
5
5
  Author-email: Wherobots <support@wherobots.com>
6
6
  License-Expression: Apache-2.0
@@ -21,7 +21,7 @@ Requires-Python: >=3.11
21
21
  Description-Content-Type: text/markdown
22
22
  License-File: LICENSE
23
23
  Requires-Dist: requests>=2.34.2
24
- Requires-Dist: boto3>=1.43.108
24
+ Requires-Dist: boto3>=1.34.0
25
25
  Dynamic: license-file
26
26
 
27
27
  # Wherobots Python SDK
@@ -118,18 +118,27 @@ job = WherobotsJob(
118
118
 
119
119
  The API key can also be set via the `WHEROBOTS_API_KEY` environment variable.
120
120
 
121
+ Inside a Wherobots notebook or job, no API key is needed. The SDK authenticates as the user who
122
+ launched the workload, using the token Wherobots mounts in the pod. An API key, when set, takes
123
+ precedence.
124
+
125
+ ```python
126
+ job = WherobotsJob(script="s3://bucket/script.py", name="my-job-001")
127
+ ```
128
+
121
129
  ## Configuration
122
130
 
123
131
  ### Environment Variables
124
132
 
125
133
  | Variable | Description | Default |
126
134
  |------------------------------------|---------------------------------|--------------------------------------|
127
- | `WHEROBOTS_API_KEY` | API key | *(required)* |
135
+ | `WHEROBOTS_API_KEY` | API key | *(required outside Wherobots)* |
128
136
  | `WHEROBOTS_REGION` | Region override (else org default) | *(none — org default)* |
129
137
  | `WHEROBOTS_API_BASE_URL` | API base URL | `https://api.cloud.wherobots.com` |
130
138
  | `WHEROBOTS_VERSION` | Wherobots version | `latest` |
131
139
  | `WHEROBOTS_REQUEST_TIMEOUT_SECONDS`| HTTP request timeout (seconds) | `30` |
132
140
  | `WHEROBOTS_CLIENT_CHAIN` | Upstream client-attribution chain (see below) | *(none)* |
141
+ | `WHEROBOTS_ID_TOKEN_PATH` | Workload identity token, set by Wherobots notebooks and jobs | *(none)* |
133
142
 
134
143
  ### WherobotsConfig
135
144
 
@@ -456,7 +465,7 @@ with wherobots.Wherobots() as wb:
456
465
  | Failure | Raised as |
457
466
  |---|---|
458
467
  | Files is not enabled for this drive in this region | `WherobotsFilesNotEnabledError` (`.drive`, `.region`) |
459
- | The API key was not accepted | `WherobotsNotAuthenticatedError` |
468
+ | The credentials were not accepted | `WherobotsNotAuthenticatedError` |
460
469
  | No such file or folder | `WherobotsFileNotFoundError` |
461
470
  | Transfer to or from storage failed | `WherobotsS3Error` |
462
471
 
@@ -673,7 +682,7 @@ All exceptions inherit from `WherobotsJobError`.
673
682
  | `WherobotsS3Error` | A transfer to or from storage failed, including Files presigned uploads and downloads. |
674
683
  | `WherobotsConfigError` | Configuration is missing or invalid. |
675
684
  | `WherobotsTimeoutError` | Job exceeded the specified timeout. |
676
- | `WherobotsNotAuthenticatedError` | Files: the API key was not accepted (HTTP 401). Subclass of `WherobotsAPIError`. |
685
+ | `WherobotsNotAuthenticatedError` | Files: the credentials were not accepted (HTTP 401). In a Wherobots pod: the workload identity token exchange failed. Subclass of `WherobotsAPIError`. |
677
686
  | `WherobotsFilesNotEnabledError` | Files is not enabled for a drive in a region. Has `drive`, `region`. Subclass of `WherobotsAPIError`. |
678
687
  | `WherobotsFileNotFoundError` | Files: no such file or folder. Subclass of `WherobotsAPIError`. |
679
688
 
@@ -92,18 +92,27 @@ job = WherobotsJob(
92
92
 
93
93
  The API key can also be set via the `WHEROBOTS_API_KEY` environment variable.
94
94
 
95
+ Inside a Wherobots notebook or job, no API key is needed. The SDK authenticates as the user who
96
+ launched the workload, using the token Wherobots mounts in the pod. An API key, when set, takes
97
+ precedence.
98
+
99
+ ```python
100
+ job = WherobotsJob(script="s3://bucket/script.py", name="my-job-001")
101
+ ```
102
+
95
103
  ## Configuration
96
104
 
97
105
  ### Environment Variables
98
106
 
99
107
  | Variable | Description | Default |
100
108
  |------------------------------------|---------------------------------|--------------------------------------|
101
- | `WHEROBOTS_API_KEY` | API key | *(required)* |
109
+ | `WHEROBOTS_API_KEY` | API key | *(required outside Wherobots)* |
102
110
  | `WHEROBOTS_REGION` | Region override (else org default) | *(none — org default)* |
103
111
  | `WHEROBOTS_API_BASE_URL` | API base URL | `https://api.cloud.wherobots.com` |
104
112
  | `WHEROBOTS_VERSION` | Wherobots version | `latest` |
105
113
  | `WHEROBOTS_REQUEST_TIMEOUT_SECONDS`| HTTP request timeout (seconds) | `30` |
106
114
  | `WHEROBOTS_CLIENT_CHAIN` | Upstream client-attribution chain (see below) | *(none)* |
115
+ | `WHEROBOTS_ID_TOKEN_PATH` | Workload identity token, set by Wherobots notebooks and jobs | *(none)* |
107
116
 
108
117
  ### WherobotsConfig
109
118
 
@@ -430,7 +439,7 @@ with wherobots.Wherobots() as wb:
430
439
  | Failure | Raised as |
431
440
  |---|---|
432
441
  | Files is not enabled for this drive in this region | `WherobotsFilesNotEnabledError` (`.drive`, `.region`) |
433
- | The API key was not accepted | `WherobotsNotAuthenticatedError` |
442
+ | The credentials were not accepted | `WherobotsNotAuthenticatedError` |
434
443
  | No such file or folder | `WherobotsFileNotFoundError` |
435
444
  | Transfer to or from storage failed | `WherobotsS3Error` |
436
445
 
@@ -647,7 +656,7 @@ All exceptions inherit from `WherobotsJobError`.
647
656
  | `WherobotsS3Error` | A transfer to or from storage failed, including Files presigned uploads and downloads. |
648
657
  | `WherobotsConfigError` | Configuration is missing or invalid. |
649
658
  | `WherobotsTimeoutError` | Job exceeded the specified timeout. |
650
- | `WherobotsNotAuthenticatedError` | Files: the API key was not accepted (HTTP 401). Subclass of `WherobotsAPIError`. |
659
+ | `WherobotsNotAuthenticatedError` | Files: the credentials were not accepted (HTTP 401). In a Wherobots pod: the workload identity token exchange failed. Subclass of `WherobotsAPIError`. |
651
660
  | `WherobotsFilesNotEnabledError` | Files is not enabled for a drive in a region. Has `drive`, `region`. Subclass of `WherobotsAPIError`. |
652
661
  | `WherobotsFileNotFoundError` | Files: no such file or folder. Subclass of `WherobotsAPIError`. |
653
662
 
@@ -25,7 +25,7 @@ classifiers = [
25
25
  keywords = ["wherobots", "spark", "sedona", "geospatial", "jobs", "sdk", "rest-api"]
26
26
  dependencies = [
27
27
  "requests>=2.34.2",
28
- "boto3>=1.43.108",
28
+ "boto3>=1.34.0",
29
29
  ]
30
30
 
31
31
  [project.urls]
@@ -39,14 +39,14 @@ Documentation = "https://docs.wherobots.com"
39
39
  dev = [
40
40
  "pytest==9.1.1",
41
41
  "pytest-cov==7.1.0",
42
- "pytest-mock==3.16.0",
43
- "ruff==0.17.0",
44
- "pyrefly==1.3.2",
42
+ "pytest-mock==3.17.0",
43
+ "ruff==0.17.3",
44
+ "pyrefly==1.4.0",
45
45
  "import-linter==2.15",
46
46
  "pre-commit==4.7.0",
47
- "hypothesis==6.170.2",
47
+ "hypothesis==6.170.12",
48
48
  "cosmic-ray==8.7.0",
49
- "ast-grep-cli==0.50.0",
49
+ "ast-grep-cli==0.51.0",
50
50
  "types-requests==2.33.0.20261006",
51
51
  ]
52
52
 
@@ -158,17 +158,70 @@ class TestClientAttributionHeader:
158
158
 
159
159
 
160
160
  class TestClientConfigValidation:
161
- def test_missing_api_key_raises(self):
161
+ def test_missing_credentials_raises(self):
162
162
  cfg = WherobotsConfig(api_key=None)
163
- with pytest.raises(WherobotsConfigError, match="API key required"):
163
+ with pytest.raises(WherobotsConfigError, match="No Wherobots credentials found"):
164
164
  BaseClient(cfg)
165
165
 
166
166
  def test_empty_api_key_raises(self):
167
- cfg = WherobotsConfig(api_key="")
168
- with pytest.raises(WherobotsConfigError, match="API key required"):
167
+ cfg = WherobotsConfig(api_key=" ")
168
+ with pytest.raises(WherobotsConfigError, match="No Wherobots credentials found"):
169
169
  BaseClient(cfg)
170
170
 
171
171
 
172
+ # ---------------------------------------------------------------------------
173
+ # Credential selection
174
+ # ---------------------------------------------------------------------------
175
+
176
+
177
+ class TestCredentialSelection:
178
+ @patch("wherobots.api.base.requests.Session")
179
+ def test_api_key_wins_over_workload_identity(self, MockSession):
180
+ session = MockSession.return_value
181
+ session.headers = {}
182
+ session.auth = None
183
+
184
+ BaseClient(
185
+ WherobotsConfig(
186
+ api_key="test-key",
187
+ id_token_path="/var/run/secrets/wherobots/id-token",
188
+ base_url="https://api.test.wherobots.com",
189
+ )
190
+ )
191
+
192
+ assert session.headers["X-API-Key"] == "test-key"
193
+ assert session.auth is None
194
+
195
+ @pytest.mark.parametrize("api_key", [None, " "])
196
+ @patch("wherobots.api.workload_auth.requests.post")
197
+ @patch("requests.adapters.HTTPAdapter.send")
198
+ def test_workload_identity_without_api_key(self, mock_send, mock_post, api_key, tmp_path):
199
+ id_token = tmp_path / "id-token"
200
+ id_token.write_text("wid-token", encoding="utf-8")
201
+ exchange = MagicMock(status_code=200)
202
+ exchange.json.return_value = {"access_token": "access-1", "expires_in": 3600}
203
+ mock_post.return_value = exchange
204
+ api_response = requests.Response()
205
+ api_response.status_code = 200
206
+ mock_send.return_value = api_response
207
+
208
+ client = BaseClient(
209
+ WherobotsConfig(
210
+ api_key=api_key,
211
+ id_token_path=str(id_token),
212
+ base_url="https://api.test.wherobots.com/",
213
+ request_timeout_seconds=12,
214
+ )
215
+ )
216
+ client.get("/runs")
217
+
218
+ sent = mock_send.call_args.args[0]
219
+ assert sent.headers["Authorization"] == "Bearer access-1"
220
+ assert "X-API-Key" not in sent.headers
221
+ assert mock_post.call_args.args == ("https://api.test.wherobots.com/ste/oauth2/token",)
222
+ assert mock_post.call_args.kwargs["timeout"] == 12
223
+
224
+
172
225
  # ---------------------------------------------------------------------------
173
226
  # Request error wrapping
174
227
  # ---------------------------------------------------------------------------
@@ -88,6 +88,7 @@ class TestFromEnv:
88
88
  "WHEROBOTS_VERSION",
89
89
  "WHEROBOTS_REQUEST_TIMEOUT_SECONDS",
90
90
  "WHEROBOTS_CLIENT_CHAIN",
91
+ "WHEROBOTS_ID_TOKEN_PATH",
91
92
  ]:
92
93
  monkeypatch.delenv(var, raising=False)
93
94
 
@@ -100,6 +101,7 @@ class TestFromEnv:
100
101
  assert cfg.version == "latest"
101
102
  assert cfg.request_timeout_seconds == 30
102
103
  assert cfg.client_chain is None
104
+ assert cfg.id_token_path is None
103
105
 
104
106
  def test_from_env_timeout_override(self, monkeypatch):
105
107
  monkeypatch.delenv("WHEROBOTS_REQUEST_TIMEOUT_SECONDS", raising=False)
@@ -116,6 +118,16 @@ class TestFromEnv:
116
118
  cfg = WherobotsConfig.from_env(client_chain="client=mcp;ver=0.9")
117
119
  assert cfg.client_chain == "client=mcp;ver=0.9"
118
120
 
121
+ def test_blank_explicit_api_key_falls_back_to_env(self, monkeypatch):
122
+ monkeypatch.setenv("WHEROBOTS_API_KEY", "env-key")
123
+ cfg = WherobotsConfig.from_env(api_key=" ")
124
+ assert cfg.api_key == "env-key"
125
+
126
+ def test_id_token_path_from_env(self, monkeypatch):
127
+ monkeypatch.setenv("WHEROBOTS_ID_TOKEN_PATH", "/var/run/secrets/wherobots/id-token")
128
+ cfg = WherobotsConfig.from_env()
129
+ assert cfg.id_token_path == "/var/run/secrets/wherobots/id-token"
130
+
119
131
 
120
132
  # ---------------------------------------------------------------------------
121
133
  # Frozen
@@ -14,7 +14,11 @@ from wherobots.api.files import (
14
14
  _validate_filename,
15
15
  )
16
16
  from wherobots.config import WherobotsConfig
17
- from wherobots.exceptions import WherobotsAPIError, WherobotsValidationError
17
+ from wherobots.exceptions import (
18
+ WherobotsAPIError,
19
+ WherobotsNotAuthenticatedError,
20
+ WherobotsValidationError,
21
+ )
18
22
  from wherobots.models import StorageIntegration, TemporaryAwsCredentials
19
23
 
20
24
  # ---------------------------------------------------------------------------
@@ -110,6 +114,15 @@ class TestGetStorageCredentials:
110
114
  with pytest.raises(WherobotsAPIError, match="Invalid JSON"):
111
115
  api.get_storage_credentials("s3://x/y")
112
116
 
117
+ def test_workload_token_exchange_failure_passes_through(self, files_api):
118
+ api, session = files_api
119
+ session.request.side_effect = WherobotsNotAuthenticatedError(
120
+ "Workload identity token exchange was rejected (status 401): Session is expired"
121
+ )
122
+
123
+ with pytest.raises(WherobotsNotAuthenticatedError, match="Session is expired"):
124
+ api.get_storage_credentials("s3://x/y")
125
+
113
126
  def test_default_access_mode_is_read(self, files_api, creds_payload):
114
127
  """Least-privilege default: public callers get read-only creds unless
115
128
  they explicitly opt into read_write. The SDK's upload paths pass
@@ -0,0 +1,151 @@
1
+ """Tests for WorkloadIdentityAuth: token exchange, caching, refresh, errors."""
2
+
3
+ from unittest.mock import MagicMock, patch
4
+
5
+ import pytest
6
+ import requests
7
+
8
+ from wherobots.api.workload_auth import WorkloadIdentityAuth, shared_workload_identity_auth
9
+ from wherobots.exceptions import WherobotsConfigError, WherobotsNotAuthenticatedError
10
+
11
+ BASE_URL = "https://api.test.wherobots.com"
12
+
13
+
14
+ @pytest.fixture
15
+ def id_token_file(tmp_path):
16
+ path = tmp_path / "id-token"
17
+ path.write_text("wid-token\n", encoding="utf-8")
18
+ return path
19
+
20
+
21
+ def _token_response(access_token, expires_in=3600):
22
+ response = MagicMock()
23
+ response.status_code = 200
24
+ response.raise_for_status.return_value = None
25
+ response.json.return_value = {
26
+ "access_token": access_token,
27
+ "expires_in": expires_in,
28
+ "token_type": "Bearer",
29
+ }
30
+ return response
31
+
32
+
33
+ def _authorize(auth):
34
+ request = requests.Request("GET", f"{BASE_URL}/runs").prepare()
35
+ return auth(request).headers["Authorization"]
36
+
37
+
38
+ @patch("wherobots.api.workload_auth.requests.post")
39
+ def test_exchanges_the_mounted_token_for_a_bearer(mock_post, id_token_file):
40
+ mock_post.return_value = _token_response("access-1")
41
+ auth = WorkloadIdentityAuth(BASE_URL, str(id_token_file), timeout=30)
42
+
43
+ assert _authorize(auth) == "Bearer access-1"
44
+ mock_post.assert_called_once_with(
45
+ f"{BASE_URL}/ste/oauth2/token",
46
+ data={"grant_type": "refresh_token", "refresh_token": "wid-token", "scope": "api"},
47
+ timeout=30,
48
+ allow_redirects=False,
49
+ )
50
+
51
+
52
+ @patch("wherobots.api.workload_auth.time.monotonic")
53
+ @patch("wherobots.api.workload_auth.requests.post")
54
+ def test_reuses_the_token_until_shortly_before_expiry(mock_post, mock_clock, id_token_file):
55
+ mock_post.side_effect = [_token_response("access-1"), _token_response("access-2")]
56
+ auth = WorkloadIdentityAuth(BASE_URL, str(id_token_file), timeout=30)
57
+
58
+ mock_clock.return_value = 1000.0
59
+ assert _authorize(auth) == "Bearer access-1"
60
+ mock_clock.return_value = 1000.0 + 3600 - 61
61
+ assert _authorize(auth) == "Bearer access-1"
62
+ mock_clock.return_value = 1000.0 + 3600 - 60
63
+ assert _authorize(auth) == "Bearer access-2"
64
+ assert mock_post.call_count == 2
65
+
66
+
67
+ @patch("wherobots.api.workload_auth.time.monotonic")
68
+ @patch("wherobots.api.workload_auth.requests.post")
69
+ def test_refresh_rereads_the_token_file(mock_post, mock_clock, id_token_file):
70
+ mock_post.side_effect = [_token_response("access-1"), _token_response("access-2")]
71
+ auth = WorkloadIdentityAuth(BASE_URL, str(id_token_file), timeout=30)
72
+
73
+ mock_clock.return_value = 0.0
74
+ _authorize(auth)
75
+ id_token_file.write_text("rotated-wid", encoding="utf-8")
76
+ mock_clock.return_value = 3600.0
77
+ _authorize(auth)
78
+
79
+ assert mock_post.call_args.kwargs["data"]["refresh_token"] == "rotated-wid"
80
+
81
+
82
+ def test_missing_token_file_raises_config_error(tmp_path):
83
+ auth = WorkloadIdentityAuth(BASE_URL, str(tmp_path / "missing"), timeout=30)
84
+
85
+ with pytest.raises(WherobotsConfigError, match="Cannot read the workload identity token"):
86
+ _authorize(auth)
87
+
88
+
89
+ def test_undecodable_token_file_raises_config_error(tmp_path):
90
+ path = tmp_path / "id-token"
91
+ path.write_bytes(b"\xff\xfe")
92
+ auth = WorkloadIdentityAuth(BASE_URL, str(path), timeout=30)
93
+
94
+ with pytest.raises(WherobotsConfigError, match="Cannot read the workload identity token"):
95
+ _authorize(auth)
96
+
97
+
98
+ @patch("wherobots.api.workload_auth.requests.post")
99
+ def test_empty_token_file_raises_config_error(mock_post, tmp_path):
100
+ path = tmp_path / "id-token"
101
+ path.write_text("\n", encoding="utf-8")
102
+ auth = WorkloadIdentityAuth(BASE_URL, str(path), timeout=30)
103
+
104
+ with pytest.raises(WherobotsConfigError, match="is empty"):
105
+ _authorize(auth)
106
+ mock_post.assert_not_called()
107
+
108
+
109
+ @patch("wherobots.api.workload_auth.requests.post")
110
+ def test_rejected_exchange_carries_the_server_reason(mock_post, id_token_file):
111
+ response = MagicMock()
112
+ response.status_code = 401
113
+ response.json.return_value = {"message": "Session is expired"}
114
+ response.raise_for_status.side_effect = requests.HTTPError(response=response)
115
+ mock_post.return_value = response
116
+ auth = WorkloadIdentityAuth(BASE_URL, str(id_token_file), timeout=30)
117
+
118
+ with pytest.raises(WherobotsNotAuthenticatedError) as exc_info:
119
+ _authorize(auth)
120
+ assert str(exc_info.value) == (
121
+ "Workload identity token exchange was rejected (status 401): Session is expired"
122
+ )
123
+ # The status belongs to the exchange, not to the request being authenticated.
124
+ assert exc_info.value.status_code is None
125
+
126
+
127
+ @patch("wherobots.api.workload_auth.requests.post")
128
+ def test_connection_failure_raises_not_authenticated(mock_post, id_token_file):
129
+ mock_post.side_effect = requests.ConnectionError("unreachable")
130
+ auth = WorkloadIdentityAuth(BASE_URL, str(id_token_file), timeout=30)
131
+
132
+ with pytest.raises(WherobotsNotAuthenticatedError, match="unreachable"):
133
+ _authorize(auth)
134
+
135
+
136
+ @patch("wherobots.api.workload_auth.requests.post")
137
+ def test_unexpected_body_raises_not_authenticated(mock_post, id_token_file):
138
+ response = _token_response("access-1")
139
+ response.json.return_value = {"access_token": "access-1"}
140
+ mock_post.return_value = response
141
+ auth = WorkloadIdentityAuth(BASE_URL, str(id_token_file), timeout=30)
142
+
143
+ with pytest.raises(WherobotsNotAuthenticatedError, match="unexpected body"):
144
+ _authorize(auth)
145
+
146
+
147
+ def test_shared_auth_is_one_instance_per_token(tmp_path):
148
+ first = shared_workload_identity_auth(BASE_URL, str(tmp_path / "a"), 30)
149
+
150
+ assert shared_workload_identity_auth(BASE_URL, str(tmp_path / "a"), 30) is first
151
+ assert shared_workload_identity_auth(BASE_URL, str(tmp_path / "b"), 30) is not first
@@ -1,5 +1,5 @@
1
1
  """Version information."""
2
2
 
3
- __version__ = "0.6.0"
3
+ __version__ = "0.7.0"
4
4
  __author__ = "Wherobots"
5
5
  __email__ = "support@wherobots.com"
@@ -1,7 +1,8 @@
1
1
  """Low-level HTTP client for the Wherobots REST API.
2
2
 
3
3
  ``BaseClient`` wraps a ``requests.Session`` with:
4
- * Auth header construction (``X-API-Key``).
4
+ * Auth: ``X-API-Key`` when an API key is configured, otherwise the
5
+ workload-identity token mounted in Wherobots notebook and job pods.
5
6
  * Client attribution via the advisory ``X-Wherobots-Client`` header,
6
7
  set once on the session so every outgoing request carries it.
7
8
  * Automatic retry with exponential back-off for transient errors
@@ -20,6 +21,7 @@ from requests.adapters import HTTPAdapter
20
21
  from urllib3.util.retry import Retry
21
22
 
22
23
  from wherobots.__version__ import __version__
24
+ from wherobots.api.workload_auth import shared_workload_identity_auth
23
25
  from wherobots.config import WherobotsConfig # noqa: TC001
24
26
  from wherobots.exceptions import WherobotsAPIError, WherobotsConfigError
25
27
  from wherobots.utils.client_header import CLIENT_HEADER, client_header_value
@@ -57,8 +59,8 @@ class BaseClient:
57
59
  back-off, and unified error mapping to ``WherobotsAPIError``.
58
60
 
59
61
  Args:
60
- config: A populated ``WherobotsConfig``. Must contain at least
61
- ``api_key`` and ``base_url``.
62
+ config: A populated ``WherobotsConfig``. Must contain ``base_url``
63
+ and either ``api_key`` or ``id_token_path``.
62
64
  max_retries: Maximum number of retries for transient failures.
63
65
  backoff_factor: Multiplier for exponential back-off between
64
66
  retries.
@@ -75,10 +77,11 @@ class BaseClient:
75
77
  # session is built.
76
78
  self._closed = False
77
79
 
78
- if not config.api_key or not config.api_key.strip():
80
+ api_key = (config.api_key or "").strip()
81
+ if not api_key and not config.id_token_path:
79
82
  raise WherobotsConfigError(
80
- "API key required. Provide via api_key parameter or "
81
- "WHEROBOTS_API_KEY environment variable."
83
+ "No Wherobots credentials found. Provide the api_key parameter, set "
84
+ "WHEROBOTS_API_KEY, or run inside a Wherobots notebook or job."
82
85
  )
83
86
 
84
87
  self._config = config
@@ -122,9 +125,14 @@ class BaseClient:
122
125
  "Content-Type": "application/json",
123
126
  "User-Agent": f"wherobots-python-sdk/{__version__}",
124
127
  CLIENT_HEADER: client_header_value(config.client_chain),
125
- **self._auth_headers(config.api_key),
126
128
  }
127
129
  )
130
+ if api_key:
131
+ self._session.headers.update(self._auth_headers(api_key))
132
+ elif config.id_token_path:
133
+ self._session.auth = shared_workload_identity_auth(
134
+ self._base_url, config.id_token_path, self._timeout
135
+ )
128
136
 
129
137
  @property
130
138
  def timeout(self) -> int:
@@ -336,7 +336,7 @@ class FilesDrive:
336
336
  """
337
337
  if err.status_code == _HTTP_UNAUTHORIZED:
338
338
  return WherobotsNotAuthenticatedError(
339
- "The Wherobots API key was not accepted (HTTP 401)",
339
+ "The Wherobots API did not accept the credentials (HTTP 401)",
340
340
  status_code=err.status_code,
341
341
  response=err.response,
342
342
  request_id=err.request_id,
@@ -29,7 +29,11 @@ from botocore.session import Session as BotocoreSession
29
29
 
30
30
  from wherobots.api.base import BaseClient
31
31
  from wherobots.config import WherobotsConfig # noqa: TC001
32
- from wherobots.exceptions import WherobotsAPIError, WherobotsValidationError
32
+ from wherobots.exceptions import (
33
+ WherobotsAPIError,
34
+ WherobotsNotAuthenticatedError,
35
+ WherobotsValidationError,
36
+ )
33
37
  from wherobots.models import StorageIntegration, TemporaryAwsCredentials
34
38
  from wherobots.utils.logger import get_logger
35
39
 
@@ -110,6 +114,9 @@ class FilesAPI:
110
114
  "/storage/credentials",
111
115
  json_body={"path": path, "access_mode": access_mode},
112
116
  )
117
+ except WherobotsNotAuthenticatedError:
118
+ # The workload token exchange failed before this request was sent (ADR-0003).
119
+ raise
113
120
  except WherobotsAPIError as exc:
114
121
  # Scrub: never let the credentials endpoint's error
115
122
  # response body bubble out — it could contain partial
@@ -0,0 +1,116 @@
1
+ """Workload-identity auth for code running inside Wherobots (ADR-0003).
2
+
3
+ Wherobots notebook and job pods mount a workload-identity token and set
4
+ ``WHEROBOTS_ID_TOKEN_PATH`` to it. The API accepts that token as an OAuth
5
+ refresh token at ``/ste/oauth2/token`` and returns a short-lived bearer token
6
+ that acts as the user who launched the workload.
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ import functools
12
+ import threading
13
+ import time
14
+ from pathlib import Path
15
+
16
+ import requests
17
+ from requests.auth import AuthBase
18
+
19
+ from wherobots.exceptions import WherobotsConfigError, WherobotsNotAuthenticatedError
20
+
21
+ _EXCHANGE_PATH = "/ste/oauth2/token"
22
+ # Descriptive only: the API does not enforce scope on these tokens.
23
+ _SCOPE = "api"
24
+ _REFRESH_MARGIN_SECONDS = 60
25
+
26
+
27
+ class WorkloadIdentityAuth(AuthBase):
28
+ """Sets ``Authorization: Bearer`` from the mounted workload-identity token.
29
+
30
+ The bearer token is cached and swapped for a new one shortly before it
31
+ expires, so long-running jobs keep working past its one-hour lifetime.
32
+
33
+ Args:
34
+ base_url: Wherobots API base URL (already validated as HTTPS).
35
+ id_token_path: Path to the mounted workload-identity token.
36
+ timeout: Timeout in seconds for the token exchange request.
37
+ """
38
+
39
+ def __init__(self, base_url: str, id_token_path: str, timeout: int) -> None:
40
+ """Build the auth; see the class docstring."""
41
+ self._token_url = f"{base_url}{_EXCHANGE_PATH}"
42
+ self._id_token_path = Path(id_token_path)
43
+ self._timeout = timeout
44
+ self._lock = threading.Lock()
45
+ self._access_token: str | None = None
46
+ self._refresh_at = 0.0
47
+
48
+ def __call__(self, request: requests.PreparedRequest) -> requests.PreparedRequest:
49
+ """Attach the current bearer token to *request*."""
50
+ request.headers["Authorization"] = f"Bearer {self._current_token()}"
51
+ return request
52
+
53
+ def _current_token(self) -> str:
54
+ with self._lock:
55
+ if self._access_token is None or time.monotonic() >= self._refresh_at:
56
+ token, expires_in = self._exchange()
57
+ self._access_token = token
58
+ self._refresh_at = time.monotonic() + max(0, expires_in - _REFRESH_MARGIN_SECONDS)
59
+ return self._access_token
60
+
61
+ def _exchange(self) -> tuple[str, int]:
62
+ try:
63
+ id_token = self._id_token_path.read_text(encoding="utf-8").strip()
64
+ except (OSError, UnicodeDecodeError) as exc:
65
+ raise WherobotsConfigError(
66
+ f"Cannot read the workload identity token at {self._id_token_path}: {exc}"
67
+ ) from exc
68
+ if not id_token:
69
+ raise WherobotsConfigError(
70
+ f"The workload identity token at {self._id_token_path} is empty"
71
+ )
72
+ # Errors carry no status_code: it would read as the status of the caller's own request.
73
+ try:
74
+ response = requests.post(
75
+ self._token_url,
76
+ data={"grant_type": "refresh_token", "refresh_token": id_token, "scope": _SCOPE},
77
+ timeout=self._timeout,
78
+ allow_redirects=False,
79
+ )
80
+ response.raise_for_status()
81
+ payload = response.json()
82
+ except requests.HTTPError as exc:
83
+ raise WherobotsNotAuthenticatedError(_rejection_message(exc.response)) from exc
84
+ except (requests.RequestException, ValueError) as exc:
85
+ raise WherobotsNotAuthenticatedError(
86
+ f"Workload identity token exchange failed: {exc}"
87
+ ) from exc
88
+
89
+ token = payload.get("access_token") if isinstance(payload, dict) else None
90
+ expires_in = payload.get("expires_in") if isinstance(payload, dict) else None
91
+ if not isinstance(token, str) or not isinstance(expires_in, int):
92
+ raise WherobotsNotAuthenticatedError(
93
+ "Workload identity token exchange returned an unexpected body"
94
+ )
95
+ return token, expires_in
96
+
97
+
98
+ @functools.cache
99
+ def shared_workload_identity_auth(
100
+ base_url: str, id_token_path: str, timeout: int
101
+ ) -> WorkloadIdentityAuth:
102
+ """Return one auth per process and token, so every client reuses one bearer token."""
103
+ return WorkloadIdentityAuth(base_url, id_token_path, timeout)
104
+
105
+
106
+ def _rejection_message(response: requests.Response | None) -> str:
107
+ message = "Workload identity token exchange was rejected"
108
+ if response is None:
109
+ return message
110
+ message = f"{message} (status {response.status_code})"
111
+ try:
112
+ body = response.json()
113
+ except ValueError:
114
+ return message
115
+ reason = body.get("message") if isinstance(body, dict) else None
116
+ return f"{message}: {reason}" if isinstance(reason, str) else message
@@ -25,6 +25,9 @@ class WherobotsConfig:
25
25
  embed this SDK (e.g. the Airflow provider). The SDK appends its
26
26
  own hop to the right of it, so the origin client stays leftmost.
27
27
  Advisory only — it never affects authentication.
28
+ id_token_path: Workload-identity token mounted in Wherobots notebook
29
+ and job pods. Used for auth only when no ``api_key`` is set
30
+ (ADR-0003).
28
31
  """
29
32
 
30
33
  api_key: str | None = None
@@ -35,6 +38,7 @@ class WherobotsConfig:
35
38
  version: str = "latest"
36
39
  request_timeout_seconds: int = 30
37
40
  client_chain: str | None = None
41
+ id_token_path: str | None = None
38
42
 
39
43
  def __repr__(self) -> str:
40
44
  """Return a string representation with the API key masked."""
@@ -54,7 +58,8 @@ class WherobotsConfig:
54
58
  f"base_url={self.base_url!r}, "
55
59
  f"version={self.version!r}, "
56
60
  f"request_timeout_seconds={self.request_timeout_seconds!r}, "
57
- f"client_chain={self.client_chain!r})"
61
+ f"client_chain={self.client_chain!r}, "
62
+ f"id_token_path={self.id_token_path!r})"
58
63
  )
59
64
 
60
65
  def __post_init__(self) -> None:
@@ -90,6 +95,7 @@ class WherobotsConfig:
90
95
  version: str | None = None,
91
96
  request_timeout_seconds: int | None = None,
92
97
  client_chain: str | None = None,
98
+ id_token_path: str | None = None,
93
99
  ) -> WherobotsConfig:
94
100
  """Build configuration from environment variables with optional overrides.
95
101
 
@@ -108,6 +114,8 @@ class WherobotsConfig:
108
114
  client_chain: Override for ``WHEROBOTS_CLIENT_CHAIN`` — the
109
115
  upstream ``X-Wherobots-Client`` chain this SDK appends its
110
116
  own hop to.
117
+ id_token_path: Override for ``WHEROBOTS_ID_TOKEN_PATH``, which
118
+ Wherobots notebook and job pods set.
111
119
 
112
120
  Returns:
113
121
  A validated ``WherobotsConfig`` instance.
@@ -163,7 +171,8 @@ class WherobotsConfig:
163
171
  )
164
172
 
165
173
  return cls(
166
- api_key=api_key or _env("WHEROBOTS_API_KEY"),
174
+ # A blank key must not hide WHEROBOTS_API_KEY behind workload identity.
175
+ api_key=(api_key or "").strip() or _env("WHEROBOTS_API_KEY"),
167
176
  region=region or _env("WHEROBOTS_REGION"),
168
177
  s3_bucket=s3_bucket or _env("WHEROBOTS_S3_BUCKET"),
169
178
  s3_prefix=s3_prefix or _env("WHEROBOTS_S3_PREFIX") or "wherobots-jobs",
@@ -173,4 +182,5 @@ class WherobotsConfig:
173
182
  version=version or _env("WHEROBOTS_VERSION") or "latest",
174
183
  request_timeout_seconds=timeout,
175
184
  client_chain=client_chain or _env("WHEROBOTS_CLIENT_CHAIN"),
185
+ id_token_path=id_token_path or _env("WHEROBOTS_ID_TOKEN_PATH"),
176
186
  )
@@ -48,7 +48,7 @@ class WherobotsTimeoutError(WherobotsJobError):
48
48
 
49
49
 
50
50
  class WherobotsNotAuthenticatedError(WherobotsAPIError):
51
- """Raised when the Wherobots API does not accept the API key (HTTP 401)."""
51
+ """Raised when the Wherobots API does not accept the credentials."""
52
52
 
53
53
 
54
54
  class WherobotsFileNotFoundError(WherobotsAPIError):
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: wherobots-python-sdk
3
- Version: 0.6.0
3
+ Version: 0.7.0
4
4
  Summary: Python SDK for Wherobots (currently covers the Jobs REST API)
5
5
  Author-email: Wherobots <support@wherobots.com>
6
6
  License-Expression: Apache-2.0
@@ -21,7 +21,7 @@ Requires-Python: >=3.11
21
21
  Description-Content-Type: text/markdown
22
22
  License-File: LICENSE
23
23
  Requires-Dist: requests>=2.34.2
24
- Requires-Dist: boto3>=1.43.108
24
+ Requires-Dist: boto3>=1.34.0
25
25
  Dynamic: license-file
26
26
 
27
27
  # Wherobots Python SDK
@@ -118,18 +118,27 @@ job = WherobotsJob(
118
118
 
119
119
  The API key can also be set via the `WHEROBOTS_API_KEY` environment variable.
120
120
 
121
+ Inside a Wherobots notebook or job, no API key is needed. The SDK authenticates as the user who
122
+ launched the workload, using the token Wherobots mounts in the pod. An API key, when set, takes
123
+ precedence.
124
+
125
+ ```python
126
+ job = WherobotsJob(script="s3://bucket/script.py", name="my-job-001")
127
+ ```
128
+
121
129
  ## Configuration
122
130
 
123
131
  ### Environment Variables
124
132
 
125
133
  | Variable | Description | Default |
126
134
  |------------------------------------|---------------------------------|--------------------------------------|
127
- | `WHEROBOTS_API_KEY` | API key | *(required)* |
135
+ | `WHEROBOTS_API_KEY` | API key | *(required outside Wherobots)* |
128
136
  | `WHEROBOTS_REGION` | Region override (else org default) | *(none — org default)* |
129
137
  | `WHEROBOTS_API_BASE_URL` | API base URL | `https://api.cloud.wherobots.com` |
130
138
  | `WHEROBOTS_VERSION` | Wherobots version | `latest` |
131
139
  | `WHEROBOTS_REQUEST_TIMEOUT_SECONDS`| HTTP request timeout (seconds) | `30` |
132
140
  | `WHEROBOTS_CLIENT_CHAIN` | Upstream client-attribution chain (see below) | *(none)* |
141
+ | `WHEROBOTS_ID_TOKEN_PATH` | Workload identity token, set by Wherobots notebooks and jobs | *(none)* |
133
142
 
134
143
  ### WherobotsConfig
135
144
 
@@ -456,7 +465,7 @@ with wherobots.Wherobots() as wb:
456
465
  | Failure | Raised as |
457
466
  |---|---|
458
467
  | Files is not enabled for this drive in this region | `WherobotsFilesNotEnabledError` (`.drive`, `.region`) |
459
- | The API key was not accepted | `WherobotsNotAuthenticatedError` |
468
+ | The credentials were not accepted | `WherobotsNotAuthenticatedError` |
460
469
  | No such file or folder | `WherobotsFileNotFoundError` |
461
470
  | Transfer to or from storage failed | `WherobotsS3Error` |
462
471
 
@@ -673,7 +682,7 @@ All exceptions inherit from `WherobotsJobError`.
673
682
  | `WherobotsS3Error` | A transfer to or from storage failed, including Files presigned uploads and downloads. |
674
683
  | `WherobotsConfigError` | Configuration is missing or invalid. |
675
684
  | `WherobotsTimeoutError` | Job exceeded the specified timeout. |
676
- | `WherobotsNotAuthenticatedError` | Files: the API key was not accepted (HTTP 401). Subclass of `WherobotsAPIError`. |
685
+ | `WherobotsNotAuthenticatedError` | Files: the credentials were not accepted (HTTP 401). In a Wherobots pod: the workload identity token exchange failed. Subclass of `WherobotsAPIError`. |
677
686
  | `WherobotsFilesNotEnabledError` | Files is not enabled for a drive in a region. Has `drive`, `region`. Subclass of `WherobotsAPIError`. |
678
687
  | `WherobotsFileNotFoundError` | Files: no such file or folder. Subclass of `WherobotsAPIError`. |
679
688
 
@@ -18,6 +18,7 @@ tests/test_regressions.py
18
18
  tests/test_security.py
19
19
  tests/test_tasks.py
20
20
  tests/test_utils.py
21
+ tests/test_workload_auth.py
21
22
  wherobots/__init__.py
22
23
  wherobots/__version__.py
23
24
  wherobots/client.py
@@ -33,6 +34,7 @@ wherobots/api/drive.py
33
34
  wherobots/api/files.py
34
35
  wherobots/api/runs.py
35
36
  wherobots/api/tasks.py
37
+ wherobots/api/workload_auth.py
36
38
  wherobots/utils/__init__.py
37
39
  wherobots/utils/client_header.py
38
40
  wherobots/utils/logger.py
@@ -1,2 +1,2 @@
1
1
  requests>=2.34.2
2
- boto3>=1.43.108
2
+ boto3>=1.34.0