Skip to content

Add per-query datasource selection to Prometheus tools #129

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
prathamesh-sonpatki merged 7 commits into from
Apr 17, 2026
Merged

Add per-query datasource selection to Prometheus tools #129

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
b7bbeb8
Add per-query datasource selection to Prometheus tools
prathamesh-sonpatki Apr 16, 2026
f2acf09
Fix typos in PromQL tool descriptions
prathamesh-sonpatki Apr 16, 2026
b007a5c
Add unit tests for datasource override and list_datasources
prathamesh-sonpatki Apr 16, 2026
ead254d
Fix resolveDatasourceCfg to propagate Region and ClusterID
prathamesh-sonpatki Apr 16, 2026
f72adc8
Merge branch 'main' into feat/prometheus-dynamic-datasource
prathamesh-sonpatki Apr 16, 2026
fcab4c8
Address CodeRabbit review: credential validation, IsDefault accuracy,…
prathamesh-sonpatki Apr 16, 2026
5e9342c
Merge remote-tracking branch 'origin/feat/prometheus-dynamic-datasour…
prathamesh-sonpatki Apr 16, 2026
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
109 changes: 100 additions & 9 deletions internal/apm/apm.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,12 +74,14 @@ type PromqlRangeQueryArgs struct {
StartTimeISO string `json:"start_time_iso,omitempty" jsonschema:"Start time in RFC3339/ISO8601 format (e.g. 2024-06-01T12:00:00Z). Optional when lookback_minutes is provided."`
EndTimeISO string `json:"end_time_iso,omitempty" jsonschema:"End time in RFC3339/ISO8601 format (e.g. 2024-06-01T13:00:00Z). Defaults to now when omitted."`
LookbackMinutes float64 `json:"lookback_minutes,omitempty" jsonschema:"Number of minutes to look back from now (default: 60, minimum: 1). Use for relative windows like last 30 minutes."`
Datasource string `json:"datasource,omitempty" jsonschema:"Name of the datasource to query. If omitted, uses the default configured datasource."`
}

type PromqlInstantQueryArgs struct {
Query string `json:"query" jsonschema:"PromQL query to execute (required)"`
TimeISO string `json:"time_iso,omitempty" jsonschema:"Evaluation time in RFC3339/ISO8601 format (e.g. 2024-06-01T12:00:00Z). If omitted, defaults to now or now-lookback_minutes."`
LookbackMinutes float64 `json:"lookback_minutes,omitempty" jsonschema:"Number of minutes to look back from now when time_iso is omitted (default: 0, minimum: 1)."`
Datasource string `json:"datasource,omitempty" jsonschema:"Name of the datasource to query. If omitted, uses the default configured datasource."`
}

type PromqlLabelValuesArgs struct {
Expand All @@ -88,13 +90,15 @@ type PromqlLabelValuesArgs struct {
StartTimeISO string `json:"start_time_iso,omitempty" jsonschema:"Start time in RFC3339/ISO8601 format (e.g. 2024-06-01T12:00:00Z). Optional when lookback_minutes is provided."`
EndTimeISO string `json:"end_time_iso,omitempty" jsonschema:"End time in RFC3339/ISO8601 format (e.g. 2024-06-01T13:00:00Z). Defaults to now when omitted."`
LookbackMinutes float64 `json:"lookback_minutes,omitempty" jsonschema:"Number of minutes to look back from now (default: 60, minimum: 1). Use for relative windows like last 30 minutes."`
Datasource string `json:"datasource,omitempty" jsonschema:"Name of the datasource to query. If omitted, uses the default configured datasource."`
}

type PromqlLabelsArgs struct {
MatchQuery string `json:"match_query,omitempty" jsonschema:"PromQL query to match series (e.g. up{job=\"prometheus\"})"`
StartTimeISO string `json:"start_time_iso,omitempty" jsonschema:"Start time in RFC3339/ISO8601 format (e.g. 2024-06-01T12:00:00Z). Optional when lookback_minutes is provided."`
EndTimeISO string `json:"end_time_iso,omitempty" jsonschema:"End time in RFC3339/ISO8601 format (e.g. 2024-06-01T13:00:00Z). Defaults to now when omitted."`
LookbackMinutes float64 `json:"lookback_minutes,omitempty" jsonschema:"Number of minutes to look back from now (default: 60, minimum: 1). Use for relative windows like last 30 minutes."`
Datasource string `json:"datasource,omitempty" jsonschema:"Name of the datasource to query. If omitted, uses the default configured datasource."`
}

func resolveTimeRange(startTimeISO, endTimeISO string, lookbackMinutes float64) (int64, int64, error) {
Expand Down Expand Up @@ -1800,6 +1804,28 @@ func NewServiceDependencyGraphHandler(client *http.Client, cfg models.Config) fu
}
}

// resolveDatasourceCfg returns a copy of cfg with Prometheus credentials overridden
// to those of the named datasource. If datasourceName is empty the original cfg is
// returned unchanged. Returns an error when the name is non-empty but not found.
func resolveDatasourceCfg(cfg models.Config, datasourceName string) (models.Config, error) {
if datasourceName == "" {
return cfg, nil
}
ds, ok := cfg.ResolveDatasource(datasourceName)
if !ok {
return cfg, fmt.Errorf("datasource %q not found", datasourceName)
}
if ds.ReadURL == "" || ds.Username == "" || ds.Password == "" || ds.Region == "" {
return cfg, fmt.Errorf("datasource %q is missing required Prometheus configuration", datasourceName)
}
cfg.PrometheusReadURL = ds.ReadURL
cfg.PrometheusUsername = ds.Username
cfg.PrometheusPassword = ds.Password
Comment thread
prathamesh-sonpatki marked this conversation as resolved.
cfg.Region = ds.Region
cfg.ClusterID = ds.ClusterID
Comment thread
rizwan2000rm marked this conversation as resolved.
return cfg, nil
Comment thread
coderabbitai[bot] marked this conversation as resolved.
}

const PromqlRangeQueryDetails = `
Perform a Prometheus range query to get metrics data.
This tool can be used to query Prometheus for metrics data over a specified time range.
Expand Down Expand Up @@ -1829,6 +1855,7 @@ const PromqlRangeQueryDetails = `
- lookback_minutes: (Optional) Number of minutes to look back from now. Defaults to 60.
- start_time_iso: (Optional) Start time of the time range in RFC3339/ISO8601 format (e.g. 2026-02-09T15:04:05Z). Overrides lookback when provided.
- end_time_iso: (Optional) End time of the time range in RFC3339/ISO8601 format (e.g. 2026-02-09T16:04:05Z). Defaults to current time.
- datasource: (Optional) Name of the datasource to query. If omitted, uses the default configured datasource.
`

func NewPromqlRangeQueryHandler(client *http.Client, cfg models.Config) func(context.Context, *mcp.CallToolRequest, PromqlRangeQueryArgs) (*mcp.CallToolResult, any, error) {
Expand All @@ -1843,7 +1870,12 @@ func NewPromqlRangeQueryHandler(client *http.Client, cfg models.Config) func(con
return nil, nil, err
}

httpResp, err := utils.MakePromRangeAPIQuery(ctx, client, query, startTimeParam, endTimeParam, cfg)
queryCfg, err := resolveDatasourceCfg(cfg, args.Datasource)
if err != nil {
return nil, nil, err
}

httpResp, err := utils.MakePromRangeAPIQuery(ctx, client, query, startTimeParam, endTimeParam, queryCfg)
if err != nil {
return nil, nil, err
}
Expand Down Expand Up @@ -1897,6 +1929,7 @@ const PromqlInstantQueryDetails = `
- query: (Required) The Prometheus query to execute.
- time_iso: (Optional) The point in time to query in RFC3339/ISO8601 format (e.g. 2026-02-09T15:04:05Z). Overrides lookback when provided.
- lookback_minutes: (Optional) Number of minutes to look back from now when time_iso is omitted.
- datasource: (Optional) Name of the datasource to query. If omitted, uses the default configured datasource.
`

func NewPromqlInstantQueryHandler(client *http.Client, cfg models.Config) func(context.Context, *mcp.CallToolRequest, PromqlInstantQueryArgs) (*mcp.CallToolResult, any, error) {
Expand All @@ -1911,7 +1944,12 @@ func NewPromqlInstantQueryHandler(client *http.Client, cfg models.Config) func(c
return nil, nil, err
}

httpResp, err := utils.MakePromInstantAPIQuery(ctx, client, query, timeParam, cfg)
queryCfg, err := resolveDatasourceCfg(cfg, args.Datasource)
if err != nil {
return nil, nil, err
}

httpResp, err := utils.MakePromInstantAPIQuery(ctx, client, query, timeParam, queryCfg)
if err != nil {
return nil, nil, err
}
Expand Down Expand Up @@ -2052,14 +2090,15 @@ const PromqlLabelValuesQueryDetails = `
It returns an array of values for the label.
Parameters:
- match_query: (Required) A valid promql filter query
- label: (Required) Name of the label to return values for
- label: (Required) Name of the label to return values for
- lookback_minutes: (Optional) Number of minutes to look back from now. Defaults to 60.
- start_time_iso: (Optional) Start time of the time range in RFC3339/ISO8601 format (e.g. 2026-02-09T15:04:05Z). Overrides lookback when provided.
- end_time_iso: (Optional) End time of the time range in RFC3339/ISO8601 format (e.g. 2026-02-09T16:04:05Z). Defaults to current time.
- datasource: (Optional) Name of the datasource to query. If omitted, uses the default configured datasource.

match_query should be a well formed, valid promql query
It is enouraged to not use default
values of start_time and end_time and use values that are appropriate for the
It is encouraged to not use default
values of start_time and end_time and use values that are appropriate for the
use case
`

Expand All @@ -2078,7 +2117,12 @@ func NewPromqlLabelValuesHandler(client *http.Client, cfg models.Config) func(co
return nil, nil, err
}

httpResp, err := utils.MakePromLabelValuesAPIQuery(ctx, client, label, query, startTimeParam, endTimeParam, cfg)
queryCfg, err := resolveDatasourceCfg(cfg, args.Datasource)
if err != nil {
return nil, nil, err
}

httpResp, err := utils.MakePromLabelValuesAPIQuery(ctx, client, label, query, startTimeParam, endTimeParam, queryCfg)
if err != nil {
return nil, nil, err
}
Expand Down Expand Up @@ -2116,10 +2160,11 @@ const PromqlLabelsQueryDetails = `
- lookback_minutes: (Optional) Number of minutes to look back from now. Defaults to 60.
- start_time_iso: (Optional) Start time of the time range in RFC3339/ISO8601 format (e.g. 2026-02-09T15:04:05Z). Overrides lookback when provided.
- end_time_iso: (Optional) End time of the time range in RFC3339/ISO8601 format (e.g. 2026-02-09T16:04:05Z). Defaults to current time.
- datasource: (Optional) Name of the datasource to query. If omitted, uses the default configured datasource.

match_query should be a well formed, valid promql query
It is enouraged to not use default
values of start_time and end_time and use values that are appropriate for the
It is encouraged to not use default
values of start_time and end_time and use values that are appropriate for the
use case
`

Expand All @@ -2134,7 +2179,12 @@ func NewPromqlLabelsHandler(client *http.Client, cfg models.Config) func(context
return nil, nil, err
}

httpResp, err := utils.MakePromLabelsAPIQuery(ctx, client, query, startTimeParam, endTimeParam, cfg)
queryCfg, err := resolveDatasourceCfg(cfg, args.Datasource)
if err != nil {
return nil, nil, err
}

httpResp, err := utils.MakePromLabelsAPIQuery(ctx, client, query, startTimeParam, endTimeParam, queryCfg)
if err != nil {
return nil, nil, err
}
Expand All @@ -2160,3 +2210,44 @@ func NewPromqlLabelsHandler(client *http.Client, cfg models.Config) func(context
}, nil, nil
}
}

// ListDatasourcesArgs has no required parameters.
type ListDatasourcesArgs struct{}

const ListDatasourcesDescription = `
List all available datasources configured for this organization.
Use this tool to discover valid datasource names before passing them to
prometheus_range_query, prometheus_instant_query, prometheus_label_values,
or prometheus_labels via the datasource parameter.

Returns an array of objects, each with:
- name: the datasource name to use in the datasource parameter
- is_default: true for the datasource that is used when no datasource is specified
`

// NewListDatasourcesHandler returns a handler that serves the datasource list from
// the in-memory cache populated at startup — no extra API call is made.
// The response is serialized once at registration time since the list never changes.
func NewListDatasourcesHandler(cfg models.Config) func(context.Context, *mcp.CallToolRequest, ListDatasourcesArgs) (*mcp.CallToolResult, any, error) {
type datasourceView struct {
Name string `json:"name"`
IsDefault bool `json:"is_default"`
}

views := make([]datasourceView, 0, len(cfg.Datasources))
for _, ds := range cfg.Datasources {
views = append(views, datasourceView{
Name: ds.Name,
IsDefault: ds.IsDefault,
Comment thread
prathamesh-sonpatki marked this conversation as resolved.
})
}
out, _ := json.Marshal(views) // slice of plain structs — cannot fail

return func(_ context.Context, _ *mcp.CallToolRequest, _ ListDatasourcesArgs) (*mcp.CallToolResult, any, error) {
return &mcp.CallToolResult{
Content: []mcp.Content{
&mcp.TextContent{Text: string(out)},
},
}, nil, nil
}
}
Loading
Loading