Skip to content
Merged
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
2 changes: 2 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,8 @@ for item in job.results:

`get_job` and `wait_for_job` are **polymorphic** — they work for both bulk and crawl `job_id`s. The SDK reads a `kind` discriminator from the API response and returns either a `BulkJob` or a `CrawlJob`. Use `isinstance(job, CrawlJob)` (or check `job.kind == "crawl"`) before reading crawl-specific fields like `job.truncated` or `item.depth`.

Both resolve a job in a **single request** to `GET /jobs/{id}`, which answers for either kind and requires neither the `bulk` nor the `crawl` scope. Earlier versions polled `/bulk/{id}` first just to read `kind`, then re-fetched `/crawl/{id}` for the crawl-only fields — so a key scoped to `crawl` alone got a `403` and could not poll its own job. Requires an API deployed on or after this endpoint's release.

Bulk submissions carry an automatically generated `Idempotency-Key` (see [Retries & idempotency](#retries--idempotency)), so an internal retry replays the original job rather than enqueuing a second one.

## Crawl
Expand Down
99 changes: 26 additions & 73 deletions tests/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -275,7 +275,7 @@ async def test_async_extract_success() -> None:
@respx.mock
async def test_async_wait_for_job_polls_until_done() -> None:
job_id = "1c4f9a02-0000-0000-0000-000000000000"
route = respx.get(f"{BASE_URL}/bulk/{job_id}")
route = respx.get(f"{BASE_URL}/jobs/{job_id}")
route.side_effect = [
httpx.Response(
200,
Expand Down Expand Up @@ -485,22 +485,10 @@ def test_crawl_negative_depth_validates_client_side() -> None:
@respx.mock
def test_get_job_surfaces_crawl_truncation_fields() -> None:
"""Regression: CrawlJob parses every truncation-related field correctly
when reached through the polymorphic get_job path. /bulk reports the
job's kind as crawl; the SDK re-fetches /crawl for the full shape."""
when reached through the polymorphic get_job path — now from the single
/jobs/{id} response rather than a /crawl re-fetch."""
job_id = "9aaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"
respx.get(f"{BASE_URL}/bulk/{job_id}").mock(
return_value=httpx.Response(
200,
json={
"job_id": job_id, "kind": "crawl",
"status": "done", "total": 1000, "completed": 1000,
"results": [],
"created_at": "2026-05-15T10:00:00+00:00",
"finished_at": "2026-05-15T10:02:30+00:00",
},
)
)
respx.get(f"{BASE_URL}/crawl/{job_id}").mock(
respx.get(f"{BASE_URL}/jobs/{job_id}").mock(
return_value=httpx.Response(
200,
json={
Expand Down Expand Up @@ -555,7 +543,7 @@ def test_get_job_surfaces_crawl_truncation_fields() -> None:
def test_get_job_returns_bulk_job_when_kind_is_bulk() -> None:
"""A response with kind='bulk' (or no kind) → BulkJob, single round-trip."""
job_id = "1c4f9a02-0000-0000-0000-000000000000"
respx.get(f"{BASE_URL}/bulk/{job_id}").mock(
respx.get(f"{BASE_URL}/jobs/{job_id}").mock(
return_value=httpx.Response(
200,
json={
Expand All @@ -579,25 +567,13 @@ def test_get_job_returns_bulk_job_when_kind_is_bulk() -> None:


@respx.mock
def test_get_job_redispatches_to_crawl_when_kind_is_crawl() -> None:
"""When /bulk reports kind='crawl', the SDK re-fetches /crawl for the
proper shape (with depths + truncated fields) and returns a CrawlJob."""
def test_get_job_returns_full_crawl_job_in_one_call() -> None:
"""The regression that matters. get_job used to hit /bulk/{id} just to
read `kind`, then re-fetch /crawl/{id} for truncated + depth — two round
trips, and a 403 for any key holding `crawl` but not `bulk`, since
/bulk/* is scope-gated. One call to /jobs/{id} now carries everything."""
job_id = "9aaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"
# First call: /bulk/{id} reports the job IS a crawl.
respx.get(f"{BASE_URL}/bulk/{job_id}").mock(
return_value=httpx.Response(
200,
json={
"job_id": job_id, "kind": "crawl",
"status": "done", "total": 1, "completed": 1,
"results": [],
"created_at": "2026-05-15T10:00:00+00:00",
"finished_at": "2026-05-15T10:00:30+00:00",
},
)
)
# Second call: /crawl/{id} returns the proper crawl shape.
respx.get(f"{BASE_URL}/crawl/{job_id}").mock(
respx.get(f"{BASE_URL}/jobs/{job_id}").mock(
return_value=httpx.Response(
200,
json={
Expand All @@ -622,31 +598,22 @@ def test_get_job_redispatches_to_crawl_when_kind_is_crawl() -> None:
assert job.truncated
assert job.truncated_reason == "page_cap_reached"
assert job.results[0].depth == 0
assert respx.calls.call_count == 1
# The scope-gated poll routes must not be touched: /bulk/* needs `bulk`
# and /crawl/* needs `crawl`, which is the 403 this endpoint removes.
assert respx.calls[0].request.url.path == f"/jobs/{job_id}"


@respx.mock
def test_wait_for_job_uses_typed_endpoint_after_first_call() -> None:
"""wait_for_job's first call goes through get_job (polymorphic). After it
detects the kind, subsequent polls hit the typed endpoint directly — so
a crawl job's polling loop doesn't re-pay the /bulk + /crawl dispatch
on every iteration."""
def test_wait_for_job_polls_crawl_jobs_with_no_dispatch_round_trip() -> None:
"""Every poll goes to /jobs/{id}, which answers for either kind. The old
loop spent three calls to reach the same result on a crawl job: a /bulk
discovery, a /crawl re-fetch, then the real poll."""
job_id = "9aaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"

# Discovery call: /bulk says crawl, /crawl returns processing.
respx.get(f"{BASE_URL}/bulk/{job_id}").mock(
return_value=httpx.Response(
200,
json={
"job_id": job_id, "kind": "crawl",
"status": "processing", "total": 2, "completed": 1,
"results": [],
"created_at": "2026-05-15T10:00:00+00:00",
},
)
)
crawl_route = respx.get(f"{BASE_URL}/crawl/{job_id}")
crawl_route.side_effect = [
# First /crawl call: completes the discovery (after /bulk re-dispatch).
jobs_route = respx.get(f"{BASE_URL}/jobs/{job_id}")
jobs_route.side_effect = [
# First call — still processing.
httpx.Response(
200,
json={
Expand All @@ -657,7 +624,7 @@ def test_wait_for_job_uses_typed_endpoint_after_first_call() -> None:
"created_at": "2026-05-15T10:00:00+00:00",
},
),
# Second /crawl call: poll iteration after the sleep — should be DONE.
# Second callpoll iteration after the sleep, now DONE.
httpx.Response(
200,
json={
Expand All @@ -679,10 +646,8 @@ def test_wait_for_job_uses_typed_endpoint_after_first_call() -> None:

assert isinstance(job, CrawlJob)
assert job.done
# /bulk called exactly once (discovery), /crawl called twice (one for
# discovery re-fetch, one for the actual poll iteration). If we'd
# routed every poll through get_job we'd see /bulk hit a 2nd time.
assert respx.calls.call_count == 3
assert respx.calls.call_count == 2
assert {c.request.url.path for c in respx.calls} == {f"/jobs/{job_id}"}


@pytest.mark.asyncio
Expand All @@ -693,19 +658,7 @@ async def test_async_wait_for_job_handles_crawl_jobs() -> None:
crawl-specific function."""
job_id = "9aaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"

respx.get(f"{BASE_URL}/bulk/{job_id}").mock(
return_value=httpx.Response(
200,
json={
"job_id": job_id, "kind": "crawl",
"status": "done", "total": 1, "completed": 1,
"results": [],
"created_at": "2026-05-15T10:00:00+00:00",
"finished_at": "2026-05-15T10:00:30+00:00",
},
)
)
respx.get(f"{BASE_URL}/crawl/{job_id}").mock(
respx.get(f"{BASE_URL}/jobs/{job_id}").mock(
return_value=httpx.Response(
200,
json={
Expand Down
31 changes: 15 additions & 16 deletions wellmarked/async_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -178,18 +178,22 @@ async def bulk(
)
return BulkJob.from_response(body)

async def get_job(self, job_id: str) -> Union[BulkJob, CrawlJob]:
"""Polymorphic job lookup — works for both bulk and crawl jobs.

See :meth:`wellmarked.WellMarked.get_job` for the full
explanation of the kind-discriminator + optional re-fetch model.
"""
body = await self._request("GET", f"/bulk/{job_id}")
@staticmethod
def _job_from_body(body: dict[str, Any]) -> Union[BulkJob, CrawlJob]:
"""Build the right job type from a ``/jobs/{id}`` body, using ``kind``."""
if body.get("kind") == "crawl":
body = await self._request("GET", f"/crawl/{job_id}")
return CrawlJob.from_response(body)
return BulkJob.from_response(body)

async def get_job(self, job_id: str) -> Union[BulkJob, CrawlJob]:
"""Polymorphic job lookup — works for both bulk and crawl jobs, in
ONE call to ``GET /jobs/{job_id}``.

See :meth:`wellmarked.WellMarked.get_job` for why this replaced the
old ``/bulk`` discovery call plus ``/crawl`` re-fetch.
"""
return self._job_from_body(await self._request("GET", f"/jobs/{job_id}"))

async def wait_for_job(
self,
job_id: str,
Expand All @@ -198,23 +202,18 @@ async def wait_for_job(
timeout: Optional[float] = 300.0,
) -> Union[BulkJob, CrawlJob]:
"""Await until a job reaches ``status="done"`` (or timeout). Works
for both bulk and crawl jobs — same polymorphic dispatch as the
sync client. See :meth:`wellmarked.WellMarked.wait_for_job`."""
for both bulk and crawl jobs — same single-call polling as the sync
client. See :meth:`wellmarked.WellMarked.wait_for_job`."""
deadline = None if timeout is None else time.monotonic() + timeout
job: Union[BulkJob, CrawlJob] = await self.get_job(job_id)
is_crawl = isinstance(job, CrawlJob)
while not job.done:
if deadline is not None and time.monotonic() >= deadline:
raise TimeoutError(
f"Job {job_id} did not finish within {timeout}s "
f"(last status: {job.status}, {job.completed}/{job.total})"
)
await asyncio.sleep(poll_interval)
# Subsequent polls hit the typed endpoint directly — skips the
# /bulk + /crawl dispatch get_job does for crawl jobs.
path = f"/crawl/{job_id}" if is_crawl else f"/bulk/{job_id}"
body = await self._request("GET", path)
job = CrawlJob.from_response(body) if is_crawl else BulkJob.from_response(body)
job = self._job_from_body(await self._request("GET", f"/jobs/{job_id}"))
return job

async def crawl(
Expand Down
49 changes: 24 additions & 25 deletions wellmarked/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -333,32 +333,37 @@ def bulk(
)
return BulkJob.from_response(body)

@staticmethod
def _job_from_body(body: dict[str, Any]) -> Union[BulkJob, CrawlJob]:
"""Build the right job type from a ``/jobs/{id}`` body, using ``kind``."""
if body.get("kind") == "crawl":
return CrawlJob.from_response(body)
return BulkJob.from_response(body)

def get_job(self, job_id: str) -> Union[BulkJob, CrawlJob]:
"""Polymorphic job lookup — works for both bulk and crawl jobs.
"""Polymorphic job lookup — works for both bulk and crawl jobs, in
ONE call.

Calls ``GET /bulk/{job_id}`` first, then inspects the response's
``kind`` discriminator field. If the job is actually a crawl, a
second request to ``GET /crawl/{job_id}`` fetches the full crawl
shape (with per-item depth and the truncated flags). Returns
``GET /jobs/{job_id}`` resolves either kind and returns the superset
shape; the ``kind`` discriminator says which you got. Returns
:class:`BulkJob` or :class:`CrawlJob` accordingly.

This used to call ``GET /bulk/{job_id}`` purely to read ``kind``,
then re-fetch ``GET /crawl/{job_id}`` for the fields the bulk shape
omits (``truncated``, ``truncated_reason``, per-item ``depth``).
That cost two round trips per crawl poll, and — because ``/bulk/*``
requires the ``bulk`` scope — a key scoped to ``crawl`` alone got a
403 on the discovery call and could never poll its own job.
``/jobs/{id}`` requires neither scope.

Use ``isinstance(job, CrawlJob)`` or ``job.kind == "crawl"`` to
branch on crawl-specific behavior. The shared interface
(``status``, ``completed``, ``total``, ``results``, ``done``)
works on either type.

Jobs are retained for 6 hours after completion.
"""
body = self._request("GET", f"/bulk/{job_id}")
# /bulk/{id} answers for any job_id today (the endpoint just
# serializes results in the bulk shape regardless of stored
# job_type — see api/routes/bulk.py). The `kind` field tells
# us whether we got a bulk-shaped response of a crawl job; if
# so, re-fetch via /crawl/{id} for the proper shape.
if body.get("kind") == "crawl":
body = self._request("GET", f"/crawl/{job_id}")
return CrawlJob.from_response(body)
return BulkJob.from_response(body)
return self._job_from_body(self._request("GET", f"/jobs/{job_id}"))

def wait_for_job(
self,
Expand All @@ -370,9 +375,9 @@ def wait_for_job(
"""Block until a job reaches ``status="done"`` (or timeout). Works
for both bulk and crawl jobs.

The first call uses the polymorphic :meth:`get_job` to discover
the job's kind. Subsequent polls go directly to the typed
endpoint, so a crawl job only pays the dispatch round-trip once.
Every poll goes to ``/jobs/{id}``, which answers for either kind, so
there is no dispatch round-trip to pay and no scope to guess — the
loop reads ``kind`` off each body.

Args:
job_id: The job id returned from :meth:`bulk` or :meth:`crawl`.
Expand All @@ -383,21 +388,15 @@ def wait_for_job(
TimeoutError: The job didn't finish before ``timeout`` elapsed.
"""
deadline = None if timeout is None else time.monotonic() + timeout
# First call: figure out which endpoint owns this job.
job: Union[BulkJob, CrawlJob] = self.get_job(job_id)
is_crawl = isinstance(job, CrawlJob)
while not job.done:
if deadline is not None and time.monotonic() >= deadline:
raise TimeoutError(
f"Job {job_id} did not finish within {timeout}s "
f"(last status: {job.status}, {job.completed}/{job.total})"
)
time.sleep(poll_interval)
# Subsequent polls hit the typed endpoint directly — skips the
# /bulk + /crawl dispatch get_job does for crawl jobs.
path = f"/crawl/{job_id}" if is_crawl else f"/bulk/{job_id}"
body = self._request("GET", path)
job = CrawlJob.from_response(body) if is_crawl else BulkJob.from_response(body)
job = self._job_from_body(self._request("GET", f"/jobs/{job_id}"))
return job

def crawl(
Expand Down