Skip to content

Commit 17dd9d3

Browse files
toppercodesampagent
toppercodes
and
ampagent
committed
fix: support edge parameter in clients
Amp-Thread-ID: https://ampcode.com/threads/T-019d7839-1ded-7439-83da-c8092b44c943 Co-authored-by: Amp <amp@ampcode.com>
1 parent a1e9f35 commit 17dd9d3

4 files changed

Lines changed: 218 additions & 52 deletions

File tree

‎src/axiom_py/client.py‎

Lines changed: 38 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -171,6 +171,7 @@ def __init__(
171171
org_id: Optional[str] = None,
172172
url: Optional[str] = None,
173173
edge_url: Optional[str] = None,
174+
edge: Optional[str] = None,
174175
):
175176
"""
176177
Initialize the Axiom client.
@@ -187,6 +188,12 @@ def __init__(
187188
requests use `/v1/ingest/{dataset}` and query requests use
188189
`/v1/query/_apl`.
189190
Must be passed explicitly (not read from environment).
191+
Takes precedence over `edge`.
192+
edge: Edge domain for ingest/query operations (e.g.,
193+
"eu-central-1.aws.edge.axiom.co"). When set, ingest
194+
requests use `https://{edge}/v1/ingest/{dataset}` and query
195+
requests use `https://{edge}/v1/query/_apl`.
196+
Must be passed explicitly (not read from environment).
190197
"""
191198
# fallback to env variables if not provided
192199
if token is None:
@@ -195,19 +202,21 @@ def __init__(
195202
org_id = os.getenv("AXIOM_ORG_ID")
196203
if url is None:
197204
url = os.getenv("AXIOM_URL")
198-
# Note: edge_url is NOT auto-read from environment.
205+
# Note: edge_url and edge are NOT auto-read from environment.
199206
# Edge configuration must be explicit to avoid accidentally routing
200-
# all requests through edge when AXIOM_EDGE_URL is set for
207+
# all requests through edge when AXIOM_EDGE_URL or AXIOM_EDGE is set for
201208
# edge-specific tests. Create a separate Client with edge_url
202209
# for edge operations.
203210

204211
# Normalize empty strings to None for edge config
205212
edge_url = edge_url or None
213+
edge = edge or None
206214

207215
# Store for building ingest/query endpoints
208216
self._token = token
209217
self._url = url
210218
self._edge_url = edge_url
219+
self._edge = edge
211220

212221
# Determine API base URL (for non-ingest/query operations)
213222
# This always uses AXIOM_URL (api.axiom.co) unless a custom url is set
@@ -250,45 +259,51 @@ def __init__(
250259

251260
def is_edge_configured(self) -> bool:
252261
"""Check if edge is configured."""
253-
return self._edge_url is not None
262+
return self._edge_url is not None or self._edge is not None
254263

255264
def _get_edge_ingest_url(self, dataset: str) -> Optional[str]:
256265
"""
257266
Get the full edge ingest URL for a dataset.
258267
Returns None if edge is not configured.
259268
"""
260-
if self._edge_url is None:
261-
return None
269+
if self._edge_url is not None:
270+
url = self._edge_url.rstrip("/")
271+
parsed = urlparse(url)
272+
path = parsed.path
273+
274+
# If path is empty or just "/", append edge ingest format
275+
if path == "" or path == "/":
276+
return f"{parsed.scheme}://{parsed.netloc}/v1/ingest/{dataset}"
262277

263-
url = self._edge_url.rstrip("/")
264-
parsed = urlparse(url)
265-
path = parsed.path
278+
# edge_url has a custom path, use as-is
279+
return url
266280

267-
# If path is empty or just "/", append edge ingest format
268-
if path == "" or path == "/":
269-
return f"{parsed.scheme}://{parsed.netloc}/v1/ingest/{dataset}"
281+
if self._edge is not None:
282+
return f"https://{self._edge.rstrip('/')}/v1/ingest/{dataset}"
270283

271-
# edge_url has a custom path, use as-is
272-
return url
284+
return None
273285

274286
def _get_edge_query_url(self) -> Optional[str]:
275287
"""
276288
Get the full edge query URL.
277289
Returns None if edge is not configured.
278290
"""
279-
if self._edge_url is None:
280-
return None
291+
if self._edge_url is not None:
292+
url = self._edge_url.rstrip("/")
293+
parsed = urlparse(url)
294+
path = parsed.path
295+
296+
# If path is empty or just "/", append edge query format
297+
if path == "" or path == "/":
298+
return f"{parsed.scheme}://{parsed.netloc}/v1/query/_apl"
281299

282-
url = self._edge_url.rstrip("/")
283-
parsed = urlparse(url)
284-
path = parsed.path
300+
# edge_url has a custom path, use as-is
301+
return url
285302

286-
# If path is empty or just "/", append edge query format
287-
if path == "" or path == "/":
288-
return f"{parsed.scheme}://{parsed.netloc}/v1/query/_apl"
303+
if self._edge is not None:
304+
return f"https://{self._edge.rstrip('/')}/v1/query/_apl"
289305

290-
# edge_url has a custom path, use as-is
291-
return url
306+
return None
292307

293308
def before_shutdown(self, func: Callable):
294309
self.before_shutdown_funcs.append(func)

‎src/axiom_py/client_async.py‎

Lines changed: 38 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ def __init__(
5050
org_id: Optional[str] = None,
5151
url: Optional[str] = None,
5252
edge_url: Optional[str] = None,
53+
edge: Optional[str] = None,
5354
):
5455
"""
5556
Initialize the async Axiom client.
@@ -65,6 +66,12 @@ def __init__(
6566
requests use `/v1/ingest/{dataset}` and query requests use
6667
`/v1/query/_apl`.
6768
Must be passed explicitly (not read from environment).
69+
Takes precedence over `edge`.
70+
edge: Edge domain for ingest/query operations (e.g.,
71+
"eu-central-1.aws.edge.axiom.co"). When set, ingest
72+
requests use `https://{edge}/v1/ingest/{dataset}` and query
73+
requests use `https://{edge}/v1/query/_apl`.
74+
Must be passed explicitly (not read from environment).
6875
6976
Example:
7077
```python
@@ -80,18 +87,20 @@ def __init__(
8087
if url is None:
8188
url = AXIOM_URL
8289

83-
# Note: edge_url is NOT auto-read from environment.
90+
# Note: edge_url and edge are NOT auto-read from environment.
8491
# Edge configuration must be explicit to avoid accidentally routing
85-
# all requests through edge when AXIOM_EDGE_URL is set for
92+
# all requests through edge when AXIOM_EDGE_URL or AXIOM_EDGE is set for
8693
# edge-specific tests. Create a separate AsyncClient with edge_url
8794
# for edge operations.
8895

8996
# Normalize empty strings to None for edge config
9097
edge_url = edge_url or None
98+
edge = edge or None
9199

92100
# Store for building ingest/query endpoints
93101
self._token = token
94102
self._edge_url = edge_url
103+
self._edge = edge
95104

96105
# Get common headers
97106
headers = get_common_headers(token, org_id)
@@ -129,45 +138,51 @@ async def close(self):
129138

130139
def is_edge_configured(self) -> bool:
131140
"""Check if edge is configured."""
132-
return self._edge_url is not None
141+
return self._edge_url is not None or self._edge is not None
133142

134143
def _get_edge_ingest_url(self, dataset: str) -> Optional[str]:
135144
"""
136145
Get the full edge ingest URL for a dataset.
137146
Returns None if edge is not configured.
138147
"""
139-
if self._edge_url is None:
140-
return None
148+
if self._edge_url is not None:
149+
url = self._edge_url.rstrip("/")
150+
parsed = urlparse(url)
151+
path = parsed.path
152+
153+
# If path is empty or just "/", append edge ingest format
154+
if path == "" or path == "/":
155+
return f"{parsed.scheme}://{parsed.netloc}/v1/ingest/{dataset}"
141156

142-
url = self._edge_url.rstrip("/")
143-
parsed = urlparse(url)
144-
path = parsed.path
157+
# edge_url has a custom path, use as-is
158+
return url
145159

146-
# If path is empty or just "/", append edge ingest format
147-
if path == "" or path == "/":
148-
return f"{parsed.scheme}://{parsed.netloc}/v1/ingest/{dataset}"
160+
if self._edge is not None:
161+
return f"https://{self._edge.rstrip('/')}/v1/ingest/{dataset}"
149162

150-
# edge_url has a custom path, use as-is
151-
return url
163+
return None
152164

153165
def _get_edge_query_url(self) -> Optional[str]:
154166
"""
155167
Get the full edge query URL.
156168
Returns None if edge is not configured.
157169
"""
158-
if self._edge_url is None:
159-
return None
170+
if self._edge_url is not None:
171+
url = self._edge_url.rstrip("/")
172+
parsed = urlparse(url)
173+
path = parsed.path
174+
175+
# If path is empty or just "/", append edge query format
176+
if path == "" or path == "/":
177+
return f"{parsed.scheme}://{parsed.netloc}/v1/query/_apl"
160178

161-
url = self._edge_url.rstrip("/")
162-
parsed = urlparse(url)
163-
path = parsed.path
179+
# edge_url has a custom path, use as-is
180+
return url
164181

165-
# If path is empty or just "/", append edge query format
166-
if path == "" or path == "/":
167-
return f"{parsed.scheme}://{parsed.netloc}/v1/query/_apl"
182+
if self._edge is not None:
183+
return f"https://{self._edge.rstrip('/')}/v1/query/_apl"
168184

169-
# edge_url has a custom path, use as-is
170-
return url
185+
return None
171186

172187
async def ingest(
173188
self,

‎tests/test_client.py‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -325,9 +325,30 @@ class TestEdgeConfiguration(unittest.TestCase):
325325

326326
def _clear_env(self):
327327
"""Helper to clear edge-related env vars."""
328+
os.environ.pop("AXIOM_EDGE", None)
328329
os.environ.pop("AXIOM_URL", None)
329330
os.environ.pop("AXIOM_EDGE_URL", None)
330331

332+
def test_edge_builds_correct_urls(self):
333+
"""Test that edge config builds correct edge ingest and query URLs."""
334+
with patch.dict(os.environ, {}, clear=False):
335+
self._clear_env()
336+
client = Client(
337+
token="xaat-test-token",
338+
org_id="test-org",
339+
edge="eu-central-1.aws.edge.axiom.co",
340+
)
341+
self.assertEqual(client._edge, "eu-central-1.aws.edge.axiom.co")
342+
self.assertEqual(
343+
client._get_edge_ingest_url("my-dataset"),
344+
"https://eu-central-1.aws.edge.axiom.co/v1/ingest/my-dataset",
345+
)
346+
self.assertEqual(
347+
client._get_edge_query_url(),
348+
"https://eu-central-1.aws.edge.axiom.co/v1/query/_apl",
349+
)
350+
self.assertTrue(client.is_edge_configured())
351+
331352
def test_edge_url_builds_correct_ingest_url(self):
332353
"""Test that edge_url config builds correct edge ingest URL."""
333354
with patch.dict(os.environ, {}, clear=False):
@@ -359,6 +380,25 @@ def test_edge_url_builds_correct_query_url(self):
359380
"https://eu-central-1.aws.edge.axiom.co/v1/query/_apl",
360381
)
361382

383+
def test_edge_url_takes_precedence_over_edge(self):
384+
"""Test that edge_url takes precedence over edge."""
385+
with patch.dict(os.environ, {}, clear=False):
386+
self._clear_env()
387+
client = Client(
388+
token="xaat-test-token",
389+
org_id="test-org",
390+
edge="ignored.aws.edge.axiom.co",
391+
edge_url="https://custom-edge.example.com",
392+
)
393+
self.assertEqual(
394+
client._get_edge_ingest_url("my-dataset"),
395+
"https://custom-edge.example.com/v1/ingest/my-dataset",
396+
)
397+
self.assertEqual(
398+
client._get_edge_query_url(),
399+
"https://custom-edge.example.com/v1/query/_apl",
400+
)
401+
362402
def test_no_edge_returns_none(self):
363403
"""Test that no edge config returns None for edge URLs."""
364404
with patch.dict(os.environ, {}, clear=False):
@@ -471,13 +511,16 @@ def test_edge_url_not_read_from_env(self):
471511
self._clear_env()
472512
# Set env var that should be ignored
473513
os.environ["AXIOM_EDGE_URL"] = "https://edge.example.com"
514+
os.environ["AXIOM_EDGE"] = "eu-central-1.aws.edge.axiom.co"
474515

475516
# Client should NOT pick up edge config from env
476517
client = Client(token="xaat-test-token", org_id="test-org")
477518
self.assertIsNone(client._edge_url)
519+
self.assertIsNone(client._edge)
478520
self.assertFalse(client.is_edge_configured())
479521

480522
os.environ.pop("AXIOM_EDGE_URL", None)
523+
os.environ.pop("AXIOM_EDGE", None)
481524

482525
def test_personal_token_rejected_for_edge_ingest(self):
483526
"""Test that personal tokens are rejected for edge ingest."""
@@ -520,3 +563,15 @@ def test_empty_string_edge_url_treated_as_none(self):
520563
)
521564
self.assertIsNone(client._edge_url)
522565
self.assertFalse(client.is_edge_configured())
566+
567+
def test_empty_string_edge_treated_as_none(self):
568+
"""Test that empty string for edge is treated as None."""
569+
with patch.dict(os.environ, {}, clear=False):
570+
self._clear_env()
571+
client = Client(
572+
token="xaat-api-token",
573+
org_id="test-org",
574+
edge="",
575+
)
576+
self.assertIsNone(client._edge)
577+
self.assertFalse(client.is_edge_configured())

0 commit comments

Comments
 (0)