Add per-query datasource selection to Prometheus tools#129
Merged
prathamesh-sonpatki merged 7 commits intomainfrom Apr 17, 2026
Merged
Add per-query datasource selection to Prometheus tools#129prathamesh-sonpatki merged 7 commits intomainfrom
prathamesh-sonpatki merged 7 commits intomainfrom
Conversation
Adds an optional `datasource` parameter to prometheus_range_query, prometheus_instant_query, prometheus_label_values, and prometheus_labels. When supplied, credentials for that datasource are resolved from the startup-cached list instead of the server-wide default, enabling an LLM agent to query multiple datasources (e.g. prod vs staging) in a single conversation without restarting the server. Also adds a `list_datasources` tool that returns all available datasource names and which one is the current default, so agents can discover valid values for the new parameter. The response is serialized once at registration time since the list is immutable after startup. Removes the now-dead GetDatasourceByName util that issued a live API call on every invocation; per-query resolution now goes through the in-memory cache. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Member
Author
|
@coderabbitai do a full review |
Contributor
|
✅ Actions performedFull review triggered. |
Contributor
📝 WalkthroughWalkthroughAdds per-request datasource selection for PromQL tools via an optional Changes
Sequence DiagramsequenceDiagram
participant Client
participant Handler as PromQL Handler
participant Resolver as resolveDatasourceCfg
participant Config as Config.ResolveDatasource
participant Prom as Prometheus API
Client->>Handler: Send PromQL request (optional datasource)
Handler->>Resolver: resolveDatasourceCfg(cfg, datasourceName)
Resolver->>Config: ResolveDatasource(datasourceName)
Config-->>Resolver: DatasourceInfo / not found
Resolver->>Resolver: Build queryCfg (override read URL, creds, region, cluster)
alt datasource resolved
Resolver-->>Handler: queryCfg
Handler->>Prom: Query Prometheus using queryCfg
Prom-->>Handler: Results
Handler-->>Client: Return results
else resolution error
Resolver-->>Handler: error
Handler-->>Client: Return error
end
Suggested reviewers
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds per-query Prometheus datasource selection by caching all datasources at startup, letting each Prometheus tool optionally override credentials by datasource name, and exposing a new discovery tool to list available datasources.
Changes:
- Added a
list_datasourcesMCP tool for discovering datasource names and the default datasource. - Added optional
datasourceargument support across all Prometheus tools and resolved credentials from the startup cache. - Cached datasource credentials in
models.ConfigduringPopulateAPICfgand removed the old per-call live lookup utility.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| tools.go | Registers the new list_datasources tool alongside existing tools. |
| internal/utils/utils.go | Caches all datasource credential info at startup in cfg.Datasources. |
| internal/models/config.go | Adds DatasourceInfo, Config.Datasources, and a ResolveDatasource lookup helper. |
| internal/apm/apm.go | Adds datasource args to Prometheus tools, resolves per-query credentials, and implements list_datasources handler. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
"enouraged" -> "encouraged" in prometheus_label_values and prometheus_labels tool descriptions (flagged by Copilot review). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Covers the two gaps flagged by Copilot review: - TestResolveDatasourceCfg: verifies empty name returns default cfg, known name overrides credentials, unknown name returns a clear error, and original cfg is never mutated. - TestPromqlRangeHandler_DatasourceOverride: end-to-end check that the overridden read_url/username/password reach the outgoing HTTP request, and that an unknown datasource errors before hitting the server. - TestNewListDatasourcesHandler: asserts correct JSON output with name + is_default fields, including the empty-list case. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
When overriding credentials for a named datasource, also copy Region and ClusterID so the resulting cfg is a self-consistent view of that datasource. Without this, handlers using queryCfg.Region (e.g. traces) or ClusterID (deep link builder) would silently use the startup-default datasource's values even when a different datasource was requested. Also removes a stale test comment that incorrectly attributed the strings import to TestNewServiceSummaryHandler_ExtraParams; it is used by the datasource tests added in the previous commit. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Contributor
✅ Actions performedFull review triggered. |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@internal/apm/apm_test.go`:
- Around line 536-542: The test using the httptest.Server with capturedReq
doesn't verify the handler was never invoked; modify the server handler to
atomically increment a request counter or set an atomic/chan flag (e.g., using a
uint32 counter via sync/atomic or a channel) when it receives any request, then
after calling the code under test assert the counter/flag is zero (or no value
received) to prove the error occurred before hitting the server; apply the same
change to the second occurrence (the block using capturedReq around lines
589-596) so both tests assert the server was not contacted.
In `@internal/apm/apm.go`:
- Around line 1810-1823: resolveDatasourceCfg currently trusts that a listed
datasource has usable connection info; validate the resolved datasource before
overwriting query config by returning an error if required fields are missing.
In resolveDatasourceCfg, after obtaining ds from cfg.ResolveDatasource, ensure
ds.ReadURL is non-empty (return fmt.Errorf("datasource %q has empty read_url",
datasourceName) if it is) and validate credentials consistently (e.g., if
ds.Username != "" require ds.Password != "" and error otherwise). Only assign
PrometheusReadURL, PrometheusUsername, PrometheusPassword, Region, and ClusterID
to cfg after these checks pass.
In `@internal/utils/utils.go`:
- Around line 470-483: The cached datasource entries in cfg.Datasources
currently use ds.IsDefault from the API, which becomes incorrect when startup
overrides the default via cfg.DatasourceName; update the loop that builds
cfg.Datasources (iterating datasourcesList and creating models.DatasourceInfo)
to compute IsDefault by comparing each datasource's Name to cfg.DatasourceName
(set true when names match, false otherwise) so the cached IsDefault reflects
the effective runtime default rather than the API-provided ds.IsDefault.
🪄 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: Repository UI (base), Organization UI (inherited)
Review profile: ASSERTIVE
Plan: Pro
Run ID: 71f33604-0d1e-41c0-addb-c7aaa3d8633f
📒 Files selected for processing (5)
internal/apm/apm.gointernal/apm/apm_test.gointernal/models/config.gointernal/utils/utils.gotools.go
… test coverage - resolveDatasourceCfg: return error when matched datasource lacks ReadURL, Username, Password, or Region (incomplete records should fail fast) - IsDefault in Datasources cache: derive from which datasource was actually selected at startup (selectedDatasource.Name == ds.Name) rather than the raw API flag, which doesn't account for DatasourceName config override - TestPromqlRangeHandler_DatasourceOverride: add atomic hitCount to assert the unknown-datasource error fires before any outbound HTTP request; expand staging fixture with Region to satisfy new validation Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…ce' into feat/prometheus-dynamic-datasource
Collaborator
|
@coderabbitai review |
Contributor
✅ Actions performedReview triggered.
|
rizwan2000rm
approved these changes
Apr 16, 2026
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@internal/apm/apm_test.go`:
- Around line 469-528: Add a new subtest inside TestResolveDatasourceCfg that
constructs a models.DatasourceInfo entry with empty Region but valid
ReadURL/Username/Password/ClusterID (e.g., name "noregion"), call
resolveDatasourceCfg(cfg, "noregion"), assert no error and that
PrometheusReadURL, PrometheusUsername, PrometheusPassword and ClusterID are
overridden to the datasource values while Region may remain empty, and also
verify the original cfg is not mutated; this ensures resolveDatasourceCfg
accepts datasources with empty Region rather than over-validating.
In `@internal/apm/apm.go`:
- Around line 1818-1825: The current validation block in internal/apm/apm.go
incorrectly requires ds.Region (and effectively ClusterID) for Prometheus
datasource overrides; change the check that currently tests ds.ReadURL,
ds.Username, ds.Password, and ds.Region to only validate ds.ReadURL,
ds.Username, and ds.Password so Prometheus-only handlers aren't rejected when
Region is empty. Keep assigning cfg.PrometheusReadURL, cfg.PrometheusUsername,
cfg.PrometheusPassword, and continue to set cfg.Region and cfg.ClusterID if
present, but do not make their presence mandatory; also update the returned
error message (the fmt.Errorf call referencing datasourceName) to reflect only
the required Prometheus configuration fields.
🪄 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: Repository UI (base), Organization UI (inherited)
Review profile: ASSERTIVE
Plan: Pro
Run ID: 50982d58-c0e5-4fbc-b250-b2968002d565
📒 Files selected for processing (3)
internal/apm/apm.gointernal/apm/apm_test.gointernal/utils/utils.go
rizwan2000rm
approved these changes
Apr 16, 2026
Member
Author
|
Datasource selection evals — 14/15 passed (7s)
Evaluated against SHA |
Member
Author
|
Datasource selection evals — 15/15 passed (8s)
Evaluated against SHA |
Member
Author
|
Datasource selection evals — 15/15 passed (10s)
Evaluated against SHA |
1 similar comment
Member
Author
|
Datasource selection evals — 15/15 passed (10s)
Evaluated against SHA |
Member
Author
|
Datasource selection evals — 14/15 passed (11s)
Evaluated against SHA |
Member
Author
|
Datasource selection evals — 15/15 passed (12s)
Evaluated against SHA |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
datasourceparameter to all 4 Prometheus tools (prometheus_range_query,prometheus_instant_query,prometheus_label_values,prometheus_labels) — when set, credentials for the named datasource are resolved from the startup-cached list instead of the server-wide defaultlist_datasourcestool so agents can discover available datasource names (and which is the default) before queryingGetDatasourceByNameutility that made a live API call on every invocation; all per-query resolution now uses the in-memory cacheHow it works
All datasources are already fetched from the API at startup in
PopulateAPICfg. We now cache them incfg.Datasources []DatasourceInfo. At query time,resolveDatasourceCfgcopiescfgby value and overwrites only the three Prometheus credential fields — safe across concurrent requests with no locking.The
list_datasourcesresponse is serialized once at handler registration since the list never changes after startup.Test plan
list_datasourcesreturns all datasources with correctis_defaultflagdatasourceparam is omitteddatasourceparam is set to a valid namedatasourceparam is set to an unknown name🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Improvements
Bug Fixes
Tests
Documentation