Skip to content

Commit fe796a3

Browse files
toppercodestoppercodes
authored
Merge pull request #183 from axiomhq/topper/axm-11747-support-edge-parameter-in-axiom-py
2 parents 9e5d3e4 + 17dd9d3 commit fe796a3

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
@@ -258,45 +267,51 @@ def __init__(
258267

259268
def is_edge_configured(self) -> bool:
260269
"""Check if edge is configured."""
261-
return self._edge_url is not None
270+
return self._edge_url is not None or self._edge is not None
262271

263272
def _get_edge_ingest_url(self, dataset: str) -> Optional[str]:
264273
"""
265274
Get the full edge ingest URL for a dataset.
266275
Returns None if edge is not configured.
267276
"""
268-
if self._edge_url is None:
269-
return None
277+
if self._edge_url is not None:
278+
url = self._edge_url.rstrip("/")
279+
parsed = urlparse(url)
280+
path = parsed.path
281+
282+
# If path is empty or just "/", append edge ingest format
283+
if path == "" or path == "/":
284+
return f"{parsed.scheme}://{parsed.netloc}/v1/ingest/{dataset}"
270285

271-
url = self._edge_url.rstrip("/")
272-
parsed = urlparse(url)
273-
path = parsed.path
286+
# edge_url has a custom path, use as-is
287+
return url
274288

275-
# If path is empty or just "/", append edge ingest format
276-
if path == "" or path == "/":
277-
return f"{parsed.scheme}://{parsed.netloc}/v1/ingest/{dataset}"
289+
if self._edge is not None:
290+
return f"https://{self._edge.rstrip('/')}/v1/ingest/{dataset}"
278291

279-
# edge_url has a custom path, use as-is
280-
return url
292+
return None
281293

282294
def _get_edge_query_url(self) -> Optional[str]:
283295
"""
284296
Get the full edge query URL.
285297
Returns None if edge is not configured.
286298
"""
287-
if self._edge_url is None:
288-
return None
299+
if self._edge_url is not None:
300+
url = self._edge_url.rstrip("/")
301+
parsed = urlparse(url)
302+
path = parsed.path
303+
304+
# If path is empty or just "/", append edge query format
305+
if path == "" or path == "/":
306+
return f"{parsed.scheme}://{parsed.netloc}/v1/query/_apl"
289307

290-
url = self._edge_url.rstrip("/")
291-
parsed = urlparse(url)
292-
path = parsed.path
308+
# edge_url has a custom path, use as-is
309+
return url
293310

294-
# If path is empty or just "/", append edge query format
295-
if path == "" or path == "/":
296-
return f"{parsed.scheme}://{parsed.netloc}/v1/query/_apl"
311+
if self._edge is not None:
312+
return f"https://{self._edge.rstrip('/')}/v1/query/_apl"
297313

298-
# edge_url has a custom path, use as-is
299-
return url
314+
return None
300315

301316
def before_shutdown(self, func: Callable):
302317
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)