Skip to content

feat: switch to API key auth for Antenna endpoints#136

Open
mihow wants to merge 2 commits intomainfrom
feat/api-key-auth
Open

feat: switch to API key auth for Antenna endpoints#136
mihow wants to merge 2 commits intomainfrom
feat/api-key-auth

Conversation

@mihow
Copy link
Copy Markdown
Collaborator

@mihow mihow commented Apr 4, 2026

Summary

Switch from user token auth (Authorization: Token) to API key auth (Authorization: Api-Key) for all Antenna API requests. The processing service is now identified by its API key rather than by name.

  • Replace antenna_api_auth_token setting with antenna_api_key (env var: AMI_ANTENNA_API_KEY)
  • Remove antenna_service_name setting (service name is now managed in Antenna admin)
  • Remove processing_service_name from pipeline registration request body
  • Add client_info metadata (hostname, software, version, platform) to pipeline registration
  • Remove get_full_service_name helper

Companion to RolnickLab/antenna#1194

Migration

  1. Generate an API key in Antenna admin (Processing Services → select service → "Generate API key") or via POST /api/v2/processing-services/{id}/generate_key/
  2. Replace AMI_ANTENNA_API_AUTH_TOKEN with AMI_ANTENNA_API_KEY in your .env file or environment
  3. Remove AMI_ANTENNA_SERVICE_NAME from your configuration (no longer used)

Test plan

  • ami worker register --project <id> registers pipelines with API key auth
  • ami worker --pipeline <slug> fetches tasks and posts results with API key auth
  • E2E test: run test_ml_job_e2e management command on Antenna side
  • Unit tests pass (pytest trapdata/antenna/tests/)

🤖 Generated with Claude Code

Summary by CodeRabbit

  • Refactor

    • Switched authentication platform-wide from token-based to API key-based
    • Updated configuration to use a single Antenna API key and updated environment variable to AMI_ANTENNA_API_KEY
    • Removed the requirement to provide a service name during pipeline registration
  • New Features

    • Pipeline registration now attaches client metadata (version, hostname, platform) for diagnostics

Replace user token auth (Authorization: Token) with API key auth
(Authorization: Api-Key) for all Antenna API requests. Remove
processing_service_name from registration (service is now identified
by its API key). Add client_info metadata to pipeline registration.

- Replace antenna_api_auth_token setting with antenna_api_key
- Remove antenna_service_name setting (managed in Antenna now)
- Update get_http_session to use Api-Key header
- Add client_info (hostname, software, version, platform) to registration
- Remove get_full_service_name helper (no longer needed)
- Update all callers, tests, and benchmark

Companion to RolnickLab/antenna#1194

Co-Authored-By: Claude <noreply@anthropic.com>
@coderabbitai
Copy link
Copy Markdown

coderabbitai bot commented Apr 4, 2026

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 2a766cdf-39ca-4421-8f2f-e23bad15bb84

📥 Commits

Reviewing files that changed from the base of the PR and between 1683543 and 1d86a46.

📒 Files selected for processing (2)
  • trapdata/antenna/registration.py
  • trapdata/antenna/result_posting.py
✅ Files skipped from review due to trivial changes (1)
  • trapdata/antenna/registration.py
🚧 Files skipped from review as they are similar to previous changes (1)
  • trapdata/antenna/result_posting.py

📝 Walkthrough

Walkthrough

Renamed Antenna authentication from auth_tokenapi_key across client, dataset, worker, benchmark, and result-posting code; removed service_name plumbing and added optional client metadata for pipeline registration; updated Settings and CLI to use antenna_api_key; updated tests to match.

Changes

Cohort / File(s) Summary
HTTP utils
trapdata/api/utils.py
Renamed get_http_session(auth_token)get_http_session(api_key) and changed Authorization header to Api-Key {api_key}.
Antenna client
trapdata/antenna/client.py
Removed socket import and get_full_service_name; renamed auth_tokenapi_key in get_jobs, post_batch_results, get_user_projects; pass api_key into get_http_session.
Result posting
trapdata/antenna/result_posting.py
Renamed ResultPoster.post_async(..., auth_token)...api_key; propagate api_key through executor/task helpers and into post_batch_results calls.
Datasets / Loader
trapdata/antenna/datasets.py
RESTDataset constructor now takes api_key (stored as self.api_key); _ensure_sessions() and get_rest_dataloader() use settings.antenna_api_key.
Benchmark & CLI
trapdata/antenna/benchmark.py, trapdata/cli/worker.py
Benchmark CLI and run_benchmark accept api_key (from AMI_ANTENNA_API_KEY) instead of auth_token; CLI register no longer reads/passes service_name.
Worker
trapdata/antenna/worker.py
Replaced checks/usage of antenna_api_auth_token and antenna_service_name with single required antenna_api_key; use api_key for job polling and result posting; removed get_full_service_name usage and related logging.
Registration & Schemas
trapdata/antenna/registration.py, trapdata/antenna/schemas.py
Removed processing_service_name field/parameter; added _get_version() and _build_client_info() and pass client_info into AsyncPipelineRegistrationRequest; registration functions now accept/propagate api_key and use settings.antenna_api_key.
Settings & Tests
trapdata/settings.py, trapdata/antenna/tests/*
Settings: removed antenna_api_auth_token and antenna_service_name, added antenna_api_key. Tests updated to use antenna_api_key / api_key values and removed service_name test args.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

Poem

🐇 I hopped through headers, keys in paw,

Old tokens shelved without a flaw,
Client-info tucked beneath my nose,
Pipelines register as the sunlight rose,
Api-key bound — the rabbit hums, all set to go.

🚥 Pre-merge checks | ✅ 3
✅ Passed checks (3 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title 'feat: switch to API key auth for Antenna endpoints' accurately summarizes the main change: replacing token-based authentication with API key authentication across the Antenna integration.
Docstring Coverage ✅ Passed Docstring coverage is 89.29% which is sufficient. The required threshold is 80.00%.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch feat/api-key-auth

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

Copy link
Copy Markdown

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
trapdata/antenna/result_posting.py (1)

63-63: ⚠️ Potential issue | 🟡 Minor

Minor docstring inconsistency: example still references auth_token.

The class docstring example should be updated to reflect the new parameter name.

📝 Suggested fix
     Example:
         poster = ResultPoster(max_pending=10)
-        poster.post_async(base_url, auth_token, job_id, results)
+        poster.post_async(base_url, api_key, job_id, results)
         metrics = poster.get_metrics()
         poster.shutdown()
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@trapdata/antenna/result_posting.py` at line 63, The class docstring example
still uses the old name auth_token; update the example so its call to
poster.post_async matches the actual parameter name in the current
poster.post_async signature (replace auth_token with the current name used in
the function signature), ensuring the example parameters (base_url,
<current_auth_param>, job_id, results) match the implementation in
result_posting.py.
🧹 Nitpick comments (1)
trapdata/antenna/schemas.py (1)

96-96: Consider using a TypedDict for client_info to improve type safety.

The _build_client_info() function in trapdata/antenna/registration.py:28-36 produces a specific shape with hostname, software, version, and platform keys. A TypedDict would provide better IDE support and catch typos.

💡 Optional improvement
from typing import TypedDict

class ClientInfo(TypedDict, total=False):
    hostname: str
    software: str
    version: str
    platform: str

class AsyncPipelineRegistrationRequest(pydantic.BaseModel):
    # ...
    client_info: ClientInfo | None = None
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@trapdata/antenna/schemas.py` at line 96, Replace the untyped client_info:
dict | None on AsyncPipelineRegistrationRequest with a TypedDict to mirror the
shape built by _build_client_info; define a ClientInfo TypedDict (total=False)
with keys hostname, software, version, platform (all str) and change the field
type to ClientInfo | None so IDEs and type checkers catch mismatches and typos
when constructing client_info in _build_client_info.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@trapdata/antenna/registration.py`:
- Around line 24-25: The except Exception in the _get_version function is too
broad; narrow it to only the expected metadata/lookup failures by catching
specific exceptions (e.g., importlib.metadata.PackageNotFoundError and other
lookup-related errors like KeyError/AttributeError) instead of Exception, and
import PackageNotFoundError from importlib.metadata (or reference it fully) so
the handler returns "unknown" only for those expected failure modes.

---

Outside diff comments:
In `@trapdata/antenna/result_posting.py`:
- Line 63: The class docstring example still uses the old name auth_token;
update the example so its call to poster.post_async matches the actual parameter
name in the current poster.post_async signature (replace auth_token with the
current name used in the function signature), ensuring the example parameters
(base_url, <current_auth_param>, job_id, results) match the implementation in
result_posting.py.

---

Nitpick comments:
In `@trapdata/antenna/schemas.py`:
- Line 96: Replace the untyped client_info: dict | None on
AsyncPipelineRegistrationRequest with a TypedDict to mirror the shape built by
_build_client_info; define a ClientInfo TypedDict (total=False) with keys
hostname, software, version, platform (all str) and change the field type to
ClientInfo | None so IDEs and type checkers catch mismatches and typos when
constructing client_info in _build_client_info.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 18a4f370-f380-4a34-965d-c9fc905943b4

📥 Commits

Reviewing files that changed from the base of the PR and between 0987099 and 1683543.

📒 Files selected for processing (12)
  • trapdata/antenna/benchmark.py
  • trapdata/antenna/client.py
  • trapdata/antenna/datasets.py
  • trapdata/antenna/registration.py
  • trapdata/antenna/result_posting.py
  • trapdata/antenna/schemas.py
  • trapdata/antenna/tests/test_memory_leak.py
  • trapdata/antenna/tests/test_worker.py
  • trapdata/antenna/worker.py
  • trapdata/api/utils.py
  • trapdata/cli/worker.py
  • trapdata/settings.py

Co-Authored-By: Claude <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant