From c04ad4629a6acdf9ac6ffed966b42a6d1619b80e Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 19:42:06 +0000 Subject: [PATCH 01/33] Add remote query Postgres match endpoint --- cmd/agent/subcommands/run/command.go | 2 + comp/dataobs/remotequeries/fx/fx.go | 20 ++ .../remotequeries/impl/postgres_match.go | 339 ++++++++++++++++++ .../remotequeries/impl/postgres_match_test.go | 239 ++++++++++++ 4 files changed, 600 insertions(+) create mode 100644 comp/dataobs/remotequeries/fx/fx.go create mode 100644 comp/dataobs/remotequeries/impl/postgres_match.go create mode 100644 comp/dataobs/remotequeries/impl/postgres_match_test.go diff --git a/cmd/agent/subcommands/run/command.go b/cmd/agent/subcommands/run/command.go index 32ce92df5900..90b4e595ffaa 100644 --- a/cmd/agent/subcommands/run/command.go +++ b/cmd/agent/subcommands/run/command.go @@ -41,6 +41,7 @@ import ( workloadselectionfx "github.com/DataDog/datadog-agent/comp/workloadselection/fx" doqueryactionsfx "github.com/DataDog/datadog-agent/comp/dataobs/queryactions/fx" + remotequeriesfx "github.com/DataDog/datadog-agent/comp/dataobs/remotequeries/fx" haagentfx "github.com/DataDog/datadog-agent/comp/haagent/fx" snmpscanfx "github.com/DataDog/datadog-agent/comp/snmpscan/fx" snmpscanmanagerfx "github.com/DataDog/datadog-agent/comp/snmpscanmanager/fx" @@ -567,6 +568,7 @@ func getSharedFxOption() fx.Option { remoteagentregistryfx.Module(), haagentfx.Module(), doqueryactionsfx.Module(), + remotequeriesfx.Module(), metricscompressorfx.Module(), diagnosefx.Module(), ipcfx.ModuleReadWrite(), diff --git a/comp/dataobs/remotequeries/fx/fx.go b/comp/dataobs/remotequeries/fx/fx.go new file mode 100644 index 000000000000..fd32dc4a4997 --- /dev/null +++ b/comp/dataobs/remotequeries/fx/fx.go @@ -0,0 +1,20 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2026-present Datadog, Inc. + +// Package fx provides the fx module for the Remote Queries POC component. +package fx + +import ( + remotequeriesimpl "github.com/DataDog/datadog-agent/comp/dataobs/remotequeries/impl" + "github.com/DataDog/datadog-agent/pkg/util/fxutil" + "go.uber.org/fx" +) + +// Module defines the fx options for this component. +func Module() fxutil.Module { + return fxutil.Component( + fx.Provide(remotequeriesimpl.NewPostgresMatchEndpointProvider), + ) +} diff --git a/comp/dataobs/remotequeries/impl/postgres_match.go b/comp/dataobs/remotequeries/impl/postgres_match.go new file mode 100644 index 000000000000..4166c3ca9312 --- /dev/null +++ b/comp/dataobs/remotequeries/impl/postgres_match.go @@ -0,0 +1,339 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2026-present Datadog, Inc. + +// Package remotequeriesimpl implements Remote Queries POC endpoints. +package remotequeriesimpl + +import ( + "encoding/json" + "fmt" + "io" + "mime" + "net/http" + "strings" + + "go.uber.org/fx" + "gopkg.in/yaml.v3" + + api "github.com/DataDog/datadog-agent/comp/api/api/def" + "github.com/DataDog/datadog-agent/comp/collector/collector" + "github.com/DataDog/datadog-agent/comp/core/config" + "github.com/DataDog/datadog-agent/pkg/collector/check" +) + +const ( + // PostgresMatchEndpointPath is mounted under /agent by the Agent command API. + PostgresMatchEndpointPath = "/remote-queries/postgres/match-check" + // PostgresMatchEnabledConfig is disabled by default when the key is absent. + PostgresMatchEnabledConfig = "remote_queries.postgres_match_check.enabled" + + statusOK = "ok" + statusTargetNotFound = "target_not_found" + statusAmbiguous = "ambiguous_target" + statusInvalidRequest = "invalid_request" + statusBridgeDisabled = "bridge_disabled" +) + +var credentialShapedFields = map[string]struct{}{ + "app_key": {}, + "apikey": {}, + "api_key": {}, + "conn_string": {}, + "connection_string": {}, + "credential": {}, + "credentials": {}, + "dsn": {}, + "pass": {}, + "password": {}, + "pwd": {}, + "secret": {}, + "sslcert": {}, + "sslkey": {}, + "token": {}, + "user": {}, + "username": {}, +} + +// Requires defines dependencies for the Remote Queries POC endpoint provider. +type Requires struct { + fx.In + + Cfg config.Component + Collector collector.Component +} + +// NewPostgresMatchEndpointProvider registers the Postgres match endpoint on the internal Agent API. +func NewPostgresMatchEndpointProvider(reqs Requires) api.AgentEndpointProvider { + h := &postgresMatchHandler{ + collector: reqs.Collector, + enabled: reqs.Cfg.GetBool(PostgresMatchEnabledConfig), + } + return api.NewAgentEndpointProvider(h.handle, PostgresMatchEndpointPath, http.MethodPost) +} + +type postgresMatchHandler struct { + collector collector.Component + enabled bool +} + +type matchResponse struct { + Status string `json:"status"` + MatchedCount int `json:"matched_count"` + Match *sanitizedMatch `json:"match,omitempty"` + Error *responseError `json:"error,omitempty"` +} + +type sanitizedMatch struct { + Integration string `json:"integration"` + Loader string `json:"loader"` + ConfigProvider string `json:"config_provider"` +} + +type responseError struct { + Code string `json:"code"` + Message string `json:"message"` +} + +type postgresTarget struct { + Host string + Port int + DBName string +} + +type postgresInstanceTarget struct { + host string + port int + dbname string +} + +func (h *postgresMatchHandler) handle(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + + if !h.enabled { + writeMatchResponse(w, http.StatusServiceUnavailable, statusBridgeDisabled, 0, nil, "remote queries bridge is disabled") + return + } + + target, err := parseMatchRequest(r) + if err != nil { + writeMatchResponse(w, http.StatusBadRequest, statusInvalidRequest, 0, nil, err.Error()) + return + } + + matches := h.findMatches(target) + switch len(matches) { + case 0: + writeMatchResponse(w, http.StatusNotFound, statusTargetNotFound, 0, nil, "no matching Postgres check found") + case 1: + writeMatchResponse(w, http.StatusOK, statusOK, 1, &matches[0], "") + default: + writeMatchResponse(w, http.StatusConflict, statusAmbiguous, len(matches), nil, "multiple matching Postgres checks found") + } +} + +func parseMatchRequest(r *http.Request) (postgresTarget, error) { + if !isJSONContentType(r.Header.Get("Content-Type")) { + return postgresTarget{}, fmt.Errorf("content-type must be application/json") + } + + defer r.Body.Close() + decoder := json.NewDecoder(r.Body) + decoder.UseNumber() + + var root map[string]json.RawMessage + if err := decoder.Decode(&root); err != nil { + return postgresTarget{}, fmt.Errorf("malformed JSON request") + } + if err := decoder.Decode(&struct{}{}); err != io.EOF { + return postgresTarget{}, fmt.Errorf("malformed JSON request") + } + + for key := range root { + if key != "target" { + if isCredentialShapedField(key) { + return postgresTarget{}, fmt.Errorf("request contains disallowed credential-shaped field") + } + return postgresTarget{}, fmt.Errorf("request contains unknown field") + } + } + + rawTarget, ok := root["target"] + if !ok { + return postgresTarget{}, fmt.Errorf("target is required") + } + + var targetFields map[string]json.RawMessage + if err := json.Unmarshal(rawTarget, &targetFields); err != nil || targetFields == nil { + return postgresTarget{}, fmt.Errorf("target must be an object") + } + + for key := range targetFields { + switch key { + case "host", "port", "dbname": + continue + default: + if isCredentialShapedField(key) { + return postgresTarget{}, fmt.Errorf("request contains disallowed credential-shaped field") + } + return postgresTarget{}, fmt.Errorf("target contains unknown field") + } + } + + host, err := parseTargetString(targetFields, "host") + if err != nil { + return postgresTarget{}, err + } + host = normalizeHost(host) + if host == "" { + return postgresTarget{}, fmt.Errorf("target.host is required") + } + + port, err := parseTargetPort(targetFields) + if err != nil { + return postgresTarget{}, err + } + + dbname, err := parseTargetString(targetFields, "dbname") + if err != nil { + return postgresTarget{}, err + } + if dbname == "" { + return postgresTarget{}, fmt.Errorf("target.dbname is required") + } + + return postgresTarget{Host: host, Port: port, DBName: dbname}, nil +} + +func isJSONContentType(contentType string) bool { + mediaType, _, err := mime.ParseMediaType(contentType) + if err != nil { + return false + } + return mediaType == "application/json" +} + +func parseTargetString(fields map[string]json.RawMessage, field string) (string, error) { + raw, ok := fields[field] + if !ok { + return "", fmt.Errorf("target.%s is required", field) + } + var value string + if err := json.Unmarshal(raw, &value); err != nil { + return "", fmt.Errorf("target.%s must be a string", field) + } + return value, nil +} + +func parseTargetPort(fields map[string]json.RawMessage) (int, error) { + raw, ok := fields["port"] + if !ok { + return 0, fmt.Errorf("target.port is required") + } + + var port int + if err := json.Unmarshal(raw, &port); err != nil { + return 0, fmt.Errorf("target.port must be an integer") + } + if port < 1 || port > 65535 { + return 0, fmt.Errorf("target.port is out of range") + } + return port, nil +} + +func normalizeHost(host string) string { + host = strings.ToLower(strings.TrimSpace(host)) + return strings.TrimSuffix(host, ".") +} + +func isCredentialShapedField(field string) bool { + _, found := credentialShapedFields[strings.ToLower(field)] + return found +} + +func (h *postgresMatchHandler) findMatches(target postgresTarget) []sanitizedMatch { + checks := h.collector.GetChecks() + matches := make([]sanitizedMatch, 0, 1) + for _, chk := range checks { + if !isPostgresCheck(chk) { + continue + } + + instanceTarget, ok := parsePostgresInstanceTarget(chk.InstanceConfig()) + if !ok { + continue + } + + if instanceTarget.host == target.Host && instanceTarget.port == target.Port && instanceTarget.dbname == target.DBName { + matches = append(matches, sanitizedMatch{ + Integration: "postgres", + Loader: chk.Loader(), + ConfigProvider: chk.ConfigProvider(), + }) + } + } + return matches +} + +func isPostgresCheck(chk check.Check) bool { + name := strings.ToLower(strings.TrimSpace(chk.String())) + return name == "postgres" || name == "postgresql" +} + +func parsePostgresInstanceTarget(instanceConfig string) (postgresInstanceTarget, bool) { + var fields map[string]any + if err := yaml.Unmarshal([]byte(instanceConfig), &fields); err != nil || fields == nil { + return postgresInstanceTarget{}, false + } + + host, ok := fields["host"].(string) + if !ok { + return postgresInstanceTarget{}, false + } + host = normalizeHost(host) + if host == "" { + return postgresInstanceTarget{}, false + } + + port, ok := yamlInt(fields["port"]) + if !ok || port < 1 || port > 65535 { + return postgresInstanceTarget{}, false + } + + dbname, ok := fields["dbname"].(string) + if !ok || dbname == "" { + return postgresInstanceTarget{}, false + } + + return postgresInstanceTarget{host: host, port: port, dbname: dbname}, true +} + +func yamlInt(value any) (int, bool) { + switch v := value.(type) { + case int: + return v, true + case int64: + return int(v), true + case uint64: + if v > uint64(^uint(0)>>1) { + return 0, false + } + return int(v), true + default: + return 0, false + } +} + +func writeMatchResponse(w http.ResponseWriter, httpStatus int, status string, matchedCount int, match *sanitizedMatch, message string) { + w.WriteHeader(httpStatus) + resp := matchResponse{ + Status: status, + MatchedCount: matchedCount, + Match: match, + } + if status != statusOK { + resp.Error = &responseError{Code: status, Message: message} + } + _ = json.NewEncoder(w).Encode(resp) +} diff --git a/comp/dataobs/remotequeries/impl/postgres_match_test.go b/comp/dataobs/remotequeries/impl/postgres_match_test.go new file mode 100644 index 000000000000..128d2a7f2627 --- /dev/null +++ b/comp/dataobs/remotequeries/impl/postgres_match_test.go @@ -0,0 +1,239 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2026-present Datadog, Inc. + +package remotequeriesimpl + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/DataDog/datadog-agent/comp/collector/collector" + "github.com/DataDog/datadog-agent/comp/core/autodiscovery/integration" + diagnose "github.com/DataDog/datadog-agent/comp/core/diagnose/def" + "github.com/DataDog/datadog-agent/pkg/aggregator/sender" + "github.com/DataDog/datadog-agent/pkg/collector/check" + checkid "github.com/DataDog/datadog-agent/pkg/collector/check/id" + "github.com/DataDog/datadog-agent/pkg/collector/check/stats" +) + +func TestParseMatchRequestValidatesStrictShape(t *testing.T) { + tests := []struct { + name string + body string + wantError string + }{ + { + name: "unknown top level field", + body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"extra":true}`, + wantError: "request contains unknown field", + }, + { + name: "unknown target field", + body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres","extra":true}}`, + wantError: "target contains unknown field", + }, + { + name: "credential-shaped top level field", + body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"password":"secret-value"}`, + wantError: "request contains disallowed credential-shaped field", + }, + { + name: "credential-shaped target field", + body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres","username":"alice"}}`, + wantError: "request contains disallowed credential-shaped field", + }, + { + name: "non-integer port", + body: `{"target":{"host":"localhost","port":5432.1,"dbname":"postgres"}}`, + wantError: "target.port must be an integer", + }, + { + name: "string port", + body: `{"target":{"host":"localhost","port":"5432","dbname":"postgres"}}`, + wantError: "target.port must be an integer", + }, + { + name: "missing dbname", + body: `{"target":{"host":"localhost","port":5432}}`, + wantError: "target.dbname is required", + }, + { + name: "malformed JSON", + body: `{"target":`, + wantError: "malformed JSON request", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, PostgresMatchEndpointPath, strings.NewReader(tt.body)) + req.Header.Set("Content-Type", "application/json") + + _, err := parseMatchRequest(req) + require.Error(t, err) + assert.Equal(t, tt.wantError, err.Error()) + assert.NotContains(t, err.Error(), "secret-value") + assert.NotContains(t, err.Error(), "alice") + }) + } +} + +func TestParseMatchRequestNormalizesTargetHost(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, PostgresMatchEndpointPath, strings.NewReader( + `{"target":{"host":" LocalHost. ","port":5432,"dbname":"Postgres"}}`, + )) + req.Header.Set("Content-Type", "application/json; charset=utf-8") + + target, err := parseMatchRequest(req) + require.NoError(t, err) + assert.Equal(t, postgresTarget{Host: "localhost", Port: 5432, DBName: "Postgres"}, target) +} + +func TestPostgresMatchHandlerDisabled(t *testing.T) { + handler := &postgresMatchHandler{enabled: false, collector: fakeCollector{}} + + recorder := callMatchHandler(handler, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"}}`) + + assert.Equal(t, http.StatusServiceUnavailable, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"status":"bridge_disabled"`) +} + +func TestPostgresMatchHandlerExactMatch(t *testing.T) { + handler := &postgresMatchHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ + fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: LOCALHOST.\nport: 5432\ndbname: postgres\nusername: alice\npassword: secret-value\n"}, + fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5433\ndbname: postgres\npassword: other-secret\n"}, + fakeCheck{name: "mysql", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: mysql-secret\n"}, + }}} + + recorder := callMatchHandler(handler, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"}}`) + + assert.Equal(t, http.StatusOK, recorder.Code) + body := recorder.Body.String() + assert.Contains(t, body, `"status":"ok"`) + assert.Contains(t, body, `"matched_count":1`) + assert.Contains(t, body, `"integration":"postgres"`) + assert.Contains(t, body, `"loader":"python"`) + assert.Contains(t, body, `"config_provider":"file"`) + assert.NotContains(t, body, "alice") + assert.NotContains(t, body, "secret-value") + assert.NotContains(t, body, "other-secret") + assert.NotContains(t, body, "mysql-secret") + assert.NotContains(t, body, "InstanceConfig") +} + +func TestPostgresMatchHandlerNoMatch(t *testing.T) { + handler := &postgresMatchHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ + fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-value\n"}, + }}} + + recorder := callMatchHandler(handler, `{"target":{"host":"localhost","port":5432,"dbname":"other"}}`) + + assert.Equal(t, http.StatusNotFound, recorder.Code) + body := recorder.Body.String() + assert.Contains(t, body, `"status":"target_not_found"`) + assert.Contains(t, body, `"matched_count":0`) + assert.NotContains(t, body, "secret-value") + assert.NotContains(t, body, "other") +} + +func TestPostgresMatchHandlerAmbiguousMatch(t *testing.T) { + handler := &postgresMatchHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ + fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-one\n"}, + fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-two\n"}, + }}} + + recorder := callMatchHandler(handler, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"}}`) + + assert.Equal(t, http.StatusConflict, recorder.Code) + body := recorder.Body.String() + assert.Contains(t, body, `"status":"ambiguous_target"`) + assert.Contains(t, body, `"matched_count":2`) + assert.NotContains(t, body, "secret-one") + assert.NotContains(t, body, "secret-two") +} + +func TestPostgresMatchHandlerCredentialRequestDoesNotEchoValue(t *testing.T) { + handler := &postgresMatchHandler{enabled: true, collector: fakeCollector{}} + + recorder := callMatchHandler(handler, `{"target":{"host":"localhost","port":5432,"dbname":"postgres","dsn":"postgres://secret-value@example/db"}}`) + + assert.Equal(t, http.StatusBadRequest, recorder.Code) + body := recorder.Body.String() + assert.Contains(t, body, `"status":"invalid_request"`) + assert.Contains(t, body, "credential-shaped field") + assert.NotContains(t, body, "postgres://secret-value@example/db") + assert.NotContains(t, body, "secret-value") +} + +func TestPostgresMatchHandlerRejectsInvalidContentType(t *testing.T) { + handler := &postgresMatchHandler{enabled: true, collector: fakeCollector{}} + req := httptest.NewRequest(http.MethodPost, PostgresMatchEndpointPath, strings.NewReader(`{"target":{"host":"localhost","port":5432,"dbname":"postgres"}}`)) + req.Header.Set("Content-Type", "text/plain") + recorder := httptest.NewRecorder() + + handler.handle(recorder, req) + + assert.Equal(t, http.StatusBadRequest, recorder.Code) + assert.Contains(t, recorder.Body.String(), "content-type must be application/json") +} + +func callMatchHandler(handler *postgresMatchHandler, body string) *httptest.ResponseRecorder { + req := httptest.NewRequest(http.MethodPost, PostgresMatchEndpointPath, strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + recorder := httptest.NewRecorder() + handler.handle(recorder, req) + return recorder +} + +type fakeCollector struct { + checks []check.Check +} + +func (f fakeCollector) RunCheck(inner check.Check) (checkid.ID, error) { return inner.ID(), nil } +func (f fakeCollector) StopCheck(checkid.ID) error { return nil } +func (f fakeCollector) MapOverChecks(cb func([]check.Info)) {} +func (f fakeCollector) GetChecks() []check.Check { return f.checks } +func (f fakeCollector) ReloadAllCheckInstances(string, []check.Check) ([]checkid.ID, error) { + return nil, nil +} +func (f fakeCollector) AddEventReceiver(collector.EventReceiver) {} + +type fakeCheck struct { + name string + loader string + provider string + instance string +} + +func (f fakeCheck) Run() error { return nil } +func (f fakeCheck) Stop() {} +func (f fakeCheck) Cancel() {} +func (f fakeCheck) String() string { + return f.name +} +func (f fakeCheck) Loader() string { return f.loader } +func (f fakeCheck) Configure(sender.SenderManager, uint64, integration.Data, integration.Data, string, string) error { + return nil +} +func (f fakeCheck) Interval() time.Duration { return 0 } +func (f fakeCheck) ID() checkid.ID { return checkid.ID(f.name) } +func (f fakeCheck) GetWarnings() []error { return nil } +func (f fakeCheck) GetSenderStats() (stats.SenderStats, error) { return stats.SenderStats{}, nil } +func (f fakeCheck) Version() string { return "" } +func (f fakeCheck) ConfigSource() string { return "" } +func (f fakeCheck) ConfigProvider() string { return f.provider } +func (f fakeCheck) IsTelemetryEnabled() bool { return false } +func (f fakeCheck) InitConfig() string { return "" } +func (f fakeCheck) InstanceConfig() string { return f.instance } +func (f fakeCheck) GetDiagnoses() ([]diagnose.Diagnosis, error) { + return nil, nil +} +func (f fakeCheck) IsHASupported() bool { return false } From f5cdad518417429a41fee5688711428560d8f85a Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 21:50:57 +0000 Subject: [PATCH 02/33] Add remote query Postgres execute bridge --- .../internal/middleware/check_wrapper.go | 5 + comp/dataobs/remotequeries/fx/fx.go | 1 + .../remotequeries/impl/postgres_execute.go | 350 ++++++++++++++++++ .../remotequeries/impl/postgres_match.go | 74 ++-- .../remotequeries/impl/postgres_match_test.go | 202 ++++++++++ pkg/collector/python/check.go | 28 ++ pkg/collector/python/check_test.go | 16 + pkg/collector/python/test_check.go | 89 +++++ rtloader/include/datadog_agent_rtloader.h | 11 + rtloader/include/rtloader.h | 8 + rtloader/rtloader/api.cpp | 5 + rtloader/three/three.cpp | 52 +++ rtloader/three/three.h | 1 + 13 files changed, 789 insertions(+), 53 deletions(-) create mode 100644 comp/dataobs/remotequeries/impl/postgres_execute.go diff --git a/comp/collector/collector/collectorimpl/internal/middleware/check_wrapper.go b/comp/collector/collector/collectorimpl/internal/middleware/check_wrapper.go index bcdb1473457b..6c9d2c33a7b1 100644 --- a/comp/collector/collector/collectorimpl/internal/middleware/check_wrapper.go +++ b/comp/collector/collector/collectorimpl/internal/middleware/check_wrapper.go @@ -43,6 +43,11 @@ func NewCheckWrapper(inner check.Check, senderManager sender.SenderManager, agen } } +// Unwrap returns the wrapped check. +func (c *CheckWrapper) Unwrap() check.Check { + return c.inner +} + // Run implements Check#Run func (c *CheckWrapper) Run() (err error) { c.runM.Lock() diff --git a/comp/dataobs/remotequeries/fx/fx.go b/comp/dataobs/remotequeries/fx/fx.go index fd32dc4a4997..6fdcf4d91a65 100644 --- a/comp/dataobs/remotequeries/fx/fx.go +++ b/comp/dataobs/remotequeries/fx/fx.go @@ -16,5 +16,6 @@ import ( func Module() fxutil.Module { return fxutil.Component( fx.Provide(remotequeriesimpl.NewPostgresMatchEndpointProvider), + fx.Provide(remotequeriesimpl.NewPostgresExecuteEndpointProvider), ) } diff --git a/comp/dataobs/remotequeries/impl/postgres_execute.go b/comp/dataobs/remotequeries/impl/postgres_execute.go new file mode 100644 index 000000000000..a6a54257ac08 --- /dev/null +++ b/comp/dataobs/remotequeries/impl/postgres_execute.go @@ -0,0 +1,350 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2026-present Datadog, Inc. + +package remotequeriesimpl + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net/http" + + api "github.com/DataDog/datadog-agent/comp/api/api/def" + "github.com/DataDog/datadog-agent/comp/collector/collector" + "github.com/DataDog/datadog-agent/pkg/collector/check" +) + +const ( + // PostgresExecuteEndpointPath is mounted under /agent by the Agent command API. + PostgresExecuteEndpointPath = "/remote-queries/postgres/execute" + // PostgresExecuteEnabledConfig is disabled by default when the key is absent. + PostgresExecuteEnabledConfig = "remote_queries.postgres_execute.enabled" + + postgresRemoteQueryProofQuery = "SELECT 1 AS value" + + statusExecutorUnavailable = "executor_unavailable" +) + +type postgresRemoteQueryRunner interface { + RunPostgresRemoteQueryJSON(requestJSON string) (string, error) +} + +type postgresCheckUnwrapper interface { + Unwrap() check.Check +} + +func postgresRemoteQueryRunnerFor(chk check.Check) (postgresRemoteQueryRunner, bool) { + for chk != nil { + if runner, ok := chk.(postgresRemoteQueryRunner); ok { + return runner, true + } + unwrapper, ok := chk.(postgresCheckUnwrapper) + if !ok { + break + } + unwrapped := unwrapper.Unwrap() + if unwrapped == chk { + break + } + chk = unwrapped + } + return nil, false +} + +// NewPostgresExecuteEndpointProvider registers the Postgres execute endpoint on the internal Agent API. +func NewPostgresExecuteEndpointProvider(reqs Requires) api.AgentEndpointProvider { + h := &postgresExecuteHandler{ + collector: reqs.Collector, + enabled: reqs.Cfg.GetBool(PostgresExecuteEnabledConfig), + } + return api.NewAgentEndpointProvider(h.handle, PostgresExecuteEndpointPath, http.MethodPost) +} + +type postgresExecuteHandler struct { + collector collector.Component + enabled bool +} + +type postgresExecuteRequest struct { + Target postgresTarget + Query string + Limits *postgresExecuteLimits +} + +type postgresExecuteLimits struct { + MaxRows int + MaxBytes int + TimeoutMs int +} + +type postgresExecuteRequestJSON struct { + Target postgresTargetJSON `json:"target"` + Query string `json:"query"` + Limits *postgresExecuteLimitsJSON `json:"limits,omitempty"` +} + +type postgresTargetJSON struct { + Host string `json:"host"` + Port int `json:"port"` + DBName string `json:"dbname"` +} + +type postgresExecuteLimitsJSON struct { + MaxRows int `json:"maxRows"` + MaxBytes int `json:"maxBytes"` + TimeoutMs int `json:"timeoutMs"` +} + +func (h *postgresExecuteHandler) handle(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + + if !h.enabled { + writeExecuteError(w, http.StatusServiceUnavailable, statusBridgeDisabled, "remote queries bridge is disabled") + return + } + + req, requestJSON, err := parseExecuteRequest(r) + if err != nil { + writeExecuteError(w, http.StatusBadRequest, statusInvalidRequest, err.Error()) + return + } + + matches := findPostgresMatches(h.collector, req.Target) + switch len(matches) { + case 0: + writeExecuteError(w, http.StatusNotFound, statusTargetNotFound, "no matching Postgres check found") + return + case 1: + // continue below + default: + writeExecuteError(w, http.StatusConflict, statusAmbiguous, "multiple matching Postgres checks found") + return + } + + runner, ok := postgresRemoteQueryRunnerFor(matches[0].check) + if !ok { + writeExecuteError(w, http.StatusFailedDependency, statusExecutorUnavailable, "matched Postgres check does not support remote query execution") + return + } + + responseJSON, err := runner.RunPostgresRemoteQueryJSON(requestJSON) + if err != nil { + writeExecuteError(w, http.StatusBadGateway, statusExecutorUnavailable, "remote query executor failed") + return + } + + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, responseJSON) +} + +func parseExecuteRequest(r *http.Request) (postgresExecuteRequest, string, error) { + if !isJSONContentType(r.Header.Get("Content-Type")) { + return postgresExecuteRequest{}, "", fmt.Errorf("content-type must be application/json") + } + + defer r.Body.Close() + decoder := json.NewDecoder(r.Body) + decoder.UseNumber() + + var root map[string]json.RawMessage + if err := decoder.Decode(&root); err != nil { + return postgresExecuteRequest{}, "", fmt.Errorf("malformed JSON request") + } + if err := decoder.Decode(&struct{}{}); err != io.EOF { + return postgresExecuteRequest{}, "", fmt.Errorf("malformed JSON request") + } + + for key := range root { + switch key { + case "target", "query", "limits": + continue + default: + if isCredentialShapedField(key) { + return postgresExecuteRequest{}, "", fmt.Errorf("request contains disallowed credential-shaped field") + } + return postgresExecuteRequest{}, "", fmt.Errorf("request contains unknown field") + } + } + + target, err := parseTargetFromRoot(root) + if err != nil { + return postgresExecuteRequest{}, "", err + } + + query, err := parseRequiredString(root, "query") + if err != nil { + return postgresExecuteRequest{}, "", err + } + if query != postgresRemoteQueryProofQuery { + return postgresExecuteRequest{}, "", fmt.Errorf("query is not allowed") + } + + limits, err := parseExecuteLimits(root) + if err != nil { + return postgresExecuteRequest{}, "", err + } + + req := postgresExecuteRequest{Target: target, Query: query, Limits: limits} + requestJSON, err := marshalExecuteRequest(req) + if err != nil { + return postgresExecuteRequest{}, "", fmt.Errorf("malformed JSON request") + } + return req, requestJSON, nil +} + +func parseTargetFromRoot(root map[string]json.RawMessage) (postgresTarget, error) { + rawTarget, ok := root["target"] + if !ok { + return postgresTarget{}, fmt.Errorf("target is required") + } + + var targetFields map[string]json.RawMessage + if err := json.Unmarshal(rawTarget, &targetFields); err != nil || targetFields == nil { + return postgresTarget{}, fmt.Errorf("target must be an object") + } + return parseTargetFields(targetFields) +} + +func parseTargetFields(targetFields map[string]json.RawMessage) (postgresTarget, error) { + for key := range targetFields { + switch key { + case "host", "port", "dbname": + continue + default: + if isCredentialShapedField(key) { + return postgresTarget{}, fmt.Errorf("request contains disallowed credential-shaped field") + } + return postgresTarget{}, fmt.Errorf("target contains unknown field") + } + } + + host, err := parseTargetString(targetFields, "host") + if err != nil { + return postgresTarget{}, err + } + host = normalizeHost(host) + if host == "" { + return postgresTarget{}, fmt.Errorf("target.host is required") + } + + port, err := parseTargetPort(targetFields) + if err != nil { + return postgresTarget{}, err + } + + dbname, err := parseTargetString(targetFields, "dbname") + if err != nil { + return postgresTarget{}, err + } + if dbname == "" { + return postgresTarget{}, fmt.Errorf("target.dbname is required") + } + + return postgresTarget{Host: host, Port: port, DBName: dbname}, nil +} + +func parseRequiredString(root map[string]json.RawMessage, field string) (string, error) { + raw, ok := root[field] + if !ok { + return "", fmt.Errorf("%s is required", field) + } + var value string + if err := json.Unmarshal(raw, &value); err != nil { + return "", fmt.Errorf("%s must be a string", field) + } + return value, nil +} + +func parseExecuteLimits(root map[string]json.RawMessage) (*postgresExecuteLimits, error) { + rawLimits, ok := root["limits"] + if !ok { + return nil, nil + } + + var limitFields map[string]json.RawMessage + if err := json.Unmarshal(rawLimits, &limitFields); err != nil || limitFields == nil { + return nil, fmt.Errorf("limits must be an object") + } + + for key := range limitFields { + switch key { + case "maxRows", "maxBytes", "timeoutMs": + continue + default: + if isCredentialShapedField(key) { + return nil, fmt.Errorf("request contains disallowed credential-shaped field") + } + return nil, fmt.Errorf("limits contains unknown field") + } + } + + maxRows, err := parsePositiveJSONInt(limitFields, "limits.maxRows", "maxRows") + if err != nil { + return nil, err + } + maxBytes, err := parsePositiveJSONInt(limitFields, "limits.maxBytes", "maxBytes") + if err != nil { + return nil, err + } + timeoutMs, err := parsePositiveJSONInt(limitFields, "limits.timeoutMs", "timeoutMs") + if err != nil { + return nil, err + } + + return &postgresExecuteLimits{MaxRows: maxRows, MaxBytes: maxBytes, TimeoutMs: timeoutMs}, nil +} + +func parsePositiveJSONInt(fields map[string]json.RawMessage, displayName string, wireName string) (int, error) { + raw, ok := fields[wireName] + if !ok { + return 0, fmt.Errorf("%s is required", displayName) + } + + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.UseNumber() + var value int + if err := decoder.Decode(&value); err != nil { + return 0, fmt.Errorf("%s must be an integer", displayName) + } + if err := decoder.Decode(&struct{}{}); err != io.EOF { + return 0, fmt.Errorf("%s must be an integer", displayName) + } + if value < 1 { + return 0, fmt.Errorf("%s must be at least 1", displayName) + } + return value, nil +} + +func marshalExecuteRequest(req postgresExecuteRequest) (string, error) { + wireReq := postgresExecuteRequestJSON{ + Target: postgresTargetJSON{Host: req.Target.Host, Port: req.Target.Port, DBName: req.Target.DBName}, + Query: req.Query, + } + if req.Limits != nil { + wireReq.Limits = &postgresExecuteLimitsJSON{ + MaxRows: req.Limits.MaxRows, + MaxBytes: req.Limits.MaxBytes, + TimeoutMs: req.Limits.TimeoutMs, + } + } + + requestJSON, err := json.Marshal(wireReq) + if err != nil { + return "", err + } + return string(requestJSON), nil +} + +func writeExecuteError(w http.ResponseWriter, httpStatus int, status string, message string) { + w.WriteHeader(httpStatus) + _ = json.NewEncoder(w).Encode(struct { + Status string `json:"status"` + Error *responseError `json:"error"` + }{ + Status: status, + Error: &responseError{Code: status, Message: message}, + }) +} diff --git a/comp/dataobs/remotequeries/impl/postgres_match.go b/comp/dataobs/remotequeries/impl/postgres_match.go index 4166c3ca9312..8b056d69202f 100644 --- a/comp/dataobs/remotequeries/impl/postgres_match.go +++ b/comp/dataobs/remotequeries/impl/postgres_match.go @@ -127,7 +127,7 @@ func (h *postgresMatchHandler) handle(w http.ResponseWriter, r *http.Request) { case 0: writeMatchResponse(w, http.StatusNotFound, statusTargetNotFound, 0, nil, "no matching Postgres check found") case 1: - writeMatchResponse(w, http.StatusOK, statusOK, 1, &matches[0], "") + writeMatchResponse(w, http.StatusOK, statusOK, 1, &matches[0].sanitized, "") default: writeMatchResponse(w, http.StatusConflict, statusAmbiguous, len(matches), nil, "multiple matching Postgres checks found") } @@ -159,51 +159,7 @@ func parseMatchRequest(r *http.Request) (postgresTarget, error) { } } - rawTarget, ok := root["target"] - if !ok { - return postgresTarget{}, fmt.Errorf("target is required") - } - - var targetFields map[string]json.RawMessage - if err := json.Unmarshal(rawTarget, &targetFields); err != nil || targetFields == nil { - return postgresTarget{}, fmt.Errorf("target must be an object") - } - - for key := range targetFields { - switch key { - case "host", "port", "dbname": - continue - default: - if isCredentialShapedField(key) { - return postgresTarget{}, fmt.Errorf("request contains disallowed credential-shaped field") - } - return postgresTarget{}, fmt.Errorf("target contains unknown field") - } - } - - host, err := parseTargetString(targetFields, "host") - if err != nil { - return postgresTarget{}, err - } - host = normalizeHost(host) - if host == "" { - return postgresTarget{}, fmt.Errorf("target.host is required") - } - - port, err := parseTargetPort(targetFields) - if err != nil { - return postgresTarget{}, err - } - - dbname, err := parseTargetString(targetFields, "dbname") - if err != nil { - return postgresTarget{}, err - } - if dbname == "" { - return postgresTarget{}, fmt.Errorf("target.dbname is required") - } - - return postgresTarget{Host: host, Port: port, DBName: dbname}, nil + return parseTargetFromRoot(root) } func isJSONContentType(contentType string) bool { @@ -252,9 +208,18 @@ func isCredentialShapedField(field string) bool { return found } -func (h *postgresMatchHandler) findMatches(target postgresTarget) []sanitizedMatch { - checks := h.collector.GetChecks() - matches := make([]sanitizedMatch, 0, 1) +type postgresCheckMatch struct { + check check.Check + sanitized sanitizedMatch +} + +func (h *postgresMatchHandler) findMatches(target postgresTarget) []postgresCheckMatch { + return findPostgresMatches(h.collector, target) +} + +func findPostgresMatches(collector collector.Component, target postgresTarget) []postgresCheckMatch { + checks := collector.GetChecks() + matches := make([]postgresCheckMatch, 0, 1) for _, chk := range checks { if !isPostgresCheck(chk) { continue @@ -266,10 +231,13 @@ func (h *postgresMatchHandler) findMatches(target postgresTarget) []sanitizedMat } if instanceTarget.host == target.Host && instanceTarget.port == target.Port && instanceTarget.dbname == target.DBName { - matches = append(matches, sanitizedMatch{ - Integration: "postgres", - Loader: chk.Loader(), - ConfigProvider: chk.ConfigProvider(), + matches = append(matches, postgresCheckMatch{ + check: chk, + sanitized: sanitizedMatch{ + Integration: "postgres", + Loader: chk.Loader(), + ConfigProvider: chk.ConfigProvider(), + }, }) } } diff --git a/comp/dataobs/remotequeries/impl/postgres_match_test.go b/comp/dataobs/remotequeries/impl/postgres_match_test.go index 128d2a7f2627..94605a81c25b 100644 --- a/comp/dataobs/remotequeries/impl/postgres_match_test.go +++ b/comp/dataobs/remotequeries/impl/postgres_match_test.go @@ -237,3 +237,205 @@ func (f fakeCheck) GetDiagnoses() ([]diagnose.Diagnosis, error) { return nil, nil } func (f fakeCheck) IsHASupported() bool { return false } + +func TestParseExecuteRequestValidatesStrictShape(t *testing.T) { + tests := []struct { + name string + body string + wantError string + }{ + { + name: "unknown top level field", + body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","extra":true}`, + wantError: "request contains unknown field", + }, + { + name: "credential-shaped top level field", + body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","token":"secret-value"}`, + wantError: "request contains disallowed credential-shaped field", + }, + { + name: "unknown target field", + body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres","extra":true},"query":"SELECT 1 AS value"}`, + wantError: "target contains unknown field", + }, + { + name: "credential-shaped target field", + body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres","password":"secret-value"},"query":"SELECT 1 AS value"}`, + wantError: "request contains disallowed credential-shaped field", + }, + { + name: "non-exact query", + body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value;"}`, + wantError: "query is not allowed", + }, + { + name: "unknown limits field", + body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":10,"maxBytes":1048576,"timeoutMs":5000,"extra":true}}`, + wantError: "limits contains unknown field", + }, + { + name: "credential-shaped limits field", + body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":10,"maxBytes":1048576,"timeoutMs":5000,"password":"secret-value"}}`, + wantError: "request contains disallowed credential-shaped field", + }, + { + name: "string maxRows", + body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":"10","maxBytes":1048576,"timeoutMs":5000}}`, + wantError: "limits.maxRows must be an integer", + }, + { + name: "zero timeout", + body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":10,"maxBytes":1048576,"timeoutMs":0}}`, + wantError: "limits.timeoutMs must be at least 1", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, PostgresExecuteEndpointPath, strings.NewReader(tt.body)) + req.Header.Set("Content-Type", "application/json") + + _, _, err := parseExecuteRequest(req) + require.Error(t, err) + assert.Equal(t, tt.wantError, err.Error()) + assert.NotContains(t, err.Error(), "secret-value") + }) + } +} + +func TestParseExecuteRequestNormalizesAndMarshalsCredentialFreeJSON(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, PostgresExecuteEndpointPath, strings.NewReader( + `{"target":{"host":" LocalHost. ","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":10,"maxBytes":1048576,"timeoutMs":5000}}`, + )) + req.Header.Set("Content-Type", "application/json; charset=utf-8") + + parsed, requestJSON, err := parseExecuteRequest(req) + require.NoError(t, err) + assert.Equal(t, postgresTarget{Host: "localhost", Port: 5432, DBName: "postgres"}, parsed.Target) + assert.JSONEq(t, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":10,"maxBytes":1048576,"timeoutMs":5000}}`, requestJSON) +} + +func TestParseExecuteRequestAllowsOmittedLimits(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, PostgresExecuteEndpointPath, strings.NewReader( + `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`, + )) + req.Header.Set("Content-Type", "application/json") + + parsed, requestJSON, err := parseExecuteRequest(req) + require.NoError(t, err) + assert.Nil(t, parsed.Limits) + assert.JSONEq(t, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`, requestJSON) +} + +func TestPostgresExecuteHandlerDisabled(t *testing.T) { + handler := &postgresExecuteHandler{enabled: false, collector: fakeCollector{}} + + recorder := callExecuteHandler(handler, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) + + assert.Equal(t, http.StatusServiceUnavailable, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"status":"bridge_disabled"`) +} + +func TestPostgresExecuteHandlerRunnerSuccess(t *testing.T) { + runner := &fakeRunnerCheck{ + fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-value\n"}, + response: `{"status":"SUCCEEDED","rows":[{"value":1}]}`, + } + handler := &postgresExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{fakeWrappedCheck{Check: runner}}}} + + recorder := callExecuteHandler(handler, `{"target":{"host":"LOCALHOST.","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) + + assert.Equal(t, http.StatusOK, recorder.Code) + assert.JSONEq(t, `{"status":"SUCCEEDED","rows":[{"value":1}]}`, recorder.Body.String()) + assert.JSONEq(t, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`, runner.seenRequest()) + assert.NotContains(t, recorder.Body.String(), "secret-value") +} + +func TestPostgresExecuteHandlerNoMatchAndAmbiguous(t *testing.T) { + t.Run("no match", func(t *testing.T) { + handler := &postgresExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ + &fakeRunnerCheck{fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-value\n"}}, + }}} + + recorder := callExecuteHandler(handler, `{"target":{"host":"localhost","port":5432,"dbname":"other"},"query":"SELECT 1 AS value"}`) + + assert.Equal(t, http.StatusNotFound, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"status":"target_not_found"`) + assert.NotContains(t, recorder.Body.String(), "secret-value") + assert.NotContains(t, recorder.Body.String(), "other") + }) + + t.Run("ambiguous", func(t *testing.T) { + handler := &postgresExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ + &fakeRunnerCheck{fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-one\n"}}, + &fakeRunnerCheck{fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-two\n"}}, + }}} + + recorder := callExecuteHandler(handler, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) + + assert.Equal(t, http.StatusConflict, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"status":"ambiguous_target"`) + assert.NotContains(t, recorder.Body.String(), "secret-one") + assert.NotContains(t, recorder.Body.String(), "secret-two") + }) +} + +func TestPostgresExecuteHandlerUnsupportedAndRunnerErrorAreSanitized(t *testing.T) { + t.Run("unsupported", func(t *testing.T) { + handler := &postgresExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ + fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-value\n"}, + }}} + + recorder := callExecuteHandler(handler, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) + + assert.Equal(t, http.StatusFailedDependency, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"status":"executor_unavailable"`) + assert.NotContains(t, recorder.Body.String(), "secret-value") + }) + + t.Run("runner error", func(t *testing.T) { + handler := &postgresExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ + &fakeRunnerCheck{fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-value\n"}, err: assert.AnError}, + }}} + + recorder := callExecuteHandler(handler, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) + + assert.Equal(t, http.StatusBadGateway, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"status":"executor_unavailable"`) + assert.NotContains(t, recorder.Body.String(), "secret-value") + assert.NotContains(t, recorder.Body.String(), assert.AnError.Error()) + }) +} + +func callExecuteHandler(handler *postgresExecuteHandler, body string) *httptest.ResponseRecorder { + req := httptest.NewRequest(http.MethodPost, PostgresExecuteEndpointPath, strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + recorder := httptest.NewRecorder() + handler.handle(recorder, req) + return recorder +} + +type fakeWrappedCheck struct { + check.Check +} + +func (f fakeWrappedCheck) Unwrap() check.Check { + return f.Check +} + +type fakeRunnerCheck struct { + fakeCheck + response string + err error + seen string +} + +func (f *fakeRunnerCheck) RunPostgresRemoteQueryJSON(requestJSON string) (string, error) { + f.seen = requestJSON + return f.response, f.err +} + +func (f *fakeRunnerCheck) seenRequest() string { + return f.seen +} diff --git a/pkg/collector/python/check.go b/pkg/collector/python/check.go index 66c608e07473..2b8840ef1f34 100644 --- a/pkg/collector/python/check.go +++ b/pkg/collector/python/check.go @@ -39,6 +39,7 @@ import ( #include "rtloader_mem.h" char *getStringAddr(char **array, unsigned int idx); +char *run_postgres_remote_query(rtloader_t *, rtloader_pyobject_t *check, const char *request_json); static inline void call_free(void* ptr) { _free(ptr); @@ -157,6 +158,33 @@ func (c *PythonCheck) RunSimple() error { return c.runCheck(false) } +// RunPostgresRemoteQueryJSON runs the fixed Postgres remote query helper for this Python check. +func (c *PythonCheck) RunPostgresRemoteQueryJSON(requestJSON string) (string, error) { + gstate, err := newStickyLock() + if err != nil { + return "", err + } + defer gstate.unlock() + + if c.cancelled { + return "", fmt.Errorf("check %s is already cancelled", c.ModuleName) + } + + cRequestJSON := C.CString(requestJSON) + defer C.free(unsafe.Pointer(cRequestJSON)) + + cResult := C.run_postgres_remote_query(rtloader, c.instance, cRequestJSON) + if cResult == nil { + if err := getRtLoaderError(); err != nil { + return "", err + } + return "", fmt.Errorf("an error occurred while running Postgres remote query") + } + defer C.rtloader_free(rtloader, unsafe.Pointer(cResult)) + + return C.GoString(cResult), nil +} + // Stop does nothing func (c *PythonCheck) Stop() {} diff --git a/pkg/collector/python/check_test.go b/pkg/collector/python/check_test.go index 9f41fc0fa4bb..30ac85a41703 100644 --- a/pkg/collector/python/check_test.go +++ b/pkg/collector/python/check_test.go @@ -70,3 +70,19 @@ func TestCheckDiagnosesDeserialization(t *testing.T) { func TestRunAfterCancel(t *testing.T) { testRunAfterCancel(t) } + +func TestRunPostgresRemoteQueryJSON(t *testing.T) { + testRunPostgresRemoteQueryJSON(t) +} + +func TestRunPostgresRemoteQueryJSONError(t *testing.T) { + testRunPostgresRemoteQueryJSONError(t) +} + +func TestRunPostgresRemoteQueryJSONWithRuntimeNotInitializedError(t *testing.T) { + testRunPostgresRemoteQueryJSONWithRuntimeNotInitializedError(t) +} + +func TestRunPostgresRemoteQueryJSONAfterCancel(t *testing.T) { + testRunPostgresRemoteQueryJSONAfterCancel(t) +} diff --git a/pkg/collector/python/test_check.go b/pkg/collector/python/test_check.go index 6ff27e7a4398..95860e728196 100644 --- a/pkg/collector/python/test_check.go +++ b/pkg/collector/python/test_check.go @@ -104,6 +104,17 @@ char *get_check_diagnoses(rtloader_t *s, rtloader_pyobject_t *check) { return get_check_diagnoses_return; } +char *run_postgres_remote_query_return = NULL; +int run_postgres_remote_query_calls = 0; +rtloader_pyobject_t *run_postgres_remote_query_instance = NULL; +const char *run_postgres_remote_query_request_json = NULL; +char *run_postgres_remote_query(rtloader_t *s, rtloader_pyobject_t *check, const char *request_json) { + run_postgres_remote_query_instance = check; + run_postgres_remote_query_request_json = strdup(request_json); + run_postgres_remote_query_calls++; + return run_postgres_remote_query_return; +} + // // get_check MOCK // @@ -203,6 +214,10 @@ void reset_check_mock() { get_check_diagnoses_return = NULL; get_check_diagnoses_calls = 0; + run_postgres_remote_query_return = NULL; + run_postgres_remote_query_calls = 0; + run_postgres_remote_query_instance = NULL; + run_postgres_remote_query_request_json = NULL; } */ import "C" @@ -668,6 +683,80 @@ func testGetDiagnoses(t *testing.T) { assert.Zero(t, len(diagnoses[1].Remediation)) } +func testRunPostgresRemoteQueryJSON(t *testing.T) { + mockRtloader(t) + + check, err := NewPythonFakeCheck(aggregator.NewNoOpSenderManager()) + require.NoError(t, err) + check.instance = newMockPyObjectPtr() + + C.reset_check_mock() + C.run_postgres_remote_query_return = C.CString(`{"status":"SUCCEEDED"}`) + + result, err := check.RunPostgresRemoteQueryJSON(`{"query":"SELECT 1 AS value"}`) + + require.NoError(t, err) + assert.JSONEq(t, `{"status":"SUCCEEDED"}`, result) + assert.Equal(t, C.int(1), C.gil_locked_calls) + assert.Equal(t, C.int(1), C.gil_unlocked_calls) + assert.Equal(t, C.int(1), C.run_postgres_remote_query_calls) + assert.Equal(t, C.int(1), C.rtloader_free_calls) + assert.Equal(t, check.instance, C.run_postgres_remote_query_instance) + assert.Equal(t, `{"query":"SELECT 1 AS value"}`, C.GoString(C.run_postgres_remote_query_request_json)) +} + +func testRunPostgresRemoteQueryJSONError(t *testing.T) { + mockRtloader(t) + + check, err := NewPythonFakeCheck(aggregator.NewNoOpSenderManager()) + require.NoError(t, err) + check.instance = newMockPyObjectPtr() + + C.reset_check_mock() + C.run_postgres_remote_query_return = nil + C.has_error_return = 1 + C.get_error_return = C.CString("rtloader helper failed") + + result, err := check.RunPostgresRemoteQueryJSON(`{"query":"SELECT 1 AS value"}`) + + assert.Empty(t, result) + require.Error(t, err) + assert.EqualError(t, err, "rtloader helper failed") + assert.Equal(t, C.int(1), C.gil_locked_calls) + assert.Equal(t, C.int(1), C.gil_unlocked_calls) + assert.Equal(t, C.int(1), C.run_postgres_remote_query_calls) + assert.Equal(t, C.int(0), C.rtloader_free_calls) +} + +func testRunPostgresRemoteQueryJSONWithRuntimeNotInitializedError(t *testing.T) { + mockRtloader(t) + check, err := NewPythonFakeCheck(aggregator.NewNoOpSenderManager()) + require.NoError(t, err) + check.instance = newMockPyObjectPtr() + + C.reset_check_mock() + rtloader = nil + + _, err = check.RunPostgresRemoteQueryJSON(`{"query":"SELECT 1 AS value"}`) + assert.ErrorIs(t, err, ErrNotInitialized) + assert.Equal(t, C.int(0), C.run_postgres_remote_query_calls) +} + +func testRunPostgresRemoteQueryJSONAfterCancel(t *testing.T) { + mockRtloader(t) + + check, err := NewPythonFakeCheck(aggregator.NewNoOpSenderManager()) + require.NoError(t, err) + check.instance = newMockPyObjectPtr() + + C.reset_check_mock() + check.Cancel() + + _, err = check.RunPostgresRemoteQueryJSON(`{"query":"SELECT 1 AS value"}`) + assert.EqualError(t, err, "check fake_check is already cancelled") + assert.Equal(t, C.int(0), C.run_postgres_remote_query_calls) +} + func testRunAfterCancel(t *testing.T) { mockRtloader(t) diff --git a/rtloader/include/datadog_agent_rtloader.h b/rtloader/include/datadog_agent_rtloader.h index 8b3b8f8036fe..6f93a576d2df 100644 --- a/rtloader/include/datadog_agent_rtloader.h +++ b/rtloader/include/datadog_agent_rtloader.h @@ -190,6 +190,17 @@ DATADOG_AGENT_RTLOADER_API int get_check_deprecated(rtloader_t *rtloader, rtload */ DATADOG_AGENT_RTLOADER_API char *run_check(rtloader_t *, rtloader_pyobject_t *check); +/*! \fn char *run_postgres_remote_query(rtloader_t *, rtloader_pyobject_t *check, const char *request_json) + \brief Runs the fixed Postgres remote query helper for a check instance. + \param rtloader_t A rtloader_t * pointer to the RtLoader instance. + \param check A rtloader_pyobject_t * pointer to the check instance we wish to use. + \param request_json A credential-free JSON request string. + \return A C-string with the JSON result. + \sa rtloader_pyobject_t, rtloader_t +*/ +DATADOG_AGENT_RTLOADER_API char *run_postgres_remote_query(rtloader_t *, rtloader_pyobject_t *check, + const char *request_json); + /*! \fn char *cancel_check(rtloader_t *, rtloader_pyobject_t *check) \brief Cancels a check instance. This allow check to be notified when they're unscheduled and can free any remaining resources. diff --git a/rtloader/include/rtloader.h b/rtloader/include/rtloader.h index 98c371e42a8c..b82f12e19b50 100644 --- a/rtloader/include/rtloader.h +++ b/rtloader/include/rtloader.h @@ -121,6 +121,14 @@ class RtLoader */ virtual char *runCheck(RtLoaderPyObject *check) = 0; + //! Pure virtual runPostgresRemoteQuery member. + /*! + \param check The python object pointer to the check we wish to use. + \param request_json A credential-free JSON request string. + \return A C-string with the JSON result. + */ + virtual char *runPostgresRemoteQuery(RtLoaderPyObject *check, const char *request_json) = 0; + //! Pure virtual cancelCheck member. /*! \param check The python object pointer to the check we wish to cancel. diff --git a/rtloader/rtloader/api.cpp b/rtloader/rtloader/api.cpp index 0c62375399d6..114b501d4556 100644 --- a/rtloader/rtloader/api.cpp +++ b/rtloader/rtloader/api.cpp @@ -278,6 +278,11 @@ char *run_check(rtloader_t *rtloader, rtloader_pyobject_t *check) return AS_TYPE(RtLoader, rtloader)->runCheck(AS_TYPE(RtLoaderPyObject, check)); } +char *run_postgres_remote_query(rtloader_t *rtloader, rtloader_pyobject_t *check, const char *request_json) +{ + return AS_TYPE(RtLoader, rtloader)->runPostgresRemoteQuery(AS_TYPE(RtLoaderPyObject, check), request_json); +} + void cancel_check(rtloader_t *rtloader, rtloader_pyobject_t *check) { AS_TYPE(RtLoader, rtloader)->cancelCheck(AS_TYPE(RtLoaderPyObject, check)); diff --git a/rtloader/three/three.cpp b/rtloader/three/three.cpp index 36906d82d5a7..96fad7a513dd 100644 --- a/rtloader/three/three.cpp +++ b/rtloader/three/three.cpp @@ -497,6 +497,58 @@ char *Three::runCheck(RtLoaderPyObject *check) return ret; } +char *Three::runPostgresRemoteQuery(RtLoaderPyObject *check, const char *request_json) +{ + if (check == NULL || request_json == NULL) { + return NULL; + } + + PyObject *py_check = reinterpret_cast(check); + char *ret = NULL; + PyObject *remote_query_module = NULL; + PyObject *execute_func = NULL; + PyObject *py_request_json = NULL; + PyObject *result = NULL; + + remote_query_module = PyImport_ImportModule("datadog_checks.postgres.remote_query"); + if (remote_query_module == NULL) { + setError("error importing Postgres remote query helper: " + _fetchPythonError()); + goto done; + } + + execute_func = PyObject_GetAttrString(remote_query_module, "execute_agent_rpc_json"); + if (execute_func == NULL || !PyCallable_Check(execute_func)) { + setError("error loading Postgres remote query helper: " + _fetchPythonError()); + goto done; + } + + py_request_json = PyUnicode_FromString(request_json); + if (py_request_json == NULL) { + setError("error converting Postgres remote query request to Python string: " + _fetchPythonError()); + goto done; + } + + result = PyObject_CallFunctionObjArgs(execute_func, py_request_json, py_check, NULL); + if (result == NULL || !PyUnicode_Check(result)) { + setError("error invoking Postgres remote query helper: " + _fetchPythonError()); + goto done; + } + + ret = as_string(result); + if (ret == NULL) { + // as_string clears the error, so we can't fetch it here + setError("error converting Postgres remote query helper result to string"); + goto done; + } + + done: + Py_XDECREF(result); + Py_XDECREF(py_request_json); + Py_XDECREF(execute_func); + Py_XDECREF(remote_query_module); + return ret; +} + void Three::cancelCheck(RtLoaderPyObject *check) { if (check == NULL) { diff --git a/rtloader/three/three.h b/rtloader/three/three.h index ffce489bc4b7..7ad49356c204 100644 --- a/rtloader/three/three.h +++ b/rtloader/three/three.h @@ -65,6 +65,7 @@ class Three : public RtLoader const char *provider_str, RtLoaderPyObject *&check); char *runCheck(RtLoaderPyObject *check); + char *runPostgresRemoteQuery(RtLoaderPyObject *check, const char *request_json); void cancelCheck(RtLoaderPyObject *check); char **getCheckWarnings(RtLoaderPyObject *check); char *getCheckDiagnoses(RtLoaderPyObject *check); From a3412f533071d084474d74061d1039cb83cc2665 Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 07:12:10 +0000 Subject: [PATCH 03/33] Add PAR-shaped remote query execute IPC proof --- .../impl/postgres_execute_par_poc.go | 106 ++++++++++++++ .../impl/postgres_execute_par_poc_test.go | 136 ++++++++++++++++++ 2 files changed, 242 insertions(+) create mode 100644 comp/dataobs/remotequeries/impl/postgres_execute_par_poc.go create mode 100644 comp/dataobs/remotequeries/impl/postgres_execute_par_poc_test.go diff --git a/comp/dataobs/remotequeries/impl/postgres_execute_par_poc.go b/comp/dataobs/remotequeries/impl/postgres_execute_par_poc.go new file mode 100644 index 000000000000..ad5f74dd7b6e --- /dev/null +++ b/comp/dataobs/remotequeries/impl/postgres_execute_par_poc.go @@ -0,0 +1,106 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2026-present Datadog, Inc. + +package remotequeriesimpl + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + + ipc "github.com/DataDog/datadog-agent/comp/core/ipc/def" + ipchttp "github.com/DataDog/datadog-agent/comp/core/ipc/httphelpers" +) + +const ( + // AgentPostgresExecuteEndpointPath is the POC-only caller path for the core-Agent command API mount. + // This deliberately documents the current dev proof shape and is not a production IPC API commitment. + AgentPostgresExecuteEndpointPath = "/agent" + PostgresExecuteEndpointPath +) + +// postgresExecutePARIPCClient is the narrow Agent IPC client surface this POC caller needs. +type postgresExecutePARIPCClient interface { + Post(url string, contentType string, body io.Reader, opts ...ipc.RequestOption) (resp []byte, err error) +} + +// PostgresExecutePARHarness is a dev-only, PAR-shaped proof caller for the remote query execute bridge. +// It accepts credential-free action-like inputs, sends them through the injected Agent IPC HTTP client, +// and decodes the execute bridge response without depending on PAR's production runner/bundle registry. +type PostgresExecutePARHarness struct { + client postgresExecutePARIPCClient + endpointURL string +} + +// PostgresExecutePARInputs is the credential-free task input shape for the PAR-shaped POC harness. +type PostgresExecutePARInputs struct { + Target postgresTargetJSON `json:"target"` + Query string `json:"query"` + Limits *postgresExecuteLimitsJSON `json:"limits,omitempty"` +} + +// PostgresExecutePARResult is the decoded execute bridge result or sanitized bridge error. +type PostgresExecutePARResult struct { + HTTPStatus int `json:"-"` + Status string `json:"status"` + Rows []map[string]any `json:"rows,omitempty"` + Error *responseError `json:"error,omitempty"` + Raw json.RawMessage `json:"-"` +} + +// NewPostgresExecutePARHarness creates a dev-only PAR-shaped IPC execute proof caller. +func NewPostgresExecutePARHarness(client postgresExecutePARIPCClient, endpointURL string) *PostgresExecutePARHarness { + return &PostgresExecutePARHarness{client: client, endpointURL: endpointURL} +} + +// Execute sends a credential-free target/query request through the injected Agent IPC HTTP client. +func (h *PostgresExecutePARHarness) Execute(ctx context.Context, inputs PostgresExecutePARInputs) (PostgresExecutePARResult, error) { + if h == nil || h.client == nil { + return PostgresExecutePARResult{}, fmt.Errorf("postgres execute PAR harness requires an IPC client") + } + if h.endpointURL == "" { + return PostgresExecutePARResult{}, fmt.Errorf("postgres execute PAR harness requires an endpoint URL") + } + + payload, err := json.Marshal(inputs) + if err != nil { + return PostgresExecutePARResult{}, fmt.Errorf("marshal postgres execute PAR inputs: %w", err) + } + + body, postErr := h.client.Post(h.endpointURL, "application/json", bytes.NewReader(payload), ipchttp.WithContext(ctx)) + result, decodeErr := decodePostgresExecutePARResponse(body) + if decodeErr == nil { + // IPC HTTPClient returns both the response body and an error for HTTP >= 400. + // The execute bridge body is the sanitized contract, so propagate that decoded body. + return result, nil + } + if postErr != nil { + if len(body) > 0 { + return PostgresExecutePARResult{}, fmt.Errorf("postgres execute IPC request failed with undecodable response") + } + return PostgresExecutePARResult{}, fmt.Errorf("postgres execute IPC request failed: %w", postErr) + } + return PostgresExecutePARResult{}, decodeErr +} + +func decodePostgresExecutePARResponse(body []byte) (PostgresExecutePARResult, error) { + if len(body) == 0 { + return PostgresExecutePARResult{}, fmt.Errorf("empty postgres execute response") + } + + decoder := json.NewDecoder(bytes.NewReader(body)) + decoder.UseNumber() + + var result PostgresExecutePARResult + if err := decoder.Decode(&result); err != nil { + return PostgresExecutePARResult{}, fmt.Errorf("decode postgres execute response: %w", err) + } + if result.Status == "" { + return PostgresExecutePARResult{}, fmt.Errorf("postgres execute response missing status") + } + result.Raw = append(result.Raw[:0], body...) + return result, nil +} diff --git a/comp/dataobs/remotequeries/impl/postgres_execute_par_poc_test.go b/comp/dataobs/remotequeries/impl/postgres_execute_par_poc_test.go new file mode 100644 index 000000000000..2d2e98ec573b --- /dev/null +++ b/comp/dataobs/remotequeries/impl/postgres_execute_par_poc_test.go @@ -0,0 +1,136 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2026-present Datadog, Inc. + +package remotequeriesimpl + +import ( + "context" + "io" + "net/http" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + ipc "github.com/DataDog/datadog-agent/comp/core/ipc/def" + ipcmock "github.com/DataDog/datadog-agent/comp/core/ipc/mock" + "github.com/DataDog/datadog-agent/pkg/collector/check" +) + +func TestPostgresExecutePARHarnessUsesCredentialFreeIPCPostShape(t *testing.T) { + client := &capturePostClient{response: []byte(`{"status":"SUCCEEDED","rows":[{"value":1}]}`)} + harness := NewPostgresExecutePARHarness(client, "https://localhost:5001"+AgentPostgresExecuteEndpointPath) + + result, err := harness.Execute(context.Background(), PostgresExecutePARInputs{ + Target: postgresTargetJSON{Host: "localhost", Port: 5432, DBName: "postgres"}, + Query: postgresRemoteQueryProofQuery, + Limits: &postgresExecuteLimitsJSON{MaxRows: 1, MaxBytes: 1024, TimeoutMs: 1000}, + }) + + require.NoError(t, err) + assert.Equal(t, "https://localhost:5001"+AgentPostgresExecuteEndpointPath, client.url) + assert.Equal(t, "application/json", client.contentType) + assert.JSONEq(t, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":1,"maxBytes":1024,"timeoutMs":1000}}`, client.body) + assert.NotContains(t, client.body, "password") + assert.NotContains(t, client.body, "secret") + assert.Equal(t, "SUCCEEDED", result.Status) + assert.JSONEq(t, `{"status":"SUCCEEDED","rows":[{"value":1}]}`, string(result.Raw)) +} + +func TestPostgresExecutePARHarnessWithRealAgentIPCClient(t *testing.T) { + runner := &fakeRunnerCheck{ + fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: datastore-secret\n"}, + response: `{"status":"SUCCEEDED","rows":[{"value":1}]}`, + } + handler := &postgresExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{fakeWrappedCheck{Check: runner}}}} + ipc := ipcmock.New(t) + mux := http.NewServeMux() + mux.HandleFunc(AgentPostgresExecuteEndpointPath, handler.handle) + server := ipc.NewMockServer(ipc.HTTPMiddleware(mux)) + harness := NewPostgresExecutePARHarness(ipc.GetClient(), server.URL+AgentPostgresExecuteEndpointPath) + + result, err := harness.Execute(context.Background(), PostgresExecutePARInputs{ + Target: postgresTargetJSON{Host: "LOCALHOST.", Port: 5432, DBName: "postgres"}, + Query: postgresRemoteQueryProofQuery, + }) + + require.NoError(t, err) + assert.Equal(t, "SUCCEEDED", result.Status) + assert.JSONEq(t, `{"status":"SUCCEEDED","rows":[{"value":1}]}`, string(result.Raw)) + assert.JSONEq(t, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`, runner.seenRequest()) + assert.NotContains(t, string(result.Raw), "datastore-secret") +} + +func TestPostgresExecutePARHarnessPropagatesSanitizedBridgeErrors(t *testing.T) { + handler := &postgresExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ + &fakeRunnerCheck{fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: datastore-secret\n"}}, + }}} + ipc := ipcmock.New(t) + mux := http.NewServeMux() + mux.HandleFunc(AgentPostgresExecuteEndpointPath, handler.handle) + server := ipc.NewMockServer(ipc.HTTPMiddleware(mux)) + harness := NewPostgresExecutePARHarness(ipc.GetClient(), server.URL+AgentPostgresExecuteEndpointPath) + + tests := []struct { + name string + inputs PostgresExecutePARInputs + wantStatus string + wantCode string + }{ + { + name: "target not found", + inputs: PostgresExecutePARInputs{ + Target: postgresTargetJSON{Host: "localhost", Port: 5432, DBName: "other"}, + Query: postgresRemoteQueryProofQuery, + }, + wantStatus: statusTargetNotFound, + wantCode: statusTargetNotFound, + }, + { + name: "invalid query", + inputs: PostgresExecutePARInputs{ + Target: postgresTargetJSON{Host: "localhost", Port: 5432, DBName: "postgres"}, + Query: "SELECT 2 AS value", + }, + wantStatus: statusInvalidRequest, + wantCode: statusInvalidRequest, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := harness.Execute(context.Background(), tt.inputs) + + require.NoError(t, err) + assert.Equal(t, tt.wantStatus, result.Status) + require.NotNil(t, result.Error) + assert.Equal(t, tt.wantCode, result.Error.Code) + assert.NotContains(t, string(result.Raw), "datastore-secret") + assert.NotContains(t, string(result.Raw), tt.inputs.Target.DBName) + assert.NotContains(t, string(result.Raw), tt.inputs.Query) + }) + } +} + +type capturePostClient struct { + response []byte + err error + url string + contentType string + body string +} + +func (c *capturePostClient) Post(url string, contentType string, body io.Reader, _ ...ipc.RequestOption) ([]byte, error) { + c.url = url + c.contentType = contentType + payload, err := io.ReadAll(body) + if err != nil { + return nil, err + } + c.body = string(payload) + return c.response, c.err +} + +var _ postgresExecutePARIPCClient = (*capturePostClient)(nil) From f4248def8082cddd549e7f6e2b0687dbc84d1deb Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 08:47:32 +0000 Subject: [PATCH 04/33] Rename remotequeries fx package --- cmd/agent/subcommands/run/command.go | 2 +- comp/{dataobs => }/remotequeries/fx/fx.go | 0 2 files changed, 1 insertion(+), 1 deletion(-) rename comp/{dataobs => }/remotequeries/fx/fx.go (100%) diff --git a/cmd/agent/subcommands/run/command.go b/cmd/agent/subcommands/run/command.go index 90b4e595ffaa..73eb97ae5a45 100644 --- a/cmd/agent/subcommands/run/command.go +++ b/cmd/agent/subcommands/run/command.go @@ -41,8 +41,8 @@ import ( workloadselectionfx "github.com/DataDog/datadog-agent/comp/workloadselection/fx" doqueryactionsfx "github.com/DataDog/datadog-agent/comp/dataobs/queryactions/fx" - remotequeriesfx "github.com/DataDog/datadog-agent/comp/dataobs/remotequeries/fx" haagentfx "github.com/DataDog/datadog-agent/comp/haagent/fx" + remotequeriesfx "github.com/DataDog/datadog-agent/comp/remotequeries/fx" snmpscanfx "github.com/DataDog/datadog-agent/comp/snmpscan/fx" snmpscanmanagerfx "github.com/DataDog/datadog-agent/comp/snmpscanmanager/fx" "github.com/DataDog/datadog-agent/pkg/aggregator" diff --git a/comp/dataobs/remotequeries/fx/fx.go b/comp/remotequeries/fx/fx.go similarity index 100% rename from comp/dataobs/remotequeries/fx/fx.go rename to comp/remotequeries/fx/fx.go From e274c2054eb97b20d76850aa405f8c8da98d5893 Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 09:09:00 +0000 Subject: [PATCH 05/33] Generalize remote query bridge integration selector --- .../impl/postgres_execute_par_poc.go | 106 ---------- comp/remotequeries/fx/fx.go | 6 +- .../impl/remote_query_execute.go} | 165 +++++++++------- .../impl/remote_query_match.go} | 141 ++++++++++--- .../impl/remote_query_par_poc.go | 107 ++++++++++ .../impl/remote_query_par_poc_test.go} | 60 +++--- .../impl/remote_query_test.go} | 187 +++++++++++------- pkg/collector/python/check.go | 14 ++ pkg/collector/python/check_test.go | 8 + pkg/collector/python/test_check.go | 36 ++++ 10 files changed, 525 insertions(+), 305 deletions(-) delete mode 100644 comp/dataobs/remotequeries/impl/postgres_execute_par_poc.go rename comp/{dataobs/remotequeries/impl/postgres_execute.go => remotequeries/impl/remote_query_execute.go} (53%) rename comp/{dataobs/remotequeries/impl/postgres_match.go => remotequeries/impl/remote_query_match.go} (61%) create mode 100644 comp/remotequeries/impl/remote_query_par_poc.go rename comp/{dataobs/remotequeries/impl/postgres_execute_par_poc_test.go => remotequeries/impl/remote_query_par_poc_test.go} (56%) rename comp/{dataobs/remotequeries/impl/postgres_match_test.go => remotequeries/impl/remote_query_test.go} (54%) diff --git a/comp/dataobs/remotequeries/impl/postgres_execute_par_poc.go b/comp/dataobs/remotequeries/impl/postgres_execute_par_poc.go deleted file mode 100644 index ad5f74dd7b6e..000000000000 --- a/comp/dataobs/remotequeries/impl/postgres_execute_par_poc.go +++ /dev/null @@ -1,106 +0,0 @@ -// Unless explicitly stated otherwise all files in this repository are licensed -// under the Apache License Version 2.0. -// This product includes software developed at Datadog (https://www.datadoghq.com/). -// Copyright 2026-present Datadog, Inc. - -package remotequeriesimpl - -import ( - "bytes" - "context" - "encoding/json" - "fmt" - "io" - - ipc "github.com/DataDog/datadog-agent/comp/core/ipc/def" - ipchttp "github.com/DataDog/datadog-agent/comp/core/ipc/httphelpers" -) - -const ( - // AgentPostgresExecuteEndpointPath is the POC-only caller path for the core-Agent command API mount. - // This deliberately documents the current dev proof shape and is not a production IPC API commitment. - AgentPostgresExecuteEndpointPath = "/agent" + PostgresExecuteEndpointPath -) - -// postgresExecutePARIPCClient is the narrow Agent IPC client surface this POC caller needs. -type postgresExecutePARIPCClient interface { - Post(url string, contentType string, body io.Reader, opts ...ipc.RequestOption) (resp []byte, err error) -} - -// PostgresExecutePARHarness is a dev-only, PAR-shaped proof caller for the remote query execute bridge. -// It accepts credential-free action-like inputs, sends them through the injected Agent IPC HTTP client, -// and decodes the execute bridge response without depending on PAR's production runner/bundle registry. -type PostgresExecutePARHarness struct { - client postgresExecutePARIPCClient - endpointURL string -} - -// PostgresExecutePARInputs is the credential-free task input shape for the PAR-shaped POC harness. -type PostgresExecutePARInputs struct { - Target postgresTargetJSON `json:"target"` - Query string `json:"query"` - Limits *postgresExecuteLimitsJSON `json:"limits,omitempty"` -} - -// PostgresExecutePARResult is the decoded execute bridge result or sanitized bridge error. -type PostgresExecutePARResult struct { - HTTPStatus int `json:"-"` - Status string `json:"status"` - Rows []map[string]any `json:"rows,omitempty"` - Error *responseError `json:"error,omitempty"` - Raw json.RawMessage `json:"-"` -} - -// NewPostgresExecutePARHarness creates a dev-only PAR-shaped IPC execute proof caller. -func NewPostgresExecutePARHarness(client postgresExecutePARIPCClient, endpointURL string) *PostgresExecutePARHarness { - return &PostgresExecutePARHarness{client: client, endpointURL: endpointURL} -} - -// Execute sends a credential-free target/query request through the injected Agent IPC HTTP client. -func (h *PostgresExecutePARHarness) Execute(ctx context.Context, inputs PostgresExecutePARInputs) (PostgresExecutePARResult, error) { - if h == nil || h.client == nil { - return PostgresExecutePARResult{}, fmt.Errorf("postgres execute PAR harness requires an IPC client") - } - if h.endpointURL == "" { - return PostgresExecutePARResult{}, fmt.Errorf("postgres execute PAR harness requires an endpoint URL") - } - - payload, err := json.Marshal(inputs) - if err != nil { - return PostgresExecutePARResult{}, fmt.Errorf("marshal postgres execute PAR inputs: %w", err) - } - - body, postErr := h.client.Post(h.endpointURL, "application/json", bytes.NewReader(payload), ipchttp.WithContext(ctx)) - result, decodeErr := decodePostgresExecutePARResponse(body) - if decodeErr == nil { - // IPC HTTPClient returns both the response body and an error for HTTP >= 400. - // The execute bridge body is the sanitized contract, so propagate that decoded body. - return result, nil - } - if postErr != nil { - if len(body) > 0 { - return PostgresExecutePARResult{}, fmt.Errorf("postgres execute IPC request failed with undecodable response") - } - return PostgresExecutePARResult{}, fmt.Errorf("postgres execute IPC request failed: %w", postErr) - } - return PostgresExecutePARResult{}, decodeErr -} - -func decodePostgresExecutePARResponse(body []byte) (PostgresExecutePARResult, error) { - if len(body) == 0 { - return PostgresExecutePARResult{}, fmt.Errorf("empty postgres execute response") - } - - decoder := json.NewDecoder(bytes.NewReader(body)) - decoder.UseNumber() - - var result PostgresExecutePARResult - if err := decoder.Decode(&result); err != nil { - return PostgresExecutePARResult{}, fmt.Errorf("decode postgres execute response: %w", err) - } - if result.Status == "" { - return PostgresExecutePARResult{}, fmt.Errorf("postgres execute response missing status") - } - result.Raw = append(result.Raw[:0], body...) - return result, nil -} diff --git a/comp/remotequeries/fx/fx.go b/comp/remotequeries/fx/fx.go index 6fdcf4d91a65..51919ed2df44 100644 --- a/comp/remotequeries/fx/fx.go +++ b/comp/remotequeries/fx/fx.go @@ -7,7 +7,7 @@ package fx import ( - remotequeriesimpl "github.com/DataDog/datadog-agent/comp/dataobs/remotequeries/impl" + remotequeriesimpl "github.com/DataDog/datadog-agent/comp/remotequeries/impl" "github.com/DataDog/datadog-agent/pkg/util/fxutil" "go.uber.org/fx" ) @@ -15,7 +15,7 @@ import ( // Module defines the fx options for this component. func Module() fxutil.Module { return fxutil.Component( - fx.Provide(remotequeriesimpl.NewPostgresMatchEndpointProvider), - fx.Provide(remotequeriesimpl.NewPostgresExecuteEndpointProvider), + fx.Provide(remotequeriesimpl.NewRemoteQueryMatchEndpointProvider), + fx.Provide(remotequeriesimpl.NewRemoteQueryExecuteEndpointProvider), ) } diff --git a/comp/dataobs/remotequeries/impl/postgres_execute.go b/comp/remotequeries/impl/remote_query_execute.go similarity index 53% rename from comp/dataobs/remotequeries/impl/postgres_execute.go rename to comp/remotequeries/impl/remote_query_execute.go index a6a54257ac08..a3e11e206ff2 100644 --- a/comp/dataobs/remotequeries/impl/postgres_execute.go +++ b/comp/remotequeries/impl/remote_query_execute.go @@ -18,30 +18,32 @@ import ( ) const ( - // PostgresExecuteEndpointPath is mounted under /agent by the Agent command API. - PostgresExecuteEndpointPath = "/remote-queries/postgres/execute" - // PostgresExecuteEnabledConfig is disabled by default when the key is absent. - PostgresExecuteEnabledConfig = "remote_queries.postgres_execute.enabled" + // RemoteQueryExecuteEndpointPath is mounted under /agent by the Agent command API. + RemoteQueryExecuteEndpointPath = "/remote-queries/execute" + // RemoteQueriesExecuteEnabledConfig is disabled by default when the key is absent. + RemoteQueriesExecuteEnabledConfig = "remote_queries.execute.enabled" + // legacyPostgresExecuteEnabledConfig preserves compatibility with the earlier POC key. + legacyPostgresExecuteEnabledConfig = "remote_queries.postgres_execute.enabled" - postgresRemoteQueryProofQuery = "SELECT 1 AS value" + remoteQueryProofQuery = "SELECT 1 AS value" statusExecutorUnavailable = "executor_unavailable" ) -type postgresRemoteQueryRunner interface { - RunPostgresRemoteQueryJSON(requestJSON string) (string, error) +type remoteQueryRunner interface { + RunRemoteQueryJSON(integration string, requestJSON string) (string, error) } -type postgresCheckUnwrapper interface { +type remoteQueryCheckUnwrapper interface { Unwrap() check.Check } -func postgresRemoteQueryRunnerFor(chk check.Check) (postgresRemoteQueryRunner, bool) { +func remoteQueryRunnerFor(chk check.Check) (remoteQueryRunner, bool) { for chk != nil { - if runner, ok := chk.(postgresRemoteQueryRunner); ok { + if runner, ok := chk.(remoteQueryRunner); ok { return runner, true } - unwrapper, ok := chk.(postgresCheckUnwrapper) + unwrapper, ok := chk.(remoteQueryCheckUnwrapper) if !ok { break } @@ -54,51 +56,53 @@ func postgresRemoteQueryRunnerFor(chk check.Check) (postgresRemoteQueryRunner, b return nil, false } -// NewPostgresExecuteEndpointProvider registers the Postgres execute endpoint on the internal Agent API. -func NewPostgresExecuteEndpointProvider(reqs Requires) api.AgentEndpointProvider { - h := &postgresExecuteHandler{ +// NewRemoteQueryExecuteEndpointProvider registers the remote query execute endpoint on the internal Agent API. +func NewRemoteQueryExecuteEndpointProvider(reqs Requires) api.AgentEndpointProvider { + h := &remoteQueryExecuteHandler{ collector: reqs.Collector, - enabled: reqs.Cfg.GetBool(PostgresExecuteEnabledConfig), + enabled: remoteQueryEnabled(reqs.Cfg, RemoteQueriesExecuteEnabledConfig, legacyPostgresExecuteEnabledConfig), } - return api.NewAgentEndpointProvider(h.handle, PostgresExecuteEndpointPath, http.MethodPost) + return api.NewAgentEndpointProvider(h.handle, RemoteQueryExecuteEndpointPath, http.MethodPost) } -type postgresExecuteHandler struct { +type remoteQueryExecuteHandler struct { collector collector.Component enabled bool } -type postgresExecuteRequest struct { - Target postgresTarget - Query string - Limits *postgresExecuteLimits +type remoteQueryExecuteRequest struct { + Integration string + Target remoteQueryTarget + Query string + Limits *remoteQueryExecuteLimits } -type postgresExecuteLimits struct { +type remoteQueryExecuteLimits struct { MaxRows int MaxBytes int TimeoutMs int } -type postgresExecuteRequestJSON struct { - Target postgresTargetJSON `json:"target"` - Query string `json:"query"` - Limits *postgresExecuteLimitsJSON `json:"limits,omitempty"` +type remoteQueryExecuteRequestJSON struct { + Integration string `json:"integration"` + Target remoteQueryTargetJSON `json:"target"` + Query string `json:"query"` + Limits *remoteQueryExecuteLimitsJSON `json:"limits,omitempty"` } -type postgresTargetJSON struct { +type remoteQueryTargetJSON struct { Host string `json:"host"` Port int `json:"port"` DBName string `json:"dbname"` } -type postgresExecuteLimitsJSON struct { +type remoteQueryExecuteLimitsJSON struct { MaxRows int `json:"maxRows"` MaxBytes int `json:"maxBytes"` TimeoutMs int `json:"timeoutMs"` } -func (h *postgresExecuteHandler) handle(w http.ResponseWriter, r *http.Request) { +func (h *remoteQueryExecuteHandler) handle(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") if !h.enabled { @@ -108,11 +112,11 @@ func (h *postgresExecuteHandler) handle(w http.ResponseWriter, r *http.Request) req, requestJSON, err := parseExecuteRequest(r) if err != nil { - writeExecuteError(w, http.StatusBadRequest, statusInvalidRequest, err.Error()) + writeExecuteParseError(w, err) return } - matches := findPostgresMatches(h.collector, req.Target) + matches := h.findMatches(req.Integration, req.Target) switch len(matches) { case 0: writeExecuteError(w, http.StatusNotFound, statusTargetNotFound, "no matching Postgres check found") @@ -124,13 +128,13 @@ func (h *postgresExecuteHandler) handle(w http.ResponseWriter, r *http.Request) return } - runner, ok := postgresRemoteQueryRunnerFor(matches[0].check) + runner, ok := remoteQueryRunnerFor(matches[0].check) if !ok { writeExecuteError(w, http.StatusFailedDependency, statusExecutorUnavailable, "matched Postgres check does not support remote query execution") return } - responseJSON, err := runner.RunPostgresRemoteQueryJSON(requestJSON) + responseJSON, err := runner.RunRemoteQueryJSON(req.Integration, requestJSON) if err != nil { writeExecuteError(w, http.StatusBadGateway, statusExecutorUnavailable, "remote query executor failed") return @@ -140,9 +144,9 @@ func (h *postgresExecuteHandler) handle(w http.ResponseWriter, r *http.Request) _, _ = io.WriteString(w, responseJSON) } -func parseExecuteRequest(r *http.Request) (postgresExecuteRequest, string, error) { +func parseExecuteRequest(r *http.Request) (remoteQueryExecuteRequest, string, error) { if !isJSONContentType(r.Header.Get("Content-Type")) { - return postgresExecuteRequest{}, "", fmt.Errorf("content-type must be application/json") + return remoteQueryExecuteRequest{}, "", invalidRequestError("content-type must be application/json") } defer r.Body.Close() @@ -151,99 +155,104 @@ func parseExecuteRequest(r *http.Request) (postgresExecuteRequest, string, error var root map[string]json.RawMessage if err := decoder.Decode(&root); err != nil { - return postgresExecuteRequest{}, "", fmt.Errorf("malformed JSON request") + return remoteQueryExecuteRequest{}, "", invalidRequestError("malformed JSON request") } if err := decoder.Decode(&struct{}{}); err != io.EOF { - return postgresExecuteRequest{}, "", fmt.Errorf("malformed JSON request") + return remoteQueryExecuteRequest{}, "", invalidRequestError("malformed JSON request") } for key := range root { switch key { - case "target", "query", "limits": + case "integration", "target", "query", "limits": continue default: if isCredentialShapedField(key) { - return postgresExecuteRequest{}, "", fmt.Errorf("request contains disallowed credential-shaped field") + return remoteQueryExecuteRequest{}, "", invalidRequestError("request contains disallowed credential-shaped field") } - return postgresExecuteRequest{}, "", fmt.Errorf("request contains unknown field") + return remoteQueryExecuteRequest{}, "", invalidRequestError("request contains unknown field") } } + integration, err := parseIntegrationFromRoot(root) + if err != nil { + return remoteQueryExecuteRequest{}, "", err + } + target, err := parseTargetFromRoot(root) if err != nil { - return postgresExecuteRequest{}, "", err + return remoteQueryExecuteRequest{}, "", err } query, err := parseRequiredString(root, "query") if err != nil { - return postgresExecuteRequest{}, "", err + return remoteQueryExecuteRequest{}, "", err } - if query != postgresRemoteQueryProofQuery { - return postgresExecuteRequest{}, "", fmt.Errorf("query is not allowed") + if query != remoteQueryProofQuery { + return remoteQueryExecuteRequest{}, "", fmt.Errorf("query is not allowed") } limits, err := parseExecuteLimits(root) if err != nil { - return postgresExecuteRequest{}, "", err + return remoteQueryExecuteRequest{}, "", err } - req := postgresExecuteRequest{Target: target, Query: query, Limits: limits} + req := remoteQueryExecuteRequest{Integration: integration, Target: target, Query: query, Limits: limits} requestJSON, err := marshalExecuteRequest(req) if err != nil { - return postgresExecuteRequest{}, "", fmt.Errorf("malformed JSON request") + return remoteQueryExecuteRequest{}, "", fmt.Errorf("malformed JSON request") } return req, requestJSON, nil } -func parseTargetFromRoot(root map[string]json.RawMessage) (postgresTarget, error) { +func parseTargetFromRoot(root map[string]json.RawMessage) (remoteQueryTarget, error) { rawTarget, ok := root["target"] if !ok { - return postgresTarget{}, fmt.Errorf("target is required") + return remoteQueryTarget{}, fmt.Errorf("target is required") } var targetFields map[string]json.RawMessage if err := json.Unmarshal(rawTarget, &targetFields); err != nil || targetFields == nil { - return postgresTarget{}, fmt.Errorf("target must be an object") + return remoteQueryTarget{}, fmt.Errorf("target must be an object") } return parseTargetFields(targetFields) } -func parseTargetFields(targetFields map[string]json.RawMessage) (postgresTarget, error) { +func parseTargetFields(targetFields map[string]json.RawMessage) (remoteQueryTarget, error) { for key := range targetFields { switch key { case "host", "port", "dbname": continue default: if isCredentialShapedField(key) { - return postgresTarget{}, fmt.Errorf("request contains disallowed credential-shaped field") + return remoteQueryTarget{}, fmt.Errorf("request contains disallowed credential-shaped field") } - return postgresTarget{}, fmt.Errorf("target contains unknown field") + return remoteQueryTarget{}, fmt.Errorf("target contains unknown field") } } host, err := parseTargetString(targetFields, "host") if err != nil { - return postgresTarget{}, err + return remoteQueryTarget{}, err } host = normalizeHost(host) if host == "" { - return postgresTarget{}, fmt.Errorf("target.host is required") + return remoteQueryTarget{}, fmt.Errorf("target.host is required") } port, err := parseTargetPort(targetFields) if err != nil { - return postgresTarget{}, err + return remoteQueryTarget{}, err } dbname, err := parseTargetString(targetFields, "dbname") if err != nil { - return postgresTarget{}, err + return remoteQueryTarget{}, err } if dbname == "" { - return postgresTarget{}, fmt.Errorf("target.dbname is required") + return remoteQueryTarget{}, fmt.Errorf("target.dbname is required") } - return postgresTarget{Host: host, Port: port, DBName: dbname}, nil + return remoteQueryTarget{Host: host, Port: port, DBName: dbname}, nil } func parseRequiredString(root map[string]json.RawMessage, field string) (string, error) { @@ -258,7 +267,7 @@ func parseRequiredString(root map[string]json.RawMessage, field string) (string, return value, nil } -func parseExecuteLimits(root map[string]json.RawMessage) (*postgresExecuteLimits, error) { +func parseExecuteLimits(root map[string]json.RawMessage) (*remoteQueryExecuteLimits, error) { rawLimits, ok := root["limits"] if !ok { return nil, nil @@ -294,7 +303,7 @@ func parseExecuteLimits(root map[string]json.RawMessage) (*postgresExecuteLimits return nil, err } - return &postgresExecuteLimits{MaxRows: maxRows, MaxBytes: maxBytes, TimeoutMs: timeoutMs}, nil + return &remoteQueryExecuteLimits{MaxRows: maxRows, MaxBytes: maxBytes, TimeoutMs: timeoutMs}, nil } func parsePositiveJSONInt(fields map[string]json.RawMessage, displayName string, wireName string) (int, error) { @@ -318,13 +327,23 @@ func parsePositiveJSONInt(fields map[string]json.RawMessage, displayName string, return value, nil } -func marshalExecuteRequest(req postgresExecuteRequest) (string, error) { - wireReq := postgresExecuteRequestJSON{ - Target: postgresTargetJSON{Host: req.Target.Host, Port: req.Target.Port, DBName: req.Target.DBName}, - Query: req.Query, +func (h *remoteQueryExecuteHandler) findMatches(integration string, target remoteQueryTarget) []postgresCheckMatch { + switch integration { + case integrationPostgres: + return findPostgresMatches(h.collector, target) + default: + return nil + } +} + +func marshalExecuteRequest(req remoteQueryExecuteRequest) (string, error) { + wireReq := remoteQueryExecuteRequestJSON{ + Integration: req.Integration, + Target: remoteQueryTargetJSON{Host: req.Target.Host, Port: req.Target.Port, DBName: req.Target.DBName}, + Query: req.Query, } if req.Limits != nil { - wireReq.Limits = &postgresExecuteLimitsJSON{ + wireReq.Limits = &remoteQueryExecuteLimitsJSON{ MaxRows: req.Limits.MaxRows, MaxBytes: req.Limits.MaxBytes, TimeoutMs: req.Limits.TimeoutMs, @@ -338,6 +357,20 @@ func marshalExecuteRequest(req postgresExecuteRequest) (string, error) { return string(requestJSON), nil } +func writeExecuteParseError(w http.ResponseWriter, err error) { + parseErr, ok := err.(requestParseError) + if !ok { + writeExecuteError(w, http.StatusBadRequest, statusInvalidRequest, err.Error()) + return + } + + httpStatus := http.StatusBadRequest + if parseErr.status == statusUnsupportedIntegration { + httpStatus = http.StatusUnprocessableEntity + } + writeExecuteError(w, httpStatus, parseErr.status, parseErr.message) +} + func writeExecuteError(w http.ResponseWriter, httpStatus int, status string, message string) { w.WriteHeader(httpStatus) _ = json.NewEncoder(w).Encode(struct { diff --git a/comp/dataobs/remotequeries/impl/postgres_match.go b/comp/remotequeries/impl/remote_query_match.go similarity index 61% rename from comp/dataobs/remotequeries/impl/postgres_match.go rename to comp/remotequeries/impl/remote_query_match.go index 8b056d69202f..0a7ad5b52bce 100644 --- a/comp/dataobs/remotequeries/impl/postgres_match.go +++ b/comp/remotequeries/impl/remote_query_match.go @@ -24,16 +24,21 @@ import ( ) const ( - // PostgresMatchEndpointPath is mounted under /agent by the Agent command API. - PostgresMatchEndpointPath = "/remote-queries/postgres/match-check" - // PostgresMatchEnabledConfig is disabled by default when the key is absent. - PostgresMatchEnabledConfig = "remote_queries.postgres_match_check.enabled" - - statusOK = "ok" - statusTargetNotFound = "target_not_found" - statusAmbiguous = "ambiguous_target" - statusInvalidRequest = "invalid_request" - statusBridgeDisabled = "bridge_disabled" + // RemoteQueryMatchEndpointPath is mounted under /agent by the Agent command API. + RemoteQueryMatchEndpointPath = "/remote-queries/match-check" + // RemoteQueriesMatchEnabledConfig is disabled by default when the key is absent. + RemoteQueriesMatchEnabledConfig = "remote_queries.match_check.enabled" + // legacyPostgresMatchEnabledConfig preserves compatibility with the earlier POC key. + legacyPostgresMatchEnabledConfig = "remote_queries.postgres_match_check.enabled" + + integrationPostgres = "postgres" + + statusOK = "ok" + statusTargetNotFound = "target_not_found" + statusAmbiguous = "ambiguous_target" + statusInvalidRequest = "invalid_request" + statusUnsupportedIntegration = "unsupported_integration" + statusBridgeDisabled = "bridge_disabled" ) var credentialShapedFields = map[string]struct{}{ @@ -64,16 +69,23 @@ type Requires struct { Collector collector.Component } -// NewPostgresMatchEndpointProvider registers the Postgres match endpoint on the internal Agent API. -func NewPostgresMatchEndpointProvider(reqs Requires) api.AgentEndpointProvider { - h := &postgresMatchHandler{ +// NewRemoteQueryMatchEndpointProvider registers the remote query match endpoint on the internal Agent API. +func NewRemoteQueryMatchEndpointProvider(reqs Requires) api.AgentEndpointProvider { + h := &remoteQueryMatchHandler{ collector: reqs.Collector, - enabled: reqs.Cfg.GetBool(PostgresMatchEnabledConfig), + enabled: remoteQueryEnabled(reqs.Cfg, RemoteQueriesMatchEnabledConfig, legacyPostgresMatchEnabledConfig), } - return api.NewAgentEndpointProvider(h.handle, PostgresMatchEndpointPath, http.MethodPost) + return api.NewAgentEndpointProvider(h.handle, RemoteQueryMatchEndpointPath, http.MethodPost) } -type postgresMatchHandler struct { +func remoteQueryEnabled(cfg config.Component, genericKey string, legacyKey string) bool { + if cfg.IsConfigured(genericKey) { + return cfg.GetBool(genericKey) + } + return cfg.GetBool(legacyKey) +} + +type remoteQueryMatchHandler struct { collector collector.Component enabled bool } @@ -96,19 +108,41 @@ type responseError struct { Message string `json:"message"` } -type postgresTarget struct { +type remoteQueryMatchRequest struct { + Integration string + Target remoteQueryTarget +} + +type remoteQueryTarget struct { Host string Port int DBName string } +type requestParseError struct { + status string + message string +} + +func (e requestParseError) Error() string { + return e.message +} + +func invalidRequestError(message string) error { + return requestParseError{status: statusInvalidRequest, message: message} +} + +func unsupportedIntegrationError() error { + return requestParseError{status: statusUnsupportedIntegration, message: "unsupported integration"} +} + type postgresInstanceTarget struct { host string port int dbname string } -func (h *postgresMatchHandler) handle(w http.ResponseWriter, r *http.Request) { +func (h *remoteQueryMatchHandler) handle(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") if !h.enabled { @@ -116,13 +150,13 @@ func (h *postgresMatchHandler) handle(w http.ResponseWriter, r *http.Request) { return } - target, err := parseMatchRequest(r) + req, err := parseMatchRequest(r) if err != nil { - writeMatchResponse(w, http.StatusBadRequest, statusInvalidRequest, 0, nil, err.Error()) + writeMatchParseError(w, err) return } - matches := h.findMatches(target) + matches := h.findMatches(req.Integration, req.Target) switch len(matches) { case 0: writeMatchResponse(w, http.StatusNotFound, statusTargetNotFound, 0, nil, "no matching Postgres check found") @@ -133,9 +167,9 @@ func (h *postgresMatchHandler) handle(w http.ResponseWriter, r *http.Request) { } } -func parseMatchRequest(r *http.Request) (postgresTarget, error) { +func parseMatchRequest(r *http.Request) (remoteQueryMatchRequest, error) { if !isJSONContentType(r.Header.Get("Content-Type")) { - return postgresTarget{}, fmt.Errorf("content-type must be application/json") + return remoteQueryMatchRequest{}, invalidRequestError("content-type must be application/json") } defer r.Body.Close() @@ -144,22 +178,46 @@ func parseMatchRequest(r *http.Request) (postgresTarget, error) { var root map[string]json.RawMessage if err := decoder.Decode(&root); err != nil { - return postgresTarget{}, fmt.Errorf("malformed JSON request") + return remoteQueryMatchRequest{}, invalidRequestError("malformed JSON request") } if err := decoder.Decode(&struct{}{}); err != io.EOF { - return postgresTarget{}, fmt.Errorf("malformed JSON request") + return remoteQueryMatchRequest{}, invalidRequestError("malformed JSON request") } for key := range root { - if key != "target" { + switch key { + case "integration", "target": + continue + default: if isCredentialShapedField(key) { - return postgresTarget{}, fmt.Errorf("request contains disallowed credential-shaped field") + return remoteQueryMatchRequest{}, invalidRequestError("request contains disallowed credential-shaped field") } - return postgresTarget{}, fmt.Errorf("request contains unknown field") + return remoteQueryMatchRequest{}, invalidRequestError("request contains unknown field") } } - return parseTargetFromRoot(root) + integration, err := parseIntegrationFromRoot(root) + if err != nil { + return remoteQueryMatchRequest{}, err + } + target, err := parseTargetFromRoot(root) + if err != nil { + return remoteQueryMatchRequest{}, err + } + return remoteQueryMatchRequest{Integration: integration, Target: target}, nil +} + +func parseIntegrationFromRoot(root map[string]json.RawMessage) (string, error) { + integration, err := parseRequiredString(root, "integration") + if err != nil { + return "", err + } + switch strings.ToLower(strings.TrimSpace(integration)) { + case integrationPostgres, "postgresql": + return integrationPostgres, nil + default: + return "", unsupportedIntegrationError() + } } func isJSONContentType(contentType string) bool { @@ -213,11 +271,16 @@ type postgresCheckMatch struct { sanitized sanitizedMatch } -func (h *postgresMatchHandler) findMatches(target postgresTarget) []postgresCheckMatch { - return findPostgresMatches(h.collector, target) +func (h *remoteQueryMatchHandler) findMatches(integration string, target remoteQueryTarget) []postgresCheckMatch { + switch integration { + case integrationPostgres: + return findPostgresMatches(h.collector, target) + default: + return nil + } } -func findPostgresMatches(collector collector.Component, target postgresTarget) []postgresCheckMatch { +func findPostgresMatches(collector collector.Component, target remoteQueryTarget) []postgresCheckMatch { checks := collector.GetChecks() matches := make([]postgresCheckMatch, 0, 1) for _, chk := range checks { @@ -293,6 +356,20 @@ func yamlInt(value any) (int, bool) { } } +func writeMatchParseError(w http.ResponseWriter, err error) { + parseErr, ok := err.(requestParseError) + if !ok { + writeMatchResponse(w, http.StatusBadRequest, statusInvalidRequest, 0, nil, err.Error()) + return + } + + httpStatus := http.StatusBadRequest + if parseErr.status == statusUnsupportedIntegration { + httpStatus = http.StatusUnprocessableEntity + } + writeMatchResponse(w, httpStatus, parseErr.status, 0, nil, parseErr.message) +} + func writeMatchResponse(w http.ResponseWriter, httpStatus int, status string, matchedCount int, match *sanitizedMatch, message string) { w.WriteHeader(httpStatus) resp := matchResponse{ diff --git a/comp/remotequeries/impl/remote_query_par_poc.go b/comp/remotequeries/impl/remote_query_par_poc.go new file mode 100644 index 000000000000..1844a2aca334 --- /dev/null +++ b/comp/remotequeries/impl/remote_query_par_poc.go @@ -0,0 +1,107 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2026-present Datadog, Inc. + +package remotequeriesimpl + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + + ipc "github.com/DataDog/datadog-agent/comp/core/ipc/def" + ipchttp "github.com/DataDog/datadog-agent/comp/core/ipc/httphelpers" +) + +const ( + // AgentRemoteQueryExecuteEndpointPath is the POC-only caller path for the core-Agent command API mount. + // This deliberately documents the current dev proof shape and is not a production IPC API commitment. + AgentRemoteQueryExecuteEndpointPath = "/agent" + RemoteQueryExecuteEndpointPath +) + +// remoteQueryPARIPCClient is the narrow Agent IPC client surface this POC caller needs. +type remoteQueryPARIPCClient interface { + Post(url string, contentType string, body io.Reader, opts ...ipc.RequestOption) (resp []byte, err error) +} + +// RemoteQueryPARHarness is a dev-only, PAR-shaped proof caller for the remote query execute bridge. +// It accepts credential-free action-like inputs, sends them through the injected Agent IPC HTTP client, +// and decodes the execute bridge response without depending on PAR's production runner/bundle registry. +type RemoteQueryPARHarness struct { + client remoteQueryPARIPCClient + endpointURL string +} + +// RemoteQueryPARInputs is the credential-free task input shape for the PAR-shaped POC harness. +type RemoteQueryPARInputs struct { + Integration string `json:"integration"` + Target remoteQueryTargetJSON `json:"target"` + Query string `json:"query"` + Limits *remoteQueryExecuteLimitsJSON `json:"limits,omitempty"` +} + +// RemoteQueryPARResult is the decoded execute bridge result or sanitized bridge error. +type RemoteQueryPARResult struct { + HTTPStatus int `json:"-"` + Status string `json:"status"` + Rows []map[string]any `json:"rows,omitempty"` + Error *responseError `json:"error,omitempty"` + Raw json.RawMessage `json:"-"` +} + +// NewRemoteQueryPARHarness creates a dev-only PAR-shaped IPC execute proof caller. +func NewRemoteQueryPARHarness(client remoteQueryPARIPCClient, endpointURL string) *RemoteQueryPARHarness { + return &RemoteQueryPARHarness{client: client, endpointURL: endpointURL} +} + +// Execute sends a credential-free target/query request through the injected Agent IPC HTTP client. +func (h *RemoteQueryPARHarness) Execute(ctx context.Context, inputs RemoteQueryPARInputs) (RemoteQueryPARResult, error) { + if h == nil || h.client == nil { + return RemoteQueryPARResult{}, fmt.Errorf("remote query PAR harness requires an IPC client") + } + if h.endpointURL == "" { + return RemoteQueryPARResult{}, fmt.Errorf("remote query PAR harness requires an endpoint URL") + } + + payload, err := json.Marshal(inputs) + if err != nil { + return RemoteQueryPARResult{}, fmt.Errorf("marshal remote query PAR inputs: %w", err) + } + + body, postErr := h.client.Post(h.endpointURL, "application/json", bytes.NewReader(payload), ipchttp.WithContext(ctx)) + result, decodeErr := decodeRemoteQueryPARResponse(body) + if decodeErr == nil { + // IPC HTTPClient returns both the response body and an error for HTTP >= 400. + // The execute bridge body is the sanitized contract, so propagate that decoded body. + return result, nil + } + if postErr != nil { + if len(body) > 0 { + return RemoteQueryPARResult{}, fmt.Errorf("remote query IPC request failed with undecodable response") + } + return RemoteQueryPARResult{}, fmt.Errorf("remote query IPC request failed: %w", postErr) + } + return RemoteQueryPARResult{}, decodeErr +} + +func decodeRemoteQueryPARResponse(body []byte) (RemoteQueryPARResult, error) { + if len(body) == 0 { + return RemoteQueryPARResult{}, fmt.Errorf("empty remote query response") + } + + decoder := json.NewDecoder(bytes.NewReader(body)) + decoder.UseNumber() + + var result RemoteQueryPARResult + if err := decoder.Decode(&result); err != nil { + return RemoteQueryPARResult{}, fmt.Errorf("decode remote query response: %w", err) + } + if result.Status == "" { + return RemoteQueryPARResult{}, fmt.Errorf("remote query response missing status") + } + result.Raw = append(result.Raw[:0], body...) + return result, nil +} diff --git a/comp/dataobs/remotequeries/impl/postgres_execute_par_poc_test.go b/comp/remotequeries/impl/remote_query_par_poc_test.go similarity index 56% rename from comp/dataobs/remotequeries/impl/postgres_execute_par_poc_test.go rename to comp/remotequeries/impl/remote_query_par_poc_test.go index 2d2e98ec573b..749c76bc825f 100644 --- a/comp/dataobs/remotequeries/impl/postgres_execute_par_poc_test.go +++ b/comp/remotequeries/impl/remote_query_par_poc_test.go @@ -19,80 +19,84 @@ import ( "github.com/DataDog/datadog-agent/pkg/collector/check" ) -func TestPostgresExecutePARHarnessUsesCredentialFreeIPCPostShape(t *testing.T) { +func TestRemoteQueryPARHarnessUsesCredentialFreeIPCPostShape(t *testing.T) { client := &capturePostClient{response: []byte(`{"status":"SUCCEEDED","rows":[{"value":1}]}`)} - harness := NewPostgresExecutePARHarness(client, "https://localhost:5001"+AgentPostgresExecuteEndpointPath) + harness := NewRemoteQueryPARHarness(client, "https://localhost:5001"+AgentRemoteQueryExecuteEndpointPath) - result, err := harness.Execute(context.Background(), PostgresExecutePARInputs{ - Target: postgresTargetJSON{Host: "localhost", Port: 5432, DBName: "postgres"}, - Query: postgresRemoteQueryProofQuery, - Limits: &postgresExecuteLimitsJSON{MaxRows: 1, MaxBytes: 1024, TimeoutMs: 1000}, + result, err := harness.Execute(context.Background(), RemoteQueryPARInputs{ + Integration: integrationPostgres, + Target: remoteQueryTargetJSON{Host: "localhost", Port: 5432, DBName: "postgres"}, + Query: remoteQueryProofQuery, + Limits: &remoteQueryExecuteLimitsJSON{MaxRows: 1, MaxBytes: 1024, TimeoutMs: 1000}, }) require.NoError(t, err) - assert.Equal(t, "https://localhost:5001"+AgentPostgresExecuteEndpointPath, client.url) + assert.Equal(t, "https://localhost:5001"+AgentRemoteQueryExecuteEndpointPath, client.url) assert.Equal(t, "application/json", client.contentType) - assert.JSONEq(t, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":1,"maxBytes":1024,"timeoutMs":1000}}`, client.body) + assert.JSONEq(t, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":1,"maxBytes":1024,"timeoutMs":1000}}`, client.body) assert.NotContains(t, client.body, "password") assert.NotContains(t, client.body, "secret") assert.Equal(t, "SUCCEEDED", result.Status) assert.JSONEq(t, `{"status":"SUCCEEDED","rows":[{"value":1}]}`, string(result.Raw)) } -func TestPostgresExecutePARHarnessWithRealAgentIPCClient(t *testing.T) { +func TestRemoteQueryPARHarnessWithRealAgentIPCClient(t *testing.T) { runner := &fakeRunnerCheck{ fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: datastore-secret\n"}, response: `{"status":"SUCCEEDED","rows":[{"value":1}]}`, } - handler := &postgresExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{fakeWrappedCheck{Check: runner}}}} + handler := &remoteQueryExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{fakeWrappedCheck{Check: runner}}}} ipc := ipcmock.New(t) mux := http.NewServeMux() - mux.HandleFunc(AgentPostgresExecuteEndpointPath, handler.handle) + mux.HandleFunc(AgentRemoteQueryExecuteEndpointPath, handler.handle) server := ipc.NewMockServer(ipc.HTTPMiddleware(mux)) - harness := NewPostgresExecutePARHarness(ipc.GetClient(), server.URL+AgentPostgresExecuteEndpointPath) + harness := NewRemoteQueryPARHarness(ipc.GetClient(), server.URL+AgentRemoteQueryExecuteEndpointPath) - result, err := harness.Execute(context.Background(), PostgresExecutePARInputs{ - Target: postgresTargetJSON{Host: "LOCALHOST.", Port: 5432, DBName: "postgres"}, - Query: postgresRemoteQueryProofQuery, + result, err := harness.Execute(context.Background(), RemoteQueryPARInputs{ + Integration: integrationPostgres, + Target: remoteQueryTargetJSON{Host: "LOCALHOST.", Port: 5432, DBName: "postgres"}, + Query: remoteQueryProofQuery, }) require.NoError(t, err) assert.Equal(t, "SUCCEEDED", result.Status) assert.JSONEq(t, `{"status":"SUCCEEDED","rows":[{"value":1}]}`, string(result.Raw)) - assert.JSONEq(t, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`, runner.seenRequest()) + assert.JSONEq(t, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`, runner.seenRequest()) assert.NotContains(t, string(result.Raw), "datastore-secret") } -func TestPostgresExecutePARHarnessPropagatesSanitizedBridgeErrors(t *testing.T) { - handler := &postgresExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ +func TestRemoteQueryPARHarnessPropagatesSanitizedBridgeErrors(t *testing.T) { + handler := &remoteQueryExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ &fakeRunnerCheck{fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: datastore-secret\n"}}, }}} ipc := ipcmock.New(t) mux := http.NewServeMux() - mux.HandleFunc(AgentPostgresExecuteEndpointPath, handler.handle) + mux.HandleFunc(AgentRemoteQueryExecuteEndpointPath, handler.handle) server := ipc.NewMockServer(ipc.HTTPMiddleware(mux)) - harness := NewPostgresExecutePARHarness(ipc.GetClient(), server.URL+AgentPostgresExecuteEndpointPath) + harness := NewRemoteQueryPARHarness(ipc.GetClient(), server.URL+AgentRemoteQueryExecuteEndpointPath) tests := []struct { name string - inputs PostgresExecutePARInputs + inputs RemoteQueryPARInputs wantStatus string wantCode string }{ { name: "target not found", - inputs: PostgresExecutePARInputs{ - Target: postgresTargetJSON{Host: "localhost", Port: 5432, DBName: "other"}, - Query: postgresRemoteQueryProofQuery, + inputs: RemoteQueryPARInputs{ + Integration: integrationPostgres, + Target: remoteQueryTargetJSON{Host: "localhost", Port: 5432, DBName: "other"}, + Query: remoteQueryProofQuery, }, wantStatus: statusTargetNotFound, wantCode: statusTargetNotFound, }, { name: "invalid query", - inputs: PostgresExecutePARInputs{ - Target: postgresTargetJSON{Host: "localhost", Port: 5432, DBName: "postgres"}, - Query: "SELECT 2 AS value", + inputs: RemoteQueryPARInputs{ + Integration: integrationPostgres, + Target: remoteQueryTargetJSON{Host: "localhost", Port: 5432, DBName: "postgres"}, + Query: "SELECT 2 AS value", }, wantStatus: statusInvalidRequest, wantCode: statusInvalidRequest, @@ -133,4 +137,4 @@ func (c *capturePostClient) Post(url string, contentType string, body io.Reader, return c.response, c.err } -var _ postgresExecutePARIPCClient = (*capturePostClient)(nil) +var _ remoteQueryPARIPCClient = (*capturePostClient)(nil) diff --git a/comp/dataobs/remotequeries/impl/postgres_match_test.go b/comp/remotequeries/impl/remote_query_test.go similarity index 54% rename from comp/dataobs/remotequeries/impl/postgres_match_test.go rename to comp/remotequeries/impl/remote_query_test.go index 94605a81c25b..18c35ec07875 100644 --- a/comp/dataobs/remotequeries/impl/postgres_match_test.go +++ b/comp/remotequeries/impl/remote_query_test.go @@ -32,49 +32,49 @@ func TestParseMatchRequestValidatesStrictShape(t *testing.T) { }{ { name: "unknown top level field", - body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"extra":true}`, + body: `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"extra":true}`, wantError: "request contains unknown field", }, { name: "unknown target field", - body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres","extra":true}}`, + body: `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres","extra":true}}`, wantError: "target contains unknown field", }, { name: "credential-shaped top level field", - body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"password":"secret-value"}`, + body: `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"password":"secret-value"}`, wantError: "request contains disallowed credential-shaped field", }, { name: "credential-shaped target field", - body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres","username":"alice"}}`, + body: `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres","username":"alice"}}`, wantError: "request contains disallowed credential-shaped field", }, { name: "non-integer port", - body: `{"target":{"host":"localhost","port":5432.1,"dbname":"postgres"}}`, + body: `{"integration":"postgres","target":{"host":"localhost","port":5432.1,"dbname":"postgres"}}`, wantError: "target.port must be an integer", }, { name: "string port", - body: `{"target":{"host":"localhost","port":"5432","dbname":"postgres"}}`, + body: `{"integration":"postgres","target":{"host":"localhost","port":"5432","dbname":"postgres"}}`, wantError: "target.port must be an integer", }, { name: "missing dbname", - body: `{"target":{"host":"localhost","port":5432}}`, + body: `{"integration":"postgres","target":{"host":"localhost","port":5432}}`, wantError: "target.dbname is required", }, { name: "malformed JSON", - body: `{"target":`, + body: `{"integration":"postgres","target":`, wantError: "malformed JSON request", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - req := httptest.NewRequest(http.MethodPost, PostgresMatchEndpointPath, strings.NewReader(tt.body)) + req := httptest.NewRequest(http.MethodPost, RemoteQueryMatchEndpointPath, strings.NewReader(tt.body)) req.Header.Set("Content-Type", "application/json") _, err := parseMatchRequest(req) @@ -87,33 +87,45 @@ func TestParseMatchRequestValidatesStrictShape(t *testing.T) { } func TestParseMatchRequestNormalizesTargetHost(t *testing.T) { - req := httptest.NewRequest(http.MethodPost, PostgresMatchEndpointPath, strings.NewReader( - `{"target":{"host":" LocalHost. ","port":5432,"dbname":"Postgres"}}`, + req := httptest.NewRequest(http.MethodPost, RemoteQueryMatchEndpointPath, strings.NewReader( + `{"integration":"postgres","target":{"host":" LocalHost. ","port":5432,"dbname":"Postgres"}}`, )) req.Header.Set("Content-Type", "application/json; charset=utf-8") - target, err := parseMatchRequest(req) + parsed, err := parseMatchRequest(req) require.NoError(t, err) - assert.Equal(t, postgresTarget{Host: "localhost", Port: 5432, DBName: "Postgres"}, target) + assert.Equal(t, integrationPostgres, parsed.Integration) + assert.Equal(t, remoteQueryTarget{Host: "localhost", Port: 5432, DBName: "Postgres"}, parsed.Target) } -func TestPostgresMatchHandlerDisabled(t *testing.T) { - handler := &postgresMatchHandler{enabled: false, collector: fakeCollector{}} +func TestParseMatchRequestRejectsUnsupportedIntegration(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, RemoteQueryMatchEndpointPath, strings.NewReader( + `{"integration":"mysql","target":{"host":"localhost","port":3306,"dbname":"mysql"}}`, + )) + req.Header.Set("Content-Type", "application/json") + + _, err := parseMatchRequest(req) + require.Error(t, err) + assert.Equal(t, "unsupported integration", err.Error()) +} + +func TestRemoteQueryMatchHandlerDisabled(t *testing.T) { + handler := &remoteQueryMatchHandler{enabled: false, collector: fakeCollector{}} - recorder := callMatchHandler(handler, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"}}`) + recorder := callMatchHandler(handler, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"}}`) assert.Equal(t, http.StatusServiceUnavailable, recorder.Code) assert.Contains(t, recorder.Body.String(), `"status":"bridge_disabled"`) } -func TestPostgresMatchHandlerExactMatch(t *testing.T) { - handler := &postgresMatchHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ +func TestRemoteQueryMatchHandlerExactMatch(t *testing.T) { + handler := &remoteQueryMatchHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: LOCALHOST.\nport: 5432\ndbname: postgres\nusername: alice\npassword: secret-value\n"}, fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5433\ndbname: postgres\npassword: other-secret\n"}, fakeCheck{name: "mysql", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: mysql-secret\n"}, }}} - recorder := callMatchHandler(handler, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"}}`) + recorder := callMatchHandler(handler, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"}}`) assert.Equal(t, http.StatusOK, recorder.Code) body := recorder.Body.String() @@ -129,12 +141,12 @@ func TestPostgresMatchHandlerExactMatch(t *testing.T) { assert.NotContains(t, body, "InstanceConfig") } -func TestPostgresMatchHandlerNoMatch(t *testing.T) { - handler := &postgresMatchHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ +func TestRemoteQueryMatchHandlerNoMatch(t *testing.T) { + handler := &remoteQueryMatchHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-value\n"}, }}} - recorder := callMatchHandler(handler, `{"target":{"host":"localhost","port":5432,"dbname":"other"}}`) + recorder := callMatchHandler(handler, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"other"}}`) assert.Equal(t, http.StatusNotFound, recorder.Code) body := recorder.Body.String() @@ -144,13 +156,13 @@ func TestPostgresMatchHandlerNoMatch(t *testing.T) { assert.NotContains(t, body, "other") } -func TestPostgresMatchHandlerAmbiguousMatch(t *testing.T) { - handler := &postgresMatchHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ +func TestRemoteQueryMatchHandlerAmbiguousMatch(t *testing.T) { + handler := &remoteQueryMatchHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-one\n"}, fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-two\n"}, }}} - recorder := callMatchHandler(handler, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"}}`) + recorder := callMatchHandler(handler, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"}}`) assert.Equal(t, http.StatusConflict, recorder.Code) body := recorder.Body.String() @@ -160,10 +172,10 @@ func TestPostgresMatchHandlerAmbiguousMatch(t *testing.T) { assert.NotContains(t, body, "secret-two") } -func TestPostgresMatchHandlerCredentialRequestDoesNotEchoValue(t *testing.T) { - handler := &postgresMatchHandler{enabled: true, collector: fakeCollector{}} +func TestRemoteQueryMatchHandlerCredentialRequestDoesNotEchoValue(t *testing.T) { + handler := &remoteQueryMatchHandler{enabled: true, collector: fakeCollector{}} - recorder := callMatchHandler(handler, `{"target":{"host":"localhost","port":5432,"dbname":"postgres","dsn":"postgres://secret-value@example/db"}}`) + recorder := callMatchHandler(handler, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres","dsn":"postgres://secret-value@example/db"}}`) assert.Equal(t, http.StatusBadRequest, recorder.Code) body := recorder.Body.String() @@ -173,9 +185,19 @@ func TestPostgresMatchHandlerCredentialRequestDoesNotEchoValue(t *testing.T) { assert.NotContains(t, body, "secret-value") } -func TestPostgresMatchHandlerRejectsInvalidContentType(t *testing.T) { - handler := &postgresMatchHandler{enabled: true, collector: fakeCollector{}} - req := httptest.NewRequest(http.MethodPost, PostgresMatchEndpointPath, strings.NewReader(`{"target":{"host":"localhost","port":5432,"dbname":"postgres"}}`)) +func TestRemoteQueryMatchHandlerRejectsUnsupportedIntegration(t *testing.T) { + handler := &remoteQueryMatchHandler{enabled: true, collector: fakeCollector{}} + + recorder := callMatchHandler(handler, `{"integration":"mysql","target":{"host":"localhost","port":3306,"dbname":"mysql"}}`) + + assert.Equal(t, http.StatusUnprocessableEntity, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"status":"unsupported_integration"`) + assert.NotContains(t, recorder.Body.String(), "mysql") +} + +func TestRemoteQueryMatchHandlerRejectsInvalidContentType(t *testing.T) { + handler := &remoteQueryMatchHandler{enabled: true, collector: fakeCollector{}} + req := httptest.NewRequest(http.MethodPost, RemoteQueryMatchEndpointPath, strings.NewReader(`{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"}}`)) req.Header.Set("Content-Type", "text/plain") recorder := httptest.NewRecorder() @@ -185,8 +207,8 @@ func TestPostgresMatchHandlerRejectsInvalidContentType(t *testing.T) { assert.Contains(t, recorder.Body.String(), "content-type must be application/json") } -func callMatchHandler(handler *postgresMatchHandler, body string) *httptest.ResponseRecorder { - req := httptest.NewRequest(http.MethodPost, PostgresMatchEndpointPath, strings.NewReader(body)) +func callMatchHandler(handler *remoteQueryMatchHandler, body string) *httptest.ResponseRecorder { + req := httptest.NewRequest(http.MethodPost, RemoteQueryMatchEndpointPath, strings.NewReader(body)) req.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() handler.handle(recorder, req) @@ -246,54 +268,54 @@ func TestParseExecuteRequestValidatesStrictShape(t *testing.T) { }{ { name: "unknown top level field", - body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","extra":true}`, + body: `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","extra":true}`, wantError: "request contains unknown field", }, { name: "credential-shaped top level field", - body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","token":"secret-value"}`, + body: `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","token":"secret-value"}`, wantError: "request contains disallowed credential-shaped field", }, { name: "unknown target field", - body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres","extra":true},"query":"SELECT 1 AS value"}`, + body: `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres","extra":true},"query":"SELECT 1 AS value"}`, wantError: "target contains unknown field", }, { name: "credential-shaped target field", - body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres","password":"secret-value"},"query":"SELECT 1 AS value"}`, + body: `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres","password":"secret-value"},"query":"SELECT 1 AS value"}`, wantError: "request contains disallowed credential-shaped field", }, { name: "non-exact query", - body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value;"}`, + body: `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value;"}`, wantError: "query is not allowed", }, { name: "unknown limits field", - body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":10,"maxBytes":1048576,"timeoutMs":5000,"extra":true}}`, + body: `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":10,"maxBytes":1048576,"timeoutMs":5000,"extra":true}}`, wantError: "limits contains unknown field", }, { name: "credential-shaped limits field", - body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":10,"maxBytes":1048576,"timeoutMs":5000,"password":"secret-value"}}`, + body: `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":10,"maxBytes":1048576,"timeoutMs":5000,"password":"secret-value"}}`, wantError: "request contains disallowed credential-shaped field", }, { name: "string maxRows", - body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":"10","maxBytes":1048576,"timeoutMs":5000}}`, + body: `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":"10","maxBytes":1048576,"timeoutMs":5000}}`, wantError: "limits.maxRows must be an integer", }, { name: "zero timeout", - body: `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":10,"maxBytes":1048576,"timeoutMs":0}}`, + body: `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":10,"maxBytes":1048576,"timeoutMs":0}}`, wantError: "limits.timeoutMs must be at least 1", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - req := httptest.NewRequest(http.MethodPost, PostgresExecuteEndpointPath, strings.NewReader(tt.body)) + req := httptest.NewRequest(http.MethodPost, RemoteQueryExecuteEndpointPath, strings.NewReader(tt.body)) req.Header.Set("Content-Type", "application/json") _, _, err := parseExecuteRequest(req) @@ -305,60 +327,82 @@ func TestParseExecuteRequestValidatesStrictShape(t *testing.T) { } func TestParseExecuteRequestNormalizesAndMarshalsCredentialFreeJSON(t *testing.T) { - req := httptest.NewRequest(http.MethodPost, PostgresExecuteEndpointPath, strings.NewReader( - `{"target":{"host":" LocalHost. ","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":10,"maxBytes":1048576,"timeoutMs":5000}}`, + req := httptest.NewRequest(http.MethodPost, RemoteQueryExecuteEndpointPath, strings.NewReader( + `{"integration":"postgres","target":{"host":" LocalHost. ","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":10,"maxBytes":1048576,"timeoutMs":5000}}`, )) req.Header.Set("Content-Type", "application/json; charset=utf-8") parsed, requestJSON, err := parseExecuteRequest(req) require.NoError(t, err) - assert.Equal(t, postgresTarget{Host: "localhost", Port: 5432, DBName: "postgres"}, parsed.Target) - assert.JSONEq(t, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":10,"maxBytes":1048576,"timeoutMs":5000}}`, requestJSON) + assert.Equal(t, integrationPostgres, parsed.Integration) + assert.Equal(t, remoteQueryTarget{Host: "localhost", Port: 5432, DBName: "postgres"}, parsed.Target) + assert.JSONEq(t, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":10,"maxBytes":1048576,"timeoutMs":5000}}`, requestJSON) +} + +func TestParseExecuteRequestRejectsUnsupportedIntegration(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, RemoteQueryExecuteEndpointPath, strings.NewReader( + `{"integration":"mysql","target":{"host":"localhost","port":3306,"dbname":"mysql"},"query":"SELECT 1 AS value"}`, + )) + req.Header.Set("Content-Type", "application/json") + + _, _, err := parseExecuteRequest(req) + require.Error(t, err) + assert.Equal(t, "unsupported integration", err.Error()) } func TestParseExecuteRequestAllowsOmittedLimits(t *testing.T) { - req := httptest.NewRequest(http.MethodPost, PostgresExecuteEndpointPath, strings.NewReader( - `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`, + req := httptest.NewRequest(http.MethodPost, RemoteQueryExecuteEndpointPath, strings.NewReader( + `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`, )) req.Header.Set("Content-Type", "application/json") parsed, requestJSON, err := parseExecuteRequest(req) require.NoError(t, err) assert.Nil(t, parsed.Limits) - assert.JSONEq(t, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`, requestJSON) + assert.JSONEq(t, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`, requestJSON) } -func TestPostgresExecuteHandlerDisabled(t *testing.T) { - handler := &postgresExecuteHandler{enabled: false, collector: fakeCollector{}} +func TestRemoteQueryExecuteHandlerDisabled(t *testing.T) { + handler := &remoteQueryExecuteHandler{enabled: false, collector: fakeCollector{}} - recorder := callExecuteHandler(handler, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) + recorder := callExecuteHandler(handler, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) assert.Equal(t, http.StatusServiceUnavailable, recorder.Code) assert.Contains(t, recorder.Body.String(), `"status":"bridge_disabled"`) } -func TestPostgresExecuteHandlerRunnerSuccess(t *testing.T) { +func TestRemoteQueryExecuteHandlerRunnerSuccess(t *testing.T) { runner := &fakeRunnerCheck{ fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-value\n"}, response: `{"status":"SUCCEEDED","rows":[{"value":1}]}`, } - handler := &postgresExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{fakeWrappedCheck{Check: runner}}}} + handler := &remoteQueryExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{fakeWrappedCheck{Check: runner}}}} - recorder := callExecuteHandler(handler, `{"target":{"host":"LOCALHOST.","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) + recorder := callExecuteHandler(handler, `{"integration":"postgres","target":{"host":"LOCALHOST.","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) assert.Equal(t, http.StatusOK, recorder.Code) assert.JSONEq(t, `{"status":"SUCCEEDED","rows":[{"value":1}]}`, recorder.Body.String()) - assert.JSONEq(t, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`, runner.seenRequest()) + assert.JSONEq(t, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`, runner.seenRequest()) assert.NotContains(t, recorder.Body.String(), "secret-value") } -func TestPostgresExecuteHandlerNoMatchAndAmbiguous(t *testing.T) { +func TestRemoteQueryExecuteHandlerRejectsUnsupportedIntegration(t *testing.T) { + handler := &remoteQueryExecuteHandler{enabled: true, collector: fakeCollector{}} + + recorder := callExecuteHandler(handler, `{"integration":"mysql","target":{"host":"localhost","port":3306,"dbname":"mysql"},"query":"SELECT 1 AS value"}`) + + assert.Equal(t, http.StatusUnprocessableEntity, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"status":"unsupported_integration"`) + assert.NotContains(t, recorder.Body.String(), "mysql") +} + +func TestRemoteQueryExecuteHandlerNoMatchAndAmbiguous(t *testing.T) { t.Run("no match", func(t *testing.T) { - handler := &postgresExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ + handler := &remoteQueryExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ &fakeRunnerCheck{fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-value\n"}}, }}} - recorder := callExecuteHandler(handler, `{"target":{"host":"localhost","port":5432,"dbname":"other"},"query":"SELECT 1 AS value"}`) + recorder := callExecuteHandler(handler, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"other"},"query":"SELECT 1 AS value"}`) assert.Equal(t, http.StatusNotFound, recorder.Code) assert.Contains(t, recorder.Body.String(), `"status":"target_not_found"`) @@ -367,12 +411,12 @@ func TestPostgresExecuteHandlerNoMatchAndAmbiguous(t *testing.T) { }) t.Run("ambiguous", func(t *testing.T) { - handler := &postgresExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ + handler := &remoteQueryExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ &fakeRunnerCheck{fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-one\n"}}, &fakeRunnerCheck{fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-two\n"}}, }}} - recorder := callExecuteHandler(handler, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) + recorder := callExecuteHandler(handler, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) assert.Equal(t, http.StatusConflict, recorder.Code) assert.Contains(t, recorder.Body.String(), `"status":"ambiguous_target"`) @@ -381,13 +425,13 @@ func TestPostgresExecuteHandlerNoMatchAndAmbiguous(t *testing.T) { }) } -func TestPostgresExecuteHandlerUnsupportedAndRunnerErrorAreSanitized(t *testing.T) { +func TestRemoteQueryExecuteHandlerUnsupportedAndRunnerErrorAreSanitized(t *testing.T) { t.Run("unsupported", func(t *testing.T) { - handler := &postgresExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ + handler := &remoteQueryExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-value\n"}, }}} - recorder := callExecuteHandler(handler, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) + recorder := callExecuteHandler(handler, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) assert.Equal(t, http.StatusFailedDependency, recorder.Code) assert.Contains(t, recorder.Body.String(), `"status":"executor_unavailable"`) @@ -395,11 +439,11 @@ func TestPostgresExecuteHandlerUnsupportedAndRunnerErrorAreSanitized(t *testing. }) t.Run("runner error", func(t *testing.T) { - handler := &postgresExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ + handler := &remoteQueryExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{ &fakeRunnerCheck{fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-value\n"}, err: assert.AnError}, }}} - recorder := callExecuteHandler(handler, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) + recorder := callExecuteHandler(handler, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) assert.Equal(t, http.StatusBadGateway, recorder.Code) assert.Contains(t, recorder.Body.String(), `"status":"executor_unavailable"`) @@ -408,8 +452,8 @@ func TestPostgresExecuteHandlerUnsupportedAndRunnerErrorAreSanitized(t *testing. }) } -func callExecuteHandler(handler *postgresExecuteHandler, body string) *httptest.ResponseRecorder { - req := httptest.NewRequest(http.MethodPost, PostgresExecuteEndpointPath, strings.NewReader(body)) +func callExecuteHandler(handler *remoteQueryExecuteHandler, body string) *httptest.ResponseRecorder { + req := httptest.NewRequest(http.MethodPost, RemoteQueryExecuteEndpointPath, strings.NewReader(body)) req.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() handler.handle(recorder, req) @@ -431,7 +475,10 @@ type fakeRunnerCheck struct { seen string } -func (f *fakeRunnerCheck) RunPostgresRemoteQueryJSON(requestJSON string) (string, error) { +func (f *fakeRunnerCheck) RunRemoteQueryJSON(integration string, requestJSON string) (string, error) { + if integration != integrationPostgres { + return "", assert.AnError + } f.seen = requestJSON return f.response, f.err } diff --git a/pkg/collector/python/check.go b/pkg/collector/python/check.go index 2b8840ef1f34..852346d90ffc 100644 --- a/pkg/collector/python/check.go +++ b/pkg/collector/python/check.go @@ -158,8 +158,22 @@ func (c *PythonCheck) RunSimple() error { return c.runCheck(false) } +// RunRemoteQueryJSON runs a remote query helper for this Python check. +func (c *PythonCheck) RunRemoteQueryJSON(integration string, requestJSON string) (string, error) { + switch strings.ToLower(strings.TrimSpace(integration)) { + case "postgres", "postgresql": + return c.runPostgresRemoteQueryJSON(requestJSON) + default: + return "", fmt.Errorf("unsupported integration") + } +} + // RunPostgresRemoteQueryJSON runs the fixed Postgres remote query helper for this Python check. func (c *PythonCheck) RunPostgresRemoteQueryJSON(requestJSON string) (string, error) { + return c.RunRemoteQueryJSON("postgres", requestJSON) +} + +func (c *PythonCheck) runPostgresRemoteQueryJSON(requestJSON string) (string, error) { gstate, err := newStickyLock() if err != nil { return "", err diff --git a/pkg/collector/python/check_test.go b/pkg/collector/python/check_test.go index 30ac85a41703..83083f1e1c6e 100644 --- a/pkg/collector/python/check_test.go +++ b/pkg/collector/python/check_test.go @@ -71,6 +71,14 @@ func TestRunAfterCancel(t *testing.T) { testRunAfterCancel(t) } +func TestRunRemoteQueryJSON(t *testing.T) { + testRunRemoteQueryJSON(t) +} + +func TestRunRemoteQueryJSONUnsupportedIntegration(t *testing.T) { + testRunRemoteQueryJSONUnsupportedIntegration(t) +} + func TestRunPostgresRemoteQueryJSON(t *testing.T) { testRunPostgresRemoteQueryJSON(t) } diff --git a/pkg/collector/python/test_check.go b/pkg/collector/python/test_check.go index 95860e728196..fd30c9093ff4 100644 --- a/pkg/collector/python/test_check.go +++ b/pkg/collector/python/test_check.go @@ -683,6 +683,42 @@ func testGetDiagnoses(t *testing.T) { assert.Zero(t, len(diagnoses[1].Remediation)) } +func testRunRemoteQueryJSON(t *testing.T) { + mockRtloader(t) + + check, err := NewPythonFakeCheck(aggregator.NewNoOpSenderManager()) + require.NoError(t, err) + check.instance = newMockPyObjectPtr() + + C.reset_check_mock() + C.run_postgres_remote_query_return = C.CString(`{"status":"SUCCEEDED"}`) + + result, err := check.RunRemoteQueryJSON("postgres", `{"integration":"postgres","query":"SELECT 1 AS value"}`) + + require.NoError(t, err) + assert.JSONEq(t, `{"status":"SUCCEEDED"}`, result) + assert.Equal(t, C.int(1), C.run_postgres_remote_query_calls) + assert.Equal(t, check.instance, C.run_postgres_remote_query_instance) + assert.JSONEq(t, `{"integration":"postgres","query":"SELECT 1 AS value"}`, C.GoString(C.run_postgres_remote_query_request_json)) +} + +func testRunRemoteQueryJSONUnsupportedIntegration(t *testing.T) { + mockRtloader(t) + + check, err := NewPythonFakeCheck(aggregator.NewNoOpSenderManager()) + require.NoError(t, err) + check.instance = newMockPyObjectPtr() + + C.reset_check_mock() + + result, err := check.RunRemoteQueryJSON("mysql", `{"integration":"mysql","query":"SELECT 1"}`) + + assert.Empty(t, result) + require.Error(t, err) + assert.EqualError(t, err, "unsupported integration") + assert.Equal(t, C.int(0), C.run_postgres_remote_query_calls) +} + func testRunPostgresRemoteQueryJSON(t *testing.T) { mockRtloader(t) From 3956a2be7076543c0033c35b8994e3e922bdf3e9 Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 09:23:24 +0000 Subject: [PATCH 06/33] Keep integration selector out of Postgres executor payload --- comp/remotequeries/impl/remote_query_execute.go | 16 +++++++--------- .../impl/remote_query_par_poc_test.go | 3 ++- comp/remotequeries/impl/remote_query_test.go | 11 +++++++---- 3 files changed, 16 insertions(+), 14 deletions(-) diff --git a/comp/remotequeries/impl/remote_query_execute.go b/comp/remotequeries/impl/remote_query_execute.go index a3e11e206ff2..c91014a2b374 100644 --- a/comp/remotequeries/impl/remote_query_execute.go +++ b/comp/remotequeries/impl/remote_query_execute.go @@ -83,11 +83,10 @@ type remoteQueryExecuteLimits struct { TimeoutMs int } -type remoteQueryExecuteRequestJSON struct { - Integration string `json:"integration"` - Target remoteQueryTargetJSON `json:"target"` - Query string `json:"query"` - Limits *remoteQueryExecuteLimitsJSON `json:"limits,omitempty"` +type remoteQueryExecutorRequestJSON struct { + Target remoteQueryTargetJSON `json:"target"` + Query string `json:"query"` + Limits *remoteQueryExecuteLimitsJSON `json:"limits,omitempty"` } type remoteQueryTargetJSON struct { @@ -337,10 +336,9 @@ func (h *remoteQueryExecuteHandler) findMatches(integration string, target remot } func marshalExecuteRequest(req remoteQueryExecuteRequest) (string, error) { - wireReq := remoteQueryExecuteRequestJSON{ - Integration: req.Integration, - Target: remoteQueryTargetJSON{Host: req.Target.Host, Port: req.Target.Port, DBName: req.Target.DBName}, - Query: req.Query, + wireReq := remoteQueryExecutorRequestJSON{ + Target: remoteQueryTargetJSON{Host: req.Target.Host, Port: req.Target.Port, DBName: req.Target.DBName}, + Query: req.Query, } if req.Limits != nil { wireReq.Limits = &remoteQueryExecuteLimitsJSON{ diff --git a/comp/remotequeries/impl/remote_query_par_poc_test.go b/comp/remotequeries/impl/remote_query_par_poc_test.go index 749c76bc825f..9a680c5fccd1 100644 --- a/comp/remotequeries/impl/remote_query_par_poc_test.go +++ b/comp/remotequeries/impl/remote_query_par_poc_test.go @@ -61,7 +61,8 @@ func TestRemoteQueryPARHarnessWithRealAgentIPCClient(t *testing.T) { require.NoError(t, err) assert.Equal(t, "SUCCEEDED", result.Status) assert.JSONEq(t, `{"status":"SUCCEEDED","rows":[{"value":1}]}`, string(result.Raw)) - assert.JSONEq(t, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`, runner.seenRequest()) + assert.JSONEq(t, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`, runner.seenRequest()) + assert.NotContains(t, runner.seenRequest(), "integration") assert.NotContains(t, string(result.Raw), "datastore-secret") } diff --git a/comp/remotequeries/impl/remote_query_test.go b/comp/remotequeries/impl/remote_query_test.go index 18c35ec07875..7368677a6abe 100644 --- a/comp/remotequeries/impl/remote_query_test.go +++ b/comp/remotequeries/impl/remote_query_test.go @@ -326,7 +326,7 @@ func TestParseExecuteRequestValidatesStrictShape(t *testing.T) { } } -func TestParseExecuteRequestNormalizesAndMarshalsCredentialFreeJSON(t *testing.T) { +func TestParseExecuteRequestNormalizesAndMarshalsExecutorJSON(t *testing.T) { req := httptest.NewRequest(http.MethodPost, RemoteQueryExecuteEndpointPath, strings.NewReader( `{"integration":"postgres","target":{"host":" LocalHost. ","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":10,"maxBytes":1048576,"timeoutMs":5000}}`, )) @@ -336,7 +336,8 @@ func TestParseExecuteRequestNormalizesAndMarshalsCredentialFreeJSON(t *testing.T require.NoError(t, err) assert.Equal(t, integrationPostgres, parsed.Integration) assert.Equal(t, remoteQueryTarget{Host: "localhost", Port: 5432, DBName: "postgres"}, parsed.Target) - assert.JSONEq(t, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":10,"maxBytes":1048576,"timeoutMs":5000}}`, requestJSON) + assert.JSONEq(t, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":10,"maxBytes":1048576,"timeoutMs":5000}}`, requestJSON) + assert.NotContains(t, requestJSON, "integration") } func TestParseExecuteRequestRejectsUnsupportedIntegration(t *testing.T) { @@ -359,7 +360,8 @@ func TestParseExecuteRequestAllowsOmittedLimits(t *testing.T) { parsed, requestJSON, err := parseExecuteRequest(req) require.NoError(t, err) assert.Nil(t, parsed.Limits) - assert.JSONEq(t, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`, requestJSON) + assert.JSONEq(t, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`, requestJSON) + assert.NotContains(t, requestJSON, "integration") } func TestRemoteQueryExecuteHandlerDisabled(t *testing.T) { @@ -382,7 +384,8 @@ func TestRemoteQueryExecuteHandlerRunnerSuccess(t *testing.T) { assert.Equal(t, http.StatusOK, recorder.Code) assert.JSONEq(t, `{"status":"SUCCEEDED","rows":[{"value":1}]}`, recorder.Body.String()) - assert.JSONEq(t, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`, runner.seenRequest()) + assert.JSONEq(t, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`, runner.seenRequest()) + assert.NotContains(t, runner.seenRequest(), "integration") assert.NotContains(t, recorder.Body.String(), "secret-value") } From 8000ca069644290392f5efb90c21e242d6b6476a Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 09:57:38 +0000 Subject: [PATCH 07/33] Remove redundant remote query credential checks --- .../impl/remote_query_execute.go | 9 ------ comp/remotequeries/impl/remote_query_match.go | 28 ------------------- comp/remotequeries/impl/remote_query_test.go | 24 ++++++++-------- 3 files changed, 12 insertions(+), 49 deletions(-) diff --git a/comp/remotequeries/impl/remote_query_execute.go b/comp/remotequeries/impl/remote_query_execute.go index c91014a2b374..403700135080 100644 --- a/comp/remotequeries/impl/remote_query_execute.go +++ b/comp/remotequeries/impl/remote_query_execute.go @@ -165,9 +165,6 @@ func parseExecuteRequest(r *http.Request) (remoteQueryExecuteRequest, string, er case "integration", "target", "query", "limits": continue default: - if isCredentialShapedField(key) { - return remoteQueryExecuteRequest{}, "", invalidRequestError("request contains disallowed credential-shaped field") - } return remoteQueryExecuteRequest{}, "", invalidRequestError("request contains unknown field") } } @@ -222,9 +219,6 @@ func parseTargetFields(targetFields map[string]json.RawMessage) (remoteQueryTarg case "host", "port", "dbname": continue default: - if isCredentialShapedField(key) { - return remoteQueryTarget{}, fmt.Errorf("request contains disallowed credential-shaped field") - } return remoteQueryTarget{}, fmt.Errorf("target contains unknown field") } } @@ -282,9 +276,6 @@ func parseExecuteLimits(root map[string]json.RawMessage) (*remoteQueryExecuteLim case "maxRows", "maxBytes", "timeoutMs": continue default: - if isCredentialShapedField(key) { - return nil, fmt.Errorf("request contains disallowed credential-shaped field") - } return nil, fmt.Errorf("limits contains unknown field") } } diff --git a/comp/remotequeries/impl/remote_query_match.go b/comp/remotequeries/impl/remote_query_match.go index 0a7ad5b52bce..cadccc62208f 100644 --- a/comp/remotequeries/impl/remote_query_match.go +++ b/comp/remotequeries/impl/remote_query_match.go @@ -41,26 +41,6 @@ const ( statusBridgeDisabled = "bridge_disabled" ) -var credentialShapedFields = map[string]struct{}{ - "app_key": {}, - "apikey": {}, - "api_key": {}, - "conn_string": {}, - "connection_string": {}, - "credential": {}, - "credentials": {}, - "dsn": {}, - "pass": {}, - "password": {}, - "pwd": {}, - "secret": {}, - "sslcert": {}, - "sslkey": {}, - "token": {}, - "user": {}, - "username": {}, -} - // Requires defines dependencies for the Remote Queries POC endpoint provider. type Requires struct { fx.In @@ -189,9 +169,6 @@ func parseMatchRequest(r *http.Request) (remoteQueryMatchRequest, error) { case "integration", "target": continue default: - if isCredentialShapedField(key) { - return remoteQueryMatchRequest{}, invalidRequestError("request contains disallowed credential-shaped field") - } return remoteQueryMatchRequest{}, invalidRequestError("request contains unknown field") } } @@ -261,11 +238,6 @@ func normalizeHost(host string) string { return strings.TrimSuffix(host, ".") } -func isCredentialShapedField(field string) bool { - _, found := credentialShapedFields[strings.ToLower(field)] - return found -} - type postgresCheckMatch struct { check check.Check sanitized sanitizedMatch diff --git a/comp/remotequeries/impl/remote_query_test.go b/comp/remotequeries/impl/remote_query_test.go index 7368677a6abe..839d3d85920c 100644 --- a/comp/remotequeries/impl/remote_query_test.go +++ b/comp/remotequeries/impl/remote_query_test.go @@ -41,14 +41,14 @@ func TestParseMatchRequestValidatesStrictShape(t *testing.T) { wantError: "target contains unknown field", }, { - name: "credential-shaped top level field", + name: "credential-like top level field is unknown", body: `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"password":"secret-value"}`, - wantError: "request contains disallowed credential-shaped field", + wantError: "request contains unknown field", }, { - name: "credential-shaped target field", + name: "credential-like target field is unknown", body: `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres","username":"alice"}}`, - wantError: "request contains disallowed credential-shaped field", + wantError: "target contains unknown field", }, { name: "non-integer port", @@ -172,7 +172,7 @@ func TestRemoteQueryMatchHandlerAmbiguousMatch(t *testing.T) { assert.NotContains(t, body, "secret-two") } -func TestRemoteQueryMatchHandlerCredentialRequestDoesNotEchoValue(t *testing.T) { +func TestRemoteQueryMatchHandlerUnknownTargetFieldDoesNotEchoValue(t *testing.T) { handler := &remoteQueryMatchHandler{enabled: true, collector: fakeCollector{}} recorder := callMatchHandler(handler, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres","dsn":"postgres://secret-value@example/db"}}`) @@ -180,7 +180,7 @@ func TestRemoteQueryMatchHandlerCredentialRequestDoesNotEchoValue(t *testing.T) assert.Equal(t, http.StatusBadRequest, recorder.Code) body := recorder.Body.String() assert.Contains(t, body, `"status":"invalid_request"`) - assert.Contains(t, body, "credential-shaped field") + assert.Contains(t, body, "target contains unknown field") assert.NotContains(t, body, "postgres://secret-value@example/db") assert.NotContains(t, body, "secret-value") } @@ -272,9 +272,9 @@ func TestParseExecuteRequestValidatesStrictShape(t *testing.T) { wantError: "request contains unknown field", }, { - name: "credential-shaped top level field", + name: "credential-like top level field is unknown", body: `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","token":"secret-value"}`, - wantError: "request contains disallowed credential-shaped field", + wantError: "request contains unknown field", }, { name: "unknown target field", @@ -282,9 +282,9 @@ func TestParseExecuteRequestValidatesStrictShape(t *testing.T) { wantError: "target contains unknown field", }, { - name: "credential-shaped target field", + name: "credential-like target field is unknown", body: `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres","password":"secret-value"},"query":"SELECT 1 AS value"}`, - wantError: "request contains disallowed credential-shaped field", + wantError: "target contains unknown field", }, { name: "non-exact query", @@ -297,9 +297,9 @@ func TestParseExecuteRequestValidatesStrictShape(t *testing.T) { wantError: "limits contains unknown field", }, { - name: "credential-shaped limits field", + name: "credential-like limits field is unknown", body: `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":10,"maxBytes":1048576,"timeoutMs":5000,"password":"secret-value"}}`, - wantError: "request contains disallowed credential-shaped field", + wantError: "limits contains unknown field", }, { name: "string maxRows", From e978b8441dd05898d8c56525e14941f093b7eff6 Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 10:33:17 +0000 Subject: [PATCH 08/33] Use strict JSON decoding for remote queries --- .../impl/remote_query_execute.go | 166 ++++++------------ comp/remotequeries/impl/remote_query_match.go | 164 ++++++++++++----- 2 files changed, 172 insertions(+), 158 deletions(-) diff --git a/comp/remotequeries/impl/remote_query_execute.go b/comp/remotequeries/impl/remote_query_execute.go index 403700135080..cf3dd2e6e155 100644 --- a/comp/remotequeries/impl/remote_query_execute.go +++ b/comp/remotequeries/impl/remote_query_execute.go @@ -8,6 +8,7 @@ package remotequeriesimpl import ( "bytes" "encoding/json" + "errors" "fmt" "io" "net/http" @@ -77,6 +78,19 @@ type remoteQueryExecuteRequest struct { Limits *remoteQueryExecuteLimits } +type remoteQueryExecuteRequestJSON struct { + Integration string `json:"integration"` + Target *remoteQueryTargetRequestJSON `json:"target"` + Query string `json:"query"` + Limits *remoteQueryExecuteLimitsRequestJSON `json:"limits,omitempty"` +} + +type remoteQueryExecuteLimitsRequestJSON struct { + MaxRows *int `json:"maxRows"` + MaxBytes *int `json:"maxBytes"` + TimeoutMs *int `json:"timeoutMs"` +} + type remoteQueryExecuteLimits struct { MaxRows int MaxBytes int @@ -149,50 +163,34 @@ func parseExecuteRequest(r *http.Request) (remoteQueryExecuteRequest, string, er } defer r.Body.Close() - decoder := json.NewDecoder(r.Body) - decoder.UseNumber() - - var root map[string]json.RawMessage - if err := decoder.Decode(&root); err != nil { - return remoteQueryExecuteRequest{}, "", invalidRequestError("malformed JSON request") - } - if err := decoder.Decode(&struct{}{}); err != io.EOF { - return remoteQueryExecuteRequest{}, "", invalidRequestError("malformed JSON request") + var wireReq remoteQueryExecuteRequestJSON + if err := decodeStrictJSON(r.Body, &wireReq); err != nil { + return remoteQueryExecuteRequest{}, "", parseJSONRequestError(err) } - for key := range root { - switch key { - case "integration", "target", "query", "limits": - continue - default: - return remoteQueryExecuteRequest{}, "", invalidRequestError("request contains unknown field") - } - } - - integration, err := parseIntegrationFromRoot(root) + integration, err := parseIntegration(wireReq.Integration) if err != nil { return remoteQueryExecuteRequest{}, "", err } - target, err := parseTargetFromRoot(root) + target, err := parseTarget(wireReq.Target) if err != nil { return remoteQueryExecuteRequest{}, "", err } - query, err := parseRequiredString(root, "query") - if err != nil { - return remoteQueryExecuteRequest{}, "", err + if wireReq.Query == "" { + return remoteQueryExecuteRequest{}, "", fmt.Errorf("query is required") } - if query != remoteQueryProofQuery { + if wireReq.Query != remoteQueryProofQuery { return remoteQueryExecuteRequest{}, "", fmt.Errorf("query is not allowed") } - limits, err := parseExecuteLimits(root) + limits, err := parseExecuteLimits(wireReq.Limits) if err != nil { return remoteQueryExecuteRequest{}, "", err } - req := remoteQueryExecuteRequest{Integration: integration, Target: target, Query: query, Limits: limits} + req := remoteQueryExecuteRequest{Integration: integration, Target: target, Query: wireReq.Query, Limits: limits} requestJSON, err := marshalExecuteRequest(req) if err != nil { return remoteQueryExecuteRequest{}, "", fmt.Errorf("malformed JSON request") @@ -200,95 +198,42 @@ func parseExecuteRequest(r *http.Request) (remoteQueryExecuteRequest, string, er return req, requestJSON, nil } -func parseTargetFromRoot(root map[string]json.RawMessage) (remoteQueryTarget, error) { - rawTarget, ok := root["target"] - if !ok { - return remoteQueryTarget{}, fmt.Errorf("target is required") - } +var ( + errLimitsUnknownField = errors.New("limits contains unknown field") + errLimitsMustBeObject = errors.New("limits must be an object") +) - var targetFields map[string]json.RawMessage - if err := json.Unmarshal(rawTarget, &targetFields); err != nil || targetFields == nil { - return remoteQueryTarget{}, fmt.Errorf("target must be an object") +func (l *remoteQueryExecuteLimitsRequestJSON) UnmarshalJSON(data []byte) error { + if !isJSONObject(data) { + return errLimitsMustBeObject } - return parseTargetFields(targetFields) -} -func parseTargetFields(targetFields map[string]json.RawMessage) (remoteQueryTarget, error) { - for key := range targetFields { - switch key { - case "host", "port", "dbname": - continue - default: - return remoteQueryTarget{}, fmt.Errorf("target contains unknown field") + type limitsAlias remoteQueryExecuteLimitsRequestJSON + var limits limitsAlias + if err := decodeStrictJSON(bytes.NewReader(data), &limits); err != nil { + if isUnknownJSONFieldError(err) { + return errLimitsUnknownField } + return err } - - host, err := parseTargetString(targetFields, "host") - if err != nil { - return remoteQueryTarget{}, err - } - host = normalizeHost(host) - if host == "" { - return remoteQueryTarget{}, fmt.Errorf("target.host is required") - } - - port, err := parseTargetPort(targetFields) - if err != nil { - return remoteQueryTarget{}, err - } - - dbname, err := parseTargetString(targetFields, "dbname") - if err != nil { - return remoteQueryTarget{}, err - } - if dbname == "" { - return remoteQueryTarget{}, fmt.Errorf("target.dbname is required") - } - - return remoteQueryTarget{Host: host, Port: port, DBName: dbname}, nil -} - -func parseRequiredString(root map[string]json.RawMessage, field string) (string, error) { - raw, ok := root[field] - if !ok { - return "", fmt.Errorf("%s is required", field) - } - var value string - if err := json.Unmarshal(raw, &value); err != nil { - return "", fmt.Errorf("%s must be a string", field) - } - return value, nil + *l = remoteQueryExecuteLimitsRequestJSON(limits) + return nil } -func parseExecuteLimits(root map[string]json.RawMessage) (*remoteQueryExecuteLimits, error) { - rawLimits, ok := root["limits"] - if !ok { +func parseExecuteLimits(limits *remoteQueryExecuteLimitsRequestJSON) (*remoteQueryExecuteLimits, error) { + if limits == nil { return nil, nil } - var limitFields map[string]json.RawMessage - if err := json.Unmarshal(rawLimits, &limitFields); err != nil || limitFields == nil { - return nil, fmt.Errorf("limits must be an object") - } - - for key := range limitFields { - switch key { - case "maxRows", "maxBytes", "timeoutMs": - continue - default: - return nil, fmt.Errorf("limits contains unknown field") - } - } - - maxRows, err := parsePositiveJSONInt(limitFields, "limits.maxRows", "maxRows") + maxRows, err := parseRequiredPositiveInt(limits.MaxRows, "limits.maxRows") if err != nil { return nil, err } - maxBytes, err := parsePositiveJSONInt(limitFields, "limits.maxBytes", "maxBytes") + maxBytes, err := parseRequiredPositiveInt(limits.MaxBytes, "limits.maxBytes") if err != nil { return nil, err } - timeoutMs, err := parsePositiveJSONInt(limitFields, "limits.timeoutMs", "timeoutMs") + timeoutMs, err := parseRequiredPositiveInt(limits.TimeoutMs, "limits.timeoutMs") if err != nil { return nil, err } @@ -296,25 +241,14 @@ func parseExecuteLimits(root map[string]json.RawMessage) (*remoteQueryExecuteLim return &remoteQueryExecuteLimits{MaxRows: maxRows, MaxBytes: maxBytes, TimeoutMs: timeoutMs}, nil } -func parsePositiveJSONInt(fields map[string]json.RawMessage, displayName string, wireName string) (int, error) { - raw, ok := fields[wireName] - if !ok { - return 0, fmt.Errorf("%s is required", displayName) - } - - decoder := json.NewDecoder(bytes.NewReader(raw)) - decoder.UseNumber() - var value int - if err := decoder.Decode(&value); err != nil { - return 0, fmt.Errorf("%s must be an integer", displayName) - } - if err := decoder.Decode(&struct{}{}); err != io.EOF { - return 0, fmt.Errorf("%s must be an integer", displayName) +func parseRequiredPositiveInt(value *int, name string) (int, error) { + if value == nil { + return 0, fmt.Errorf("%s is required", name) } - if value < 1 { - return 0, fmt.Errorf("%s must be at least 1", displayName) + if *value < 1 { + return 0, fmt.Errorf("%s must be at least 1", name) } - return value, nil + return *value, nil } func (h *remoteQueryExecuteHandler) findMatches(integration string, target remoteQueryTarget) []postgresCheckMatch { diff --git a/comp/remotequeries/impl/remote_query_match.go b/comp/remotequeries/impl/remote_query_match.go index cadccc62208f..b275517772ca 100644 --- a/comp/remotequeries/impl/remote_query_match.go +++ b/comp/remotequeries/impl/remote_query_match.go @@ -7,7 +7,9 @@ package remotequeriesimpl import ( + "bytes" "encoding/json" + "errors" "fmt" "io" "mime" @@ -93,6 +95,17 @@ type remoteQueryMatchRequest struct { Target remoteQueryTarget } +type remoteQueryMatchRequestJSON struct { + Integration string `json:"integration"` + Target *remoteQueryTargetRequestJSON `json:"target"` +} + +type remoteQueryTargetRequestJSON struct { + Host string `json:"host"` + Port *int `json:"port"` + DBName string `json:"dbname"` +} + type remoteQueryTarget struct { Host string Port int @@ -153,43 +166,26 @@ func parseMatchRequest(r *http.Request) (remoteQueryMatchRequest, error) { } defer r.Body.Close() - decoder := json.NewDecoder(r.Body) - decoder.UseNumber() - - var root map[string]json.RawMessage - if err := decoder.Decode(&root); err != nil { - return remoteQueryMatchRequest{}, invalidRequestError("malformed JSON request") - } - if err := decoder.Decode(&struct{}{}); err != io.EOF { - return remoteQueryMatchRequest{}, invalidRequestError("malformed JSON request") - } - - for key := range root { - switch key { - case "integration", "target": - continue - default: - return remoteQueryMatchRequest{}, invalidRequestError("request contains unknown field") - } + var wireReq remoteQueryMatchRequestJSON + if err := decodeStrictJSON(r.Body, &wireReq); err != nil { + return remoteQueryMatchRequest{}, parseJSONRequestError(err) } - integration, err := parseIntegrationFromRoot(root) + integration, err := parseIntegration(wireReq.Integration) if err != nil { return remoteQueryMatchRequest{}, err } - target, err := parseTargetFromRoot(root) + target, err := parseTarget(wireReq.Target) if err != nil { return remoteQueryMatchRequest{}, err } return remoteQueryMatchRequest{Integration: integration, Target: target}, nil } -func parseIntegrationFromRoot(root map[string]json.RawMessage) (string, error) { - integration, err := parseRequiredString(root, "integration") - if err != nil { - return "", err - } +func parseIntegration(integration string) (string, error) { switch strings.ToLower(strings.TrimSpace(integration)) { + case "": + return "", fmt.Errorf("integration is required") case integrationPostgres, "postgresql": return integrationPostgres, nil default: @@ -205,32 +201,116 @@ func isJSONContentType(contentType string) bool { return mediaType == "application/json" } -func parseTargetString(fields map[string]json.RawMessage, field string) (string, error) { - raw, ok := fields[field] - if !ok { - return "", fmt.Errorf("target.%s is required", field) +var ( + errMultipleJSONValues = errors.New("multiple JSON values") + errTargetUnknownField = errors.New("target contains unknown field") + errTargetMustBeObject = errors.New("target must be an object") +) + +func (t *remoteQueryTargetRequestJSON) UnmarshalJSON(data []byte) error { + if !isJSONObject(data) { + return errTargetMustBeObject } - var value string - if err := json.Unmarshal(raw, &value); err != nil { - return "", fmt.Errorf("target.%s must be a string", field) + + type targetAlias remoteQueryTargetRequestJSON + var target targetAlias + if err := decodeStrictJSON(bytes.NewReader(data), &target); err != nil { + if isUnknownJSONFieldError(err) { + return errTargetUnknownField + } + return err } - return value, nil + *t = remoteQueryTargetRequestJSON(target) + return nil } -func parseTargetPort(fields map[string]json.RawMessage) (int, error) { - raw, ok := fields["port"] - if !ok { - return 0, fmt.Errorf("target.port is required") +func parseTarget(target *remoteQueryTargetRequestJSON) (remoteQueryTarget, error) { + if target == nil { + return remoteQueryTarget{}, fmt.Errorf("target is required") + } + + host := normalizeHost(target.Host) + if host == "" { + return remoteQueryTarget{}, fmt.Errorf("target.host is required") } - var port int - if err := json.Unmarshal(raw, &port); err != nil { - return 0, fmt.Errorf("target.port must be an integer") + port, err := parseRequiredPort(target.Port) + if err != nil { + return remoteQueryTarget{}, err } - if port < 1 || port > 65535 { + + if target.DBName == "" { + return remoteQueryTarget{}, fmt.Errorf("target.dbname is required") + } + + return remoteQueryTarget{Host: host, Port: port, DBName: target.DBName}, nil +} + +func parseRequiredPort(port *int) (int, error) { + if port == nil { + return 0, fmt.Errorf("target.port is required") + } + if *port < 1 || *port > 65535 { return 0, fmt.Errorf("target.port is out of range") } - return port, nil + return *port, nil +} + +func decodeStrictJSON(r io.Reader, value any) error { + decoder := json.NewDecoder(r) + decoder.DisallowUnknownFields() + if err := decoder.Decode(value); err != nil { + return err + } + if err := decoder.Decode(&struct{}{}); err != io.EOF { + return errMultipleJSONValues + } + return nil +} + +func parseJSONRequestError(err error) error { + switch { + case errors.Is(err, errMultipleJSONValues): + return invalidRequestError("malformed JSON request") + case errors.Is(err, errTargetUnknownField): + return errTargetUnknownField + case errors.Is(err, errTargetMustBeObject): + return errTargetMustBeObject + case errors.Is(err, errLimitsUnknownField): + return errLimitsUnknownField + case errors.Is(err, errLimitsMustBeObject): + return errLimitsMustBeObject + case isUnknownJSONFieldError(err): + return invalidRequestError("request contains unknown field") + } + + var typeErr *json.UnmarshalTypeError + if errors.As(err, &typeErr) { + switch typeErr.Field { + case "port", "target.port": + return fmt.Errorf("target.port must be an integer") + case "target": + return errTargetMustBeObject + case "maxRows", "limits.maxRows": + return fmt.Errorf("limits.maxRows must be an integer") + case "maxBytes", "limits.maxBytes": + return fmt.Errorf("limits.maxBytes must be an integer") + case "timeoutMs", "limits.timeoutMs": + return fmt.Errorf("limits.timeoutMs must be an integer") + case "limits": + return errLimitsMustBeObject + } + } + + return invalidRequestError("malformed JSON request") +} + +func isUnknownJSONFieldError(err error) bool { + return strings.HasPrefix(err.Error(), "json: unknown field ") +} + +func isJSONObject(data []byte) bool { + return bytes.HasPrefix(bytes.TrimSpace(data), []byte("{")) } func normalizeHost(host string) string { From bed8a6f0d5698295c3b68e5e7f394db79e1ca3a3 Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 10:36:24 +0000 Subject: [PATCH 09/33] Remove Postgres remote query wrapper --- pkg/collector/python/check.go | 5 ---- pkg/collector/python/check_test.go | 16 +++++-------- pkg/collector/python/test_check.go | 37 ++++++++---------------------- 3 files changed, 15 insertions(+), 43 deletions(-) diff --git a/pkg/collector/python/check.go b/pkg/collector/python/check.go index 852346d90ffc..b569146c5714 100644 --- a/pkg/collector/python/check.go +++ b/pkg/collector/python/check.go @@ -168,11 +168,6 @@ func (c *PythonCheck) RunRemoteQueryJSON(integration string, requestJSON string) } } -// RunPostgresRemoteQueryJSON runs the fixed Postgres remote query helper for this Python check. -func (c *PythonCheck) RunPostgresRemoteQueryJSON(requestJSON string) (string, error) { - return c.RunRemoteQueryJSON("postgres", requestJSON) -} - func (c *PythonCheck) runPostgresRemoteQueryJSON(requestJSON string) (string, error) { gstate, err := newStickyLock() if err != nil { diff --git a/pkg/collector/python/check_test.go b/pkg/collector/python/check_test.go index 83083f1e1c6e..9c9e8175c6a2 100644 --- a/pkg/collector/python/check_test.go +++ b/pkg/collector/python/check_test.go @@ -79,18 +79,14 @@ func TestRunRemoteQueryJSONUnsupportedIntegration(t *testing.T) { testRunRemoteQueryJSONUnsupportedIntegration(t) } -func TestRunPostgresRemoteQueryJSON(t *testing.T) { - testRunPostgresRemoteQueryJSON(t) +func TestRunRemoteQueryJSONPostgresError(t *testing.T) { + testRunRemoteQueryJSONPostgresError(t) } -func TestRunPostgresRemoteQueryJSONError(t *testing.T) { - testRunPostgresRemoteQueryJSONError(t) +func TestRunRemoteQueryJSONWithRuntimeNotInitializedError(t *testing.T) { + testRunRemoteQueryJSONWithRuntimeNotInitializedError(t) } -func TestRunPostgresRemoteQueryJSONWithRuntimeNotInitializedError(t *testing.T) { - testRunPostgresRemoteQueryJSONWithRuntimeNotInitializedError(t) -} - -func TestRunPostgresRemoteQueryJSONAfterCancel(t *testing.T) { - testRunPostgresRemoteQueryJSONAfterCancel(t) +func TestRunRemoteQueryJSONAfterCancel(t *testing.T) { + testRunRemoteQueryJSONAfterCancel(t) } diff --git a/pkg/collector/python/test_check.go b/pkg/collector/python/test_check.go index fd30c9093ff4..765a5c5a36d5 100644 --- a/pkg/collector/python/test_check.go +++ b/pkg/collector/python/test_check.go @@ -697,7 +697,10 @@ func testRunRemoteQueryJSON(t *testing.T) { require.NoError(t, err) assert.JSONEq(t, `{"status":"SUCCEEDED"}`, result) + assert.Equal(t, C.int(1), C.gil_locked_calls) + assert.Equal(t, C.int(1), C.gil_unlocked_calls) assert.Equal(t, C.int(1), C.run_postgres_remote_query_calls) + assert.Equal(t, C.int(1), C.rtloader_free_calls) assert.Equal(t, check.instance, C.run_postgres_remote_query_instance) assert.JSONEq(t, `{"integration":"postgres","query":"SELECT 1 AS value"}`, C.GoString(C.run_postgres_remote_query_request_json)) } @@ -719,29 +722,7 @@ func testRunRemoteQueryJSONUnsupportedIntegration(t *testing.T) { assert.Equal(t, C.int(0), C.run_postgres_remote_query_calls) } -func testRunPostgresRemoteQueryJSON(t *testing.T) { - mockRtloader(t) - - check, err := NewPythonFakeCheck(aggregator.NewNoOpSenderManager()) - require.NoError(t, err) - check.instance = newMockPyObjectPtr() - - C.reset_check_mock() - C.run_postgres_remote_query_return = C.CString(`{"status":"SUCCEEDED"}`) - - result, err := check.RunPostgresRemoteQueryJSON(`{"query":"SELECT 1 AS value"}`) - - require.NoError(t, err) - assert.JSONEq(t, `{"status":"SUCCEEDED"}`, result) - assert.Equal(t, C.int(1), C.gil_locked_calls) - assert.Equal(t, C.int(1), C.gil_unlocked_calls) - assert.Equal(t, C.int(1), C.run_postgres_remote_query_calls) - assert.Equal(t, C.int(1), C.rtloader_free_calls) - assert.Equal(t, check.instance, C.run_postgres_remote_query_instance) - assert.Equal(t, `{"query":"SELECT 1 AS value"}`, C.GoString(C.run_postgres_remote_query_request_json)) -} - -func testRunPostgresRemoteQueryJSONError(t *testing.T) { +func testRunRemoteQueryJSONPostgresError(t *testing.T) { mockRtloader(t) check, err := NewPythonFakeCheck(aggregator.NewNoOpSenderManager()) @@ -753,7 +734,7 @@ func testRunPostgresRemoteQueryJSONError(t *testing.T) { C.has_error_return = 1 C.get_error_return = C.CString("rtloader helper failed") - result, err := check.RunPostgresRemoteQueryJSON(`{"query":"SELECT 1 AS value"}`) + result, err := check.RunRemoteQueryJSON("postgres", `{"query":"SELECT 1 AS value"}`) assert.Empty(t, result) require.Error(t, err) @@ -764,7 +745,7 @@ func testRunPostgresRemoteQueryJSONError(t *testing.T) { assert.Equal(t, C.int(0), C.rtloader_free_calls) } -func testRunPostgresRemoteQueryJSONWithRuntimeNotInitializedError(t *testing.T) { +func testRunRemoteQueryJSONWithRuntimeNotInitializedError(t *testing.T) { mockRtloader(t) check, err := NewPythonFakeCheck(aggregator.NewNoOpSenderManager()) require.NoError(t, err) @@ -773,12 +754,12 @@ func testRunPostgresRemoteQueryJSONWithRuntimeNotInitializedError(t *testing.T) C.reset_check_mock() rtloader = nil - _, err = check.RunPostgresRemoteQueryJSON(`{"query":"SELECT 1 AS value"}`) + _, err = check.RunRemoteQueryJSON("postgres", `{"query":"SELECT 1 AS value"}`) assert.ErrorIs(t, err, ErrNotInitialized) assert.Equal(t, C.int(0), C.run_postgres_remote_query_calls) } -func testRunPostgresRemoteQueryJSONAfterCancel(t *testing.T) { +func testRunRemoteQueryJSONAfterCancel(t *testing.T) { mockRtloader(t) check, err := NewPythonFakeCheck(aggregator.NewNoOpSenderManager()) @@ -788,7 +769,7 @@ func testRunPostgresRemoteQueryJSONAfterCancel(t *testing.T) { C.reset_check_mock() check.Cancel() - _, err = check.RunPostgresRemoteQueryJSON(`{"query":"SELECT 1 AS value"}`) + _, err = check.RunRemoteQueryJSON("postgres", `{"query":"SELECT 1 AS value"}`) assert.EqualError(t, err, "check fake_check is already cancelled") assert.Equal(t, C.int(0), C.run_postgres_remote_query_calls) } From 2839cb0c2e2f2e22dc5a4c1135acab5047a06b4f Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 10:55:26 +0000 Subject: [PATCH 10/33] Remove hardcoded Postgres remote query bridge --- .../impl/remote_query_execute.go | 25 ++--- comp/remotequeries/impl/remote_query_match.go | 95 +++++++------------ .../impl/remote_query_par_poc_test.go | 8 +- comp/remotequeries/impl/remote_query_test.go | 36 +++---- pkg/collector/python/check.go | 18 ++-- pkg/collector/python/check_test.go | 8 +- pkg/collector/python/test_check.go | 61 ++++++------ rtloader/include/datadog_agent_rtloader.h | 9 +- rtloader/include/rtloader.h | 5 +- rtloader/rtloader/api.cpp | 4 +- rtloader/three/three.cpp | 54 +++++++++-- rtloader/three/three.h | 2 +- 12 files changed, 169 insertions(+), 156 deletions(-) diff --git a/comp/remotequeries/impl/remote_query_execute.go b/comp/remotequeries/impl/remote_query_execute.go index cf3dd2e6e155..c16ceb07d293 100644 --- a/comp/remotequeries/impl/remote_query_execute.go +++ b/comp/remotequeries/impl/remote_query_execute.go @@ -23,8 +23,6 @@ const ( RemoteQueryExecuteEndpointPath = "/remote-queries/execute" // RemoteQueriesExecuteEnabledConfig is disabled by default when the key is absent. RemoteQueriesExecuteEnabledConfig = "remote_queries.execute.enabled" - // legacyPostgresExecuteEnabledConfig preserves compatibility with the earlier POC key. - legacyPostgresExecuteEnabledConfig = "remote_queries.postgres_execute.enabled" remoteQueryProofQuery = "SELECT 1 AS value" @@ -61,7 +59,7 @@ func remoteQueryRunnerFor(chk check.Check) (remoteQueryRunner, bool) { func NewRemoteQueryExecuteEndpointProvider(reqs Requires) api.AgentEndpointProvider { h := &remoteQueryExecuteHandler{ collector: reqs.Collector, - enabled: remoteQueryEnabled(reqs.Cfg, RemoteQueriesExecuteEnabledConfig, legacyPostgresExecuteEnabledConfig), + enabled: reqs.Cfg.GetBool(RemoteQueriesExecuteEnabledConfig), } return api.NewAgentEndpointProvider(h.handle, RemoteQueryExecuteEndpointPath, http.MethodPost) } @@ -132,18 +130,18 @@ func (h *remoteQueryExecuteHandler) handle(w http.ResponseWriter, r *http.Reques matches := h.findMatches(req.Integration, req.Target) switch len(matches) { case 0: - writeExecuteError(w, http.StatusNotFound, statusTargetNotFound, "no matching Postgres check found") + writeExecuteError(w, http.StatusNotFound, statusTargetNotFound, "no matching integration check found") return case 1: // continue below default: - writeExecuteError(w, http.StatusConflict, statusAmbiguous, "multiple matching Postgres checks found") + writeExecuteError(w, http.StatusConflict, statusAmbiguous, "multiple matching integration checks found") return } runner, ok := remoteQueryRunnerFor(matches[0].check) if !ok { - writeExecuteError(w, http.StatusFailedDependency, statusExecutorUnavailable, "matched Postgres check does not support remote query execution") + writeExecuteError(w, http.StatusFailedDependency, statusExecutorUnavailable, "matched integration check does not support remote query execution") return } @@ -251,13 +249,8 @@ func parseRequiredPositiveInt(value *int, name string) (int, error) { return *value, nil } -func (h *remoteQueryExecuteHandler) findMatches(integration string, target remoteQueryTarget) []postgresCheckMatch { - switch integration { - case integrationPostgres: - return findPostgresMatches(h.collector, target) - default: - return nil - } +func (h *remoteQueryExecuteHandler) findMatches(integration string, target remoteQueryTarget) []integrationCheckMatch { + return findIntegrationMatches(h.collector, integration, target) } func marshalExecuteRequest(req remoteQueryExecuteRequest) (string, error) { @@ -287,11 +280,7 @@ func writeExecuteParseError(w http.ResponseWriter, err error) { return } - httpStatus := http.StatusBadRequest - if parseErr.status == statusUnsupportedIntegration { - httpStatus = http.StatusUnprocessableEntity - } - writeExecuteError(w, httpStatus, parseErr.status, parseErr.message) + writeExecuteError(w, http.StatusBadRequest, parseErr.status, parseErr.message) } func writeExecuteError(w http.ResponseWriter, httpStatus int, status string, message string) { diff --git a/comp/remotequeries/impl/remote_query_match.go b/comp/remotequeries/impl/remote_query_match.go index b275517772ca..078757ceae04 100644 --- a/comp/remotequeries/impl/remote_query_match.go +++ b/comp/remotequeries/impl/remote_query_match.go @@ -14,6 +14,7 @@ import ( "io" "mime" "net/http" + "regexp" "strings" "go.uber.org/fx" @@ -30,17 +31,12 @@ const ( RemoteQueryMatchEndpointPath = "/remote-queries/match-check" // RemoteQueriesMatchEnabledConfig is disabled by default when the key is absent. RemoteQueriesMatchEnabledConfig = "remote_queries.match_check.enabled" - // legacyPostgresMatchEnabledConfig preserves compatibility with the earlier POC key. - legacyPostgresMatchEnabledConfig = "remote_queries.postgres_match_check.enabled" - integrationPostgres = "postgres" - - statusOK = "ok" - statusTargetNotFound = "target_not_found" - statusAmbiguous = "ambiguous_target" - statusInvalidRequest = "invalid_request" - statusUnsupportedIntegration = "unsupported_integration" - statusBridgeDisabled = "bridge_disabled" + statusOK = "ok" + statusTargetNotFound = "target_not_found" + statusAmbiguous = "ambiguous_target" + statusInvalidRequest = "invalid_request" + statusBridgeDisabled = "bridge_disabled" ) // Requires defines dependencies for the Remote Queries POC endpoint provider. @@ -55,18 +51,11 @@ type Requires struct { func NewRemoteQueryMatchEndpointProvider(reqs Requires) api.AgentEndpointProvider { h := &remoteQueryMatchHandler{ collector: reqs.Collector, - enabled: remoteQueryEnabled(reqs.Cfg, RemoteQueriesMatchEnabledConfig, legacyPostgresMatchEnabledConfig), + enabled: reqs.Cfg.GetBool(RemoteQueriesMatchEnabledConfig), } return api.NewAgentEndpointProvider(h.handle, RemoteQueryMatchEndpointPath, http.MethodPost) } -func remoteQueryEnabled(cfg config.Component, genericKey string, legacyKey string) bool { - if cfg.IsConfigured(genericKey) { - return cfg.GetBool(genericKey) - } - return cfg.GetBool(legacyKey) -} - type remoteQueryMatchHandler struct { collector collector.Component enabled bool @@ -125,11 +114,9 @@ func invalidRequestError(message string) error { return requestParseError{status: statusInvalidRequest, message: message} } -func unsupportedIntegrationError() error { - return requestParseError{status: statusUnsupportedIntegration, message: "unsupported integration"} -} +var integrationNamePattern = regexp.MustCompile(`^[a-z0-9_]+$`) -type postgresInstanceTarget struct { +type integrationInstanceTarget struct { host string port int dbname string @@ -152,11 +139,11 @@ func (h *remoteQueryMatchHandler) handle(w http.ResponseWriter, r *http.Request) matches := h.findMatches(req.Integration, req.Target) switch len(matches) { case 0: - writeMatchResponse(w, http.StatusNotFound, statusTargetNotFound, 0, nil, "no matching Postgres check found") + writeMatchResponse(w, http.StatusNotFound, statusTargetNotFound, 0, nil, "no matching integration check found") case 1: writeMatchResponse(w, http.StatusOK, statusOK, 1, &matches[0].sanitized, "") default: - writeMatchResponse(w, http.StatusConflict, statusAmbiguous, len(matches), nil, "multiple matching Postgres checks found") + writeMatchResponse(w, http.StatusConflict, statusAmbiguous, len(matches), nil, "multiple matching integration checks found") } } @@ -183,14 +170,14 @@ func parseMatchRequest(r *http.Request) (remoteQueryMatchRequest, error) { } func parseIntegration(integration string) (string, error) { - switch strings.ToLower(strings.TrimSpace(integration)) { - case "": + integration = strings.ToLower(strings.TrimSpace(integration)) + if integration == "" { return "", fmt.Errorf("integration is required") - case integrationPostgres, "postgresql": - return integrationPostgres, nil - default: - return "", unsupportedIntegrationError() } + if !integrationNamePattern.MatchString(integration) { + return "", invalidRequestError("integration contains invalid characters") + } + return integration, nil } func isJSONContentType(contentType string) bool { @@ -318,38 +305,33 @@ func normalizeHost(host string) string { return strings.TrimSuffix(host, ".") } -type postgresCheckMatch struct { +type integrationCheckMatch struct { check check.Check sanitized sanitizedMatch } -func (h *remoteQueryMatchHandler) findMatches(integration string, target remoteQueryTarget) []postgresCheckMatch { - switch integration { - case integrationPostgres: - return findPostgresMatches(h.collector, target) - default: - return nil - } +func (h *remoteQueryMatchHandler) findMatches(integration string, target remoteQueryTarget) []integrationCheckMatch { + return findIntegrationMatches(h.collector, integration, target) } -func findPostgresMatches(collector collector.Component, target remoteQueryTarget) []postgresCheckMatch { +func findIntegrationMatches(collector collector.Component, integration string, target remoteQueryTarget) []integrationCheckMatch { checks := collector.GetChecks() - matches := make([]postgresCheckMatch, 0, 1) + matches := make([]integrationCheckMatch, 0, 1) for _, chk := range checks { - if !isPostgresCheck(chk) { + if normalizeIntegrationName(chk.String()) != integration { continue } - instanceTarget, ok := parsePostgresInstanceTarget(chk.InstanceConfig()) + instanceTarget, ok := parseIntegrationInstanceTarget(chk.InstanceConfig()) if !ok { continue } if instanceTarget.host == target.Host && instanceTarget.port == target.Port && instanceTarget.dbname == target.DBName { - matches = append(matches, postgresCheckMatch{ + matches = append(matches, integrationCheckMatch{ check: chk, sanitized: sanitizedMatch{ - Integration: "postgres", + Integration: integration, Loader: chk.Loader(), ConfigProvider: chk.ConfigProvider(), }, @@ -359,37 +341,36 @@ func findPostgresMatches(collector collector.Component, target remoteQueryTarget return matches } -func isPostgresCheck(chk check.Check) bool { - name := strings.ToLower(strings.TrimSpace(chk.String())) - return name == "postgres" || name == "postgresql" +func normalizeIntegrationName(name string) string { + return strings.ToLower(strings.TrimSpace(name)) } -func parsePostgresInstanceTarget(instanceConfig string) (postgresInstanceTarget, bool) { +func parseIntegrationInstanceTarget(instanceConfig string) (integrationInstanceTarget, bool) { var fields map[string]any if err := yaml.Unmarshal([]byte(instanceConfig), &fields); err != nil || fields == nil { - return postgresInstanceTarget{}, false + return integrationInstanceTarget{}, false } host, ok := fields["host"].(string) if !ok { - return postgresInstanceTarget{}, false + return integrationInstanceTarget{}, false } host = normalizeHost(host) if host == "" { - return postgresInstanceTarget{}, false + return integrationInstanceTarget{}, false } port, ok := yamlInt(fields["port"]) if !ok || port < 1 || port > 65535 { - return postgresInstanceTarget{}, false + return integrationInstanceTarget{}, false } dbname, ok := fields["dbname"].(string) if !ok || dbname == "" { - return postgresInstanceTarget{}, false + return integrationInstanceTarget{}, false } - return postgresInstanceTarget{host: host, port: port, dbname: dbname}, true + return integrationInstanceTarget{host: host, port: port, dbname: dbname}, true } func yamlInt(value any) (int, bool) { @@ -415,11 +396,7 @@ func writeMatchParseError(w http.ResponseWriter, err error) { return } - httpStatus := http.StatusBadRequest - if parseErr.status == statusUnsupportedIntegration { - httpStatus = http.StatusUnprocessableEntity - } - writeMatchResponse(w, httpStatus, parseErr.status, 0, nil, parseErr.message) + writeMatchResponse(w, http.StatusBadRequest, parseErr.status, 0, nil, parseErr.message) } func writeMatchResponse(w http.ResponseWriter, httpStatus int, status string, matchedCount int, match *sanitizedMatch, message string) { diff --git a/comp/remotequeries/impl/remote_query_par_poc_test.go b/comp/remotequeries/impl/remote_query_par_poc_test.go index 9a680c5fccd1..7525784c885c 100644 --- a/comp/remotequeries/impl/remote_query_par_poc_test.go +++ b/comp/remotequeries/impl/remote_query_par_poc_test.go @@ -24,7 +24,7 @@ func TestRemoteQueryPARHarnessUsesCredentialFreeIPCPostShape(t *testing.T) { harness := NewRemoteQueryPARHarness(client, "https://localhost:5001"+AgentRemoteQueryExecuteEndpointPath) result, err := harness.Execute(context.Background(), RemoteQueryPARInputs{ - Integration: integrationPostgres, + Integration: "postgres", Target: remoteQueryTargetJSON{Host: "localhost", Port: 5432, DBName: "postgres"}, Query: remoteQueryProofQuery, Limits: &remoteQueryExecuteLimitsJSON{MaxRows: 1, MaxBytes: 1024, TimeoutMs: 1000}, @@ -53,7 +53,7 @@ func TestRemoteQueryPARHarnessWithRealAgentIPCClient(t *testing.T) { harness := NewRemoteQueryPARHarness(ipc.GetClient(), server.URL+AgentRemoteQueryExecuteEndpointPath) result, err := harness.Execute(context.Background(), RemoteQueryPARInputs{ - Integration: integrationPostgres, + Integration: "postgres", Target: remoteQueryTargetJSON{Host: "LOCALHOST.", Port: 5432, DBName: "postgres"}, Query: remoteQueryProofQuery, }) @@ -85,7 +85,7 @@ func TestRemoteQueryPARHarnessPropagatesSanitizedBridgeErrors(t *testing.T) { { name: "target not found", inputs: RemoteQueryPARInputs{ - Integration: integrationPostgres, + Integration: "postgres", Target: remoteQueryTargetJSON{Host: "localhost", Port: 5432, DBName: "other"}, Query: remoteQueryProofQuery, }, @@ -95,7 +95,7 @@ func TestRemoteQueryPARHarnessPropagatesSanitizedBridgeErrors(t *testing.T) { { name: "invalid query", inputs: RemoteQueryPARInputs{ - Integration: integrationPostgres, + Integration: "postgres", Target: remoteQueryTargetJSON{Host: "localhost", Port: 5432, DBName: "postgres"}, Query: "SELECT 2 AS value", }, diff --git a/comp/remotequeries/impl/remote_query_test.go b/comp/remotequeries/impl/remote_query_test.go index 839d3d85920c..78e69e3f2b87 100644 --- a/comp/remotequeries/impl/remote_query_test.go +++ b/comp/remotequeries/impl/remote_query_test.go @@ -94,19 +94,19 @@ func TestParseMatchRequestNormalizesTargetHost(t *testing.T) { parsed, err := parseMatchRequest(req) require.NoError(t, err) - assert.Equal(t, integrationPostgres, parsed.Integration) + assert.Equal(t, "postgres", parsed.Integration) assert.Equal(t, remoteQueryTarget{Host: "localhost", Port: 5432, DBName: "Postgres"}, parsed.Target) } -func TestParseMatchRequestRejectsUnsupportedIntegration(t *testing.T) { +func TestParseMatchRequestRejectsInvalidIntegration(t *testing.T) { req := httptest.NewRequest(http.MethodPost, RemoteQueryMatchEndpointPath, strings.NewReader( - `{"integration":"mysql","target":{"host":"localhost","port":3306,"dbname":"mysql"}}`, + `{"integration":"my-sql","target":{"host":"localhost","port":3306,"dbname":"mysql"}}`, )) req.Header.Set("Content-Type", "application/json") _, err := parseMatchRequest(req) require.Error(t, err) - assert.Equal(t, "unsupported integration", err.Error()) + assert.Equal(t, "integration contains invalid characters", err.Error()) } func TestRemoteQueryMatchHandlerDisabled(t *testing.T) { @@ -185,13 +185,14 @@ func TestRemoteQueryMatchHandlerUnknownTargetFieldDoesNotEchoValue(t *testing.T) assert.NotContains(t, body, "secret-value") } -func TestRemoteQueryMatchHandlerRejectsUnsupportedIntegration(t *testing.T) { +func TestRemoteQueryMatchHandlerRejectsInvalidIntegration(t *testing.T) { handler := &remoteQueryMatchHandler{enabled: true, collector: fakeCollector{}} - recorder := callMatchHandler(handler, `{"integration":"mysql","target":{"host":"localhost","port":3306,"dbname":"mysql"}}`) + recorder := callMatchHandler(handler, `{"integration":"my-sql","target":{"host":"localhost","port":3306,"dbname":"mysql"}}`) - assert.Equal(t, http.StatusUnprocessableEntity, recorder.Code) - assert.Contains(t, recorder.Body.String(), `"status":"unsupported_integration"`) + assert.Equal(t, http.StatusBadRequest, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"status":"invalid_request"`) + assert.Contains(t, recorder.Body.String(), "integration contains invalid characters") assert.NotContains(t, recorder.Body.String(), "mysql") } @@ -334,21 +335,21 @@ func TestParseExecuteRequestNormalizesAndMarshalsExecutorJSON(t *testing.T) { parsed, requestJSON, err := parseExecuteRequest(req) require.NoError(t, err) - assert.Equal(t, integrationPostgres, parsed.Integration) + assert.Equal(t, "postgres", parsed.Integration) assert.Equal(t, remoteQueryTarget{Host: "localhost", Port: 5432, DBName: "postgres"}, parsed.Target) assert.JSONEq(t, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":10,"maxBytes":1048576,"timeoutMs":5000}}`, requestJSON) assert.NotContains(t, requestJSON, "integration") } -func TestParseExecuteRequestRejectsUnsupportedIntegration(t *testing.T) { +func TestParseExecuteRequestRejectsInvalidIntegration(t *testing.T) { req := httptest.NewRequest(http.MethodPost, RemoteQueryExecuteEndpointPath, strings.NewReader( - `{"integration":"mysql","target":{"host":"localhost","port":3306,"dbname":"mysql"},"query":"SELECT 1 AS value"}`, + `{"integration":"my-sql","target":{"host":"localhost","port":3306,"dbname":"mysql"},"query":"SELECT 1 AS value"}`, )) req.Header.Set("Content-Type", "application/json") _, _, err := parseExecuteRequest(req) require.Error(t, err) - assert.Equal(t, "unsupported integration", err.Error()) + assert.Equal(t, "integration contains invalid characters", err.Error()) } func TestParseExecuteRequestAllowsOmittedLimits(t *testing.T) { @@ -389,13 +390,14 @@ func TestRemoteQueryExecuteHandlerRunnerSuccess(t *testing.T) { assert.NotContains(t, recorder.Body.String(), "secret-value") } -func TestRemoteQueryExecuteHandlerRejectsUnsupportedIntegration(t *testing.T) { +func TestRemoteQueryExecuteHandlerRejectsInvalidIntegration(t *testing.T) { handler := &remoteQueryExecuteHandler{enabled: true, collector: fakeCollector{}} - recorder := callExecuteHandler(handler, `{"integration":"mysql","target":{"host":"localhost","port":3306,"dbname":"mysql"},"query":"SELECT 1 AS value"}`) + recorder := callExecuteHandler(handler, `{"integration":"my-sql","target":{"host":"localhost","port":3306,"dbname":"mysql"},"query":"SELECT 1 AS value"}`) - assert.Equal(t, http.StatusUnprocessableEntity, recorder.Code) - assert.Contains(t, recorder.Body.String(), `"status":"unsupported_integration"`) + assert.Equal(t, http.StatusBadRequest, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"status":"invalid_request"`) + assert.Contains(t, recorder.Body.String(), "integration contains invalid characters") assert.NotContains(t, recorder.Body.String(), "mysql") } @@ -479,7 +481,7 @@ type fakeRunnerCheck struct { } func (f *fakeRunnerCheck) RunRemoteQueryJSON(integration string, requestJSON string) (string, error) { - if integration != integrationPostgres { + if integration != "postgres" { return "", assert.AnError } f.seen = requestJSON diff --git a/pkg/collector/python/check.go b/pkg/collector/python/check.go index b569146c5714..7f13184fc8aa 100644 --- a/pkg/collector/python/check.go +++ b/pkg/collector/python/check.go @@ -39,7 +39,7 @@ import ( #include "rtloader_mem.h" char *getStringAddr(char **array, unsigned int idx); -char *run_postgres_remote_query(rtloader_t *, rtloader_pyobject_t *check, const char *request_json); +char *run_remote_query(rtloader_t *, rtloader_pyobject_t *check, const char *integration, const char *request_json); static inline void call_free(void* ptr) { _free(ptr); @@ -160,15 +160,11 @@ func (c *PythonCheck) RunSimple() error { // RunRemoteQueryJSON runs a remote query helper for this Python check. func (c *PythonCheck) RunRemoteQueryJSON(integration string, requestJSON string) (string, error) { - switch strings.ToLower(strings.TrimSpace(integration)) { - case "postgres", "postgresql": - return c.runPostgresRemoteQueryJSON(requestJSON) - default: - return "", fmt.Errorf("unsupported integration") + integration = strings.ToLower(strings.TrimSpace(integration)) + if integration == "" { + return "", fmt.Errorf("integration is required") } -} -func (c *PythonCheck) runPostgresRemoteQueryJSON(requestJSON string) (string, error) { gstate, err := newStickyLock() if err != nil { return "", err @@ -179,15 +175,17 @@ func (c *PythonCheck) runPostgresRemoteQueryJSON(requestJSON string) (string, er return "", fmt.Errorf("check %s is already cancelled", c.ModuleName) } + cIntegration := C.CString(integration) + defer C.free(unsafe.Pointer(cIntegration)) cRequestJSON := C.CString(requestJSON) defer C.free(unsafe.Pointer(cRequestJSON)) - cResult := C.run_postgres_remote_query(rtloader, c.instance, cRequestJSON) + cResult := C.run_remote_query(rtloader, c.instance, cIntegration, cRequestJSON) if cResult == nil { if err := getRtLoaderError(); err != nil { return "", err } - return "", fmt.Errorf("an error occurred while running Postgres remote query") + return "", fmt.Errorf("an error occurred while running remote query") } defer C.rtloader_free(rtloader, unsafe.Pointer(cResult)) diff --git a/pkg/collector/python/check_test.go b/pkg/collector/python/check_test.go index 9c9e8175c6a2..8a27336aa5e9 100644 --- a/pkg/collector/python/check_test.go +++ b/pkg/collector/python/check_test.go @@ -75,12 +75,12 @@ func TestRunRemoteQueryJSON(t *testing.T) { testRunRemoteQueryJSON(t) } -func TestRunRemoteQueryJSONUnsupportedIntegration(t *testing.T) { - testRunRemoteQueryJSONUnsupportedIntegration(t) +func TestRunRemoteQueryJSONNormalizesIntegration(t *testing.T) { + testRunRemoteQueryJSONNormalizesIntegration(t) } -func TestRunRemoteQueryJSONPostgresError(t *testing.T) { - testRunRemoteQueryJSONPostgresError(t) +func TestRunRemoteQueryJSONError(t *testing.T) { + testRunRemoteQueryJSONError(t) } func TestRunRemoteQueryJSONWithRuntimeNotInitializedError(t *testing.T) { diff --git a/pkg/collector/python/test_check.go b/pkg/collector/python/test_check.go index 765a5c5a36d5..c704da3aa67d 100644 --- a/pkg/collector/python/test_check.go +++ b/pkg/collector/python/test_check.go @@ -104,15 +104,17 @@ char *get_check_diagnoses(rtloader_t *s, rtloader_pyobject_t *check) { return get_check_diagnoses_return; } -char *run_postgres_remote_query_return = NULL; -int run_postgres_remote_query_calls = 0; -rtloader_pyobject_t *run_postgres_remote_query_instance = NULL; -const char *run_postgres_remote_query_request_json = NULL; -char *run_postgres_remote_query(rtloader_t *s, rtloader_pyobject_t *check, const char *request_json) { - run_postgres_remote_query_instance = check; - run_postgres_remote_query_request_json = strdup(request_json); - run_postgres_remote_query_calls++; - return run_postgres_remote_query_return; +char *run_remote_query_return = NULL; +int run_remote_query_calls = 0; +rtloader_pyobject_t *run_remote_query_instance = NULL; +const char *run_remote_query_integration = NULL; +const char *run_remote_query_request_json = NULL; +char *run_remote_query(rtloader_t *s, rtloader_pyobject_t *check, const char *integration, const char *request_json) { + run_remote_query_instance = check; + run_remote_query_integration = strdup(integration); + run_remote_query_request_json = strdup(request_json); + run_remote_query_calls++; + return run_remote_query_return; } // @@ -214,10 +216,11 @@ void reset_check_mock() { get_check_diagnoses_return = NULL; get_check_diagnoses_calls = 0; - run_postgres_remote_query_return = NULL; - run_postgres_remote_query_calls = 0; - run_postgres_remote_query_instance = NULL; - run_postgres_remote_query_request_json = NULL; + run_remote_query_return = NULL; + run_remote_query_calls = 0; + run_remote_query_instance = NULL; + run_remote_query_integration = NULL; + run_remote_query_request_json = NULL; } */ import "C" @@ -691,7 +694,7 @@ func testRunRemoteQueryJSON(t *testing.T) { check.instance = newMockPyObjectPtr() C.reset_check_mock() - C.run_postgres_remote_query_return = C.CString(`{"status":"SUCCEEDED"}`) + C.run_remote_query_return = C.CString(`{"status":"SUCCEEDED"}`) result, err := check.RunRemoteQueryJSON("postgres", `{"integration":"postgres","query":"SELECT 1 AS value"}`) @@ -699,13 +702,14 @@ func testRunRemoteQueryJSON(t *testing.T) { assert.JSONEq(t, `{"status":"SUCCEEDED"}`, result) assert.Equal(t, C.int(1), C.gil_locked_calls) assert.Equal(t, C.int(1), C.gil_unlocked_calls) - assert.Equal(t, C.int(1), C.run_postgres_remote_query_calls) + assert.Equal(t, C.int(1), C.run_remote_query_calls) assert.Equal(t, C.int(1), C.rtloader_free_calls) - assert.Equal(t, check.instance, C.run_postgres_remote_query_instance) - assert.JSONEq(t, `{"integration":"postgres","query":"SELECT 1 AS value"}`, C.GoString(C.run_postgres_remote_query_request_json)) + assert.Equal(t, check.instance, C.run_remote_query_instance) + assert.Equal(t, "postgres", C.GoString(C.run_remote_query_integration)) + assert.JSONEq(t, `{"integration":"postgres","query":"SELECT 1 AS value"}`, C.GoString(C.run_remote_query_request_json)) } -func testRunRemoteQueryJSONUnsupportedIntegration(t *testing.T) { +func testRunRemoteQueryJSONNormalizesIntegration(t *testing.T) { mockRtloader(t) check, err := NewPythonFakeCheck(aggregator.NewNoOpSenderManager()) @@ -713,16 +717,17 @@ func testRunRemoteQueryJSONUnsupportedIntegration(t *testing.T) { check.instance = newMockPyObjectPtr() C.reset_check_mock() + C.run_remote_query_return = C.CString(`{"status":"SUCCEEDED"}`) - result, err := check.RunRemoteQueryJSON("mysql", `{"integration":"mysql","query":"SELECT 1"}`) + result, err := check.RunRemoteQueryJSON(" MySQL ", `{"integration":"mysql","query":"SELECT 1"}`) - assert.Empty(t, result) - require.Error(t, err) - assert.EqualError(t, err, "unsupported integration") - assert.Equal(t, C.int(0), C.run_postgres_remote_query_calls) + require.NoError(t, err) + assert.JSONEq(t, `{"status":"SUCCEEDED"}`, result) + assert.Equal(t, C.int(1), C.run_remote_query_calls) + assert.Equal(t, "mysql", C.GoString(C.run_remote_query_integration)) } -func testRunRemoteQueryJSONPostgresError(t *testing.T) { +func testRunRemoteQueryJSONError(t *testing.T) { mockRtloader(t) check, err := NewPythonFakeCheck(aggregator.NewNoOpSenderManager()) @@ -730,7 +735,7 @@ func testRunRemoteQueryJSONPostgresError(t *testing.T) { check.instance = newMockPyObjectPtr() C.reset_check_mock() - C.run_postgres_remote_query_return = nil + C.run_remote_query_return = nil C.has_error_return = 1 C.get_error_return = C.CString("rtloader helper failed") @@ -741,7 +746,7 @@ func testRunRemoteQueryJSONPostgresError(t *testing.T) { assert.EqualError(t, err, "rtloader helper failed") assert.Equal(t, C.int(1), C.gil_locked_calls) assert.Equal(t, C.int(1), C.gil_unlocked_calls) - assert.Equal(t, C.int(1), C.run_postgres_remote_query_calls) + assert.Equal(t, C.int(1), C.run_remote_query_calls) assert.Equal(t, C.int(0), C.rtloader_free_calls) } @@ -756,7 +761,7 @@ func testRunRemoteQueryJSONWithRuntimeNotInitializedError(t *testing.T) { _, err = check.RunRemoteQueryJSON("postgres", `{"query":"SELECT 1 AS value"}`) assert.ErrorIs(t, err, ErrNotInitialized) - assert.Equal(t, C.int(0), C.run_postgres_remote_query_calls) + assert.Equal(t, C.int(0), C.run_remote_query_calls) } func testRunRemoteQueryJSONAfterCancel(t *testing.T) { @@ -771,7 +776,7 @@ func testRunRemoteQueryJSONAfterCancel(t *testing.T) { _, err = check.RunRemoteQueryJSON("postgres", `{"query":"SELECT 1 AS value"}`) assert.EqualError(t, err, "check fake_check is already cancelled") - assert.Equal(t, C.int(0), C.run_postgres_remote_query_calls) + assert.Equal(t, C.int(0), C.run_remote_query_calls) } func testRunAfterCancel(t *testing.T) { diff --git a/rtloader/include/datadog_agent_rtloader.h b/rtloader/include/datadog_agent_rtloader.h index 6f93a576d2df..2efbf61f641e 100644 --- a/rtloader/include/datadog_agent_rtloader.h +++ b/rtloader/include/datadog_agent_rtloader.h @@ -190,16 +190,17 @@ DATADOG_AGENT_RTLOADER_API int get_check_deprecated(rtloader_t *rtloader, rtload */ DATADOG_AGENT_RTLOADER_API char *run_check(rtloader_t *, rtloader_pyobject_t *check); -/*! \fn char *run_postgres_remote_query(rtloader_t *, rtloader_pyobject_t *check, const char *request_json) - \brief Runs the fixed Postgres remote query helper for a check instance. +/*! \fn char *run_remote_query(rtloader_t *, rtloader_pyobject_t *check, const char *integration, const char *request_json) + \brief Runs the integration remote query helper for a check instance. \param rtloader_t A rtloader_t * pointer to the RtLoader instance. \param check A rtloader_pyobject_t * pointer to the check instance we wish to use. + \param integration The integration helper module name. \param request_json A credential-free JSON request string. \return A C-string with the JSON result. \sa rtloader_pyobject_t, rtloader_t */ -DATADOG_AGENT_RTLOADER_API char *run_postgres_remote_query(rtloader_t *, rtloader_pyobject_t *check, - const char *request_json); +DATADOG_AGENT_RTLOADER_API char *run_remote_query(rtloader_t *, rtloader_pyobject_t *check, const char *integration, + const char *request_json); /*! \fn char *cancel_check(rtloader_t *, rtloader_pyobject_t *check) \brief Cancels a check instance. This allow check to be notified when diff --git a/rtloader/include/rtloader.h b/rtloader/include/rtloader.h index b82f12e19b50..37c10cca3e9d 100644 --- a/rtloader/include/rtloader.h +++ b/rtloader/include/rtloader.h @@ -121,13 +121,14 @@ class RtLoader */ virtual char *runCheck(RtLoaderPyObject *check) = 0; - //! Pure virtual runPostgresRemoteQuery member. + //! Pure virtual runRemoteQuery member. /*! \param check The python object pointer to the check we wish to use. + \param integration The integration helper module name. \param request_json A credential-free JSON request string. \return A C-string with the JSON result. */ - virtual char *runPostgresRemoteQuery(RtLoaderPyObject *check, const char *request_json) = 0; + virtual char *runRemoteQuery(RtLoaderPyObject *check, const char *integration, const char *request_json) = 0; //! Pure virtual cancelCheck member. /*! diff --git a/rtloader/rtloader/api.cpp b/rtloader/rtloader/api.cpp index 114b501d4556..6aa8fd0f909b 100644 --- a/rtloader/rtloader/api.cpp +++ b/rtloader/rtloader/api.cpp @@ -278,9 +278,9 @@ char *run_check(rtloader_t *rtloader, rtloader_pyobject_t *check) return AS_TYPE(RtLoader, rtloader)->runCheck(AS_TYPE(RtLoaderPyObject, check)); } -char *run_postgres_remote_query(rtloader_t *rtloader, rtloader_pyobject_t *check, const char *request_json) +char *run_remote_query(rtloader_t *rtloader, rtloader_pyobject_t *check, const char *integration, const char *request_json) { - return AS_TYPE(RtLoader, rtloader)->runPostgresRemoteQuery(AS_TYPE(RtLoaderPyObject, check), request_json); + return AS_TYPE(RtLoader, rtloader)->runRemoteQuery(AS_TYPE(RtLoaderPyObject, check), integration, request_json); } void cancel_check(rtloader_t *rtloader, rtloader_pyobject_t *check) diff --git a/rtloader/three/three.cpp b/rtloader/three/three.cpp index 96fad7a513dd..e6f891560688 100644 --- a/rtloader/three/three.cpp +++ b/rtloader/three/three.cpp @@ -19,6 +19,7 @@ #include "util.h" #include +#include #include #include @@ -497,47 +498,86 @@ char *Three::runCheck(RtLoaderPyObject *check) return ret; } -char *Three::runPostgresRemoteQuery(RtLoaderPyObject *check, const char *request_json) +namespace { +std::string normalizeRemoteQueryIntegration(const char *integration) +{ + if (integration == NULL) { + return ""; + } + + std::string normalized(integration); + normalized.erase(normalized.begin(), std::find_if(normalized.begin(), normalized.end(), [](unsigned char ch) { + return !std::isspace(ch); + })); + normalized.erase(std::find_if(normalized.rbegin(), normalized.rend(), [](unsigned char ch) { + return !std::isspace(ch); + }).base(), + normalized.end()); + std::transform(normalized.begin(), normalized.end(), normalized.begin(), [](unsigned char ch) { + return static_cast(std::tolower(ch)); + }); + return normalized; +} + +bool isValidRemoteQueryIntegration(const std::string &integration) +{ + if (integration.empty()) { + return false; + } + return std::all_of(integration.begin(), integration.end(), [](unsigned char ch) { + return std::islower(ch) || std::isdigit(ch) || ch == '_'; + }); +} +} // namespace + +char *Three::runRemoteQuery(RtLoaderPyObject *check, const char *integration, const char *request_json) { if (check == NULL || request_json == NULL) { return NULL; } + std::string normalized_integration = normalizeRemoteQueryIntegration(integration); + if (!isValidRemoteQueryIntegration(normalized_integration)) { + setError("invalid remote query integration name"); + return NULL; + } + PyObject *py_check = reinterpret_cast(check); char *ret = NULL; PyObject *remote_query_module = NULL; PyObject *execute_func = NULL; PyObject *py_request_json = NULL; PyObject *result = NULL; + std::string module_name = "datadog_checks." + normalized_integration + ".remote_query"; - remote_query_module = PyImport_ImportModule("datadog_checks.postgres.remote_query"); + remote_query_module = PyImport_ImportModule(module_name.c_str()); if (remote_query_module == NULL) { - setError("error importing Postgres remote query helper: " + _fetchPythonError()); + setError("error importing remote query helper: " + _fetchPythonError()); goto done; } execute_func = PyObject_GetAttrString(remote_query_module, "execute_agent_rpc_json"); if (execute_func == NULL || !PyCallable_Check(execute_func)) { - setError("error loading Postgres remote query helper: " + _fetchPythonError()); + setError("error loading remote query helper: " + _fetchPythonError()); goto done; } py_request_json = PyUnicode_FromString(request_json); if (py_request_json == NULL) { - setError("error converting Postgres remote query request to Python string: " + _fetchPythonError()); + setError("error converting remote query request to Python string: " + _fetchPythonError()); goto done; } result = PyObject_CallFunctionObjArgs(execute_func, py_request_json, py_check, NULL); if (result == NULL || !PyUnicode_Check(result)) { - setError("error invoking Postgres remote query helper: " + _fetchPythonError()); + setError("error invoking remote query helper: " + _fetchPythonError()); goto done; } ret = as_string(result); if (ret == NULL) { // as_string clears the error, so we can't fetch it here - setError("error converting Postgres remote query helper result to string"); + setError("error converting remote query helper result to string"); goto done; } diff --git a/rtloader/three/three.h b/rtloader/three/three.h index 7ad49356c204..f946886ef16a 100644 --- a/rtloader/three/three.h +++ b/rtloader/three/three.h @@ -65,7 +65,7 @@ class Three : public RtLoader const char *provider_str, RtLoaderPyObject *&check); char *runCheck(RtLoaderPyObject *check); - char *runPostgresRemoteQuery(RtLoaderPyObject *check, const char *request_json); + char *runRemoteQuery(RtLoaderPyObject *check, const char *integration, const char *request_json); void cancelCheck(RtLoaderPyObject *check); char **getCheckWarnings(RtLoaderPyObject *check); char *getCheckDiagnoses(RtLoaderPyObject *check); From 972934ac44c1f93c474e5114df8a63c89bfe4a85 Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 11:00:32 +0000 Subject: [PATCH 11/33] Fix rtloader test path precedence --- rtloader/test/rtloader/rtloader.go | 5 ++--- rtloader/three/three.cpp | 9 ++++++--- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/rtloader/test/rtloader/rtloader.go b/rtloader/test/rtloader/rtloader.go index d0d3a3b82595..b46e38066d5d 100644 --- a/rtloader/test/rtloader/rtloader.go +++ b/rtloader/test/rtloader/rtloader.go @@ -143,7 +143,9 @@ func getFakeCheck() (string, error) { var version *C.char runtime.LockOSThread() + defer runtime.UnlockOSThread() state := C.ensure_gil(rtloader) + defer C.release_gil(rtloader, state) // class classStr := helpers.TrackedCString("fake_check") @@ -181,9 +183,6 @@ func getFakeCheck() (string, error) { return "", errors.New(C.GoString(C.get_error(rtloader))) } - C.release_gil(rtloader, state) - runtime.UnlockOSThread() - return C.GoString(version), fetchError() } diff --git a/rtloader/three/three.cpp b/rtloader/three/three.cpp index e6f891560688..90667d7b4e28 100644 --- a/rtloader/three/three.cpp +++ b/rtloader/three/three.cpp @@ -141,16 +141,19 @@ bool Three::init() setError("could not access sys.path"); goto done; } - for (PyPaths::iterator pit = _pythonPaths.begin(); pit != _pythonPaths.end(); ++pit) { + // Explicit rtloader paths must take precedence over ambient site-packages. + // This keeps tests on their stubs even when a developer has datadog_checks installed locally. + Py_ssize_t pythonPathIndex = 0; + for (PyPaths::iterator pit = _pythonPaths.begin(); pit != _pythonPaths.end(); ++pit, ++pythonPathIndex) { PyObject *p = PyUnicode_FromString(pit->c_str()); if (p == NULL) { setError("could not set pythonPath: " + _fetchPythonError()); goto done; } - int retval = PyList_Append(path, p); + int retval = PyList_Insert(path, pythonPathIndex, p); Py_XDECREF(p); if (retval == -1) { - setError("could not append path to pythonPath: " + _fetchPythonError()); + setError("could not add path to pythonPath: " + _fetchPythonError()); goto done; } } From aadc440459a79a531a2d17346d18ade7890ca9ba Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 11:40:16 +0000 Subject: [PATCH 12/33] Add local Remote Queries PAR action proof --- pkg/privateactionrunner/bundles/registry.go | 2 + .../bundles/registry_kubeapiserver.go | 2 + .../bundles/remotequeries/entrypoint.go | 32 +++++ .../bundles/remotequeries/execute.go | 131 ++++++++++++++++++ .../bundles/remotequeries/execute_test.go | 131 ++++++++++++++++++ .../bundles/remotequeries/ipc_client.go | 59 ++++++++ 6 files changed, 357 insertions(+) create mode 100644 pkg/privateactionrunner/bundles/remotequeries/entrypoint.go create mode 100644 pkg/privateactionrunner/bundles/remotequeries/execute.go create mode 100644 pkg/privateactionrunner/bundles/remotequeries/execute_test.go create mode 100644 pkg/privateactionrunner/bundles/remotequeries/ipc_client.go diff --git a/pkg/privateactionrunner/bundles/registry.go b/pkg/privateactionrunner/bundles/registry.go index 1e3307de1c09..5a20b6598ebb 100644 --- a/pkg/privateactionrunner/bundles/registry.go +++ b/pkg/privateactionrunner/bundles/registry.go @@ -43,6 +43,7 @@ import ( com_datadoghq_remoteaction "github.com/DataDog/datadog-agent/pkg/privateactionrunner/bundles/remoteaction" com_datadoghq_remoteaction_networks "github.com/DataDog/datadog-agent/pkg/privateactionrunner/bundles/remoteaction/networks" com_datadoghq_remoteaction_rshell "github.com/DataDog/datadog-agent/pkg/privateactionrunner/bundles/remoteaction/rshell" + com_datadoghq_remotequeries "github.com/DataDog/datadog-agent/pkg/privateactionrunner/bundles/remotequeries" com_datadoghq_script "github.com/DataDog/datadog-agent/pkg/privateactionrunner/bundles/script" com_datadoghq_temporal "github.com/DataDog/datadog-agent/pkg/privateactionrunner/bundles/temporal" "github.com/DataDog/datadog-agent/pkg/privateactionrunner/types" @@ -87,6 +88,7 @@ func NewRegistry(configuration *config.Config, traceroute traceroute.Component, "com.datadoghq.remoteaction": com_datadoghq_remoteaction.NewRemoteAction(), "com.datadoghq.remoteaction.networks": com_datadoghq_remoteaction_networks.NewNetworks(traceroute, eventPlatform), "com.datadoghq.remoteaction.rshell": com_datadoghq_remoteaction_rshell.NewRshellBundle(configuration), + "com.datadoghq.remotequeries": com_datadoghq_remotequeries.NewRemoteQueriesBundle(), "com.datadoghq.script": com_datadoghq_script.NewScript(), "com.datadoghq.temporal": com_datadoghq_temporal.NewTemporal(), }, diff --git a/pkg/privateactionrunner/bundles/registry_kubeapiserver.go b/pkg/privateactionrunner/bundles/registry_kubeapiserver.go index c7d6a5127d95..be3ad99af70a 100644 --- a/pkg/privateactionrunner/bundles/registry_kubeapiserver.go +++ b/pkg/privateactionrunner/bundles/registry_kubeapiserver.go @@ -43,6 +43,7 @@ import ( com_datadoghq_mongodb "github.com/DataDog/datadog-agent/pkg/privateactionrunner/bundles/mongodb" com_datadoghq_remoteaction "github.com/DataDog/datadog-agent/pkg/privateactionrunner/bundles/remoteaction" com_datadoghq_remoteaction_networks "github.com/DataDog/datadog-agent/pkg/privateactionrunner/bundles/remoteaction/networks" + com_datadoghq_remotequeries "github.com/DataDog/datadog-agent/pkg/privateactionrunner/bundles/remotequeries" com_datadoghq_script "github.com/DataDog/datadog-agent/pkg/privateactionrunner/bundles/script" com_datadoghq_temporal "github.com/DataDog/datadog-agent/pkg/privateactionrunner/bundles/temporal" "github.com/DataDog/datadog-agent/pkg/privateactionrunner/types" @@ -86,6 +87,7 @@ func NewRegistry(configuration *config.Config, traceroute traceroute.Component, "com.datadoghq.kubernetes.customresources": com_datadoghq_kubernetes_customresources.NewKubernetesCustomResources(), "com.datadoghq.kubernetes.discovery": com_datadoghq_kubernetes_discovery.NewKubernetesDiscovery(), "com.datadoghq.mongodb": com_datadoghq_mongodb.NewMongoDB(), + "com.datadoghq.remotequeries": com_datadoghq_remotequeries.NewRemoteQueriesBundle(), "com.datadoghq.script": com_datadoghq_script.NewScript(), "com.datadoghq.temporal": com_datadoghq_temporal.NewTemporal(), }, diff --git a/pkg/privateactionrunner/bundles/remotequeries/entrypoint.go b/pkg/privateactionrunner/bundles/remotequeries/entrypoint.go new file mode 100644 index 000000000000..39b8fbffdd4d --- /dev/null +++ b/pkg/privateactionrunner/bundles/remotequeries/entrypoint.go @@ -0,0 +1,32 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2026-present Datadog, Inc. + +package com_datadoghq_remotequeries + +import "github.com/DataDog/datadog-agent/pkg/privateactionrunner/types" + +const ( + // BundleID is the local-only Remote Queries PAR bundle FQN. + BundleID = "com.datadoghq.remotequeries" + + // ExecuteActionName is the action part of com.datadoghq.remotequeries.execute. + ExecuteActionName = "execute" +) + +type RemoteQueriesBundle struct { + actions map[string]types.Action +} + +func NewRemoteQueriesBundle() *RemoteQueriesBundle { + return &RemoteQueriesBundle{ + actions: map[string]types.Action{ + ExecuteActionName: NewExecuteAction(NewDefaultBridgeClient), + }, + } +} + +func (b *RemoteQueriesBundle) GetAction(actionName string) types.Action { + return b.actions[actionName] +} diff --git a/pkg/privateactionrunner/bundles/remotequeries/execute.go b/pkg/privateactionrunner/bundles/remotequeries/execute.go new file mode 100644 index 000000000000..9f4594969d7e --- /dev/null +++ b/pkg/privateactionrunner/bundles/remotequeries/execute.go @@ -0,0 +1,131 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2026-present Datadog, Inc. + +package com_datadoghq_remotequeries + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + + ipc "github.com/DataDog/datadog-agent/comp/core/ipc/def" + ipchttp "github.com/DataDog/datadog-agent/comp/core/ipc/httphelpers" + "github.com/DataDog/datadog-agent/pkg/privateactionrunner/libs/privateconnection" + "github.com/DataDog/datadog-agent/pkg/privateactionrunner/types" + "github.com/DataDog/datadog-agent/pkg/privateactionrunner/util" +) + +const ( + // AgentRemoteQueryExecuteEndpointPath is the POC-only Agent command API path this PAR action calls. + // It is intentionally local-only and is not a production API/IPC commitment. + AgentRemoteQueryExecuteEndpointPath = "/agent/remote-queries/execute" +) + +// BridgeClient is the narrow Agent IPC HTTP client surface required by this action. +type BridgeClient interface { + Post(url string, contentType string, body io.Reader, opts ...ipc.RequestOption) (resp []byte, err error) +} + +// BridgeClientFactory returns an IPC client and fully-qualified local Agent endpoint URL. +type BridgeClientFactory func() (BridgeClient, string, error) + +type ExecuteAction struct { + newBridgeClient BridgeClientFactory +} + +func NewExecuteAction(newBridgeClient BridgeClientFactory) *ExecuteAction { + return &ExecuteAction{newBridgeClient: newBridgeClient} +} + +type ExecuteInputs struct { + Integration string `json:"integration"` + Target TargetInputs `json:"target"` + Query string `json:"query"` + Limits *LimitsInputs `json:"limits,omitempty"` +} + +type TargetInputs struct { + Host string `json:"host"` + Port int `json:"port"` + DBName string `json:"dbname"` +} + +type LimitsInputs struct { + MaxRows int `json:"maxRows"` + MaxBytes int `json:"maxBytes"` + TimeoutMs int `json:"timeoutMs"` +} + +func (a *ExecuteAction) Run( + ctx context.Context, + task *types.Task, + _ *privateconnection.PrivateCredentials, +) (interface{}, error) { + inputs, err := types.ExtractInputs[ExecuteInputs](task) + if err != nil { + return nil, util.DefaultActionErrorWithDisplayError( + fmt.Errorf("invalid remote query action inputs"), + "invalid remote query action inputs", + ) + } + + payload, err := json.Marshal(inputs) + if err != nil { + return nil, util.DefaultActionErrorWithDisplayError( + fmt.Errorf("marshal remote query action inputs"), + "invalid remote query action inputs", + ) + } + + if a == nil || a.newBridgeClient == nil { + return nil, util.DefaultActionError(fmt.Errorf("remote query action requires an Agent IPC client")) + } + client, endpointURL, err := a.newBridgeClient() + if err != nil { + return nil, util.DefaultActionErrorWithDisplayError(err, "remote query action could not create an Agent IPC client") + } + if client == nil || endpointURL == "" { + return nil, util.DefaultActionError(fmt.Errorf("remote query action requires an Agent IPC client and endpoint URL")) + } + + body, postErr := client.Post(endpointURL, "application/json", bytes.NewReader(payload), ipchttp.WithContext(ctx)) + output, decodeErr := decodeBridgeResponse(body) + if decodeErr == nil { + // IPC HTTPClient returns both the response body and an error for HTTP >= 400. + // The bridge body is already sanitized, so preserve its status/error payload as the action output. + return output, nil + } + if postErr != nil { + if len(body) > 0 { + return nil, util.DefaultActionErrorWithDisplayError( + fmt.Errorf("remote query IPC request failed with undecodable response"), + "remote query IPC request failed with undecodable response", + ) + } + return nil, util.DefaultActionErrorWithDisplayError(postErr, "remote query IPC request failed") + } + return nil, util.DefaultActionErrorWithDisplayError(decodeErr, "remote query IPC response was invalid") +} + +func decodeBridgeResponse(body []byte) (map[string]interface{}, error) { + if len(body) == 0 { + return nil, fmt.Errorf("empty remote query response") + } + + decoder := json.NewDecoder(bytes.NewReader(body)) + decoder.UseNumber() + + var output map[string]interface{} + if err := decoder.Decode(&output); err != nil { + return nil, fmt.Errorf("decode remote query response: %w", err) + } + status, ok := output["status"].(string) + if !ok || status == "" { + return nil, fmt.Errorf("remote query response missing status") + } + return output, nil +} diff --git a/pkg/privateactionrunner/bundles/remotequeries/execute_test.go b/pkg/privateactionrunner/bundles/remotequeries/execute_test.go new file mode 100644 index 000000000000..ada8916c5b49 --- /dev/null +++ b/pkg/privateactionrunner/bundles/remotequeries/execute_test.go @@ -0,0 +1,131 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2026-present Datadog, Inc. + +package com_datadoghq_remotequeries + +import ( + "context" + "encoding/json" + "errors" + "io" + "testing" + + ipc "github.com/DataDog/datadog-agent/comp/core/ipc/def" + "github.com/DataDog/datadog-agent/pkg/privateactionrunner/libs/privateconnection" + "github.com/DataDog/datadog-agent/pkg/privateactionrunner/types" + "github.com/DataDog/datadog-agent/pkg/privateactionrunner/util" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestExecuteActionUsesCredentialFreeIPCPostShape(t *testing.T) { + client := &captureBridgeClient{response: []byte(`{"status":"SUCCEEDED","rows":[{"value":1}]}`)} + action := NewExecuteAction(func() (BridgeClient, string, error) { + return client, "https://localhost:5001" + AgentRemoteQueryExecuteEndpointPath, nil + }) + + output, err := action.Run(context.Background(), taskWithInputs(map[string]interface{}{ + "integration": "postgres", + "target": map[string]interface{}{ + "host": "localhost", + "port": 5432, + "dbname": "postgres", + }, + "query": "SELECT 1 AS value", + "limits": map[string]interface{}{ + "maxRows": 1, + "maxBytes": 1024, + "timeoutMs": 1000, + }, + }), &privateconnection.PrivateCredentials{Tokens: []privateconnection.PrivateCredentialsToken{{Name: "password", Value: "secret-value"}}}) + + require.NoError(t, err) + assert.Equal(t, "https://localhost:5001"+AgentRemoteQueryExecuteEndpointPath, client.url) + assert.Equal(t, "application/json", client.contentType) + assert.JSONEq(t, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":1,"maxBytes":1024,"timeoutMs":1000}}`, client.body) + assert.NotContains(t, client.body, "secret-value") + assert.Equal(t, map[string]interface{}{ + "status": "SUCCEEDED", + "rows": []interface{}{ + map[string]interface{}{"value": json.Number("1")}, + }, + }, output) +} + +func TestExecuteActionPreservesSanitizedBridgeErrorBody(t *testing.T) { + client := &captureBridgeClient{ + response: []byte(`{"status":"target_not_found","error":{"code":"target_not_found","message":"no matching integration check found"}}`), + err: errors.New("status 404"), + } + action := NewExecuteAction(func() (BridgeClient, string, error) { + return client, "https://localhost:5001" + AgentRemoteQueryExecuteEndpointPath, nil + }) + + output, err := action.Run(context.Background(), taskWithInputs(map[string]interface{}{ + "integration": "postgres", + "target": map[string]interface{}{"host": "localhost", "port": 5432, "dbname": "secret-db"}, + "query": "SELECT 1 AS value", + }), nil) + + require.NoError(t, err) + assert.Equal(t, map[string]interface{}{ + "status": "target_not_found", + "error": map[string]interface{}{ + "code": "target_not_found", + "message": "no matching integration check found", + }, + }, output) +} + +func TestExecuteActionSanitizesInputExtractionErrors(t *testing.T) { + action := NewExecuteAction(func() (BridgeClient, string, error) { + require.Fail(t, "bridge client should not be created for invalid inputs") + return nil, "", nil + }) + + _, err := action.Run(context.Background(), taskWithInputs(map[string]interface{}{ + "integration": "postgres", + "target": map[string]interface{}{"host": "localhost", "port": 5432, "dbname": "secret-db"}, + "query": "SELECT secret FROM private_table", + "bad": make(chan struct{}), + }), nil) + + require.Error(t, err) + var parErr util.PARError + require.ErrorAs(t, err, &parErr) + assert.Equal(t, "invalid remote query action inputs", parErr.Message) + assert.Equal(t, "invalid remote query action inputs", parErr.ExternalMessage) + assert.NotContains(t, err.Error(), "secret-db") + assert.NotContains(t, err.Error(), "SELECT secret") +} + +func taskWithInputs(inputs map[string]interface{}) *types.Task { + task := &types.Task{} + task.Data.Attributes = &types.Attributes{ + BundleID: BundleID, + Name: ExecuteActionName, + Inputs: inputs, + } + return task +} + +type captureBridgeClient struct { + url string + contentType string + body string + response []byte + err error +} + +func (c *captureBridgeClient) Post(url string, contentType string, body io.Reader, _ ...ipc.RequestOption) ([]byte, error) { + c.url = url + c.contentType = contentType + payload, err := io.ReadAll(body) + if err != nil { + return nil, err + } + c.body = string(payload) + return c.response, c.err +} diff --git a/pkg/privateactionrunner/bundles/remotequeries/ipc_client.go b/pkg/privateactionrunner/bundles/remotequeries/ipc_client.go new file mode 100644 index 000000000000..7646301c1b23 --- /dev/null +++ b/pkg/privateactionrunner/bundles/remotequeries/ipc_client.go @@ -0,0 +1,59 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2026-present Datadog, Inc. + +package com_datadoghq_remotequeries + +import ( + "fmt" + "net" + "net/url" + "strconv" + + ipchttp "github.com/DataDog/datadog-agent/comp/core/ipc/httphelpers" + pkgtoken "github.com/DataDog/datadog-agent/pkg/api/security" + "github.com/DataDog/datadog-agent/pkg/api/security/cert" + pkgconfigmodel "github.com/DataDog/datadog-agent/pkg/config/model" + pkgconfigsetup "github.com/DataDog/datadog-agent/pkg/config/setup" + "github.com/DataDog/datadog-agent/pkg/util/system" +) + +// NewDefaultBridgeClient creates the local Agent IPC client used by the registered PAR action. +func NewDefaultBridgeClient() (BridgeClient, string, error) { + cfg := pkgconfigsetup.Datadog() + + token, err := pkgtoken.FetchAuthToken(cfg) + if err != nil { + return nil, "", fmt.Errorf("fetch Agent IPC auth token: %w", err) + } + clientTLSConfig, _, _, err := cert.FetchIPCCert(cfg) + if err != nil { + return nil, "", fmt.Errorf("fetch Agent IPC certificate: %w", err) + } + + endpointURL, err := agentIPCURL(cfg, AgentRemoteQueryExecuteEndpointPath) + if err != nil { + return nil, "", err + } + return ipchttp.NewClient(token, clientTLSConfig, cfg), endpointURL, nil +} + +func agentIPCURL(cfg pkgconfigmodel.Reader, endpointPath string) (string, error) { + cmdHostKey := "cmd_host" + if cfg.IsConfigured("ipc_address") { + cmdHostKey = "ipc_address" + } + + ipcHost, err := system.IsLocalAddress(cfg.GetString(cmdHostKey)) + if err != nil { + return "", fmt.Errorf("%s: %w", cmdHostKey, err) + } + + endpointURL := url.URL{ + Scheme: "https", + Host: net.JoinHostPort(ipcHost, strconv.Itoa(cfg.GetInt("cmd_port"))), + Path: endpointPath, + } + return endpointURL.String(), nil +} From 239e718f7cd077054dd962ea453ccb2ca41992f0 Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 12:19:39 +0000 Subject: [PATCH 13/33] Add local Remote Queries live PAR proof --- .../bundles/remotequeries/entrypoint.go | 15 +- .../remotequeries/live_par_loop_test.go | 175 ++++++++++++++++++ 2 files changed, 189 insertions(+), 1 deletion(-) create mode 100644 pkg/privateactionrunner/bundles/remotequeries/live_par_loop_test.go diff --git a/pkg/privateactionrunner/bundles/remotequeries/entrypoint.go b/pkg/privateactionrunner/bundles/remotequeries/entrypoint.go index 39b8fbffdd4d..f7b61b70a180 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/entrypoint.go +++ b/pkg/privateactionrunner/bundles/remotequeries/entrypoint.go @@ -19,14 +19,27 @@ type RemoteQueriesBundle struct { actions map[string]types.Action } +var defaultBridgeClientFactory BridgeClientFactory = NewDefaultBridgeClient + func NewRemoteQueriesBundle() *RemoteQueriesBundle { return &RemoteQueriesBundle{ actions: map[string]types.Action{ - ExecuteActionName: NewExecuteAction(NewDefaultBridgeClient), + ExecuteActionName: NewExecuteAction(defaultBridgeClientFactory), }, } } +// SetBridgeClientFactoryForTest overrides the bridge client factory used by newly-created +// Remote Queries bundles. It is intended for tests that exercise the registered PAR +// registry/runner path without depending on a live Agent IPC server. +func SetBridgeClientFactoryForTest(factory BridgeClientFactory) func() { + previousFactory := defaultBridgeClientFactory + defaultBridgeClientFactory = factory + return func() { + defaultBridgeClientFactory = previousFactory + } +} + func (b *RemoteQueriesBundle) GetAction(actionName string) types.Action { return b.actions[actionName] } diff --git a/pkg/privateactionrunner/bundles/remotequeries/live_par_loop_test.go b/pkg/privateactionrunner/bundles/remotequeries/live_par_loop_test.go new file mode 100644 index 000000000000..957ecd76cb31 --- /dev/null +++ b/pkg/privateactionrunner/bundles/remotequeries/live_par_loop_test.go @@ -0,0 +1,175 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2026-present Datadog, Inc. + +//go:build !windows + +package com_datadoghq_remotequeries_test + +import ( + "bytes" + "context" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "testing" + "time" + + ipc "github.com/DataDog/datadog-agent/comp/core/ipc/def" + parconfig "github.com/DataDog/datadog-agent/pkg/privateactionrunner/adapters/config" + app "github.com/DataDog/datadog-agent/pkg/privateactionrunner/adapters/constants" + com_datadoghq_remotequeries "github.com/DataDog/datadog-agent/pkg/privateactionrunner/bundles/remotequeries" + "github.com/DataDog/datadog-agent/pkg/privateactionrunner/observability" + "github.com/DataDog/datadog-agent/pkg/privateactionrunner/opms" + "github.com/DataDog/datadog-agent/pkg/privateactionrunner/runners" + taskverifier "github.com/DataDog/datadog-agent/pkg/privateactionrunner/task-verifier" + fakeintakeclient "github.com/DataDog/datadog-agent/test/fakeintake/client" + fakeintakeserver "github.com/DataDog/datadog-agent/test/fakeintake/server" + "github.com/DataDog/datadog-go/v5/statsd" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "k8s.io/apimachinery/pkg/util/sets" +) + +func TestRemoteQueriesActionRunsThroughLivePARLoopAndFakeintake(t *testing.T) { + t.Setenv(app.InternalSkipTaskVerificationEnvVar, "true") + + bridgeRequests := make(chan []byte, 1) + bridgeServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, http.MethodPost, r.Method) + require.Equal(t, com_datadoghq_remotequeries.AgentRemoteQueryExecuteEndpointPath, r.URL.Path) + require.Contains(t, r.Header.Get("Content-Type"), "application/json") + + body, err := io.ReadAll(r.Body) + require.NoError(t, err) + bridgeRequests <- body + + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"status":"SUCCEEDED","rows":[{"value":1}]}`)) + })) + defer bridgeServer.Close() + + restoreFactory := com_datadoghq_remotequeries.SetBridgeClientFactoryForTest(func() (com_datadoghq_remotequeries.BridgeClient, string, error) { + return httpBridgeClient{}, bridgeServer.URL + com_datadoghq_remotequeries.AgentRemoteQueryExecuteEndpointPath, nil + }) + defer restoreFactory() + + fakeintake, _ := fakeintakeserver.InitialiseForTests(t) + defer func() { require.NoError(t, fakeintake.Stop()) }() + fakeintakeClient := fakeintakeclient.NewClient(fakeintake.URL()) + require.NoError(t, fakeintakeClient.FlushPAR()) + + cfg := newLivePARTestConfig(t, fakeintake.URL()) + keysManager := taskverifier.NewKeyManager(nil) + verifier := taskverifier.NewTaskVerifier(keysManager, cfg) + workflowRunner, err := runners.NewWorkflowRunner(cfg, keysManager, verifier, opms.NewClient(cfg), nil, nil) + require.NoError(t, err) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + require.NoError(t, workflowRunner.Start(ctx)) + defer func() { + stopCtx, stopCancel := context.WithTimeout(context.Background(), 2*time.Second) + defer stopCancel() + require.NoError(t, workflowRunner.Stop(stopCtx)) + }() + + taskID := "remotequeries-live-par-local-proof" + require.NoError(t, fakeintakeClient.EnqueuePARTask(taskID, com_datadoghq_remotequeries.BundleID+"."+com_datadoghq_remotequeries.ExecuteActionName, map[string]interface{}{ + "integration": "postgres", + "target": map[string]interface{}{ + "host": "localhost", + "port": 5432, + "dbname": "postgres", + }, + "query": "SELECT 1 AS value", + "limits": map[string]interface{}{ + "maxRows": 1, + "maxBytes": 1024, + "timeoutMs": 1000, + }, + })) + + result, err := fakeintakeClient.GetPARTaskResult(taskID, 10*time.Second) + require.NoError(t, err) + require.True(t, result.Success) + require.Equal(t, taskID, result.TaskID) + assert.Equal(t, "SUCCEEDED", result.Outputs["status"]) + require.Contains(t, result.Outputs, "rows") + + var bridgeRequest map[string]interface{} + select { + case body := <-bridgeRequests: + require.NotContains(t, string(body), "password") + require.NotContains(t, string(body), "token") + require.NotContains(t, string(body), "secret") + require.NoError(t, json.Unmarshal(body, &bridgeRequest)) + case <-time.After(2 * time.Second): + require.FailNow(t, "remote query action did not call the local bridge") + } + assert.Equal(t, "postgres", bridgeRequest["integration"]) + assert.Equal(t, "SELECT 1 AS value", bridgeRequest["query"]) + assert.Equal(t, map[string]interface{}{"host": "localhost", "port": float64(5432), "dbname": "postgres"}, bridgeRequest["target"]) + + dequeueCalls, err := fakeintakeClient.GetPARDequeueCount() + require.NoError(t, err) + assert.GreaterOrEqual(t, dequeueCalls, 1) +} + +func newLivePARTestConfig(t *testing.T, fakeintakeURL string) *parconfig.Config { + t.Helper() + privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + require.NoError(t, err) + taskTimeoutSeconds := int32(5) + + return &parconfig.Config{ + ActionsAllowlist: map[string]sets.Set[string]{ + com_datadoghq_remotequeries.BundleID: sets.New[string](com_datadoghq_remotequeries.ExecuteActionName), + }, + DDHost: fakeintakeURL, + DDApiHost: "unused.local", + DatadogSite: "local", + OrgId: 123456, + PrivateKey: privateKey, + RunnerId: "remotequeries-live-par-local-proof-runner", + Urn: "urn:dd:apps:on-prem-runner:us1:123456:remotequeries-live-par-local-proof-runner", + LoopInterval: 10 * time.Millisecond, + MinBackoff: 10 * time.Millisecond, + MaxBackoff: 50 * time.Millisecond, + WaitBeforeRetry: 50 * time.Millisecond, + MaxAttempts: 3, + OpmsRequestTimeout: 1000, + RunnerPoolSize: 1, + HeartbeatInterval: time.Hour, + TaskTimeoutSeconds: &taskTimeoutSeconds, + MetricsClient: &statsd.NoOpClient{}, + Tags: []observability.Tag{}, + } +} + +type httpBridgeClient struct{} + +func (httpBridgeClient) Post(url string, contentType string, body io.Reader, _ ...ipc.RequestOption) ([]byte, error) { + payload, err := io.ReadAll(body) + if err != nil { + return nil, err + } + resp, err := http.Post(url, contentType, bytes.NewReader(payload)) + if err != nil { + return nil, err + } + defer resp.Body.Close() + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + if resp.StatusCode >= 400 { + return respBody, assert.AnError + } + return respBody, nil +} From 36decf0f899585974a52088c0dcf75709813fe08 Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 14:15:16 +0000 Subject: [PATCH 14/33] Add fused local Remote Queries PAR proof --- .../live_agent_ipc_par_loop_test.go | 157 +++++++++ .../fused-local-par-agent-postgres-proof.sh | 320 ++++++++++++++++++ 2 files changed, 477 insertions(+) create mode 100644 pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go create mode 100755 test/remotequeries/fused-local-par-agent-postgres-proof.sh diff --git a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go new file mode 100644 index 000000000000..0f0d709a1663 --- /dev/null +++ b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go @@ -0,0 +1,157 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2026-present Datadog, Inc. + +//go:build !windows + +package com_datadoghq_remotequeries_test + +import ( + "context" + "encoding/json" + "fmt" + "os" + "strconv" + "testing" + "time" + + pkgconfigsetup "github.com/DataDog/datadog-agent/pkg/config/setup" + parapp "github.com/DataDog/datadog-agent/pkg/privateactionrunner/adapters/constants" + com_datadoghq_remotequeries "github.com/DataDog/datadog-agent/pkg/privateactionrunner/bundles/remotequeries" + "github.com/DataDog/datadog-agent/pkg/privateactionrunner/opms" + "github.com/DataDog/datadog-agent/pkg/privateactionrunner/runners" + taskverifier "github.com/DataDog/datadog-agent/pkg/privateactionrunner/task-verifier" + fakeintakeclient "github.com/DataDog/datadog-agent/test/fakeintake/client" + fakeintakeserver "github.com/DataDog/datadog-agent/test/fakeintake/server" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const fusedLocalProofEnv = "RQ_FUSED_PROOF" + +func TestRemoteQueriesActionRunsThroughLivePARLoopWithRealAgentIPC(t *testing.T) { + if os.Getenv(fusedLocalProofEnv) != "1" { + t.Skipf("set %s=1 and start a local Agent with a loaded Postgres check to run the fused local proof", fusedLocalProofEnv) + } + + cmdPort := getenvRequired(t, "RQ_FUSED_AGENT_CMD_PORT") + authTokenFile := getenvRequired(t, "RQ_FUSED_AGENT_AUTH_TOKEN_FILE") + ipcCertFile := getenvRequired(t, "RQ_FUSED_AGENT_IPC_CERT_FILE") + cmdPortInt, err := strconv.Atoi(cmdPort) + require.NoError(t, err) + + // NewDefaultBridgeClient reads the process-wide Datadog config. Point it at the + // separate local Agent process started by the fused proof harness so the PAR + // action uses the real Agent IPC HTTP endpoint, not an httptest bridge. + cfg := pkgconfigsetup.Datadog() + cfg.SetWithoutSource("cmd_host", "127.0.0.1") + cfg.SetWithoutSource("cmd_port", cmdPortInt) + cfg.SetWithoutSource("auth_token_file_path", authTokenFile) + cfg.SetWithoutSource("ipc_cert_file_path", ipcCertFile) + + t.Setenv(parapp.InternalSkipTaskVerificationEnvVar, "true") + + fakeintake, _ := fakeintakeserver.InitialiseForTests(t) + defer func() { require.NoError(t, fakeintake.Stop()) }() + fakeintakeClient := fakeintakeclient.NewClient(fakeintake.URL()) + require.NoError(t, fakeintakeClient.FlushPAR()) + + cfgPAR := newLivePARTestConfig(t, fakeintake.URL()) + keysManager := taskverifier.NewKeyManager(nil) + verifier := taskverifier.NewTaskVerifier(keysManager, cfgPAR) + workflowRunner, err := runners.NewWorkflowRunner(cfgPAR, keysManager, verifier, opms.NewClient(cfgPAR), nil, nil) + require.NoError(t, err) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + require.NoError(t, workflowRunner.Start(ctx)) + defer func() { + stopCtx, stopCancel := context.WithTimeout(context.Background(), 2*time.Second) + defer stopCancel() + require.NoError(t, workflowRunner.Stop(stopCtx)) + }() + + taskID := fmt.Sprintf("remotequeries-fused-local-proof-%d", time.Now().UnixNano()) + inputs := map[string]interface{}{ + "integration": "postgres", + "target": map[string]interface{}{ + "host": "localhost", + "port": 5432, + "dbname": "postgres", + }, + "query": "SELECT 1 AS value", + "limits": map[string]interface{}{ + "maxRows": 1, + "maxBytes": 1024, + "timeoutMs": 1000, + }, + } + requestEvidence, err := json.Marshal(inputs) + require.NoError(t, err) + require.NotContains(t, string(requestEvidence), "password") + require.NotContains(t, string(requestEvidence), "token") + require.NotContains(t, string(requestEvidence), "secret") + + fqn := com_datadoghq_remotequeries.BundleID + "." + com_datadoghq_remotequeries.ExecuteActionName + t.Logf("fakeintake task enqueued: task_id=%s action_fqn=%s inputs=%s", taskID, fqn, requestEvidence) + t.Logf("real Agent IPC endpoint configured: https://127.0.0.1:%d%s", cmdPortInt, com_datadoghq_remotequeries.AgentRemoteQueryExecuteEndpointPath) + require.NoError(t, fakeintakeClient.EnqueuePARTask(taskID, fqn, inputs)) + + result, err := fakeintakeClient.GetPARTaskResult(taskID, 20*time.Second) + require.NoError(t, err) + require.True(t, result.Success) + require.Equal(t, taskID, result.TaskID) + assert.Equal(t, "SUCCEEDED", result.Outputs["status"]) + require.Contains(t, result.Outputs, "rows") + + rows, ok := result.Outputs["rows"].([]interface{}) + require.True(t, ok) + require.Len(t, rows, 1) + firstRow, ok := rows[0].(map[string]interface{}) + require.True(t, ok) + assert.Equal(t, float64(1), firstRow["value"]) + + resultEvidence, err := json.Marshal(result.Outputs) + require.NoError(t, err) + require.NotContains(t, string(resultEvidence), "password") + require.NotContains(t, string(resultEvidence), "token") + require.NotContains(t, string(resultEvidence), "secret") + t.Logf("fakeintake captured successful PAR task result: %s", resultEvidence) + + dequeueCalls, err := fakeintakeClient.GetPARDequeueCount() + require.NoError(t, err) + assert.GreaterOrEqual(t, dequeueCalls, 1) + t.Logf("live PAR loop dequeued from fakeintake: dequeue_calls=%d", dequeueCalls) + writeFusedEvidence(t, getenvOptional("RQ_FUSED_EVIDENCE_FILE"), []string{ + fmt.Sprintf("fakeintake task enqueued: task_id=%s action_fqn=%s inputs=%s", taskID, fqn, requestEvidence), + "live PAR loop dequeued the fakeintake OPMS task and invoked the registered action", + fmt.Sprintf("real Agent IPC endpoint called via NewDefaultBridgeClient: https://127.0.0.1:%d%s", cmdPortInt, com_datadoghq_remotequeries.AgentRemoteQueryExecuteEndpointPath), + fmt.Sprintf("fakeintake captured successful PAR task result: %s", resultEvidence), + fmt.Sprintf("dequeue_calls=%d", dequeueCalls), + "task verification skipped locally with DD_INTERNAL_PAR_SKIP_TASK_VERIFICATION=true", + }) +} + +func getenvOptional(name string) string { + return os.Getenv(name) +} + +func writeFusedEvidence(t *testing.T, path string, lines []string) { + t.Helper() + if path == "" { + return + } + payload := "" + for _, line := range lines { + payload += line + "\n" + } + require.NoError(t, os.WriteFile(path, []byte(payload), 0o600)) +} + +func getenvRequired(t *testing.T, name string) string { + t.Helper() + value := os.Getenv(name) + require.NotEmptyf(t, value, "%s is required", name) + return value +} diff --git a/test/remotequeries/fused-local-par-agent-postgres-proof.sh b/test/remotequeries/fused-local-par-agent-postgres-proof.sh new file mode 100755 index 000000000000..3486255e4479 --- /dev/null +++ b/test/remotequeries/fused-local-par-agent-postgres-proof.sh @@ -0,0 +1,320 @@ +#!/usr/bin/env bash +# Runs the fused local-only Remote Queries proof: +# fakeintake -> live WorkflowRunner PAR loop -> com.datadoghq.remotequeries.execute +# -> real local Agent IPC /agent/remote-queries/execute -> loaded Postgres check +# -> SELECT 1 AS value -> fakeintake publish. +# +# Defaults assume the remote-queries-poc worktree layout. Override AGENT_REPO, +# INTEGRATIONS_CORE, TMP_ROOT, CMD_PORT, POSTGRES_IMAGE, or AGENT_PYTHON_VERSION / +# AGENT_PYTHON_ABI if needed. The proof is local-only and intentionally sets +# DD_INTERNAL_PAR_SKIP_TASK_VERIFICATION=true inside the Go proof test. + +set -euo pipefail + +AGENT_REPO=${AGENT_REPO:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)} +INTEGRATIONS_CORE=${INTEGRATIONS_CORE:-/home/bits/dd/tasks/remote-queries-poc/worktrees/integrations-core} +TMP_ROOT=${TMP_ROOT:-/tmp/rq-fused-local-par-agent-postgres} +CMD_PORT=${CMD_PORT:-55003} +POSTGRES_IMAGE=${POSTGRES_IMAGE:-postgres:11-alpine} +POSTGRES_CONTAINER=${POSTGRES_CONTAINER:-rq-fused-local-par-agent-postgres-$$} +PIP_PLATFORM=${PIP_PLATFORM:-manylinux2014_x86_64} +AGENT_PYTHON_VERSION=${AGENT_PYTHON_VERSION:-} +AGENT_PYTHON_ABI=${AGENT_PYTHON_ABI:-} + +AGENT_PID="" +POSTGRES_STARTED=0 + +log() { + printf '\n[%s] %s\n' "$(date -u +%H:%M:%S)" "$*" +} + +cleanup_agent() { + if [[ -n "${AGENT_PID:-}" ]] && kill -0 "$AGENT_PID" 2>/dev/null; then + kill "$AGENT_PID" 2>/dev/null || true + wait "$AGENT_PID" 2>/dev/null || true + fi + AGENT_PID="" +} + +cleanup_all() { + cleanup_agent + if [[ "$POSTGRES_STARTED" == "1" ]]; then + docker rm -f "$POSTGRES_CONTAINER" >/dev/null 2>&1 || true + fi +} +trap cleanup_all EXIT + +require_cmd() { + command -v "$1" >/dev/null 2>&1 || { + echo "missing required command: $1" >&2 + exit 1 + } +} + +python_minor_from_version() { + local version=$1 + awk -F. '{print $1 "." $2}' <<<"$version" +} + +abi_from_python_minor() { + local minor=$1 + printf 'cp%s\n' "${minor//./}" +} + +write_agent_config() { + cat > "$TMP_ROOT/datadog.yaml" < "$TMP_ROOT/agent.log" + "$AGENT_REPO/bin/agent/agent" run -c "$TMP_ROOT" \ + > "$TMP_ROOT/python-detect-stdout.log" 2> "$TMP_ROOT/python-detect-stderr.log" & + AGENT_PID=$! + + local detected="" + for _ in $(seq 1 80); do + detected=$(grep -aoE '"pythonV":"[0-9]+\.[0-9]+(\.[0-9]+)?' "$TMP_ROOT/agent.log" 2>/dev/null | head -1 | cut -d: -f2 | tr -d '"' || true) + if [[ -n "$detected" ]]; then + break + fi + if ! kill -0 "$AGENT_PID" 2>/dev/null; then + break + fi + sleep 0.5 + done + cleanup_agent + + if [[ -z "$detected" && -f "$AGENT_REPO/omnibus/config/software/python3.rb" ]]; then + detected=$(sed -nE 's/^default_version "([0-9]+\.[0-9]+)(\.[0-9]+)?"/\1/p' "$AGENT_REPO/omnibus/config/software/python3.rb" | head -1) + if [[ -n "$detected" ]]; then + log "Could not detect runtime Python from Agent logs; falling back to source config: $detected" + fi + fi + + if [[ -z "$detected" ]]; then + echo "Unable to detect Agent Python version. Set AGENT_PYTHON_VERSION=3.12 or AGENT_PYTHON_ABI=cp312." >&2 + tail -80 "$TMP_ROOT/python-detect-stderr.log" >&2 || true + exit 1 + fi + + AGENT_PYTHON_VERSION=$(python_minor_from_version "$detected") + AGENT_PYTHON_ABI=$(abi_from_python_minor "$AGENT_PYTHON_VERSION") + log "Detected Agent Python runtime: $detected; installing wheels for $AGENT_PYTHON_VERSION ($AGENT_PYTHON_ABI)" +} + +install_python_deps() { + local py_digits + py_digits=${AGENT_PYTHON_VERSION//./} + + log "Installing temporary Python deps into $TMP_ROOT/pydeps for $AGENT_PYTHON_ABI on $PIP_PLATFORM" + python3 -m pip install --quiet --target "$TMP_ROOT/pydeps" \ + --only-binary=:all: --platform "$PIP_PLATFORM" \ + --implementation cp --python-version "$py_digits" --abi "$AGENT_PYTHON_ABI" \ + 'psycopg[binary,pool]' cachetools packaging semver 'pydantic<3' python-dateutil mmh3 + + for dep in "$TMP_ROOT"/pydeps/*; do + local base + base=$(basename "$dep") + [[ -e "$TMP_ROOT/checks.d/$base" ]] || ln -s "$dep" "$TMP_ROOT/checks.d/$base" + done +} + +setup_tmp_tree() { + rm -rf "$TMP_ROOT" + mkdir -p "$TMP_ROOT/conf.d/postgres.d" "$TMP_ROOT/run" "$TMP_ROOT/checks.d/datadog_checks" "$TMP_ROOT/pydeps" "$TMP_ROOT/results" + + ln -s "$INTEGRATIONS_CORE/datadog_checks_base/datadog_checks/base" "$TMP_ROOT/checks.d/datadog_checks/base" + ln -s "$INTEGRATIONS_CORE/datadog_checks_base/datadog_checks/checks" "$TMP_ROOT/checks.d/datadog_checks/checks" + ln -s "$INTEGRATIONS_CORE/datadog_checks_base/datadog_checks/config.py" "$TMP_ROOT/checks.d/datadog_checks/config.py" + ln -s "$INTEGRATIONS_CORE/datadog_checks_base/datadog_checks/errors.py" "$TMP_ROOT/checks.d/datadog_checks/errors.py" + ln -s "$INTEGRATIONS_CORE/datadog_checks_base/datadog_checks/log.py" "$TMP_ROOT/checks.d/datadog_checks/log.py" + ln -s "$INTEGRATIONS_CORE/postgres/datadog_checks/postgres" "$TMP_ROOT/checks.d/datadog_checks/postgres" + cat > "$TMP_ROOT/checks.d/datadog_checks/__init__.py" <<'PY' +__path__ = __import__('pkgutil').extend_path(__path__, __name__) +PY + + write_agent_config + detect_agent_python + install_python_deps +} + +start_postgres_fixture() { + if psql 'postgresql://postgres@localhost:5432/postgres' -c 'select 1' >/dev/null 2>&1; then + log "Using existing local Postgres on localhost:5432" + return + fi + + log "Starting disposable Postgres fixture $POSTGRES_CONTAINER from $POSTGRES_IMAGE" + docker run --rm --name "$POSTGRES_CONTAINER" \ + -e POSTGRES_HOST_AUTH_METHOD=trust \ + -p 5432:5432 \ + -d "$POSTGRES_IMAGE" >/dev/null + POSTGRES_STARTED=1 + + for _ in $(seq 1 60); do + if psql 'postgresql://postgres@localhost:5432/postgres' -c 'select 1' >/dev/null 2>&1; then + log "Postgres fixture is ready" + return + fi + sleep 0.5 + done + + docker logs "$POSTGRES_CONTAINER" | tail -80 >&2 || true + echo "Postgres fixture did not become ready" >&2 + exit 1 +} + +write_postgres_config() { + cat > "$TMP_ROOT/conf.d/postgres.d/conf.yaml" <<'YAML' +init_config: {} +instances: + - host: localhost + port: 5432 + dbname: postgres + username: postgres +YAML +} + +start_agent_and_wait_for_postgres_check() { + : > "$TMP_ROOT/agent.log" + PYTHONPATH="$TMP_ROOT/checks.d:$TMP_ROOT/pydeps" \ + "$AGENT_REPO/bin/agent/agent" run -c "$TMP_ROOT" \ + > "$TMP_ROOT/live-stdout.log" 2> "$TMP_ROOT/live-stderr.log" & + AGENT_PID=$! + + for _ in $(seq 1 80); do + if grep -q "successfully loaded check 'postgres'" "$TMP_ROOT/agent.log" && [[ -f "$TMP_ROOT/run/auth_token" && -f "$TMP_ROOT/run/ipc_cert.pem" ]]; then + log "Agent loaded the Postgres check and exposed IPC artifacts" + grep -n "successfully loaded check 'postgres'\|Scheduling check postgres" "$TMP_ROOT/agent.log" | tail -20 + return + fi + if ! kill -0 "$AGENT_PID" 2>/dev/null; then + echo "Agent exited early; stderr follows:" >&2 + tail -80 "$TMP_ROOT/live-stderr.log" >&2 || true + exit 1 + fi + sleep 0.5 + done + + echo "Timed out waiting for loaded Postgres check" >&2 + grep -ni 'postgres\|remote_queries\|ModuleNotFound\|ImportError\|unable to load check\|AttributeError' "$TMP_ROOT/agent.log" >&2 || true + exit 1 +} + +call_agent_execute_preflight() { + local payload='{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":1,"maxBytes":1024,"timeoutMs":1000}}' + local token + token=$(cat "$TMP_ROOT/run/auth_token") + + log "Preflight real Agent IPC execute endpoint" + local status + status=$(curl -sS -k -o "$TMP_ROOT/results/agent-execute-preflight.body" -w '%{http_code}' \ + -H "Authorization: Bearer ${token}" \ + -H 'Content-Type: application/json' \ + --data "$payload" \ + "https://127.0.0.1:${CMD_PORT}/agent/remote-queries/execute") + printf 'agent_execute_http_status=%s\n' "$status" | tee "$TMP_ROOT/results/agent-execute-preflight.status" + cat "$TMP_ROOT/results/agent-execute-preflight.body" + printf '\n' + + if [[ "$status" != "200" ]]; then + echo "FAIL: expected Agent execute preflight HTTP 200, got $status" >&2 + exit 1 + fi + if ! grep -Eq '"status"[[:space:]]*:[[:space:]]*"SUCCEEDED".*"value"[[:space:]]*:[[:space:]]*1|"value"[[:space:]]*:[[:space:]]*1.*"status"[[:space:]]*:[[:space:]]*"SUCCEEDED"' "$TMP_ROOT/results/agent-execute-preflight.body"; then + echo "FAIL: Agent execute preflight response did not contain SUCCEEDED row value=1" >&2 + exit 1 + fi + if grep -Eq 'password|token|secret' "$TMP_ROOT/results/agent-execute-preflight.body"; then + echo "FAIL: Agent execute preflight response contained credential-shaped text" >&2 + exit 1 + fi +} + +run_fused_go_proof() { + log "Running fused PAR -> real Agent IPC -> Postgres -> fakeintake proof test" + ( + cd "$AGENT_REPO" + RQ_FUSED_PROOF=1 \ + RQ_FUSED_AGENT_CMD_PORT="$CMD_PORT" \ + RQ_FUSED_AGENT_AUTH_TOKEN_FILE="$TMP_ROOT/run/auth_token" \ + RQ_FUSED_AGENT_IPC_CERT_FILE="$TMP_ROOT/run/ipc_cert.pem" \ + RQ_FUSED_EVIDENCE_FILE="$TMP_ROOT/results/fused-proof-evidence.txt" \ + dda inv test --targets=./pkg/privateactionrunner/bundles/remotequeries \ + --extra-args='-run TestRemoteQueriesActionRunsThroughLivePARLoopWithRealAgentIPC -count=1 -v' + ) | tee "$TMP_ROOT/results/fused-proof-test.log" +} + +main() { + require_cmd docker + require_cmd psql + require_cmd python3 + require_cmd curl + require_cmd dda + + [[ -x "$AGENT_REPO/bin/agent/agent" ]] || { + echo "Agent binary not found/executable: $AGENT_REPO/bin/agent/agent" >&2 + echo "Build it first with: dda inv agent.build --build-exclude=systemd" >&2 + exit 1 + } + [[ -d "$INTEGRATIONS_CORE/postgres/datadog_checks/postgres" ]] || { + echo "Postgres integration not found under: $INTEGRATIONS_CORE" >&2 + exit 1 + } + + log "Preparing temporary harness at $TMP_ROOT" + setup_tmp_tree + start_postgres_fixture + write_postgres_config + start_agent_and_wait_for_postgres_check + call_agent_execute_preflight + run_fused_go_proof + + log "Sanitized fused proof evidence" + cat "$TMP_ROOT/results/fused-proof-evidence.txt" + + log "Done. Sanitized artifacts left in $TMP_ROOT" + log "Key evidence: fakeintake enqueue/dequeue/publish and real Agent IPC endpoint evidence are in $TMP_ROOT/results/fused-proof-evidence.txt" +} + +main "$@" From da2339390d0c6cea7ffd0bb1dc0d5021652c0019 Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 14:19:20 +0000 Subject: [PATCH 15/33] Make fused proof Postgres fixture reusable --- .../fused-local-par-agent-postgres-proof.sh | 26 +++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/test/remotequeries/fused-local-par-agent-postgres-proof.sh b/test/remotequeries/fused-local-par-agent-postgres-proof.sh index 3486255e4479..03348e793672 100755 --- a/test/remotequeries/fused-local-par-agent-postgres-proof.sh +++ b/test/remotequeries/fused-local-par-agent-postgres-proof.sh @@ -184,11 +184,33 @@ start_postgres_fixture() { return fi + local published_containers proof_containers other_containers + published_containers=$(docker ps --filter publish=5432 --format '{{.Names}}' || true) + if [[ -n "$published_containers" ]]; then + proof_containers=$(grep -E '^rq-fused-local-par-agent-postgres-' <<<"$published_containers" || true) + other_containers=$(grep -Ev '^rq-fused-local-par-agent-postgres-' <<<"$published_containers" || true) + + if [[ -n "$proof_containers" ]]; then + log "Removing stale proof Postgres container(s) bound to localhost:5432 but not accepting psql" + xargs -r docker rm -f <<<"$proof_containers" >/dev/null + fi + if [[ -n "$other_containers" ]]; then + echo "Port 5432 is published by non-proof Docker container(s), and psql select 1 failed:" >&2 + sed 's/^/ /' <<<"$other_containers" >&2 + echo "Refusing to remove unrelated containers. Stop them or point this proof at a reachable local Postgres." >&2 + exit 1 + fi + fi + log "Starting disposable Postgres fixture $POSTGRES_CONTAINER from $POSTGRES_IMAGE" - docker run --rm --name "$POSTGRES_CONTAINER" \ + if ! docker run --rm --name "$POSTGRES_CONTAINER" \ -e POSTGRES_HOST_AUTH_METHOD=trust \ -p 5432:5432 \ - -d "$POSTGRES_IMAGE" >/dev/null + -d "$POSTGRES_IMAGE" >/dev/null; then + echo "Failed to start disposable Postgres fixture on port 5432." >&2 + echo "If another local Postgres is intended, ensure this succeeds: psql postgresql://postgres@localhost:5432/postgres -c 'select 1'" >&2 + exit 1 + fi POSTGRES_STARTED=1 for _ in $(seq 1 60); do From 977d61f6d4bdbb0eb641122dac109dcb280d2f8c Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 16:08:45 +0000 Subject: [PATCH 16/33] Add Remote Queries AgentSecure RPC bridge --- comp/api/grpcserver/impl-agent/grpc.go | 5 + .../impl-agent/remote_query_execute_test.go | 38 ++ comp/api/grpcserver/impl-agent/server.go | 139 +++++ .../impl/remote_query_execute.go | 200 ++++++- .../bundles/remotequeries/execute.go | 116 ++-- .../bundles/remotequeries/execute_test.go | 62 +- .../bundles/remotequeries/ipc_client.go | 39 +- .../live_agent_ipc_par_loop_test.go | 7 +- .../remotequeries/live_par_loop_test.go | 84 +-- pkg/proto/datadog/api/v1/api.proto | 37 ++ pkg/proto/pbgo/core/api.pb.go | 545 +++++++++++++++--- pkg/proto/pbgo/core/api_grpc.pb.go | 40 ++ pkg/proto/pbgo/mocks/core/api_mockgen.pb.go | 35 ++ .../fused-local-par-agent-postgres-proof.sh | 11 +- 14 files changed, 1083 insertions(+), 275 deletions(-) create mode 100644 comp/api/grpcserver/impl-agent/remote_query_execute_test.go diff --git a/comp/api/grpcserver/impl-agent/grpc.go b/comp/api/grpcserver/impl-agent/grpc.go index 134ef7111b92..ab7532809b89 100644 --- a/comp/api/grpcserver/impl-agent/grpc.go +++ b/comp/api/grpcserver/impl-agent/grpc.go @@ -32,6 +32,7 @@ import ( dogstatsdServer "github.com/DataDog/datadog-agent/comp/dogstatsd/server" rcservice "github.com/DataDog/datadog-agent/comp/remote-config/rcservice/def" rcservicemrf "github.com/DataDog/datadog-agent/comp/remote-config/rcservicemrf/def" + remotequeriesimpl "github.com/DataDog/datadog-agent/comp/remotequeries/impl" pb "github.com/DataDog/datadog-agent/pkg/proto/pbgo/core" grpcutil "github.com/DataDog/datadog-agent/pkg/util/grpc" "github.com/DataDog/datadog-agent/pkg/util/option" @@ -81,6 +82,7 @@ type server struct { telemetry telemetry.Component hostname hostnameinterface.Component configStream configstream.Component + remoteQueries *remotequeriesimpl.RemoteQueryExecuteService } func (s *server) BuildServer() http.Handler { @@ -122,6 +124,7 @@ func (s *server) BuildServer() http.Handler { autodiscovery: s.autodiscovery, configComp: s.configComp, configStreamServer: configstreamServer.NewServer(s.configComp, s.configStream, s.remoteAgentRegistry), + remoteQueries: s.remoteQueries, }) return grpcServer @@ -134,6 +137,7 @@ type Provides struct { // NewComponent creates a new grpc component func NewComponent(reqs Requires) (Provides, error) { + collector, _ := reqs.Collector.Get() provides := Provides{ Comp: &server{ IPC: reqs.IPC, @@ -152,6 +156,7 @@ func NewComponent(reqs Requires) (Provides, error) { telemetry: reqs.Telemetry, hostname: reqs.Hostname, configStream: reqs.ConfigStream, + remoteQueries: remotequeriesimpl.NewRemoteQueryExecuteService(collector, reqs.Cfg.GetBool(remotequeriesimpl.RemoteQueriesExecuteEnabledConfig)), }, } return provides, nil diff --git a/comp/api/grpcserver/impl-agent/remote_query_execute_test.go b/comp/api/grpcserver/impl-agent/remote_query_execute_test.go new file mode 100644 index 000000000000..7bf21f7ec4d9 --- /dev/null +++ b/comp/api/grpcserver/impl-agent/remote_query_execute_test.go @@ -0,0 +1,38 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2026-present Datadog, Inc. + +package agentimpl + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + pb "github.com/DataDog/datadog-agent/pkg/proto/pbgo/core" +) + +func TestRemoteQueryExecuteResponseFromJSONMapsStructuredRows(t *testing.T) { + resp, err := remoteQueryExecuteResponseFromJSON(`{"status":"SUCCEEDED","columns":[{"name":"value","type":"integer"}],"rows":[{"value":1}],"stats":{"elapsed_ms":2},"truncated":true}`) + + require.NoError(t, err) + assert.Equal(t, "SUCCEEDED", resp.GetStatus()) + require.Len(t, resp.GetColumns(), 1) + assert.Equal(t, "value", resp.GetColumns()[0].AsMap()["name"]) + require.Len(t, resp.GetRows(), 1) + assert.Equal(t, float64(1), resp.GetRows()[0].AsMap()["value"]) + assert.True(t, resp.GetTruncated()) + assert.Equal(t, float64(2), resp.GetStats().AsMap()["elapsed_ms"]) +} + +func TestRemoteQueryExecuteReturnsSanitizedUnavailableWhenServiceMissing(t *testing.T) { + resp, err := (&serverSecure{}).RemoteQueryExecute(context.Background(), &pb.RemoteQueryExecuteRequest{}) + + require.NoError(t, err) + assert.Equal(t, "executor_unavailable", resp.GetStatus()) + require.NotNil(t, resp.GetError()) + assert.Equal(t, "executor_unavailable", resp.GetError().GetCode()) +} diff --git a/comp/api/grpcserver/impl-agent/server.go b/comp/api/grpcserver/impl-agent/server.go index 22f91a3d174d..67a57ebe99b3 100644 --- a/comp/api/grpcserver/impl-agent/server.go +++ b/comp/api/grpcserver/impl-agent/server.go @@ -7,13 +7,17 @@ package agentimpl import ( "context" + "encoding/json" "errors" + "strconv" + "strings" "time" "google.golang.org/grpc/codes" "google.golang.org/grpc/grpclog" "google.golang.org/grpc/status" "google.golang.org/protobuf/types/known/emptypb" + "google.golang.org/protobuf/types/known/structpb" "github.com/DataDog/datadog-agent/comp/core/autodiscovery" autodiscoverystream "github.com/DataDog/datadog-agent/comp/core/autodiscovery/stream" @@ -33,6 +37,7 @@ import ( "github.com/DataDog/datadog-agent/comp/metadata/host/hostimpl/hosttags" rcservice "github.com/DataDog/datadog-agent/comp/remote-config/rcservice/def" rcservicemrf "github.com/DataDog/datadog-agent/comp/remote-config/rcservicemrf/def" + remotequeriesimpl "github.com/DataDog/datadog-agent/comp/remotequeries/impl" pb "github.com/DataDog/datadog-agent/pkg/proto/pbgo/core" "github.com/DataDog/datadog-agent/pkg/util/grpc" "github.com/DataDog/datadog-agent/pkg/util/log" @@ -60,6 +65,7 @@ type serverSecure struct { autodiscovery autodiscovery.Component configComp config.Component configStreamServer *configstreamServer.Server + remoteQueries *remotequeriesimpl.RemoteQueryExecuteService } func (s *agentServer) GetHostname(ctx context.Context, _ *pb.HostnameRequest) (*pb.HostnameReply, error) { @@ -274,3 +280,136 @@ func (s *serverSecure) CreateConfigSubscription(stream pb.AgentSecure_CreateConf func (s *serverSecure) WorkloadFilterEvaluate(ctx context.Context, req *pb.WorkloadFilterEvaluateRequest) (*pb.WorkloadFilterEvaluateResponse, error) { return s.workloadfilterServer.WorkloadFilterEvaluate(ctx, req) } + +func (s *serverSecure) RemoteQueryExecute(_ context.Context, req *pb.RemoteQueryExecuteRequest) (*pb.RemoteQueryExecuteResponse, error) { + if s.remoteQueries == nil { + return remoteQueryExecuteErrorResponse(remotequeriesimpl.RemoteQueryStatusExecutorUnavailable, "remote query executor is unavailable"), nil + } + + limits := remoteQueryLimitsFromProto(req.GetLimits()) + execReq, err := remotequeriesimpl.NewRemoteQueryExecuteRequest( + req.GetIntegration(), + remotequeriesimpl.RemoteQueryExecuteTarget{ + Host: req.GetTarget().GetHost(), + Port: int(req.GetTarget().GetPort()), + DBName: req.GetTarget().GetDbname(), + }, + req.GetQuery(), + limits, + ) + if err != nil { + return remoteQueryExecuteErrorResponse(remotequeriesimpl.RemoteQueryStatusInvalidRequest, err.Error()), nil + } + + result := s.remoteQueries.Execute(execReq) + if result.Error != nil { + return remoteQueryExecuteErrorResponse(result.Error.Code, result.Error.Message), nil + } + + return remoteQueryExecuteResponseFromJSON(result.ResponseJSON) +} + +func remoteQueryLimitsFromProto(limits *pb.RemoteQueryExecuteLimits) *remotequeriesimpl.RemoteQueryExecuteLimits { + if limits == nil { + return nil + } + return &remotequeriesimpl.RemoteQueryExecuteLimits{ + MaxRows: int(limits.GetMaxRows()), + MaxBytes: int(limits.GetMaxBytes()), + TimeoutMs: int(limits.GetTimeoutMs()), + } +} + +func remoteQueryExecuteErrorResponse(code string, message string) *pb.RemoteQueryExecuteResponse { + return &pb.RemoteQueryExecuteResponse{ + Status: code, + Error: &pb.RemoteQueryExecuteError{Code: code, Message: message}, + } +} + +type remoteQueryExecuteJSONResponse struct { + Status string `json:"status"` + Error *remoteQueryExecuteError `json:"error,omitempty"` + Columns []map[string]interface{} `json:"columns,omitempty"` + Rows []map[string]interface{} `json:"rows,omitempty"` + Truncated bool `json:"truncated,omitempty"` + Stats map[string]interface{} `json:"stats,omitempty"` +} + +type remoteQueryExecuteError struct { + Code string `json:"code"` + Message string `json:"message"` +} + +func remoteQueryExecuteResponseFromJSON(responseJSON string) (*pb.RemoteQueryExecuteResponse, error) { + var payload remoteQueryExecuteJSONResponse + decoder := json.NewDecoder(strings.NewReader(responseJSON)) + decoder.UseNumber() + if err := decoder.Decode(&payload); err != nil { + return nil, status.Error(codes.Internal, "remote query executor returned invalid JSON") + } + if payload.Status == "" { + return nil, status.Error(codes.Internal, "remote query executor response missing status") + } + + out := &pb.RemoteQueryExecuteResponse{ + Status: payload.Status, + Truncated: payload.Truncated, + } + if payload.Error != nil { + out.Error = &pb.RemoteQueryExecuteError{Code: payload.Error.Code, Message: payload.Error.Message} + } + for _, column := range payload.Columns { + pbColumn, err := structpb.NewStruct(normalizeRemoteQueryStruct(column)) + if err != nil { + return nil, status.Error(codes.Internal, "remote query executor returned invalid column data") + } + out.Columns = append(out.Columns, pbColumn) + } + for _, row := range payload.Rows { + pbRow, err := structpb.NewStruct(normalizeRemoteQueryStruct(row)) + if err != nil { + return nil, status.Error(codes.Internal, "remote query executor returned invalid row data") + } + out.Rows = append(out.Rows, pbRow) + } + if payload.Stats != nil { + stats, err := structpb.NewStruct(normalizeRemoteQueryStruct(payload.Stats)) + if err != nil { + return nil, status.Error(codes.Internal, "remote query executor returned invalid stats data") + } + out.Stats = stats + } + return out, nil +} + +func normalizeRemoteQueryStruct(in map[string]interface{}) map[string]interface{} { + out := make(map[string]interface{}, len(in)) + for key, value := range in { + out[key] = normalizeRemoteQueryValue(value) + } + return out +} + +func normalizeRemoteQueryValue(value interface{}) interface{} { + switch v := value.(type) { + case json.Number: + if i, err := strconv.ParseInt(v.String(), 10, 64); err == nil { + return i + } + if f, err := strconv.ParseFloat(v.String(), 64); err == nil { + return f + } + return v.String() + case map[string]interface{}: + return normalizeRemoteQueryStruct(v) + case []interface{}: + out := make([]interface{}, len(v)) + for i, item := range v { + out[i] = normalizeRemoteQueryValue(item) + } + return out + default: + return v + } +} diff --git a/comp/remotequeries/impl/remote_query_execute.go b/comp/remotequeries/impl/remote_query_execute.go index c16ceb07d293..bf2fcc59008c 100644 --- a/comp/remotequeries/impl/remote_query_execute.go +++ b/comp/remotequeries/impl/remote_query_execute.go @@ -58,17 +58,110 @@ func remoteQueryRunnerFor(chk check.Check) (remoteQueryRunner, bool) { // NewRemoteQueryExecuteEndpointProvider registers the remote query execute endpoint on the internal Agent API. func NewRemoteQueryExecuteEndpointProvider(reqs Requires) api.AgentEndpointProvider { h := &remoteQueryExecuteHandler{ - collector: reqs.Collector, - enabled: reqs.Cfg.GetBool(RemoteQueriesExecuteEnabledConfig), + service: NewRemoteQueryExecuteService(reqs.Collector, reqs.Cfg.GetBool(RemoteQueriesExecuteEnabledConfig)), } return api.NewAgentEndpointProvider(h.handle, RemoteQueryExecuteEndpointPath, http.MethodPost) } type remoteQueryExecuteHandler struct { + service *RemoteQueryExecuteService collector collector.Component enabled bool } +// RemoteQueryExecuteService executes credential-free Remote Queries requests through loaded checks. +type RemoteQueryExecuteService struct { + collector collector.Component + enabled bool +} + +// NewRemoteQueryExecuteService creates the shared executor used by the HTTP POC endpoint and AgentSecure RPC. +func NewRemoteQueryExecuteService(collector collector.Component, enabled bool) *RemoteQueryExecuteService { + return &RemoteQueryExecuteService{collector: collector, enabled: enabled} +} + +// RemoteQueryExecuteTarget identifies the datastore target without carrying credentials. +type RemoteQueryExecuteTarget struct { + Host string + Port int + DBName string +} + +// RemoteQueryExecuteLimits contains optional execution limits for a remote query. +type RemoteQueryExecuteLimits struct { + MaxRows int + MaxBytes int + TimeoutMs int +} + +// RemoteQueryExecuteRequest is the typed internal request shape shared by HTTP and gRPC callers. +type RemoteQueryExecuteRequest struct { + Integration string + Target RemoteQueryExecuteTarget + Query string + Limits *RemoteQueryExecuteLimits +} + +// RemoteQueryExecuteError is a sanitized remote query bridge error. +type RemoteQueryExecuteError struct { + Code string + Message string +} + +// RemoteQueryExecuteResult is the service result. ResponseJSON is set only for successful executor responses. +type RemoteQueryExecuteResult struct { + HTTPStatus int + Status string + Error *RemoteQueryExecuteError + ResponseJSON string +} + +const ( + // RemoteQueryStatusInvalidRequest reports a malformed or disallowed request. + RemoteQueryStatusInvalidRequest = statusInvalidRequest + // RemoteQueryStatusExecutorUnavailable reports an unavailable matched executor or bridge dependency. + RemoteQueryStatusExecutorUnavailable = statusExecutorUnavailable +) + +// NewRemoteQueryExecuteRequest validates and normalizes a typed Remote Queries execute request. +func NewRemoteQueryExecuteRequest(integration string, target RemoteQueryExecuteTarget, query string, limits *RemoteQueryExecuteLimits) (RemoteQueryExecuteRequest, error) { + parsedIntegration, err := parseIntegration(integration) + if err != nil { + return RemoteQueryExecuteRequest{}, err + } + + parsedTarget, err := parseTarget(&remoteQueryTargetRequestJSON{Host: target.Host, Port: &target.Port, DBName: target.DBName}) + if err != nil { + return RemoteQueryExecuteRequest{}, err + } + + if query == "" { + return RemoteQueryExecuteRequest{}, fmt.Errorf("query is required") + } + if query != remoteQueryProofQuery { + return RemoteQueryExecuteRequest{}, fmt.Errorf("query is not allowed") + } + + var parsedLimits *remoteQueryExecuteLimits + if limits != nil { + parsedLimits, err = parseExecuteLimits(&remoteQueryExecuteLimitsRequestJSON{ + MaxRows: &limits.MaxRows, + MaxBytes: &limits.MaxBytes, + TimeoutMs: &limits.TimeoutMs, + }) + if err != nil { + return RemoteQueryExecuteRequest{}, err + } + } + + return remoteQueryExecuteRequestFromInternal(remoteQueryExecuteRequest{ + Integration: parsedIntegration, + Target: parsedTarget, + Query: query, + Limits: parsedLimits, + }), nil +} + type remoteQueryExecuteRequest struct { Integration string Target remoteQueryTarget @@ -116,43 +209,29 @@ type remoteQueryExecuteLimitsJSON struct { func (h *remoteQueryExecuteHandler) handle(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") - if !h.enabled { + service := h.service + if service == nil { + service = NewRemoteQueryExecuteService(h.collector, h.enabled) + } + if service == nil || !service.enabled { writeExecuteError(w, http.StatusServiceUnavailable, statusBridgeDisabled, "remote queries bridge is disabled") return } - req, requestJSON, err := parseExecuteRequest(r) + req, _, err := parseExecuteRequest(r) if err != nil { writeExecuteParseError(w, err) return } - matches := h.findMatches(req.Integration, req.Target) - switch len(matches) { - case 0: - writeExecuteError(w, http.StatusNotFound, statusTargetNotFound, "no matching integration check found") - return - case 1: - // continue below - default: - writeExecuteError(w, http.StatusConflict, statusAmbiguous, "multiple matching integration checks found") - return - } - - runner, ok := remoteQueryRunnerFor(matches[0].check) - if !ok { - writeExecuteError(w, http.StatusFailedDependency, statusExecutorUnavailable, "matched integration check does not support remote query execution") - return - } - - responseJSON, err := runner.RunRemoteQueryJSON(req.Integration, requestJSON) - if err != nil { - writeExecuteError(w, http.StatusBadGateway, statusExecutorUnavailable, "remote query executor failed") + result := service.Execute(remoteQueryExecuteRequestFromInternal(req)) + if result.Error != nil { + writeExecuteError(w, result.HTTPStatus, result.Error.Code, result.Error.Message) return } w.WriteHeader(http.StatusOK) - _, _ = io.WriteString(w, responseJSON) + _, _ = io.WriteString(w, result.ResponseJSON) } func parseExecuteRequest(r *http.Request) (remoteQueryExecuteRequest, string, error) { @@ -249,8 +328,73 @@ func parseRequiredPositiveInt(value *int, name string) (int, error) { return *value, nil } -func (h *remoteQueryExecuteHandler) findMatches(integration string, target remoteQueryTarget) []integrationCheckMatch { - return findIntegrationMatches(h.collector, integration, target) +func (s *RemoteQueryExecuteService) Execute(req RemoteQueryExecuteRequest) RemoteQueryExecuteResult { + if s == nil || !s.enabled { + return remoteQueryExecuteErrorResult(http.StatusServiceUnavailable, statusBridgeDisabled, "remote queries bridge is disabled") + } + if s.collector == nil { + return remoteQueryExecuteErrorResult(http.StatusFailedDependency, statusExecutorUnavailable, "remote query executor is unavailable") + } + + internal := req.internal() + matches := findIntegrationMatches(s.collector, internal.Integration, internal.Target) + switch len(matches) { + case 0: + return remoteQueryExecuteErrorResult(http.StatusNotFound, statusTargetNotFound, "no matching integration check found") + case 1: + // continue below + default: + return remoteQueryExecuteErrorResult(http.StatusConflict, statusAmbiguous, "multiple matching integration checks found") + } + + runner, ok := remoteQueryRunnerFor(matches[0].check) + if !ok { + return remoteQueryExecuteErrorResult(http.StatusFailedDependency, statusExecutorUnavailable, "matched integration check does not support remote query execution") + } + + requestJSON, err := marshalExecuteRequest(internal) + if err != nil { + return remoteQueryExecuteErrorResult(http.StatusBadRequest, statusInvalidRequest, "malformed JSON request") + } + + responseJSON, err := runner.RunRemoteQueryJSON(internal.Integration, requestJSON) + if err != nil { + return remoteQueryExecuteErrorResult(http.StatusBadGateway, statusExecutorUnavailable, "remote query executor failed") + } + + return RemoteQueryExecuteResult{HTTPStatus: http.StatusOK, ResponseJSON: responseJSON} +} + +func remoteQueryExecuteErrorResult(httpStatus int, status string, message string) RemoteQueryExecuteResult { + return RemoteQueryExecuteResult{ + HTTPStatus: httpStatus, + Status: status, + Error: &RemoteQueryExecuteError{Code: status, Message: message}, + } +} + +func (r RemoteQueryExecuteRequest) internal() remoteQueryExecuteRequest { + internal := remoteQueryExecuteRequest{ + Integration: r.Integration, + Target: remoteQueryTarget{Host: r.Target.Host, Port: r.Target.Port, DBName: r.Target.DBName}, + Query: r.Query, + } + if r.Limits != nil { + internal.Limits = &remoteQueryExecuteLimits{MaxRows: r.Limits.MaxRows, MaxBytes: r.Limits.MaxBytes, TimeoutMs: r.Limits.TimeoutMs} + } + return internal +} + +func remoteQueryExecuteRequestFromInternal(req remoteQueryExecuteRequest) RemoteQueryExecuteRequest { + out := RemoteQueryExecuteRequest{ + Integration: req.Integration, + Target: RemoteQueryExecuteTarget{Host: req.Target.Host, Port: req.Target.Port, DBName: req.Target.DBName}, + Query: req.Query, + } + if req.Limits != nil { + out.Limits = &RemoteQueryExecuteLimits{MaxRows: req.Limits.MaxRows, MaxBytes: req.Limits.MaxBytes, TimeoutMs: req.Limits.TimeoutMs} + } + return out } func marshalExecuteRequest(req remoteQueryExecuteRequest) (string, error) { diff --git a/pkg/privateactionrunner/bundles/remotequeries/execute.go b/pkg/privateactionrunner/bundles/remotequeries/execute.go index 9f4594969d7e..59b347b52030 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/execute.go +++ b/pkg/privateactionrunner/bundles/remotequeries/execute.go @@ -6,32 +6,24 @@ package com_datadoghq_remotequeries import ( - "bytes" "context" - "encoding/json" "fmt" - "io" - ipc "github.com/DataDog/datadog-agent/comp/core/ipc/def" - ipchttp "github.com/DataDog/datadog-agent/comp/core/ipc/httphelpers" + "google.golang.org/grpc" + "github.com/DataDog/datadog-agent/pkg/privateactionrunner/libs/privateconnection" "github.com/DataDog/datadog-agent/pkg/privateactionrunner/types" "github.com/DataDog/datadog-agent/pkg/privateactionrunner/util" + pb "github.com/DataDog/datadog-agent/pkg/proto/pbgo/core" ) -const ( - // AgentRemoteQueryExecuteEndpointPath is the POC-only Agent command API path this PAR action calls. - // It is intentionally local-only and is not a production API/IPC commitment. - AgentRemoteQueryExecuteEndpointPath = "/agent/remote-queries/execute" -) - -// BridgeClient is the narrow Agent IPC HTTP client surface required by this action. +// BridgeClient is the narrow AgentSecure gRPC client surface required by this action. type BridgeClient interface { - Post(url string, contentType string, body io.Reader, opts ...ipc.RequestOption) (resp []byte, err error) + RemoteQueryExecute(ctx context.Context, in *pb.RemoteQueryExecuteRequest, opts ...grpc.CallOption) (*pb.RemoteQueryExecuteResponse, error) } -// BridgeClientFactory returns an IPC client and fully-qualified local Agent endpoint URL. -type BridgeClientFactory func() (BridgeClient, string, error) +// BridgeClientFactory returns an authenticated AgentSecure client over the local Agent IPC channel. +type BridgeClientFactory func() (BridgeClient, error) type ExecuteAction struct { newBridgeClient BridgeClientFactory @@ -73,59 +65,79 @@ func (a *ExecuteAction) Run( ) } - payload, err := json.Marshal(inputs) - if err != nil { - return nil, util.DefaultActionErrorWithDisplayError( - fmt.Errorf("marshal remote query action inputs"), - "invalid remote query action inputs", - ) - } - if a == nil || a.newBridgeClient == nil { return nil, util.DefaultActionError(fmt.Errorf("remote query action requires an Agent IPC client")) } - client, endpointURL, err := a.newBridgeClient() + client, err := a.newBridgeClient() if err != nil { return nil, util.DefaultActionErrorWithDisplayError(err, "remote query action could not create an Agent IPC client") } - if client == nil || endpointURL == "" { - return nil, util.DefaultActionError(fmt.Errorf("remote query action requires an Agent IPC client and endpoint URL")) + if client == nil { + return nil, util.DefaultActionError(fmt.Errorf("remote query action requires an AgentSecure client")) } - body, postErr := client.Post(endpointURL, "application/json", bytes.NewReader(payload), ipchttp.WithContext(ctx)) - output, decodeErr := decodeBridgeResponse(body) - if decodeErr == nil { - // IPC HTTPClient returns both the response body and an error for HTTP >= 400. - // The bridge body is already sanitized, so preserve its status/error payload as the action output. - return output, nil + resp, err := client.RemoteQueryExecute(ctx, remoteQueryExecuteRequestFromInputs(inputs)) + if err != nil { + return nil, util.DefaultActionErrorWithDisplayError(err, "remote query AgentSecure RPC failed") } - if postErr != nil { - if len(body) > 0 { - return nil, util.DefaultActionErrorWithDisplayError( - fmt.Errorf("remote query IPC request failed with undecodable response"), - "remote query IPC request failed with undecodable response", - ) - } - return nil, util.DefaultActionErrorWithDisplayError(postErr, "remote query IPC request failed") + output, err := remoteQueryExecuteOutputFromProto(resp) + if err != nil { + return nil, util.DefaultActionErrorWithDisplayError(err, "remote query AgentSecure RPC response was invalid") } - return nil, util.DefaultActionErrorWithDisplayError(decodeErr, "remote query IPC response was invalid") + return output, nil } -func decodeBridgeResponse(body []byte) (map[string]interface{}, error) { - if len(body) == 0 { - return nil, fmt.Errorf("empty remote query response") +func remoteQueryExecuteRequestFromInputs(inputs ExecuteInputs) *pb.RemoteQueryExecuteRequest { + req := &pb.RemoteQueryExecuteRequest{ + Integration: inputs.Integration, + Target: &pb.RemoteQueryTarget{ + Host: inputs.Target.Host, + Port: int32(inputs.Target.Port), + Dbname: inputs.Target.DBName, + }, + Query: inputs.Query, } + if inputs.Limits != nil { + req.Limits = &pb.RemoteQueryExecuteLimits{ + MaxRows: int32(inputs.Limits.MaxRows), + MaxBytes: int32(inputs.Limits.MaxBytes), + TimeoutMs: int32(inputs.Limits.TimeoutMs), + } + } + return req +} - decoder := json.NewDecoder(bytes.NewReader(body)) - decoder.UseNumber() +func remoteQueryExecuteOutputFromProto(resp *pb.RemoteQueryExecuteResponse) (map[string]interface{}, error) { + if resp == nil || resp.GetStatus() == "" { + return nil, fmt.Errorf("remote query response missing status") + } - var output map[string]interface{} - if err := decoder.Decode(&output); err != nil { - return nil, fmt.Errorf("decode remote query response: %w", err) + output := map[string]interface{}{"status": resp.GetStatus()} + if resp.GetError() != nil { + output["error"] = map[string]interface{}{ + "code": resp.GetError().GetCode(), + "message": resp.GetError().GetMessage(), + } } - status, ok := output["status"].(string) - if !ok || status == "" { - return nil, fmt.Errorf("remote query response missing status") + if len(resp.GetColumns()) > 0 { + columns := make([]interface{}, 0, len(resp.GetColumns())) + for _, column := range resp.GetColumns() { + columns = append(columns, column.AsMap()) + } + output["columns"] = columns + } + if len(resp.GetRows()) > 0 { + rows := make([]interface{}, 0, len(resp.GetRows())) + for _, row := range resp.GetRows() { + rows = append(rows, row.AsMap()) + } + output["rows"] = rows + } + if resp.GetTruncated() { + output["truncated"] = true + } + if resp.GetStats() != nil { + output["stats"] = resp.GetStats().AsMap() } return output, nil } diff --git a/pkg/privateactionrunner/bundles/remotequeries/execute_test.go b/pkg/privateactionrunner/bundles/remotequeries/execute_test.go index ada8916c5b49..9c2e27711084 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/execute_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/execute_test.go @@ -8,22 +8,25 @@ package com_datadoghq_remotequeries import ( "context" "encoding/json" - "errors" - "io" "testing" - ipc "github.com/DataDog/datadog-agent/comp/core/ipc/def" + "google.golang.org/grpc" + "google.golang.org/protobuf/types/known/structpb" + "github.com/DataDog/datadog-agent/pkg/privateactionrunner/libs/privateconnection" "github.com/DataDog/datadog-agent/pkg/privateactionrunner/types" "github.com/DataDog/datadog-agent/pkg/privateactionrunner/util" + pb "github.com/DataDog/datadog-agent/pkg/proto/pbgo/core" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) -func TestExecuteActionUsesCredentialFreeIPCPostShape(t *testing.T) { - client := &captureBridgeClient{response: []byte(`{"status":"SUCCEEDED","rows":[{"value":1}]}`)} - action := NewExecuteAction(func() (BridgeClient, string, error) { - return client, "https://localhost:5001" + AgentRemoteQueryExecuteEndpointPath, nil +func TestExecuteActionUsesCredentialFreeAgentSecureRequestShape(t *testing.T) { + row, err := structpb.NewStruct(map[string]interface{}{"value": 1}) + require.NoError(t, err) + client := &captureBridgeClient{response: &pb.RemoteQueryExecuteResponse{Status: "SUCCEEDED", Rows: []*structpb.Struct{row}}} + action := NewExecuteAction(func() (BridgeClient, error) { + return client, nil }) output, err := action.Run(context.Background(), taskWithInputs(map[string]interface{}{ @@ -42,25 +45,30 @@ func TestExecuteActionUsesCredentialFreeIPCPostShape(t *testing.T) { }), &privateconnection.PrivateCredentials{Tokens: []privateconnection.PrivateCredentialsToken{{Name: "password", Value: "secret-value"}}}) require.NoError(t, err) - assert.Equal(t, "https://localhost:5001"+AgentRemoteQueryExecuteEndpointPath, client.url) - assert.Equal(t, "application/json", client.contentType) - assert.JSONEq(t, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":1,"maxBytes":1024,"timeoutMs":1000}}`, client.body) - assert.NotContains(t, client.body, "secret-value") + require.NotNil(t, client.request) + assert.Equal(t, "postgres", client.request.GetIntegration()) + assert.Equal(t, "localhost", client.request.GetTarget().GetHost()) + assert.Equal(t, int32(5432), client.request.GetTarget().GetPort()) + assert.Equal(t, "postgres", client.request.GetTarget().GetDbname()) + assert.Equal(t, "SELECT 1 AS value", client.request.GetQuery()) + assert.Equal(t, int32(1), client.request.GetLimits().GetMaxRows()) + requestEvidence, err := json.Marshal(client.request) + require.NoError(t, err) + assert.NotContains(t, string(requestEvidence), "secret-value") assert.Equal(t, map[string]interface{}{ "status": "SUCCEEDED", "rows": []interface{}{ - map[string]interface{}{"value": json.Number("1")}, + map[string]interface{}{"value": float64(1)}, }, }, output) } func TestExecuteActionPreservesSanitizedBridgeErrorBody(t *testing.T) { client := &captureBridgeClient{ - response: []byte(`{"status":"target_not_found","error":{"code":"target_not_found","message":"no matching integration check found"}}`), - err: errors.New("status 404"), + response: &pb.RemoteQueryExecuteResponse{Status: "target_not_found", Error: &pb.RemoteQueryExecuteError{Code: "target_not_found", Message: "no matching integration check found"}}, } - action := NewExecuteAction(func() (BridgeClient, string, error) { - return client, "https://localhost:5001" + AgentRemoteQueryExecuteEndpointPath, nil + action := NewExecuteAction(func() (BridgeClient, error) { + return client, nil }) output, err := action.Run(context.Background(), taskWithInputs(map[string]interface{}{ @@ -80,9 +88,9 @@ func TestExecuteActionPreservesSanitizedBridgeErrorBody(t *testing.T) { } func TestExecuteActionSanitizesInputExtractionErrors(t *testing.T) { - action := NewExecuteAction(func() (BridgeClient, string, error) { + action := NewExecuteAction(func() (BridgeClient, error) { require.Fail(t, "bridge client should not be created for invalid inputs") - return nil, "", nil + return nil, nil }) _, err := action.Run(context.Background(), taskWithInputs(map[string]interface{}{ @@ -112,20 +120,12 @@ func taskWithInputs(inputs map[string]interface{}) *types.Task { } type captureBridgeClient struct { - url string - contentType string - body string - response []byte - err error + request *pb.RemoteQueryExecuteRequest + response *pb.RemoteQueryExecuteResponse + err error } -func (c *captureBridgeClient) Post(url string, contentType string, body io.Reader, _ ...ipc.RequestOption) ([]byte, error) { - c.url = url - c.contentType = contentType - payload, err := io.ReadAll(body) - if err != nil { - return nil, err - } - c.body = string(payload) +func (c *captureBridgeClient) RemoteQueryExecute(_ context.Context, req *pb.RemoteQueryExecuteRequest, _ ...grpc.CallOption) (*pb.RemoteQueryExecuteResponse, error) { + c.request = req return c.response, c.err } diff --git a/pkg/privateactionrunner/bundles/remotequeries/ipc_client.go b/pkg/privateactionrunner/bundles/remotequeries/ipc_client.go index 7646301c1b23..14cdab6181db 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/ipc_client.go +++ b/pkg/privateactionrunner/bundles/remotequeries/ipc_client.go @@ -6,40 +6,41 @@ package com_datadoghq_remotequeries import ( + "context" "fmt" - "net" - "net/url" "strconv" + "time" - ipchttp "github.com/DataDog/datadog-agent/comp/core/ipc/httphelpers" - pkgtoken "github.com/DataDog/datadog-agent/pkg/api/security" "github.com/DataDog/datadog-agent/pkg/api/security/cert" - pkgconfigmodel "github.com/DataDog/datadog-agent/pkg/config/model" pkgconfigsetup "github.com/DataDog/datadog-agent/pkg/config/setup" + agentgrpc "github.com/DataDog/datadog-agent/pkg/util/grpc" "github.com/DataDog/datadog-agent/pkg/util/system" ) // NewDefaultBridgeClient creates the local Agent IPC client used by the registered PAR action. -func NewDefaultBridgeClient() (BridgeClient, string, error) { +func NewDefaultBridgeClient() (BridgeClient, error) { cfg := pkgconfigsetup.Datadog() - token, err := pkgtoken.FetchAuthToken(cfg) - if err != nil { - return nil, "", fmt.Errorf("fetch Agent IPC auth token: %w", err) - } clientTLSConfig, _, _, err := cert.FetchIPCCert(cfg) if err != nil { - return nil, "", fmt.Errorf("fetch Agent IPC certificate: %w", err) + return nil, fmt.Errorf("fetch Agent IPC certificate: %w", err) } - endpointURL, err := agentIPCURL(cfg, AgentRemoteQueryExecuteEndpointPath) + ipcHost, err := agentIPCHost() + if err != nil { + return nil, err + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + client, err := agentgrpc.GetDDAgentSecureClient(ctx, ipcHost, strconv.Itoa(cfg.GetInt("cmd_port")), clientTLSConfig) if err != nil { - return nil, "", err + return nil, fmt.Errorf("create AgentSecure client: %w", err) } - return ipchttp.NewClient(token, clientTLSConfig, cfg), endpointURL, nil + return client, nil } -func agentIPCURL(cfg pkgconfigmodel.Reader, endpointPath string) (string, error) { +func agentIPCHost() (string, error) { + cfg := pkgconfigsetup.Datadog() cmdHostKey := "cmd_host" if cfg.IsConfigured("ipc_address") { cmdHostKey = "ipc_address" @@ -49,11 +50,5 @@ func agentIPCURL(cfg pkgconfigmodel.Reader, endpointPath string) (string, error) if err != nil { return "", fmt.Errorf("%s: %w", cmdHostKey, err) } - - endpointURL := url.URL{ - Scheme: "https", - Host: net.JoinHostPort(ipcHost, strconv.Itoa(cfg.GetInt("cmd_port"))), - Path: endpointPath, - } - return endpointURL.String(), nil + return ipcHost, nil } diff --git a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go index 0f0d709a1663..2e2eea43dbc8 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go @@ -95,11 +95,14 @@ func TestRemoteQueriesActionRunsThroughLivePARLoopWithRealAgentIPC(t *testing.T) fqn := com_datadoghq_remotequeries.BundleID + "." + com_datadoghq_remotequeries.ExecuteActionName t.Logf("fakeintake task enqueued: task_id=%s action_fqn=%s inputs=%s", taskID, fqn, requestEvidence) - t.Logf("real Agent IPC endpoint configured: https://127.0.0.1:%d%s", cmdPortInt, com_datadoghq_remotequeries.AgentRemoteQueryExecuteEndpointPath) + t.Logf("real AgentSecure IPC configured: 127.0.0.1:%d RemoteQueryExecute", cmdPortInt) require.NoError(t, fakeintakeClient.EnqueuePARTask(taskID, fqn, inputs)) result, err := fakeintakeClient.GetPARTaskResult(taskID, 20*time.Second) require.NoError(t, err) + if !result.Success { + t.Logf("failed PAR task result: %+v", result) + } require.True(t, result.Success) require.Equal(t, taskID, result.TaskID) assert.Equal(t, "SUCCEEDED", result.Outputs["status"]) @@ -126,7 +129,7 @@ func TestRemoteQueriesActionRunsThroughLivePARLoopWithRealAgentIPC(t *testing.T) writeFusedEvidence(t, getenvOptional("RQ_FUSED_EVIDENCE_FILE"), []string{ fmt.Sprintf("fakeintake task enqueued: task_id=%s action_fqn=%s inputs=%s", taskID, fqn, requestEvidence), "live PAR loop dequeued the fakeintake OPMS task and invoked the registered action", - fmt.Sprintf("real Agent IPC endpoint called via NewDefaultBridgeClient: https://127.0.0.1:%d%s", cmdPortInt, com_datadoghq_remotequeries.AgentRemoteQueryExecuteEndpointPath), + fmt.Sprintf("real AgentSecure IPC called via NewDefaultBridgeClient: 127.0.0.1:%d RemoteQueryExecute", cmdPortInt), fmt.Sprintf("fakeintake captured successful PAR task result: %s", resultEvidence), fmt.Sprintf("dequeue_calls=%d", dequeueCalls), "task verification skipped locally with DD_INTERNAL_PAR_SKIP_TASK_VERIFICATION=true", diff --git a/pkg/privateactionrunner/bundles/remotequeries/live_par_loop_test.go b/pkg/privateactionrunner/bundles/remotequeries/live_par_loop_test.go index 957ecd76cb31..6ccb94d4f669 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/live_par_loop_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/live_par_loop_test.go @@ -8,19 +8,18 @@ package com_datadoghq_remotequeries_test import ( - "bytes" "context" "crypto/ecdsa" "crypto/elliptic" "crypto/rand" "encoding/json" - "io" - "net/http" - "net/http/httptest" "testing" "time" - ipc "github.com/DataDog/datadog-agent/comp/core/ipc/def" + "google.golang.org/grpc" + "google.golang.org/protobuf/types/known/structpb" + "k8s.io/apimachinery/pkg/util/sets" + parconfig "github.com/DataDog/datadog-agent/pkg/privateactionrunner/adapters/config" app "github.com/DataDog/datadog-agent/pkg/privateactionrunner/adapters/constants" com_datadoghq_remotequeries "github.com/DataDog/datadog-agent/pkg/privateactionrunner/bundles/remotequeries" @@ -28,34 +27,25 @@ import ( "github.com/DataDog/datadog-agent/pkg/privateactionrunner/opms" "github.com/DataDog/datadog-agent/pkg/privateactionrunner/runners" taskverifier "github.com/DataDog/datadog-agent/pkg/privateactionrunner/task-verifier" + pb "github.com/DataDog/datadog-agent/pkg/proto/pbgo/core" fakeintakeclient "github.com/DataDog/datadog-agent/test/fakeintake/client" fakeintakeserver "github.com/DataDog/datadog-agent/test/fakeintake/server" "github.com/DataDog/datadog-go/v5/statsd" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - "k8s.io/apimachinery/pkg/util/sets" ) func TestRemoteQueriesActionRunsThroughLivePARLoopAndFakeintake(t *testing.T) { t.Setenv(app.InternalSkipTaskVerificationEnvVar, "true") - bridgeRequests := make(chan []byte, 1) - bridgeServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - require.Equal(t, http.MethodPost, r.Method) - require.Equal(t, com_datadoghq_remotequeries.AgentRemoteQueryExecuteEndpointPath, r.URL.Path) - require.Contains(t, r.Header.Get("Content-Type"), "application/json") - - body, err := io.ReadAll(r.Body) - require.NoError(t, err) - bridgeRequests <- body - - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{"status":"SUCCEEDED","rows":[{"value":1}]}`)) - })) - defer bridgeServer.Close() - - restoreFactory := com_datadoghq_remotequeries.SetBridgeClientFactoryForTest(func() (com_datadoghq_remotequeries.BridgeClient, string, error) { - return httpBridgeClient{}, bridgeServer.URL + com_datadoghq_remotequeries.AgentRemoteQueryExecuteEndpointPath, nil + row, err := structpb.NewStruct(map[string]interface{}{"value": 1}) + require.NoError(t, err) + bridgeRequests := make(chan *pb.RemoteQueryExecuteRequest, 1) + restoreFactory := com_datadoghq_remotequeries.SetBridgeClientFactoryForTest(func() (com_datadoghq_remotequeries.BridgeClient, error) { + return &captureAgentSecureClient{ + requests: bridgeRequests, + response: &pb.RemoteQueryExecuteResponse{Status: "SUCCEEDED", Rows: []*structpb.Struct{row}}, + }, nil }) defer restoreFactory() @@ -102,19 +92,21 @@ func TestRemoteQueriesActionRunsThroughLivePARLoopAndFakeintake(t *testing.T) { assert.Equal(t, "SUCCEEDED", result.Outputs["status"]) require.Contains(t, result.Outputs, "rows") - var bridgeRequest map[string]interface{} select { - case body := <-bridgeRequests: - require.NotContains(t, string(body), "password") - require.NotContains(t, string(body), "token") - require.NotContains(t, string(body), "secret") - require.NoError(t, json.Unmarshal(body, &bridgeRequest)) + case req := <-bridgeRequests: + requestEvidence, err := json.Marshal(req) + require.NoError(t, err) + require.NotContains(t, string(requestEvidence), "password") + require.NotContains(t, string(requestEvidence), "token") + require.NotContains(t, string(requestEvidence), "secret") + assert.Equal(t, "postgres", req.GetIntegration()) + assert.Equal(t, "SELECT 1 AS value", req.GetQuery()) + assert.Equal(t, "localhost", req.GetTarget().GetHost()) + assert.Equal(t, int32(5432), req.GetTarget().GetPort()) + assert.Equal(t, "postgres", req.GetTarget().GetDbname()) case <-time.After(2 * time.Second): - require.FailNow(t, "remote query action did not call the local bridge") + require.FailNow(t, "remote query action did not call the AgentSecure client") } - assert.Equal(t, "postgres", bridgeRequest["integration"]) - assert.Equal(t, "SELECT 1 AS value", bridgeRequest["query"]) - assert.Equal(t, map[string]interface{}{"host": "localhost", "port": float64(5432), "dbname": "postgres"}, bridgeRequest["target"]) dequeueCalls, err := fakeintakeClient.GetPARDequeueCount() require.NoError(t, err) @@ -152,24 +144,12 @@ func newLivePARTestConfig(t *testing.T, fakeintakeURL string) *parconfig.Config } } -type httpBridgeClient struct{} +type captureAgentSecureClient struct { + requests chan<- *pb.RemoteQueryExecuteRequest + response *pb.RemoteQueryExecuteResponse +} -func (httpBridgeClient) Post(url string, contentType string, body io.Reader, _ ...ipc.RequestOption) ([]byte, error) { - payload, err := io.ReadAll(body) - if err != nil { - return nil, err - } - resp, err := http.Post(url, contentType, bytes.NewReader(payload)) - if err != nil { - return nil, err - } - defer resp.Body.Close() - respBody, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - if resp.StatusCode >= 400 { - return respBody, assert.AnError - } - return respBody, nil +func (c *captureAgentSecureClient) RemoteQueryExecute(_ context.Context, req *pb.RemoteQueryExecuteRequest, _ ...grpc.CallOption) (*pb.RemoteQueryExecuteResponse, error) { + c.requests <- req + return c.response, nil } diff --git a/pkg/proto/datadog/api/v1/api.proto b/pkg/proto/datadog/api/v1/api.proto index 91882e0fc173..bc0856677cd8 100644 --- a/pkg/proto/datadog/api/v1/api.proto +++ b/pkg/proto/datadog/api/v1/api.proto @@ -10,6 +10,7 @@ import "datadog/workloadfilter/workloadfilter.proto"; import "datadog/autodiscovery/autodiscovery.proto"; import "datadog/kubemetadata/kubemetadata.proto"; import "google/protobuf/empty.proto"; +import "google/protobuf/struct.proto"; option go_package = "pkg/proto/pbgo/core"; // golang @@ -20,6 +21,39 @@ service Agent { rpc GetHostname (datadog.model.v1.HostnameRequest) returns (datadog.model.v1.HostnameReply); } +message RemoteQueryTarget { + string host = 1; + int32 port = 2; + string dbname = 3; +} + +message RemoteQueryExecuteLimits { + int32 max_rows = 1; + int32 max_bytes = 2; + int32 timeout_ms = 3; +} + +message RemoteQueryExecuteRequest { + string integration = 1; + RemoteQueryTarget target = 2; + string query = 3; + RemoteQueryExecuteLimits limits = 4; +} + +message RemoteQueryExecuteError { + string code = 1; + string message = 2; +} + +message RemoteQueryExecuteResponse { + string status = 1; + RemoteQueryExecuteError error = 2; + repeated google.protobuf.Struct columns = 3; + repeated google.protobuf.Struct rows = 4; + bool truncated = 5; + google.protobuf.Struct stats = 6; +} + service AgentSecure { // subscribes to added, removed, or changed entities in the Tagger // and streams them to clients as events. @@ -71,6 +105,9 @@ service AgentSecure { // Evaluates a workloadfilter rule on behalf of remote agents. rpc WorkloadFilterEvaluate(datadog.workloadfilter.WorkloadFilterEvaluateRequest) returns (datadog.workloadfilter.WorkloadFilterEvaluateResponse); + // Executes an Agent-local Remote Queries request through a matched integration check. + rpc RemoteQueryExecute(RemoteQueryExecuteRequest) returns (RemoteQueryExecuteResponse); + // Streams pod-to-service metadata for a specific node. rpc StreamKubeMetadata(datadog.kubemetadata.KubeMetadataStreamRequest) returns (stream datadog.kubemetadata.KubeMetadataStreamResponse); } diff --git a/pkg/proto/pbgo/core/api.pb.go b/pkg/proto/pbgo/core/api.pb.go index 8660eeab5541..a438caa58b51 100644 --- a/pkg/proto/pbgo/core/api.pb.go +++ b/pkg/proto/pbgo/core/api.pb.go @@ -10,7 +10,9 @@ import ( protoreflect "google.golang.org/protobuf/reflect/protoreflect" protoimpl "google.golang.org/protobuf/runtime/protoimpl" emptypb "google.golang.org/protobuf/types/known/emptypb" + structpb "google.golang.org/protobuf/types/known/structpb" reflect "reflect" + sync "sync" unsafe "unsafe" ) @@ -21,13 +23,361 @@ const ( _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) ) +type RemoteQueryTarget struct { + state protoimpl.MessageState `protogen:"open.v1"` + Host string `protobuf:"bytes,1,opt,name=host,proto3" json:"host,omitempty"` + Port int32 `protobuf:"varint,2,opt,name=port,proto3" json:"port,omitempty"` + Dbname string `protobuf:"bytes,3,opt,name=dbname,proto3" json:"dbname,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RemoteQueryTarget) Reset() { + *x = RemoteQueryTarget{} + mi := &file_datadog_api_v1_api_proto_msgTypes[0] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RemoteQueryTarget) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RemoteQueryTarget) ProtoMessage() {} + +func (x *RemoteQueryTarget) ProtoReflect() protoreflect.Message { + mi := &file_datadog_api_v1_api_proto_msgTypes[0] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RemoteQueryTarget.ProtoReflect.Descriptor instead. +func (*RemoteQueryTarget) Descriptor() ([]byte, []int) { + return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{0} +} + +func (x *RemoteQueryTarget) GetHost() string { + if x != nil { + return x.Host + } + return "" +} + +func (x *RemoteQueryTarget) GetPort() int32 { + if x != nil { + return x.Port + } + return 0 +} + +func (x *RemoteQueryTarget) GetDbname() string { + if x != nil { + return x.Dbname + } + return "" +} + +type RemoteQueryExecuteLimits struct { + state protoimpl.MessageState `protogen:"open.v1"` + MaxRows int32 `protobuf:"varint,1,opt,name=max_rows,json=maxRows,proto3" json:"max_rows,omitempty"` + MaxBytes int32 `protobuf:"varint,2,opt,name=max_bytes,json=maxBytes,proto3" json:"max_bytes,omitempty"` + TimeoutMs int32 `protobuf:"varint,3,opt,name=timeout_ms,json=timeoutMs,proto3" json:"timeout_ms,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RemoteQueryExecuteLimits) Reset() { + *x = RemoteQueryExecuteLimits{} + mi := &file_datadog_api_v1_api_proto_msgTypes[1] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RemoteQueryExecuteLimits) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RemoteQueryExecuteLimits) ProtoMessage() {} + +func (x *RemoteQueryExecuteLimits) ProtoReflect() protoreflect.Message { + mi := &file_datadog_api_v1_api_proto_msgTypes[1] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RemoteQueryExecuteLimits.ProtoReflect.Descriptor instead. +func (*RemoteQueryExecuteLimits) Descriptor() ([]byte, []int) { + return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{1} +} + +func (x *RemoteQueryExecuteLimits) GetMaxRows() int32 { + if x != nil { + return x.MaxRows + } + return 0 +} + +func (x *RemoteQueryExecuteLimits) GetMaxBytes() int32 { + if x != nil { + return x.MaxBytes + } + return 0 +} + +func (x *RemoteQueryExecuteLimits) GetTimeoutMs() int32 { + if x != nil { + return x.TimeoutMs + } + return 0 +} + +type RemoteQueryExecuteRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + Integration string `protobuf:"bytes,1,opt,name=integration,proto3" json:"integration,omitempty"` + Target *RemoteQueryTarget `protobuf:"bytes,2,opt,name=target,proto3" json:"target,omitempty"` + Query string `protobuf:"bytes,3,opt,name=query,proto3" json:"query,omitempty"` + Limits *RemoteQueryExecuteLimits `protobuf:"bytes,4,opt,name=limits,proto3" json:"limits,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RemoteQueryExecuteRequest) Reset() { + *x = RemoteQueryExecuteRequest{} + mi := &file_datadog_api_v1_api_proto_msgTypes[2] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RemoteQueryExecuteRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RemoteQueryExecuteRequest) ProtoMessage() {} + +func (x *RemoteQueryExecuteRequest) ProtoReflect() protoreflect.Message { + mi := &file_datadog_api_v1_api_proto_msgTypes[2] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RemoteQueryExecuteRequest.ProtoReflect.Descriptor instead. +func (*RemoteQueryExecuteRequest) Descriptor() ([]byte, []int) { + return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{2} +} + +func (x *RemoteQueryExecuteRequest) GetIntegration() string { + if x != nil { + return x.Integration + } + return "" +} + +func (x *RemoteQueryExecuteRequest) GetTarget() *RemoteQueryTarget { + if x != nil { + return x.Target + } + return nil +} + +func (x *RemoteQueryExecuteRequest) GetQuery() string { + if x != nil { + return x.Query + } + return "" +} + +func (x *RemoteQueryExecuteRequest) GetLimits() *RemoteQueryExecuteLimits { + if x != nil { + return x.Limits + } + return nil +} + +type RemoteQueryExecuteError struct { + state protoimpl.MessageState `protogen:"open.v1"` + Code string `protobuf:"bytes,1,opt,name=code,proto3" json:"code,omitempty"` + Message string `protobuf:"bytes,2,opt,name=message,proto3" json:"message,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RemoteQueryExecuteError) Reset() { + *x = RemoteQueryExecuteError{} + mi := &file_datadog_api_v1_api_proto_msgTypes[3] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RemoteQueryExecuteError) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RemoteQueryExecuteError) ProtoMessage() {} + +func (x *RemoteQueryExecuteError) ProtoReflect() protoreflect.Message { + mi := &file_datadog_api_v1_api_proto_msgTypes[3] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RemoteQueryExecuteError.ProtoReflect.Descriptor instead. +func (*RemoteQueryExecuteError) Descriptor() ([]byte, []int) { + return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{3} +} + +func (x *RemoteQueryExecuteError) GetCode() string { + if x != nil { + return x.Code + } + return "" +} + +func (x *RemoteQueryExecuteError) GetMessage() string { + if x != nil { + return x.Message + } + return "" +} + +type RemoteQueryExecuteResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Status string `protobuf:"bytes,1,opt,name=status,proto3" json:"status,omitempty"` + Error *RemoteQueryExecuteError `protobuf:"bytes,2,opt,name=error,proto3" json:"error,omitempty"` + Columns []*structpb.Struct `protobuf:"bytes,3,rep,name=columns,proto3" json:"columns,omitempty"` + Rows []*structpb.Struct `protobuf:"bytes,4,rep,name=rows,proto3" json:"rows,omitempty"` + Truncated bool `protobuf:"varint,5,opt,name=truncated,proto3" json:"truncated,omitempty"` + Stats *structpb.Struct `protobuf:"bytes,6,opt,name=stats,proto3" json:"stats,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RemoteQueryExecuteResponse) Reset() { + *x = RemoteQueryExecuteResponse{} + mi := &file_datadog_api_v1_api_proto_msgTypes[4] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RemoteQueryExecuteResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RemoteQueryExecuteResponse) ProtoMessage() {} + +func (x *RemoteQueryExecuteResponse) ProtoReflect() protoreflect.Message { + mi := &file_datadog_api_v1_api_proto_msgTypes[4] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RemoteQueryExecuteResponse.ProtoReflect.Descriptor instead. +func (*RemoteQueryExecuteResponse) Descriptor() ([]byte, []int) { + return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{4} +} + +func (x *RemoteQueryExecuteResponse) GetStatus() string { + if x != nil { + return x.Status + } + return "" +} + +func (x *RemoteQueryExecuteResponse) GetError() *RemoteQueryExecuteError { + if x != nil { + return x.Error + } + return nil +} + +func (x *RemoteQueryExecuteResponse) GetColumns() []*structpb.Struct { + if x != nil { + return x.Columns + } + return nil +} + +func (x *RemoteQueryExecuteResponse) GetRows() []*structpb.Struct { + if x != nil { + return x.Rows + } + return nil +} + +func (x *RemoteQueryExecuteResponse) GetTruncated() bool { + if x != nil { + return x.Truncated + } + return false +} + +func (x *RemoteQueryExecuteResponse) GetStats() *structpb.Struct { + if x != nil { + return x.Stats + } + return nil +} + var File_datadog_api_v1_api_proto protoreflect.FileDescriptor const file_datadog_api_v1_api_proto_rawDesc = "" + "\n" + - "\x18datadog/api/v1/api.proto\x12\x0edatadog.api.v1\x1a\x1cdatadog/model/v1/model.proto\x1a%datadog/remoteagent/remoteagent.proto\x1a'datadog/remoteconfig/remoteconfig.proto\x1a'datadog/workloadmeta/workloadmeta.proto\x1a+datadog/workloadfilter/workloadfilter.proto\x1a)datadog/autodiscovery/autodiscovery.proto\x1a'datadog/kubemetadata/kubemetadata.proto\x1a\x1bgoogle/protobuf/empty.proto2Z\n" + + "\x18datadog/api/v1/api.proto\x12\x0edatadog.api.v1\x1a\x1cdatadog/model/v1/model.proto\x1a%datadog/remoteagent/remoteagent.proto\x1a'datadog/remoteconfig/remoteconfig.proto\x1a'datadog/workloadmeta/workloadmeta.proto\x1a+datadog/workloadfilter/workloadfilter.proto\x1a)datadog/autodiscovery/autodiscovery.proto\x1a'datadog/kubemetadata/kubemetadata.proto\x1a\x1bgoogle/protobuf/empty.proto\x1a\x1cgoogle/protobuf/struct.proto\"S\n" + + "\x11RemoteQueryTarget\x12\x12\n" + + "\x04host\x18\x01 \x01(\tR\x04host\x12\x12\n" + + "\x04port\x18\x02 \x01(\x05R\x04port\x12\x16\n" + + "\x06dbname\x18\x03 \x01(\tR\x06dbname\"q\n" + + "\x18RemoteQueryExecuteLimits\x12\x19\n" + + "\bmax_rows\x18\x01 \x01(\x05R\amaxRows\x12\x1b\n" + + "\tmax_bytes\x18\x02 \x01(\x05R\bmaxBytes\x12\x1d\n" + + "\n" + + "timeout_ms\x18\x03 \x01(\x05R\ttimeoutMs\"\xd0\x01\n" + + "\x19RemoteQueryExecuteRequest\x12 \n" + + "\vintegration\x18\x01 \x01(\tR\vintegration\x129\n" + + "\x06target\x18\x02 \x01(\v2!.datadog.api.v1.RemoteQueryTargetR\x06target\x12\x14\n" + + "\x05query\x18\x03 \x01(\tR\x05query\x12@\n" + + "\x06limits\x18\x04 \x01(\v2(.datadog.api.v1.RemoteQueryExecuteLimitsR\x06limits\"G\n" + + "\x17RemoteQueryExecuteError\x12\x12\n" + + "\x04code\x18\x01 \x01(\tR\x04code\x12\x18\n" + + "\amessage\x18\x02 \x01(\tR\amessage\"\xa0\x02\n" + + "\x1aRemoteQueryExecuteResponse\x12\x16\n" + + "\x06status\x18\x01 \x01(\tR\x06status\x12=\n" + + "\x05error\x18\x02 \x01(\v2'.datadog.api.v1.RemoteQueryExecuteErrorR\x05error\x121\n" + + "\acolumns\x18\x03 \x03(\v2\x17.google.protobuf.StructR\acolumns\x12+\n" + + "\x04rows\x18\x04 \x03(\v2\x17.google.protobuf.StructR\x04rows\x12\x1c\n" + + "\ttruncated\x18\x05 \x01(\bR\ttruncated\x12-\n" + + "\x05stats\x18\x06 \x01(\v2\x17.google.protobuf.StructR\x05stats2Z\n" + "\x05Agent\x12Q\n" + - "\vGetHostname\x12!.datadog.model.v1.HostnameRequest\x1a\x1f.datadog.model.v1.HostnameReply2\xab\x10\n" + + "\vGetHostname\x12!.datadog.model.v1.HostnameRequest\x1a\x1f.datadog.model.v1.HostnameReply2\x98\x11\n" + "\vAgentSecure\x12c\n" + "\x14TaggerStreamEntities\x12#.datadog.model.v1.StreamTagsRequest\x1a$.datadog.model.v1.StreamTagsResponse0\x01\x12\xa2\x01\n" + "'TaggerGenerateContainerIDFromOriginInfo\x12:.datadog.model.v1.GenerateContainerIDFromOriginInfoRequest\x1a;.datadog.model.v1.GenerateContainerIDFromOriginInfoResponse\x12`\n" + @@ -46,91 +396,119 @@ const file_datadog_api_v1_api_proto_rawDesc = "" + "\x19AutodiscoveryStreamConfig\x12\x16.google.protobuf.Empty\x1a2.datadog.autodiscovery.AutodiscoveryStreamResponse0\x01\x12O\n" + "\vGetHostTags\x12 .datadog.model.v1.HostTagRequest\x1a\x1e.datadog.model.v1.HostTagReply\x12\\\n" + "\x12StreamConfigEvents\x12%.datadog.model.v1.ConfigStreamRequest\x1a\x1d.datadog.model.v1.ConfigEvent0\x01\x12\x87\x01\n" + - "\x16WorkloadFilterEvaluate\x125.datadog.workloadfilter.WorkloadFilterEvaluateRequest\x1a6.datadog.workloadfilter.WorkloadFilterEvaluateResponse\x12y\n" + + "\x16WorkloadFilterEvaluate\x125.datadog.workloadfilter.WorkloadFilterEvaluateRequest\x1a6.datadog.workloadfilter.WorkloadFilterEvaluateResponse\x12k\n" + + "\x12RemoteQueryExecute\x12).datadog.api.v1.RemoteQueryExecuteRequest\x1a*.datadog.api.v1.RemoteQueryExecuteResponse\x12y\n" + "\x12StreamKubeMetadata\x12/.datadog.kubemetadata.KubeMetadataStreamRequest\x1a0.datadog.kubemetadata.KubeMetadataStreamResponse0\x01B\x15Z\x13pkg/proto/pbgo/coreb\x06proto3" +var ( + file_datadog_api_v1_api_proto_rawDescOnce sync.Once + file_datadog_api_v1_api_proto_rawDescData []byte +) + +func file_datadog_api_v1_api_proto_rawDescGZIP() []byte { + file_datadog_api_v1_api_proto_rawDescOnce.Do(func() { + file_datadog_api_v1_api_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_datadog_api_v1_api_proto_rawDesc), len(file_datadog_api_v1_api_proto_rawDesc))) + }) + return file_datadog_api_v1_api_proto_rawDescData +} + +var file_datadog_api_v1_api_proto_msgTypes = make([]protoimpl.MessageInfo, 5) var file_datadog_api_v1_api_proto_goTypes = []any{ - (*HostnameRequest)(nil), // 0: datadog.model.v1.HostnameRequest - (*StreamTagsRequest)(nil), // 1: datadog.model.v1.StreamTagsRequest - (*GenerateContainerIDFromOriginInfoRequest)(nil), // 2: datadog.model.v1.GenerateContainerIDFromOriginInfoRequest - (*FetchEntityRequest)(nil), // 3: datadog.model.v1.FetchEntityRequest - (*CaptureTriggerRequest)(nil), // 4: datadog.model.v1.CaptureTriggerRequest - (*TaggerState)(nil), // 5: datadog.model.v1.TaggerState - (*ClientGetConfigsRequest)(nil), // 6: datadog.config.ClientGetConfigsRequest - (*emptypb.Empty)(nil), // 7: google.protobuf.Empty - (*ConfigSubscriptionRequest)(nil), // 8: datadog.config.ConfigSubscriptionRequest - (*WorkloadmetaStreamRequest)(nil), // 9: datadog.workloadmeta.WorkloadmetaStreamRequest - (*RegisterRemoteAgentRequest)(nil), // 10: datadog.remoteagent.v1.RegisterRemoteAgentRequest - (*RefreshRemoteAgentRequest)(nil), // 11: datadog.remoteagent.v1.RefreshRemoteAgentRequest - (*HostTagRequest)(nil), // 12: datadog.model.v1.HostTagRequest - (*ConfigStreamRequest)(nil), // 13: datadog.model.v1.ConfigStreamRequest - (*WorkloadFilterEvaluateRequest)(nil), // 14: datadog.workloadfilter.WorkloadFilterEvaluateRequest - (*KubeMetadataStreamRequest)(nil), // 15: datadog.kubemetadata.KubeMetadataStreamRequest - (*HostnameReply)(nil), // 16: datadog.model.v1.HostnameReply - (*StreamTagsResponse)(nil), // 17: datadog.model.v1.StreamTagsResponse - (*GenerateContainerIDFromOriginInfoResponse)(nil), // 18: datadog.model.v1.GenerateContainerIDFromOriginInfoResponse - (*FetchEntityResponse)(nil), // 19: datadog.model.v1.FetchEntityResponse - (*CaptureTriggerResponse)(nil), // 20: datadog.model.v1.CaptureTriggerResponse - (*TaggerStateResponse)(nil), // 21: datadog.model.v1.TaggerStateResponse - (*ClientGetConfigsResponse)(nil), // 22: datadog.config.ClientGetConfigsResponse - (*GetStateConfigResponse)(nil), // 23: datadog.config.GetStateConfigResponse - (*ConfigSubscriptionResponse)(nil), // 24: datadog.config.ConfigSubscriptionResponse - (*ResetStateConfigResponse)(nil), // 25: datadog.config.ResetStateConfigResponse - (*WorkloadmetaStreamResponse)(nil), // 26: datadog.workloadmeta.WorkloadmetaStreamResponse - (*RegisterRemoteAgentResponse)(nil), // 27: datadog.remoteagent.v1.RegisterRemoteAgentResponse - (*RefreshRemoteAgentResponse)(nil), // 28: datadog.remoteagent.v1.RefreshRemoteAgentResponse - (*AutodiscoveryStreamResponse)(nil), // 29: datadog.autodiscovery.AutodiscoveryStreamResponse - (*HostTagReply)(nil), // 30: datadog.model.v1.HostTagReply - (*ConfigEvent)(nil), // 31: datadog.model.v1.ConfigEvent - (*WorkloadFilterEvaluateResponse)(nil), // 32: datadog.workloadfilter.WorkloadFilterEvaluateResponse - (*KubeMetadataStreamResponse)(nil), // 33: datadog.kubemetadata.KubeMetadataStreamResponse + (*RemoteQueryTarget)(nil), // 0: datadog.api.v1.RemoteQueryTarget + (*RemoteQueryExecuteLimits)(nil), // 1: datadog.api.v1.RemoteQueryExecuteLimits + (*RemoteQueryExecuteRequest)(nil), // 2: datadog.api.v1.RemoteQueryExecuteRequest + (*RemoteQueryExecuteError)(nil), // 3: datadog.api.v1.RemoteQueryExecuteError + (*RemoteQueryExecuteResponse)(nil), // 4: datadog.api.v1.RemoteQueryExecuteResponse + (*structpb.Struct)(nil), // 5: google.protobuf.Struct + (*HostnameRequest)(nil), // 6: datadog.model.v1.HostnameRequest + (*StreamTagsRequest)(nil), // 7: datadog.model.v1.StreamTagsRequest + (*GenerateContainerIDFromOriginInfoRequest)(nil), // 8: datadog.model.v1.GenerateContainerIDFromOriginInfoRequest + (*FetchEntityRequest)(nil), // 9: datadog.model.v1.FetchEntityRequest + (*CaptureTriggerRequest)(nil), // 10: datadog.model.v1.CaptureTriggerRequest + (*TaggerState)(nil), // 11: datadog.model.v1.TaggerState + (*ClientGetConfigsRequest)(nil), // 12: datadog.config.ClientGetConfigsRequest + (*emptypb.Empty)(nil), // 13: google.protobuf.Empty + (*ConfigSubscriptionRequest)(nil), // 14: datadog.config.ConfigSubscriptionRequest + (*WorkloadmetaStreamRequest)(nil), // 15: datadog.workloadmeta.WorkloadmetaStreamRequest + (*RegisterRemoteAgentRequest)(nil), // 16: datadog.remoteagent.v1.RegisterRemoteAgentRequest + (*RefreshRemoteAgentRequest)(nil), // 17: datadog.remoteagent.v1.RefreshRemoteAgentRequest + (*HostTagRequest)(nil), // 18: datadog.model.v1.HostTagRequest + (*ConfigStreamRequest)(nil), // 19: datadog.model.v1.ConfigStreamRequest + (*WorkloadFilterEvaluateRequest)(nil), // 20: datadog.workloadfilter.WorkloadFilterEvaluateRequest + (*KubeMetadataStreamRequest)(nil), // 21: datadog.kubemetadata.KubeMetadataStreamRequest + (*HostnameReply)(nil), // 22: datadog.model.v1.HostnameReply + (*StreamTagsResponse)(nil), // 23: datadog.model.v1.StreamTagsResponse + (*GenerateContainerIDFromOriginInfoResponse)(nil), // 24: datadog.model.v1.GenerateContainerIDFromOriginInfoResponse + (*FetchEntityResponse)(nil), // 25: datadog.model.v1.FetchEntityResponse + (*CaptureTriggerResponse)(nil), // 26: datadog.model.v1.CaptureTriggerResponse + (*TaggerStateResponse)(nil), // 27: datadog.model.v1.TaggerStateResponse + (*ClientGetConfigsResponse)(nil), // 28: datadog.config.ClientGetConfigsResponse + (*GetStateConfigResponse)(nil), // 29: datadog.config.GetStateConfigResponse + (*ConfigSubscriptionResponse)(nil), // 30: datadog.config.ConfigSubscriptionResponse + (*ResetStateConfigResponse)(nil), // 31: datadog.config.ResetStateConfigResponse + (*WorkloadmetaStreamResponse)(nil), // 32: datadog.workloadmeta.WorkloadmetaStreamResponse + (*RegisterRemoteAgentResponse)(nil), // 33: datadog.remoteagent.v1.RegisterRemoteAgentResponse + (*RefreshRemoteAgentResponse)(nil), // 34: datadog.remoteagent.v1.RefreshRemoteAgentResponse + (*AutodiscoveryStreamResponse)(nil), // 35: datadog.autodiscovery.AutodiscoveryStreamResponse + (*HostTagReply)(nil), // 36: datadog.model.v1.HostTagReply + (*ConfigEvent)(nil), // 37: datadog.model.v1.ConfigEvent + (*WorkloadFilterEvaluateResponse)(nil), // 38: datadog.workloadfilter.WorkloadFilterEvaluateResponse + (*KubeMetadataStreamResponse)(nil), // 39: datadog.kubemetadata.KubeMetadataStreamResponse } var file_datadog_api_v1_api_proto_depIdxs = []int32{ - 0, // 0: datadog.api.v1.Agent.GetHostname:input_type -> datadog.model.v1.HostnameRequest - 1, // 1: datadog.api.v1.AgentSecure.TaggerStreamEntities:input_type -> datadog.model.v1.StreamTagsRequest - 2, // 2: datadog.api.v1.AgentSecure.TaggerGenerateContainerIDFromOriginInfo:input_type -> datadog.model.v1.GenerateContainerIDFromOriginInfoRequest - 3, // 3: datadog.api.v1.AgentSecure.TaggerFetchEntity:input_type -> datadog.model.v1.FetchEntityRequest - 4, // 4: datadog.api.v1.AgentSecure.DogstatsdCaptureTrigger:input_type -> datadog.model.v1.CaptureTriggerRequest - 5, // 5: datadog.api.v1.AgentSecure.DogstatsdSetTaggerState:input_type -> datadog.model.v1.TaggerState - 6, // 6: datadog.api.v1.AgentSecure.ClientGetConfigs:input_type -> datadog.config.ClientGetConfigsRequest - 7, // 7: datadog.api.v1.AgentSecure.GetConfigState:input_type -> google.protobuf.Empty - 6, // 8: datadog.api.v1.AgentSecure.ClientGetConfigsHA:input_type -> datadog.config.ClientGetConfigsRequest - 7, // 9: datadog.api.v1.AgentSecure.GetConfigStateHA:input_type -> google.protobuf.Empty - 8, // 10: datadog.api.v1.AgentSecure.CreateConfigSubscription:input_type -> datadog.config.ConfigSubscriptionRequest - 7, // 11: datadog.api.v1.AgentSecure.ResetConfigState:input_type -> google.protobuf.Empty - 9, // 12: datadog.api.v1.AgentSecure.WorkloadmetaStreamEntities:input_type -> datadog.workloadmeta.WorkloadmetaStreamRequest - 10, // 13: datadog.api.v1.AgentSecure.RegisterRemoteAgent:input_type -> datadog.remoteagent.v1.RegisterRemoteAgentRequest - 11, // 14: datadog.api.v1.AgentSecure.RefreshRemoteAgent:input_type -> datadog.remoteagent.v1.RefreshRemoteAgentRequest - 7, // 15: datadog.api.v1.AgentSecure.AutodiscoveryStreamConfig:input_type -> google.protobuf.Empty - 12, // 16: datadog.api.v1.AgentSecure.GetHostTags:input_type -> datadog.model.v1.HostTagRequest - 13, // 17: datadog.api.v1.AgentSecure.StreamConfigEvents:input_type -> datadog.model.v1.ConfigStreamRequest - 14, // 18: datadog.api.v1.AgentSecure.WorkloadFilterEvaluate:input_type -> datadog.workloadfilter.WorkloadFilterEvaluateRequest - 15, // 19: datadog.api.v1.AgentSecure.StreamKubeMetadata:input_type -> datadog.kubemetadata.KubeMetadataStreamRequest - 16, // 20: datadog.api.v1.Agent.GetHostname:output_type -> datadog.model.v1.HostnameReply - 17, // 21: datadog.api.v1.AgentSecure.TaggerStreamEntities:output_type -> datadog.model.v1.StreamTagsResponse - 18, // 22: datadog.api.v1.AgentSecure.TaggerGenerateContainerIDFromOriginInfo:output_type -> datadog.model.v1.GenerateContainerIDFromOriginInfoResponse - 19, // 23: datadog.api.v1.AgentSecure.TaggerFetchEntity:output_type -> datadog.model.v1.FetchEntityResponse - 20, // 24: datadog.api.v1.AgentSecure.DogstatsdCaptureTrigger:output_type -> datadog.model.v1.CaptureTriggerResponse - 21, // 25: datadog.api.v1.AgentSecure.DogstatsdSetTaggerState:output_type -> datadog.model.v1.TaggerStateResponse - 22, // 26: datadog.api.v1.AgentSecure.ClientGetConfigs:output_type -> datadog.config.ClientGetConfigsResponse - 23, // 27: datadog.api.v1.AgentSecure.GetConfigState:output_type -> datadog.config.GetStateConfigResponse - 22, // 28: datadog.api.v1.AgentSecure.ClientGetConfigsHA:output_type -> datadog.config.ClientGetConfigsResponse - 23, // 29: datadog.api.v1.AgentSecure.GetConfigStateHA:output_type -> datadog.config.GetStateConfigResponse - 24, // 30: datadog.api.v1.AgentSecure.CreateConfigSubscription:output_type -> datadog.config.ConfigSubscriptionResponse - 25, // 31: datadog.api.v1.AgentSecure.ResetConfigState:output_type -> datadog.config.ResetStateConfigResponse - 26, // 32: datadog.api.v1.AgentSecure.WorkloadmetaStreamEntities:output_type -> datadog.workloadmeta.WorkloadmetaStreamResponse - 27, // 33: datadog.api.v1.AgentSecure.RegisterRemoteAgent:output_type -> datadog.remoteagent.v1.RegisterRemoteAgentResponse - 28, // 34: datadog.api.v1.AgentSecure.RefreshRemoteAgent:output_type -> datadog.remoteagent.v1.RefreshRemoteAgentResponse - 29, // 35: datadog.api.v1.AgentSecure.AutodiscoveryStreamConfig:output_type -> datadog.autodiscovery.AutodiscoveryStreamResponse - 30, // 36: datadog.api.v1.AgentSecure.GetHostTags:output_type -> datadog.model.v1.HostTagReply - 31, // 37: datadog.api.v1.AgentSecure.StreamConfigEvents:output_type -> datadog.model.v1.ConfigEvent - 32, // 38: datadog.api.v1.AgentSecure.WorkloadFilterEvaluate:output_type -> datadog.workloadfilter.WorkloadFilterEvaluateResponse - 33, // 39: datadog.api.v1.AgentSecure.StreamKubeMetadata:output_type -> datadog.kubemetadata.KubeMetadataStreamResponse - 20, // [20:40] is the sub-list for method output_type - 0, // [0:20] is the sub-list for method input_type - 0, // [0:0] is the sub-list for extension type_name - 0, // [0:0] is the sub-list for extension extendee - 0, // [0:0] is the sub-list for field type_name + 0, // 0: datadog.api.v1.RemoteQueryExecuteRequest.target:type_name -> datadog.api.v1.RemoteQueryTarget + 1, // 1: datadog.api.v1.RemoteQueryExecuteRequest.limits:type_name -> datadog.api.v1.RemoteQueryExecuteLimits + 3, // 2: datadog.api.v1.RemoteQueryExecuteResponse.error:type_name -> datadog.api.v1.RemoteQueryExecuteError + 5, // 3: datadog.api.v1.RemoteQueryExecuteResponse.columns:type_name -> google.protobuf.Struct + 5, // 4: datadog.api.v1.RemoteQueryExecuteResponse.rows:type_name -> google.protobuf.Struct + 5, // 5: datadog.api.v1.RemoteQueryExecuteResponse.stats:type_name -> google.protobuf.Struct + 6, // 6: datadog.api.v1.Agent.GetHostname:input_type -> datadog.model.v1.HostnameRequest + 7, // 7: datadog.api.v1.AgentSecure.TaggerStreamEntities:input_type -> datadog.model.v1.StreamTagsRequest + 8, // 8: datadog.api.v1.AgentSecure.TaggerGenerateContainerIDFromOriginInfo:input_type -> datadog.model.v1.GenerateContainerIDFromOriginInfoRequest + 9, // 9: datadog.api.v1.AgentSecure.TaggerFetchEntity:input_type -> datadog.model.v1.FetchEntityRequest + 10, // 10: datadog.api.v1.AgentSecure.DogstatsdCaptureTrigger:input_type -> datadog.model.v1.CaptureTriggerRequest + 11, // 11: datadog.api.v1.AgentSecure.DogstatsdSetTaggerState:input_type -> datadog.model.v1.TaggerState + 12, // 12: datadog.api.v1.AgentSecure.ClientGetConfigs:input_type -> datadog.config.ClientGetConfigsRequest + 13, // 13: datadog.api.v1.AgentSecure.GetConfigState:input_type -> google.protobuf.Empty + 12, // 14: datadog.api.v1.AgentSecure.ClientGetConfigsHA:input_type -> datadog.config.ClientGetConfigsRequest + 13, // 15: datadog.api.v1.AgentSecure.GetConfigStateHA:input_type -> google.protobuf.Empty + 14, // 16: datadog.api.v1.AgentSecure.CreateConfigSubscription:input_type -> datadog.config.ConfigSubscriptionRequest + 13, // 17: datadog.api.v1.AgentSecure.ResetConfigState:input_type -> google.protobuf.Empty + 15, // 18: datadog.api.v1.AgentSecure.WorkloadmetaStreamEntities:input_type -> datadog.workloadmeta.WorkloadmetaStreamRequest + 16, // 19: datadog.api.v1.AgentSecure.RegisterRemoteAgent:input_type -> datadog.remoteagent.v1.RegisterRemoteAgentRequest + 17, // 20: datadog.api.v1.AgentSecure.RefreshRemoteAgent:input_type -> datadog.remoteagent.v1.RefreshRemoteAgentRequest + 13, // 21: datadog.api.v1.AgentSecure.AutodiscoveryStreamConfig:input_type -> google.protobuf.Empty + 18, // 22: datadog.api.v1.AgentSecure.GetHostTags:input_type -> datadog.model.v1.HostTagRequest + 19, // 23: datadog.api.v1.AgentSecure.StreamConfigEvents:input_type -> datadog.model.v1.ConfigStreamRequest + 20, // 24: datadog.api.v1.AgentSecure.WorkloadFilterEvaluate:input_type -> datadog.workloadfilter.WorkloadFilterEvaluateRequest + 2, // 25: datadog.api.v1.AgentSecure.RemoteQueryExecute:input_type -> datadog.api.v1.RemoteQueryExecuteRequest + 21, // 26: datadog.api.v1.AgentSecure.StreamKubeMetadata:input_type -> datadog.kubemetadata.KubeMetadataStreamRequest + 22, // 27: datadog.api.v1.Agent.GetHostname:output_type -> datadog.model.v1.HostnameReply + 23, // 28: datadog.api.v1.AgentSecure.TaggerStreamEntities:output_type -> datadog.model.v1.StreamTagsResponse + 24, // 29: datadog.api.v1.AgentSecure.TaggerGenerateContainerIDFromOriginInfo:output_type -> datadog.model.v1.GenerateContainerIDFromOriginInfoResponse + 25, // 30: datadog.api.v1.AgentSecure.TaggerFetchEntity:output_type -> datadog.model.v1.FetchEntityResponse + 26, // 31: datadog.api.v1.AgentSecure.DogstatsdCaptureTrigger:output_type -> datadog.model.v1.CaptureTriggerResponse + 27, // 32: datadog.api.v1.AgentSecure.DogstatsdSetTaggerState:output_type -> datadog.model.v1.TaggerStateResponse + 28, // 33: datadog.api.v1.AgentSecure.ClientGetConfigs:output_type -> datadog.config.ClientGetConfigsResponse + 29, // 34: datadog.api.v1.AgentSecure.GetConfigState:output_type -> datadog.config.GetStateConfigResponse + 28, // 35: datadog.api.v1.AgentSecure.ClientGetConfigsHA:output_type -> datadog.config.ClientGetConfigsResponse + 29, // 36: datadog.api.v1.AgentSecure.GetConfigStateHA:output_type -> datadog.config.GetStateConfigResponse + 30, // 37: datadog.api.v1.AgentSecure.CreateConfigSubscription:output_type -> datadog.config.ConfigSubscriptionResponse + 31, // 38: datadog.api.v1.AgentSecure.ResetConfigState:output_type -> datadog.config.ResetStateConfigResponse + 32, // 39: datadog.api.v1.AgentSecure.WorkloadmetaStreamEntities:output_type -> datadog.workloadmeta.WorkloadmetaStreamResponse + 33, // 40: datadog.api.v1.AgentSecure.RegisterRemoteAgent:output_type -> datadog.remoteagent.v1.RegisterRemoteAgentResponse + 34, // 41: datadog.api.v1.AgentSecure.RefreshRemoteAgent:output_type -> datadog.remoteagent.v1.RefreshRemoteAgentResponse + 35, // 42: datadog.api.v1.AgentSecure.AutodiscoveryStreamConfig:output_type -> datadog.autodiscovery.AutodiscoveryStreamResponse + 36, // 43: datadog.api.v1.AgentSecure.GetHostTags:output_type -> datadog.model.v1.HostTagReply + 37, // 44: datadog.api.v1.AgentSecure.StreamConfigEvents:output_type -> datadog.model.v1.ConfigEvent + 38, // 45: datadog.api.v1.AgentSecure.WorkloadFilterEvaluate:output_type -> datadog.workloadfilter.WorkloadFilterEvaluateResponse + 4, // 46: datadog.api.v1.AgentSecure.RemoteQueryExecute:output_type -> datadog.api.v1.RemoteQueryExecuteResponse + 39, // 47: datadog.api.v1.AgentSecure.StreamKubeMetadata:output_type -> datadog.kubemetadata.KubeMetadataStreamResponse + 27, // [27:48] is the sub-list for method output_type + 6, // [6:27] is the sub-list for method input_type + 6, // [6:6] is the sub-list for extension type_name + 6, // [6:6] is the sub-list for extension extendee + 0, // [0:6] is the sub-list for field type_name } func init() { file_datadog_api_v1_api_proto_init() } @@ -151,12 +529,13 @@ func file_datadog_api_v1_api_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_datadog_api_v1_api_proto_rawDesc), len(file_datadog_api_v1_api_proto_rawDesc)), NumEnums: 0, - NumMessages: 0, + NumMessages: 5, NumExtensions: 0, NumServices: 2, }, GoTypes: file_datadog_api_v1_api_proto_goTypes, DependencyIndexes: file_datadog_api_v1_api_proto_depIdxs, + MessageInfos: file_datadog_api_v1_api_proto_msgTypes, }.Build() File_datadog_api_v1_api_proto = out.File file_datadog_api_v1_api_proto_goTypes = nil diff --git a/pkg/proto/pbgo/core/api_grpc.pb.go b/pkg/proto/pbgo/core/api_grpc.pb.go index 6f7d4ee8f419..891713c46a26 100644 --- a/pkg/proto/pbgo/core/api_grpc.pb.go +++ b/pkg/proto/pbgo/core/api_grpc.pb.go @@ -146,6 +146,7 @@ const ( AgentSecure_GetHostTags_FullMethodName = "/datadog.api.v1.AgentSecure/GetHostTags" AgentSecure_StreamConfigEvents_FullMethodName = "/datadog.api.v1.AgentSecure/StreamConfigEvents" AgentSecure_WorkloadFilterEvaluate_FullMethodName = "/datadog.api.v1.AgentSecure/WorkloadFilterEvaluate" + AgentSecure_RemoteQueryExecute_FullMethodName = "/datadog.api.v1.AgentSecure/RemoteQueryExecute" AgentSecure_StreamKubeMetadata_FullMethodName = "/datadog.api.v1.AgentSecure/StreamKubeMetadata" ) @@ -185,6 +186,8 @@ type AgentSecureClient interface { StreamConfigEvents(ctx context.Context, in *ConfigStreamRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[ConfigEvent], error) // Evaluates a workloadfilter rule on behalf of remote agents. WorkloadFilterEvaluate(ctx context.Context, in *WorkloadFilterEvaluateRequest, opts ...grpc.CallOption) (*WorkloadFilterEvaluateResponse, error) + // Executes an Agent-local Remote Queries request through a matched integration check. + RemoteQueryExecute(ctx context.Context, in *RemoteQueryExecuteRequest, opts ...grpc.CallOption) (*RemoteQueryExecuteResponse, error) // Streams pod-to-service metadata for a specific node. StreamKubeMetadata(ctx context.Context, in *KubeMetadataStreamRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[KubeMetadataStreamResponse], error) } @@ -416,6 +419,16 @@ func (c *agentSecureClient) WorkloadFilterEvaluate(ctx context.Context, in *Work return out, nil } +func (c *agentSecureClient) RemoteQueryExecute(ctx context.Context, in *RemoteQueryExecuteRequest, opts ...grpc.CallOption) (*RemoteQueryExecuteResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(RemoteQueryExecuteResponse) + err := c.cc.Invoke(ctx, AgentSecure_RemoteQueryExecute_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + func (c *agentSecureClient) StreamKubeMetadata(ctx context.Context, in *KubeMetadataStreamRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[KubeMetadataStreamResponse], error) { cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) stream, err := c.cc.NewStream(ctx, &AgentSecure_ServiceDesc.Streams[5], AgentSecure_StreamKubeMetadata_FullMethodName, cOpts...) @@ -471,6 +484,8 @@ type AgentSecureServer interface { StreamConfigEvents(*ConfigStreamRequest, grpc.ServerStreamingServer[ConfigEvent]) error // Evaluates a workloadfilter rule on behalf of remote agents. WorkloadFilterEvaluate(context.Context, *WorkloadFilterEvaluateRequest) (*WorkloadFilterEvaluateResponse, error) + // Executes an Agent-local Remote Queries request through a matched integration check. + RemoteQueryExecute(context.Context, *RemoteQueryExecuteRequest) (*RemoteQueryExecuteResponse, error) // Streams pod-to-service metadata for a specific node. StreamKubeMetadata(*KubeMetadataStreamRequest, grpc.ServerStreamingServer[KubeMetadataStreamResponse]) error mustEmbedUnimplementedAgentSecureServer() @@ -537,6 +552,9 @@ func (UnimplementedAgentSecureServer) StreamConfigEvents(*ConfigStreamRequest, g func (UnimplementedAgentSecureServer) WorkloadFilterEvaluate(context.Context, *WorkloadFilterEvaluateRequest) (*WorkloadFilterEvaluateResponse, error) { return nil, status.Errorf(codes.Unimplemented, "method WorkloadFilterEvaluate not implemented") } +func (UnimplementedAgentSecureServer) RemoteQueryExecute(context.Context, *RemoteQueryExecuteRequest) (*RemoteQueryExecuteResponse, error) { + return nil, status.Errorf(codes.Unimplemented, "method RemoteQueryExecute not implemented") +} func (UnimplementedAgentSecureServer) StreamKubeMetadata(*KubeMetadataStreamRequest, grpc.ServerStreamingServer[KubeMetadataStreamResponse]) error { return status.Errorf(codes.Unimplemented, "method StreamKubeMetadata not implemented") } @@ -846,6 +864,24 @@ func _AgentSecure_WorkloadFilterEvaluate_Handler(srv interface{}, ctx context.Co return interceptor(ctx, in, info, handler) } +func _AgentSecure_RemoteQueryExecute_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(RemoteQueryExecuteRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(AgentSecureServer).RemoteQueryExecute(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: AgentSecure_RemoteQueryExecute_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(AgentSecureServer).RemoteQueryExecute(ctx, req.(*RemoteQueryExecuteRequest)) + } + return interceptor(ctx, in, info, handler) +} + func _AgentSecure_StreamKubeMetadata_Handler(srv interface{}, stream grpc.ServerStream) error { m := new(KubeMetadataStreamRequest) if err := stream.RecvMsg(m); err != nil { @@ -916,6 +952,10 @@ var AgentSecure_ServiceDesc = grpc.ServiceDesc{ MethodName: "WorkloadFilterEvaluate", Handler: _AgentSecure_WorkloadFilterEvaluate_Handler, }, + { + MethodName: "RemoteQueryExecute", + Handler: _AgentSecure_RemoteQueryExecute_Handler, + }, }, Streams: []grpc.StreamDesc{ { diff --git a/pkg/proto/pbgo/mocks/core/api_mockgen.pb.go b/pkg/proto/pbgo/mocks/core/api_mockgen.pb.go index 1c09277b913a..defe22225aa0 100644 --- a/pkg/proto/pbgo/mocks/core/api_mockgen.pb.go +++ b/pkg/proto/pbgo/mocks/core/api_mockgen.pb.go @@ -385,6 +385,26 @@ func (mr *MockAgentSecureClientMockRecorder) RegisterRemoteAgent(ctx, in interfa return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RegisterRemoteAgent", reflect.TypeOf((*MockAgentSecureClient)(nil).RegisterRemoteAgent), varargs...) } +// RemoteQueryExecute mocks base method. +func (m *MockAgentSecureClient) RemoteQueryExecute(ctx context.Context, in *core.RemoteQueryExecuteRequest, opts ...grpc.CallOption) (*core.RemoteQueryExecuteResponse, error) { + m.ctrl.T.Helper() + varargs := []interface{}{ctx, in} + for _, a := range opts { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "RemoteQueryExecute", varargs...) + ret0, _ := ret[0].(*core.RemoteQueryExecuteResponse) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// RemoteQueryExecute indicates an expected call of RemoteQueryExecute. +func (mr *MockAgentSecureClientMockRecorder) RemoteQueryExecute(ctx, in interface{}, opts ...interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]interface{}{ctx, in}, opts...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RemoteQueryExecute", reflect.TypeOf((*MockAgentSecureClient)(nil).RemoteQueryExecute), varargs...) +} + // ResetConfigState mocks base method. func (m *MockAgentSecureClient) ResetConfigState(ctx context.Context, in *emptypb.Empty, opts ...grpc.CallOption) (*core.ResetStateConfigResponse, error) { m.ctrl.T.Helper() @@ -731,6 +751,21 @@ func (mr *MockAgentSecureServerMockRecorder) RegisterRemoteAgent(arg0, arg1 inte return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RegisterRemoteAgent", reflect.TypeOf((*MockAgentSecureServer)(nil).RegisterRemoteAgent), arg0, arg1) } +// RemoteQueryExecute mocks base method. +func (m *MockAgentSecureServer) RemoteQueryExecute(arg0 context.Context, arg1 *core.RemoteQueryExecuteRequest) (*core.RemoteQueryExecuteResponse, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "RemoteQueryExecute", arg0, arg1) + ret0, _ := ret[0].(*core.RemoteQueryExecuteResponse) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// RemoteQueryExecute indicates an expected call of RemoteQueryExecute. +func (mr *MockAgentSecureServerMockRecorder) RemoteQueryExecute(arg0, arg1 interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RemoteQueryExecute", reflect.TypeOf((*MockAgentSecureServer)(nil).RemoteQueryExecute), arg0, arg1) +} + // ResetConfigState mocks base method. func (m *MockAgentSecureServer) ResetConfigState(arg0 context.Context, arg1 *emptypb.Empty) (*core.ResetStateConfigResponse, error) { m.ctrl.T.Helper() diff --git a/test/remotequeries/fused-local-par-agent-postgres-proof.sh b/test/remotequeries/fused-local-par-agent-postgres-proof.sh index 03348e793672..94dc692bb751 100755 --- a/test/remotequeries/fused-local-par-agent-postgres-proof.sh +++ b/test/remotequeries/fused-local-par-agent-postgres-proof.sh @@ -1,8 +1,9 @@ #!/usr/bin/env bash # Runs the fused local-only Remote Queries proof: # fakeintake -> live WorkflowRunner PAR loop -> com.datadoghq.remotequeries.execute -# -> real local Agent IPC /agent/remote-queries/execute -> loaded Postgres check -# -> SELECT 1 AS value -> fakeintake publish. +# -> real local AgentSecure gRPC RemoteQueryExecute over Agent IPC TLS/auth +# -> loaded Postgres check -> SELECT 1 AS value -> fakeintake publish. +# The HTTP execute endpoint remains as a dev preflight for local evidence only. # # Defaults assume the remote-queries-poc worktree layout. Override AGENT_REPO, # INTEGRATIONS_CORE, TMP_ROOT, CMD_PORT, POSTGRES_IMAGE, or AGENT_PYTHON_VERSION / @@ -268,7 +269,7 @@ call_agent_execute_preflight() { local token token=$(cat "$TMP_ROOT/run/auth_token") - log "Preflight real Agent IPC execute endpoint" + log "Preflight real Agent IPC HTTP execute endpoint (dev evidence only)" local status status=$(curl -sS -k -o "$TMP_ROOT/results/agent-execute-preflight.body" -w '%{http_code}' \ -H "Authorization: Bearer ${token}" \ @@ -294,7 +295,7 @@ call_agent_execute_preflight() { } run_fused_go_proof() { - log "Running fused PAR -> real Agent IPC -> Postgres -> fakeintake proof test" + log "Running fused PAR -> real AgentSecure gRPC IPC -> Postgres -> fakeintake proof test" ( cd "$AGENT_REPO" RQ_FUSED_PROOF=1 \ @@ -336,7 +337,7 @@ main() { cat "$TMP_ROOT/results/fused-proof-evidence.txt" log "Done. Sanitized artifacts left in $TMP_ROOT" - log "Key evidence: fakeintake enqueue/dequeue/publish and real Agent IPC endpoint evidence are in $TMP_ROOT/results/fused-proof-evidence.txt" + log "Key evidence: fakeintake enqueue/dequeue/publish and real AgentSecure IPC evidence are in $TMP_ROOT/results/fused-proof-evidence.txt" } main "$@" From 3b530eb29b00a71212ae75cd123a7486133ef4e5 Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 16:28:50 +0000 Subject: [PATCH 17/33] Add standalone Remote Queries PAR proof --- .../standalone_par_process_proof_test.go | 214 +++++++++++ ...andalone-par-agentsecure-postgres-proof.sh | 352 ++++++++++++++++++ 2 files changed, 566 insertions(+) create mode 100644 pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go create mode 100755 test/remotequeries/standalone-par-agentsecure-postgres-proof.sh diff --git a/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go b/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go new file mode 100644 index 000000000000..fbc4ac404b1b --- /dev/null +++ b/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go @@ -0,0 +1,214 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2026-present Datadog, Inc. + +//go:build !windows + +package com_datadoghq_remotequeries_test + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "fmt" + "os" + "os/exec" + "path/filepath" + "strconv" + "strings" + "testing" + "time" + + "github.com/DataDog/datadog-agent/pkg/privateactionrunner/adapters/constants" + com_datadoghq_remotequeries "github.com/DataDog/datadog-agent/pkg/privateactionrunner/bundles/remotequeries" + "github.com/DataDog/datadog-agent/pkg/privateactionrunner/util" + fakeintakeclient "github.com/DataDog/datadog-agent/test/fakeintake/client" + fakeintakeserver "github.com/DataDog/datadog-agent/test/fakeintake/server" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const standaloneLocalProofEnv = "RQ_STANDALONE_PROOF" + +func TestRemoteQueriesActionRunsThroughStandalonePARProcessWithRealAgentIPC(t *testing.T) { + if os.Getenv(standaloneLocalProofEnv) != "1" { + t.Skipf("set %s=1 and start a local Agent with a loaded Postgres check to run the standalone PAR process proof", standaloneLocalProofEnv) + } + + parBin := getenvRequired(t, "RQ_STANDALONE_PAR_BIN") + cmdPort := getenvRequired(t, "RQ_STANDALONE_AGENT_CMD_PORT") + authTokenFile := getenvRequired(t, "RQ_STANDALONE_AGENT_AUTH_TOKEN_FILE") + ipcCertFile := getenvRequired(t, "RQ_STANDALONE_AGENT_IPC_CERT_FILE") + agentPID := getenvOptional("RQ_STANDALONE_AGENT_PID") + cmdPortInt, err := strconv.Atoi(cmdPort) + require.NoError(t, err) + + fakeintake, _ := fakeintakeserver.InitialiseForTests(t) + defer func() { require.NoError(t, fakeintake.Stop()) }() + fakeintakeClient := fakeintakeclient.NewClient(fakeintake.URL()) + require.NoError(t, fakeintakeClient.FlushPAR()) + + parDir := t.TempDir() + parLog := filepath.Join(parDir, "private-action-runner.log") + writeStandalonePARConfig(t, parDir, parLog, fakeintake.URL(), cmdPortInt, authTokenFile, ipcCertFile) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + var stdout, stderr bytes.Buffer + cmd := exec.CommandContext(ctx, parBin, "run", "-c", parDir) + cmd.Stdout = &stdout + cmd.Stderr = &stderr + cmd.Env = append(os.Environ(), constants.InternalSkipTaskVerificationEnvVar+"=true") + require.NoError(t, cmd.Start()) + defer func() { + cancel() + _ = cmd.Wait() + }() + require.NotNil(t, cmd.Process) + parPID := cmd.Process.Pid + t.Logf("standalone private-action-runner process started: pid=%d bin=%s cfg=%s", parPID, parBin, parDir) + if agentPID != "" { + parsedAgentPID, err := strconv.Atoi(agentPID) + require.NoError(t, err) + require.NotEqual(t, parsedAgentPID, parPID, "PAR must be a separate OS process from the Agent") + } + + waitForStandalonePARPolling(t, fakeintakeClient, cmd, parLog, &stdout, &stderr) + + taskID := fmt.Sprintf("remotequeries-standalone-par-proof-%d", time.Now().UnixNano()) + inputs := map[string]interface{}{ + "integration": "postgres", + "target": map[string]interface{}{ + "host": "localhost", + "port": 5432, + "dbname": "postgres", + }, + "query": "SELECT 1 AS value", + "limits": map[string]interface{}{ + "maxRows": 1, + "maxBytes": 1024, + "timeoutMs": 1000, + }, + } + requestEvidence, err := json.Marshal(inputs) + require.NoError(t, err) + requireNoCredentialShape(t, requestEvidence) + + fqn := com_datadoghq_remotequeries.BundleID + "." + com_datadoghq_remotequeries.ExecuteActionName + t.Logf("fakeintake task enqueued: task_id=%s action_fqn=%s inputs=%s", taskID, fqn, requestEvidence) + t.Logf("real AgentSecure IPC configured for standalone PAR: 127.0.0.1:%d RemoteQueryExecute", cmdPortInt) + require.NoError(t, fakeintakeClient.EnqueuePARTask(taskID, fqn, inputs)) + + result, err := fakeintakeClient.GetPARTaskResult(taskID, 30*time.Second) + require.NoError(t, err) + if !result.Success { + t.Logf("failed PAR task result: %+v", result) + t.Logf("PAR log tail:\n%s", readTail(parLog, 120)) + } + require.True(t, result.Success) + require.Equal(t, taskID, result.TaskID) + assert.Equal(t, "SUCCEEDED", result.Outputs["status"]) + require.Contains(t, result.Outputs, "rows") + + rows, ok := result.Outputs["rows"].([]interface{}) + require.True(t, ok) + require.Len(t, rows, 1) + firstRow, ok := rows[0].(map[string]interface{}) + require.True(t, ok) + assert.Equal(t, float64(1), firstRow["value"]) + + resultEvidence, err := json.Marshal(result.Outputs) + require.NoError(t, err) + requireNoCredentialShape(t, resultEvidence) + t.Logf("fakeintake captured successful PAR task result: %s", resultEvidence) + + dequeueCalls, err := fakeintakeClient.GetPARDequeueCount() + require.NoError(t, err) + assert.GreaterOrEqual(t, dequeueCalls, 1) + t.Logf("standalone PAR process dequeued from fakeintake: dequeue_calls=%d", dequeueCalls) + + writeFusedEvidence(t, getenvOptional("RQ_STANDALONE_EVIDENCE_FILE"), []string{ + fmt.Sprintf("standalone private-action-runner process pid=%d", parPID), + fmt.Sprintf("separate Agent process pid=%s", agentPID), + fmt.Sprintf("fakeintake task enqueued: task_id=%s action_fqn=%s inputs=%s", taskID, fqn, requestEvidence), + "standalone PAR process dequeued the fakeintake OPMS task and invoked the registered action", + fmt.Sprintf("real AgentSecure IPC called by standalone PAR: 127.0.0.1:%d RemoteQueryExecute", cmdPortInt), + fmt.Sprintf("fakeintake captured successful PAR task result: %s", resultEvidence), + fmt.Sprintf("dequeue_calls=%d", dequeueCalls), + "task verification skipped for this standalone tracer bullet with DD_INTERNAL_PAR_SKIP_TASK_VERIFICATION=true", + }) +} + +func writeStandalonePARConfig(t *testing.T, dir, logFile, fakeintakeURL string, cmdPort int, authTokenFile, ipcCertFile string) { + t.Helper() + privateJWK, _, err := util.GenerateKeys() + require.NoError(t, err) + privateJWKJSON, err := json.Marshal(privateJWK) + require.NoError(t, err) + privateKeyB64 := base64.RawURLEncoding.EncodeToString(privateJWKJSON) + + cfg := fmt.Sprintf(`api_key: '00000000000000000000000000000000' +dd_url: %q +hostname: rq-standalone-par-proof +cmd_host: 127.0.0.1 +cmd_port: %d +auth_token_file_path: %q +ipc_cert_file_path: %q +log_level: debug +telemetry.enabled: false +inventories_enabled: false +process_config.enabled: 'false' +logs_enabled: false +apm_config.enabled: false +private_action_runner: + enabled: true + self_enroll: false + urn: "urn:dd:apps:on-prem-runner:us1:123456:remotequeries-standalone-par-local-proof-runner" + private_key: %q + log_file: %q + default_actions_enabled: false + actions_allowlist: + - "com.datadoghq.remotequeries.execute" + task_concurrency: 1 + task_timeout_seconds: 10 +`, fakeintakeURL, cmdPort, authTokenFile, ipcCertFile, privateKeyB64, logFile) + require.NoError(t, os.WriteFile(filepath.Join(dir, "datadog.yaml"), []byte(cfg), 0o600)) +} + +func waitForStandalonePARPolling(t *testing.T, client *fakeintakeclient.Client, cmd *exec.Cmd, parLog string, stdout, stderr *bytes.Buffer) { + t.Helper() + deadline := time.Now().Add(20 * time.Second) + for time.Now().Before(deadline) { + if cmd.ProcessState != nil && cmd.ProcessState.Exited() { + require.FailNowf(t, "standalone PAR process exited before polling fakeintake", "stdout:\n%s\nstderr:\n%s\nlog tail:\n%s", stdout.String(), stderr.String(), readTail(parLog, 120)) + } + if count, err := client.GetPARDequeueCount(); err == nil && count > 0 { + t.Logf("standalone PAR process is polling fakeintake: dequeue_calls=%d", count) + return + } + time.Sleep(250 * time.Millisecond) + } + require.FailNowf(t, "timed out waiting for standalone PAR process to poll fakeintake", "stdout:\n%s\nstderr:\n%s\nlog tail:\n%s", stdout.String(), stderr.String(), readTail(parLog, 120)) +} + +func requireNoCredentialShape(t *testing.T, payload []byte) { + t.Helper() + lower := strings.ToLower(string(payload)) + require.NotContains(t, lower, "password") + require.NotContains(t, lower, "token") + require.NotContains(t, lower, "secret") +} + +func readTail(path string, maxLines int) string { + content, err := os.ReadFile(path) + if err != nil { + return fmt.Sprintf("", path, err) + } + lines := strings.Split(string(content), "\n") + if len(lines) <= maxLines { + return string(content) + } + return strings.Join(lines[len(lines)-maxLines:], "\n") +} diff --git a/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh b/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh new file mode 100755 index 000000000000..063f921bce9a --- /dev/null +++ b/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh @@ -0,0 +1,352 @@ +#!/usr/bin/env bash +# Runs the standalone local-only Remote Queries proof: +# fakeintake -> standalone OS private-action-runner process -> com.datadoghq.remotequeries.execute +# -> real local AgentSecure gRPC RemoteQueryExecute over Agent IPC TLS/auth +# -> loaded Postgres check -> SELECT 1 AS value -> fakeintake publish. +# The HTTP execute endpoint remains as a dev preflight for local evidence only. +# +# Defaults assume the remote-queries-poc worktree layout. Override AGENT_REPO, +# INTEGRATIONS_CORE, TMP_ROOT, CMD_PORT, POSTGRES_IMAGE, or AGENT_PYTHON_VERSION / +# AGENT_PYTHON_ABI if needed. This proof intentionally follows the repository's +# fakeintake/OPMS precedent and sets DD_INTERNAL_PAR_SKIP_TASK_VERIFICATION=true +# for the standalone-process tracer bullet. Signed task verification is postponed +# to backend/AP/RC work. + +set -euo pipefail + +AGENT_REPO=${AGENT_REPO:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)} +INTEGRATIONS_CORE=${INTEGRATIONS_CORE:-/home/bits/dd/tasks/remote-queries-poc/worktrees/integrations-core} +TMP_ROOT=${TMP_ROOT:-/tmp/rq-standalone-par-agent-postgres} +CMD_PORT=${CMD_PORT:-55003} +POSTGRES_IMAGE=${POSTGRES_IMAGE:-postgres:11-alpine} +POSTGRES_CONTAINER=${POSTGRES_CONTAINER:-rq-standalone-par-agent-postgres-$$} +PIP_PLATFORM=${PIP_PLATFORM:-manylinux2014_x86_64} +AGENT_PYTHON_VERSION=${AGENT_PYTHON_VERSION:-} +AGENT_PYTHON_ABI=${AGENT_PYTHON_ABI:-} + +AGENT_PID="" +POSTGRES_STARTED=0 + +log() { + printf '\n[%s] %s\n' "$(date -u +%H:%M:%S)" "$*" +} + +cleanup_agent() { + if [[ -n "${AGENT_PID:-}" ]] && kill -0 "$AGENT_PID" 2>/dev/null; then + kill "$AGENT_PID" 2>/dev/null || true + wait "$AGENT_PID" 2>/dev/null || true + fi + AGENT_PID="" +} + +cleanup_all() { + cleanup_agent + if [[ "$POSTGRES_STARTED" == "1" ]]; then + docker rm -f "$POSTGRES_CONTAINER" >/dev/null 2>&1 || true + fi +} +trap cleanup_all EXIT + +require_cmd() { + command -v "$1" >/dev/null 2>&1 || { + echo "missing required command: $1" >&2 + exit 1 + } +} + +python_minor_from_version() { + local version=$1 + awk -F. '{print $1 "." $2}' <<<"$version" +} + +abi_from_python_minor() { + local minor=$1 + printf 'cp%s\n' "${minor//./}" +} + +write_agent_config() { + cat > "$TMP_ROOT/datadog.yaml" < "$TMP_ROOT/agent.log" + "$AGENT_REPO/bin/agent/agent" run -c "$TMP_ROOT" \ + > "$TMP_ROOT/python-detect-stdout.log" 2> "$TMP_ROOT/python-detect-stderr.log" & + AGENT_PID=$! + + local detected="" + for _ in $(seq 1 80); do + detected=$(grep -aoE '"pythonV":"[0-9]+\.[0-9]+(\.[0-9]+)?' "$TMP_ROOT/agent.log" 2>/dev/null | head -1 | cut -d: -f2 | tr -d '"' || true) + if [[ -n "$detected" ]]; then + break + fi + if ! kill -0 "$AGENT_PID" 2>/dev/null; then + break + fi + sleep 0.5 + done + cleanup_agent + + if [[ -z "$detected" && -f "$AGENT_REPO/omnibus/config/software/python3.rb" ]]; then + detected=$(sed -nE 's/^default_version "([0-9]+\.[0-9]+)(\.[0-9]+)?"/\1/p' "$AGENT_REPO/omnibus/config/software/python3.rb" | head -1) + if [[ -n "$detected" ]]; then + log "Could not detect runtime Python from Agent logs; falling back to source config: $detected" + fi + fi + + if [[ -z "$detected" ]]; then + echo "Unable to detect Agent Python version. Set AGENT_PYTHON_VERSION=3.12 or AGENT_PYTHON_ABI=cp312." >&2 + tail -80 "$TMP_ROOT/python-detect-stderr.log" >&2 || true + exit 1 + fi + + AGENT_PYTHON_VERSION=$(python_minor_from_version "$detected") + AGENT_PYTHON_ABI=$(abi_from_python_minor "$AGENT_PYTHON_VERSION") + log "Detected Agent Python runtime: $detected; installing wheels for $AGENT_PYTHON_VERSION ($AGENT_PYTHON_ABI)" +} + +install_python_deps() { + local py_digits + py_digits=${AGENT_PYTHON_VERSION//./} + + log "Installing temporary Python deps into $TMP_ROOT/pydeps for $AGENT_PYTHON_ABI on $PIP_PLATFORM" + python3 -m pip install --quiet --target "$TMP_ROOT/pydeps" \ + --only-binary=:all: --platform "$PIP_PLATFORM" \ + --implementation cp --python-version "$py_digits" --abi "$AGENT_PYTHON_ABI" \ + 'psycopg[binary,pool]' cachetools packaging semver 'pydantic<3' python-dateutil mmh3 + + for dep in "$TMP_ROOT"/pydeps/*; do + local base + base=$(basename "$dep") + [[ -e "$TMP_ROOT/checks.d/$base" ]] || ln -s "$dep" "$TMP_ROOT/checks.d/$base" + done +} + +setup_tmp_tree() { + rm -rf "$TMP_ROOT" + mkdir -p "$TMP_ROOT/conf.d/postgres.d" "$TMP_ROOT/run" "$TMP_ROOT/checks.d/datadog_checks" "$TMP_ROOT/pydeps" "$TMP_ROOT/results" + + ln -s "$INTEGRATIONS_CORE/datadog_checks_base/datadog_checks/base" "$TMP_ROOT/checks.d/datadog_checks/base" + ln -s "$INTEGRATIONS_CORE/datadog_checks_base/datadog_checks/checks" "$TMP_ROOT/checks.d/datadog_checks/checks" + ln -s "$INTEGRATIONS_CORE/datadog_checks_base/datadog_checks/config.py" "$TMP_ROOT/checks.d/datadog_checks/config.py" + ln -s "$INTEGRATIONS_CORE/datadog_checks_base/datadog_checks/errors.py" "$TMP_ROOT/checks.d/datadog_checks/errors.py" + ln -s "$INTEGRATIONS_CORE/datadog_checks_base/datadog_checks/log.py" "$TMP_ROOT/checks.d/datadog_checks/log.py" + ln -s "$INTEGRATIONS_CORE/postgres/datadog_checks/postgres" "$TMP_ROOT/checks.d/datadog_checks/postgres" + cat > "$TMP_ROOT/checks.d/datadog_checks/__init__.py" <<'PY' +__path__ = __import__('pkgutil').extend_path(__path__, __name__) +PY + + write_agent_config + detect_agent_python + install_python_deps +} + +start_postgres_fixture() { + if psql 'postgresql://postgres@localhost:5432/postgres' -c 'select 1' >/dev/null 2>&1; then + log "Using existing local Postgres on localhost:5432" + return + fi + + local published_containers proof_containers other_containers + published_containers=$(docker ps --filter publish=5432 --format '{{.Names}}' || true) + if [[ -n "$published_containers" ]]; then + proof_containers=$(grep -E '^rq-standalone-par-agent-postgres-' <<<"$published_containers" || true) + other_containers=$(grep -Ev '^rq-standalone-par-agent-postgres-' <<<"$published_containers" || true) + + if [[ -n "$proof_containers" ]]; then + log "Removing stale proof Postgres container(s) bound to localhost:5432 but not accepting psql" + xargs -r docker rm -f <<<"$proof_containers" >/dev/null + fi + if [[ -n "$other_containers" ]]; then + echo "Port 5432 is published by non-proof Docker container(s), and psql select 1 failed:" >&2 + sed 's/^/ /' <<<"$other_containers" >&2 + echo "Refusing to remove unrelated containers. Stop them or point this proof at a reachable local Postgres." >&2 + exit 1 + fi + fi + + log "Starting disposable Postgres fixture $POSTGRES_CONTAINER from $POSTGRES_IMAGE" + if ! docker run --rm --name "$POSTGRES_CONTAINER" \ + -e POSTGRES_HOST_AUTH_METHOD=trust \ + -p 5432:5432 \ + -d "$POSTGRES_IMAGE" >/dev/null; then + echo "Failed to start disposable Postgres fixture on port 5432." >&2 + echo "If another local Postgres is intended, ensure this succeeds: psql postgresql://postgres@localhost:5432/postgres -c 'select 1'" >&2 + exit 1 + fi + POSTGRES_STARTED=1 + + for _ in $(seq 1 60); do + if psql 'postgresql://postgres@localhost:5432/postgres' -c 'select 1' >/dev/null 2>&1; then + log "Postgres fixture is ready" + return + fi + sleep 0.5 + done + + docker logs "$POSTGRES_CONTAINER" | tail -80 >&2 || true + echo "Postgres fixture did not become ready" >&2 + exit 1 +} + +write_postgres_config() { + cat > "$TMP_ROOT/conf.d/postgres.d/conf.yaml" <<'YAML' +init_config: {} +instances: + - host: localhost + port: 5432 + dbname: postgres + username: postgres +YAML +} + +start_agent_and_wait_for_postgres_check() { + : > "$TMP_ROOT/agent.log" + PYTHONPATH="$TMP_ROOT/checks.d:$TMP_ROOT/pydeps" \ + "$AGENT_REPO/bin/agent/agent" run -c "$TMP_ROOT" \ + > "$TMP_ROOT/live-stdout.log" 2> "$TMP_ROOT/live-stderr.log" & + AGENT_PID=$! + + for _ in $(seq 1 80); do + if grep -q "successfully loaded check 'postgres'" "$TMP_ROOT/agent.log" && [[ -f "$TMP_ROOT/run/auth_token" && -f "$TMP_ROOT/run/ipc_cert.pem" ]]; then + log "Agent loaded the Postgres check and exposed IPC artifacts" + grep -n "successfully loaded check 'postgres'\|Scheduling check postgres" "$TMP_ROOT/agent.log" | tail -20 + return + fi + if ! kill -0 "$AGENT_PID" 2>/dev/null; then + echo "Agent exited early; stderr follows:" >&2 + tail -80 "$TMP_ROOT/live-stderr.log" >&2 || true + exit 1 + fi + sleep 0.5 + done + + echo "Timed out waiting for loaded Postgres check" >&2 + grep -ni 'postgres\|remote_queries\|ModuleNotFound\|ImportError\|unable to load check\|AttributeError' "$TMP_ROOT/agent.log" >&2 || true + exit 1 +} + +call_agent_execute_preflight() { + local payload='{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":1,"maxBytes":1024,"timeoutMs":1000}}' + local token + token=$(cat "$TMP_ROOT/run/auth_token") + + log "Preflight real Agent IPC HTTP execute endpoint (dev evidence only)" + local status + status=$(curl -sS -k -o "$TMP_ROOT/results/agent-execute-preflight.body" -w '%{http_code}' \ + -H "Authorization: Bearer ${token}" \ + -H 'Content-Type: application/json' \ + --data "$payload" \ + "https://127.0.0.1:${CMD_PORT}/agent/remote-queries/execute") + printf 'agent_execute_http_status=%s\n' "$status" | tee "$TMP_ROOT/results/agent-execute-preflight.status" + cat "$TMP_ROOT/results/agent-execute-preflight.body" + printf '\n' + + if [[ "$status" != "200" ]]; then + echo "FAIL: expected Agent execute preflight HTTP 200, got $status" >&2 + exit 1 + fi + if ! grep -Eq '"status"[[:space:]]*:[[:space:]]*"SUCCEEDED".*"value"[[:space:]]*:[[:space:]]*1|"value"[[:space:]]*:[[:space:]]*1.*"status"[[:space:]]*:[[:space:]]*"SUCCEEDED"' "$TMP_ROOT/results/agent-execute-preflight.body"; then + echo "FAIL: Agent execute preflight response did not contain SUCCEEDED row value=1" >&2 + exit 1 + fi + if grep -Eq 'password|token|secret' "$TMP_ROOT/results/agent-execute-preflight.body"; then + echo "FAIL: Agent execute preflight response contained credential-shaped text" >&2 + exit 1 + fi +} + +run_standalone_go_proof() { + log "Running standalone PAR process -> real AgentSecure gRPC IPC -> Postgres -> fakeintake proof test" + ( + cd "$AGENT_REPO" + RQ_STANDALONE_PROOF=1 \ + RQ_STANDALONE_PAR_BIN="$AGENT_REPO/bin/privateactionrunner/privateactionrunner" \ + RQ_STANDALONE_AGENT_PID="$AGENT_PID" \ + RQ_STANDALONE_AGENT_CMD_PORT="$CMD_PORT" \ + RQ_STANDALONE_AGENT_AUTH_TOKEN_FILE="$TMP_ROOT/run/auth_token" \ + RQ_STANDALONE_AGENT_IPC_CERT_FILE="$TMP_ROOT/run/ipc_cert.pem" \ + RQ_STANDALONE_EVIDENCE_FILE="$TMP_ROOT/results/standalone-proof-evidence.txt" \ + dda inv test --targets=./pkg/privateactionrunner/bundles/remotequeries \ + --extra-args='-run TestRemoteQueriesActionRunsThroughStandalonePARProcessWithRealAgentIPC -count=1 -v' + ) | tee "$TMP_ROOT/results/standalone-proof-test.log" +} + +main() { + require_cmd docker + require_cmd psql + require_cmd python3 + require_cmd curl + require_cmd dda + + [[ -x "$AGENT_REPO/bin/agent/agent" ]] || { + echo "Agent binary not found/executable: $AGENT_REPO/bin/agent/agent" >&2 + echo "Build it first with: dda inv agent.build --build-exclude=systemd" >&2 + exit 1 + } + [[ -x "$AGENT_REPO/bin/privateactionrunner/privateactionrunner" ]] || { + echo "Private Action Runner binary not found/executable: $AGENT_REPO/bin/privateactionrunner/privateactionrunner" >&2 + echo "Build it first with: dda inv privateactionrunner.build" >&2 + exit 1 + } + [[ -d "$INTEGRATIONS_CORE/postgres/datadog_checks/postgres" ]] || { + echo "Postgres integration not found under: $INTEGRATIONS_CORE" >&2 + exit 1 + } + + log "Preparing temporary harness at $TMP_ROOT" + setup_tmp_tree + start_postgres_fixture + write_postgres_config + start_agent_and_wait_for_postgres_check + call_agent_execute_preflight + run_standalone_go_proof + + log "Sanitized standalone proof evidence" + cat "$TMP_ROOT/results/standalone-proof-evidence.txt" + + log "Done. Sanitized artifacts left in $TMP_ROOT" + log "Key evidence: fakeintake enqueue/dequeue/publish, standalone PAR PID, and real AgentSecure IPC evidence are in $TMP_ROOT/results/standalone-proof-evidence.txt" +} + +main "$@" From 3b4a1ce5d0de018cb118c6d836e5702e1135ee0a Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 16:49:16 +0000 Subject: [PATCH 18/33] Use Postgres integration fixture in Remote Queries proof --- .../live_agent_ipc_par_loop_test.go | 34 ++++- .../standalone_par_process_proof_test.go | 8 +- .../fused-local-par-agent-postgres-proof.sh | 133 ++++++++++++------ ...andalone-par-agentsecure-postgres-proof.sh | 133 ++++++++++++------ 4 files changed, 204 insertions(+), 104 deletions(-) diff --git a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go index 2e2eea43dbc8..b0af7c3350ce 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go @@ -75,12 +75,8 @@ func TestRemoteQueriesActionRunsThroughLivePARLoopWithRealAgentIPC(t *testing.T) taskID := fmt.Sprintf("remotequeries-fused-local-proof-%d", time.Now().UnixNano()) inputs := map[string]interface{}{ "integration": "postgres", - "target": map[string]interface{}{ - "host": "localhost", - "port": 5432, - "dbname": "postgres", - }, - "query": "SELECT 1 AS value", + "target": remoteQueriesPostgresTargetFromEnv(t), + "query": "SELECT 1 AS value", "limits": map[string]interface{}{ "maxRows": 1, "maxBytes": 1024, @@ -140,6 +136,32 @@ func getenvOptional(name string) string { return os.Getenv(name) } +func remoteQueriesPostgresTargetFromEnv(t *testing.T) map[string]interface{} { + t.Helper() + + port := 5432 + if value := os.Getenv("RQ_POSTGRES_PORT"); value != "" { + parsed, err := strconv.Atoi(value) + require.NoError(t, err) + port = parsed + } + + host := os.Getenv("RQ_POSTGRES_HOST") + if host == "" { + host = "localhost" + } + dbname := os.Getenv("RQ_POSTGRES_DBNAME") + if dbname == "" { + dbname = "datadog_test" + } + + return map[string]interface{}{ + "host": host, + "port": port, + "dbname": dbname, + } +} + func writeFusedEvidence(t *testing.T, path string, lines []string) { t.Helper() if path == "" { diff --git a/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go b/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go index fbc4ac404b1b..927af14899f5 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go @@ -80,12 +80,8 @@ func TestRemoteQueriesActionRunsThroughStandalonePARProcessWithRealAgentIPC(t *t taskID := fmt.Sprintf("remotequeries-standalone-par-proof-%d", time.Now().UnixNano()) inputs := map[string]interface{}{ "integration": "postgres", - "target": map[string]interface{}{ - "host": "localhost", - "port": 5432, - "dbname": "postgres", - }, - "query": "SELECT 1 AS value", + "target": remoteQueriesPostgresTargetFromEnv(t), + "query": "SELECT 1 AS value", "limits": map[string]interface{}{ "maxRows": 1, "maxBytes": 1024, diff --git a/test/remotequeries/fused-local-par-agent-postgres-proof.sh b/test/remotequeries/fused-local-par-agent-postgres-proof.sh index 94dc692bb751..1fb72a2a170e 100755 --- a/test/remotequeries/fused-local-par-agent-postgres-proof.sh +++ b/test/remotequeries/fused-local-par-agent-postgres-proof.sh @@ -5,9 +5,11 @@ # -> loaded Postgres check -> SELECT 1 AS value -> fakeintake publish. # The HTTP execute endpoint remains as a dev preflight for local evidence only. # -# Defaults assume the remote-queries-poc worktree layout. Override AGENT_REPO, -# INTEGRATIONS_CORE, TMP_ROOT, CMD_PORT, POSTGRES_IMAGE, or AGENT_PYTHON_VERSION / -# AGENT_PYTHON_ABI if needed. The proof is local-only and intentionally sets +# Defaults assume the remote-queries-poc worktree layout and reuse the +# integrations-core Postgres integration test compose fixture. Override +# AGENT_REPO, INTEGRATIONS_CORE, TMP_ROOT, CMD_PORT, POSTGRES_COMPOSE_FILE, +# POSTGRES_COMPOSE_PROJECT, POSTGRES_IMAGE, RQ_POSTGRES_*, or +# AGENT_PYTHON_VERSION / AGENT_PYTHON_ABI if needed. The proof is local-only and intentionally sets # DD_INTERNAL_PAR_SKIP_TASK_VERIFICATION=true inside the Go proof test. set -euo pipefail @@ -16,14 +18,21 @@ AGENT_REPO=${AGENT_REPO:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)} INTEGRATIONS_CORE=${INTEGRATIONS_CORE:-/home/bits/dd/tasks/remote-queries-poc/worktrees/integrations-core} TMP_ROOT=${TMP_ROOT:-/tmp/rq-fused-local-par-agent-postgres} CMD_PORT=${CMD_PORT:-55003} -POSTGRES_IMAGE=${POSTGRES_IMAGE:-postgres:11-alpine} -POSTGRES_CONTAINER=${POSTGRES_CONTAINER:-rq-fused-local-par-agent-postgres-$$} +POSTGRES_COMPOSE_FILE=${POSTGRES_COMPOSE_FILE:-$INTEGRATIONS_CORE/postgres/tests/compose/docker-compose.yaml} +POSTGRES_COMPOSE_PROJECT=${POSTGRES_COMPOSE_PROJECT:-rq-fused-local-par-agent-postgres-$$} +POSTGRES_IMAGE=${POSTGRES_IMAGE:-13-alpine} +POSTGRES_LOCALE=${POSTGRES_LOCALE:-UTF8} +RQ_POSTGRES_HOST=${RQ_POSTGRES_HOST:-localhost} +RQ_POSTGRES_PORT=${RQ_POSTGRES_PORT:-5432} +RQ_POSTGRES_DBNAME=${RQ_POSTGRES_DBNAME:-datadog_test} +RQ_POSTGRES_USERNAME=${RQ_POSTGRES_USERNAME:-datadog} +RQ_POSTGRES_PASSWORD=${RQ_POSTGRES_PASSWORD:-datadog} PIP_PLATFORM=${PIP_PLATFORM:-manylinux2014_x86_64} AGENT_PYTHON_VERSION=${AGENT_PYTHON_VERSION:-} AGENT_PYTHON_ABI=${AGENT_PYTHON_ABI:-} AGENT_PID="" -POSTGRES_STARTED=0 +POSTGRES_COMPOSE_STARTED=0 log() { printf '\n[%s] %s\n' "$(date -u +%H:%M:%S)" "$*" @@ -37,10 +46,14 @@ cleanup_agent() { AGENT_PID="" } +docker_compose() { + docker compose -f "$POSTGRES_COMPOSE_FILE" -p "$POSTGRES_COMPOSE_PROJECT" "$@" +} + cleanup_all() { cleanup_agent - if [[ "$POSTGRES_STARTED" == "1" ]]; then - docker rm -f "$POSTGRES_CONTAINER" >/dev/null 2>&1 || true + if [[ "$POSTGRES_COMPOSE_STARTED" == "1" ]]; then + docker_compose down --remove-orphans >/dev/null 2>&1 || true fi } trap cleanup_all EXIT @@ -179,62 +192,72 @@ PY install_python_deps } +postgres_is_ready() { + PGPASSWORD="$RQ_POSTGRES_PASSWORD" psql \ + -h "$RQ_POSTGRES_HOST" \ + -p "$RQ_POSTGRES_PORT" \ + -U "$RQ_POSTGRES_USERNAME" \ + -d "$RQ_POSTGRES_DBNAME" \ + -c 'select 1' >/dev/null 2>&1 +} + +postgres_target_is_default_fixture() { + [[ "$RQ_POSTGRES_HOST" == "localhost" && \ + "$RQ_POSTGRES_PORT" == "5432" && \ + "$RQ_POSTGRES_DBNAME" == "datadog_test" && \ + "$RQ_POSTGRES_USERNAME" == "datadog" && \ + "$RQ_POSTGRES_PASSWORD" == "datadog" ]] +} + start_postgres_fixture() { - if psql 'postgresql://postgres@localhost:5432/postgres' -c 'select 1' >/dev/null 2>&1; then - log "Using existing local Postgres on localhost:5432" + if postgres_is_ready; then + log "Using existing compatible Postgres fixture at $RQ_POSTGRES_HOST:$RQ_POSTGRES_PORT/$RQ_POSTGRES_DBNAME" return fi - local published_containers proof_containers other_containers - published_containers=$(docker ps --filter publish=5432 --format '{{.Names}}' || true) - if [[ -n "$published_containers" ]]; then - proof_containers=$(grep -E '^rq-fused-local-par-agent-postgres-' <<<"$published_containers" || true) - other_containers=$(grep -Ev '^rq-fused-local-par-agent-postgres-' <<<"$published_containers" || true) - - if [[ -n "$proof_containers" ]]; then - log "Removing stale proof Postgres container(s) bound to localhost:5432 but not accepting psql" - xargs -r docker rm -f <<<"$proof_containers" >/dev/null - fi - if [[ -n "$other_containers" ]]; then - echo "Port 5432 is published by non-proof Docker container(s), and psql select 1 failed:" >&2 - sed 's/^/ /' <<<"$other_containers" >&2 - echo "Refusing to remove unrelated containers. Stop them or point this proof at a reachable local Postgres." >&2 - exit 1 - fi + if ! postgres_target_is_default_fixture; then + echo "Overridden RQ_POSTGRES_* target is not reachable; refusing to start the default compose fixture for a different target." >&2 + exit 1 fi - log "Starting disposable Postgres fixture $POSTGRES_CONTAINER from $POSTGRES_IMAGE" - if ! docker run --rm --name "$POSTGRES_CONTAINER" \ - -e POSTGRES_HOST_AUTH_METHOD=trust \ - -p 5432:5432 \ - -d "$POSTGRES_IMAGE" >/dev/null; then - echo "Failed to start disposable Postgres fixture on port 5432." >&2 - echo "If another local Postgres is intended, ensure this succeeds: psql postgresql://postgres@localhost:5432/postgres -c 'select 1'" >&2 + [[ -f "$POSTGRES_COMPOSE_FILE" ]] || { + echo "Postgres compose fixture not found: $POSTGRES_COMPOSE_FILE" >&2 exit 1 - fi - POSTGRES_STARTED=1 + } + docker compose version >/dev/null 2>&1 || { + echo "docker compose is required to start the integrations-core Postgres fixture" >&2 + exit 1 + } - for _ in $(seq 1 60); do - if psql 'postgresql://postgres@localhost:5432/postgres' -c 'select 1' >/dev/null 2>&1; then - log "Postgres fixture is ready" + log "Starting integrations-core Postgres fixture project=$POSTGRES_COMPOSE_PROJECT image=postgres:$POSTGRES_IMAGE" + export POSTGRES_IMAGE POSTGRES_LOCALE + docker_compose down --remove-orphans >/dev/null 2>&1 || true + docker_compose up -d postgres + POSTGRES_COMPOSE_STARTED=1 + + for _ in $(seq 1 120); do + if postgres_is_ready && docker_compose exec -T postgres test -e /tmp/container_ready.txt >/dev/null 2>&1; then + log "Integrations-core Postgres fixture is ready at $RQ_POSTGRES_HOST:$RQ_POSTGRES_PORT/$RQ_POSTGRES_DBNAME" return fi sleep 0.5 done - docker logs "$POSTGRES_CONTAINER" | tail -80 >&2 || true - echo "Postgres fixture did not become ready" >&2 + docker_compose ps >&2 || true + docker_compose logs --tail=80 postgres >&2 || true + echo "Integrations-core Postgres fixture did not become ready" >&2 exit 1 } write_postgres_config() { - cat > "$TMP_ROOT/conf.d/postgres.d/conf.yaml" <<'YAML' + cat > "$TMP_ROOT/conf.d/postgres.d/conf.yaml" < loaded Postgres check -> SELECT 1 AS value -> fakeintake publish. # The HTTP execute endpoint remains as a dev preflight for local evidence only. # -# Defaults assume the remote-queries-poc worktree layout. Override AGENT_REPO, -# INTEGRATIONS_CORE, TMP_ROOT, CMD_PORT, POSTGRES_IMAGE, or AGENT_PYTHON_VERSION / -# AGENT_PYTHON_ABI if needed. This proof intentionally follows the repository's +# Defaults assume the remote-queries-poc worktree layout and reuse the +# integrations-core Postgres integration test compose fixture. Override +# AGENT_REPO, INTEGRATIONS_CORE, TMP_ROOT, CMD_PORT, POSTGRES_COMPOSE_FILE, +# POSTGRES_COMPOSE_PROJECT, POSTGRES_IMAGE, RQ_POSTGRES_*, or +# AGENT_PYTHON_VERSION / AGENT_PYTHON_ABI if needed. This proof intentionally follows the repository's # fakeintake/OPMS precedent and sets DD_INTERNAL_PAR_SKIP_TASK_VERIFICATION=true # for the standalone-process tracer bullet. Signed task verification is postponed # to backend/AP/RC work. @@ -18,14 +20,21 @@ AGENT_REPO=${AGENT_REPO:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)} INTEGRATIONS_CORE=${INTEGRATIONS_CORE:-/home/bits/dd/tasks/remote-queries-poc/worktrees/integrations-core} TMP_ROOT=${TMP_ROOT:-/tmp/rq-standalone-par-agent-postgres} CMD_PORT=${CMD_PORT:-55003} -POSTGRES_IMAGE=${POSTGRES_IMAGE:-postgres:11-alpine} -POSTGRES_CONTAINER=${POSTGRES_CONTAINER:-rq-standalone-par-agent-postgres-$$} +POSTGRES_COMPOSE_FILE=${POSTGRES_COMPOSE_FILE:-$INTEGRATIONS_CORE/postgres/tests/compose/docker-compose.yaml} +POSTGRES_COMPOSE_PROJECT=${POSTGRES_COMPOSE_PROJECT:-rq-standalone-par-agent-postgres-$$} +POSTGRES_IMAGE=${POSTGRES_IMAGE:-13-alpine} +POSTGRES_LOCALE=${POSTGRES_LOCALE:-UTF8} +RQ_POSTGRES_HOST=${RQ_POSTGRES_HOST:-localhost} +RQ_POSTGRES_PORT=${RQ_POSTGRES_PORT:-5432} +RQ_POSTGRES_DBNAME=${RQ_POSTGRES_DBNAME:-datadog_test} +RQ_POSTGRES_USERNAME=${RQ_POSTGRES_USERNAME:-datadog} +RQ_POSTGRES_PASSWORD=${RQ_POSTGRES_PASSWORD:-datadog} PIP_PLATFORM=${PIP_PLATFORM:-manylinux2014_x86_64} AGENT_PYTHON_VERSION=${AGENT_PYTHON_VERSION:-} AGENT_PYTHON_ABI=${AGENT_PYTHON_ABI:-} AGENT_PID="" -POSTGRES_STARTED=0 +POSTGRES_COMPOSE_STARTED=0 log() { printf '\n[%s] %s\n' "$(date -u +%H:%M:%S)" "$*" @@ -39,10 +48,14 @@ cleanup_agent() { AGENT_PID="" } +docker_compose() { + docker compose -f "$POSTGRES_COMPOSE_FILE" -p "$POSTGRES_COMPOSE_PROJECT" "$@" +} + cleanup_all() { cleanup_agent - if [[ "$POSTGRES_STARTED" == "1" ]]; then - docker rm -f "$POSTGRES_CONTAINER" >/dev/null 2>&1 || true + if [[ "$POSTGRES_COMPOSE_STARTED" == "1" ]]; then + docker_compose down --remove-orphans >/dev/null 2>&1 || true fi } trap cleanup_all EXIT @@ -181,62 +194,72 @@ PY install_python_deps } +postgres_is_ready() { + PGPASSWORD="$RQ_POSTGRES_PASSWORD" psql \ + -h "$RQ_POSTGRES_HOST" \ + -p "$RQ_POSTGRES_PORT" \ + -U "$RQ_POSTGRES_USERNAME" \ + -d "$RQ_POSTGRES_DBNAME" \ + -c 'select 1' >/dev/null 2>&1 +} + +postgres_target_is_default_fixture() { + [[ "$RQ_POSTGRES_HOST" == "localhost" && \ + "$RQ_POSTGRES_PORT" == "5432" && \ + "$RQ_POSTGRES_DBNAME" == "datadog_test" && \ + "$RQ_POSTGRES_USERNAME" == "datadog" && \ + "$RQ_POSTGRES_PASSWORD" == "datadog" ]] +} + start_postgres_fixture() { - if psql 'postgresql://postgres@localhost:5432/postgres' -c 'select 1' >/dev/null 2>&1; then - log "Using existing local Postgres on localhost:5432" + if postgres_is_ready; then + log "Using existing compatible Postgres fixture at $RQ_POSTGRES_HOST:$RQ_POSTGRES_PORT/$RQ_POSTGRES_DBNAME" return fi - local published_containers proof_containers other_containers - published_containers=$(docker ps --filter publish=5432 --format '{{.Names}}' || true) - if [[ -n "$published_containers" ]]; then - proof_containers=$(grep -E '^rq-standalone-par-agent-postgres-' <<<"$published_containers" || true) - other_containers=$(grep -Ev '^rq-standalone-par-agent-postgres-' <<<"$published_containers" || true) - - if [[ -n "$proof_containers" ]]; then - log "Removing stale proof Postgres container(s) bound to localhost:5432 but not accepting psql" - xargs -r docker rm -f <<<"$proof_containers" >/dev/null - fi - if [[ -n "$other_containers" ]]; then - echo "Port 5432 is published by non-proof Docker container(s), and psql select 1 failed:" >&2 - sed 's/^/ /' <<<"$other_containers" >&2 - echo "Refusing to remove unrelated containers. Stop them or point this proof at a reachable local Postgres." >&2 - exit 1 - fi + if ! postgres_target_is_default_fixture; then + echo "Overridden RQ_POSTGRES_* target is not reachable; refusing to start the default compose fixture for a different target." >&2 + exit 1 fi - log "Starting disposable Postgres fixture $POSTGRES_CONTAINER from $POSTGRES_IMAGE" - if ! docker run --rm --name "$POSTGRES_CONTAINER" \ - -e POSTGRES_HOST_AUTH_METHOD=trust \ - -p 5432:5432 \ - -d "$POSTGRES_IMAGE" >/dev/null; then - echo "Failed to start disposable Postgres fixture on port 5432." >&2 - echo "If another local Postgres is intended, ensure this succeeds: psql postgresql://postgres@localhost:5432/postgres -c 'select 1'" >&2 + [[ -f "$POSTGRES_COMPOSE_FILE" ]] || { + echo "Postgres compose fixture not found: $POSTGRES_COMPOSE_FILE" >&2 exit 1 - fi - POSTGRES_STARTED=1 + } + docker compose version >/dev/null 2>&1 || { + echo "docker compose is required to start the integrations-core Postgres fixture" >&2 + exit 1 + } - for _ in $(seq 1 60); do - if psql 'postgresql://postgres@localhost:5432/postgres' -c 'select 1' >/dev/null 2>&1; then - log "Postgres fixture is ready" + log "Starting integrations-core Postgres fixture project=$POSTGRES_COMPOSE_PROJECT image=postgres:$POSTGRES_IMAGE" + export POSTGRES_IMAGE POSTGRES_LOCALE + docker_compose down --remove-orphans >/dev/null 2>&1 || true + docker_compose up -d postgres + POSTGRES_COMPOSE_STARTED=1 + + for _ in $(seq 1 120); do + if postgres_is_ready && docker_compose exec -T postgres test -e /tmp/container_ready.txt >/dev/null 2>&1; then + log "Integrations-core Postgres fixture is ready at $RQ_POSTGRES_HOST:$RQ_POSTGRES_PORT/$RQ_POSTGRES_DBNAME" return fi sleep 0.5 done - docker logs "$POSTGRES_CONTAINER" | tail -80 >&2 || true - echo "Postgres fixture did not become ready" >&2 + docker_compose ps >&2 || true + docker_compose logs --tail=80 postgres >&2 || true + echo "Integrations-core Postgres fixture did not become ready" >&2 exit 1 } write_postgres_config() { - cat > "$TMP_ROOT/conf.d/postgres.d/conf.yaml" <<'YAML' + cat > "$TMP_ROOT/conf.d/postgres.d/conf.yaml" < Date: Thu, 21 May 2026 18:21:23 +0000 Subject: [PATCH 19/33] Query fixture table in Remote Queries proof --- .../impl/remote_query_execute.go | 16 ++++- .../impl/remote_query_par_poc_test.go | 6 +- comp/remotequeries/impl/remote_query_test.go | 37 +++++++++++ .../bundles/remotequeries/execute_test.go | 17 +++-- .../live_agent_ipc_par_loop_test.go | 64 +++++++++++++++---- .../standalone_par_process_proof_test.go | 14 ++-- .../fused-local-par-agent-postgres-proof.sh | 47 ++++++++++---- ...andalone-par-agentsecure-postgres-proof.sh | 47 ++++++++++---- 8 files changed, 188 insertions(+), 60 deletions(-) diff --git a/comp/remotequeries/impl/remote_query_execute.go b/comp/remotequeries/impl/remote_query_execute.go index bf2fcc59008c..b98db6a7f42c 100644 --- a/comp/remotequeries/impl/remote_query_execute.go +++ b/comp/remotequeries/impl/remote_query_execute.go @@ -24,7 +24,8 @@ const ( // RemoteQueriesExecuteEnabledConfig is disabled by default when the key is absent. RemoteQueriesExecuteEnabledConfig = "remote_queries.execute.enabled" - remoteQueryProofQuery = "SELECT 1 AS value" + remoteQueryProofSeedQuery = "SELECT 1 AS value" + remoteQueryFixtureTableProofQuery = "SELECT city, country FROM cities ORDER BY city" statusExecutorUnavailable = "executor_unavailable" ) @@ -33,6 +34,15 @@ type remoteQueryRunner interface { RunRemoteQueryJSON(integration string, requestJSON string) (string, error) } +func isRemoteQueryAllowedProofQuery(query string) bool { + switch query { + case remoteQueryProofSeedQuery, remoteQueryFixtureTableProofQuery: + return true + default: + return false + } +} + type remoteQueryCheckUnwrapper interface { Unwrap() check.Check } @@ -138,7 +148,7 @@ func NewRemoteQueryExecuteRequest(integration string, target RemoteQueryExecuteT if query == "" { return RemoteQueryExecuteRequest{}, fmt.Errorf("query is required") } - if query != remoteQueryProofQuery { + if !isRemoteQueryAllowedProofQuery(query) { return RemoteQueryExecuteRequest{}, fmt.Errorf("query is not allowed") } @@ -258,7 +268,7 @@ func parseExecuteRequest(r *http.Request) (remoteQueryExecuteRequest, string, er if wireReq.Query == "" { return remoteQueryExecuteRequest{}, "", fmt.Errorf("query is required") } - if wireReq.Query != remoteQueryProofQuery { + if !isRemoteQueryAllowedProofQuery(wireReq.Query) { return remoteQueryExecuteRequest{}, "", fmt.Errorf("query is not allowed") } diff --git a/comp/remotequeries/impl/remote_query_par_poc_test.go b/comp/remotequeries/impl/remote_query_par_poc_test.go index 7525784c885c..f4269b4f8d41 100644 --- a/comp/remotequeries/impl/remote_query_par_poc_test.go +++ b/comp/remotequeries/impl/remote_query_par_poc_test.go @@ -26,7 +26,7 @@ func TestRemoteQueryPARHarnessUsesCredentialFreeIPCPostShape(t *testing.T) { result, err := harness.Execute(context.Background(), RemoteQueryPARInputs{ Integration: "postgres", Target: remoteQueryTargetJSON{Host: "localhost", Port: 5432, DBName: "postgres"}, - Query: remoteQueryProofQuery, + Query: remoteQueryProofSeedQuery, Limits: &remoteQueryExecuteLimitsJSON{MaxRows: 1, MaxBytes: 1024, TimeoutMs: 1000}, }) @@ -55,7 +55,7 @@ func TestRemoteQueryPARHarnessWithRealAgentIPCClient(t *testing.T) { result, err := harness.Execute(context.Background(), RemoteQueryPARInputs{ Integration: "postgres", Target: remoteQueryTargetJSON{Host: "LOCALHOST.", Port: 5432, DBName: "postgres"}, - Query: remoteQueryProofQuery, + Query: remoteQueryProofSeedQuery, }) require.NoError(t, err) @@ -87,7 +87,7 @@ func TestRemoteQueryPARHarnessPropagatesSanitizedBridgeErrors(t *testing.T) { inputs: RemoteQueryPARInputs{ Integration: "postgres", Target: remoteQueryTargetJSON{Host: "localhost", Port: 5432, DBName: "other"}, - Query: remoteQueryProofQuery, + Query: remoteQueryProofSeedQuery, }, wantStatus: statusTargetNotFound, wantCode: statusTargetNotFound, diff --git a/comp/remotequeries/impl/remote_query_test.go b/comp/remotequeries/impl/remote_query_test.go index 78e69e3f2b87..19982d213a98 100644 --- a/comp/remotequeries/impl/remote_query_test.go +++ b/comp/remotequeries/impl/remote_query_test.go @@ -365,6 +365,27 @@ func TestParseExecuteRequestAllowsOmittedLimits(t *testing.T) { assert.NotContains(t, requestJSON, "integration") } +func TestParseExecuteRequestAllowsFixtureTableProofQuery(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, RemoteQueryExecuteEndpointPath, strings.NewReader( + `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT city, country FROM cities ORDER BY city","limits":{"maxRows":2,"maxBytes":1024,"timeoutMs":1000}}`, + )) + req.Header.Set("Content-Type", "application/json") + + parsed, requestJSON, err := parseExecuteRequest(req) + require.NoError(t, err) + assert.Equal(t, remoteQueryFixtureTableProofQuery, parsed.Query) + assert.JSONEq(t, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT city, country FROM cities ORDER BY city","limits":{"maxRows":2,"maxBytes":1024,"timeoutMs":1000}}`, requestJSON) + assert.NotContains(t, requestJSON, "integration") +} + +func TestNewRemoteQueryExecuteRequestAllowsFixtureTableProofQuery(t *testing.T) { + req, err := NewRemoteQueryExecuteRequest("postgres", RemoteQueryExecuteTarget{Host: " LocalHost. ", Port: 5432, DBName: "postgres"}, remoteQueryFixtureTableProofQuery, &RemoteQueryExecuteLimits{MaxRows: 2, MaxBytes: 1024, TimeoutMs: 1000}) + require.NoError(t, err) + assert.Equal(t, "postgres", req.Integration) + assert.Equal(t, RemoteQueryExecuteTarget{Host: "localhost", Port: 5432, DBName: "postgres"}, req.Target) + assert.Equal(t, remoteQueryFixtureTableProofQuery, req.Query) +} + func TestRemoteQueryExecuteHandlerDisabled(t *testing.T) { handler := &remoteQueryExecuteHandler{enabled: false, collector: fakeCollector{}} @@ -390,6 +411,22 @@ func TestRemoteQueryExecuteHandlerRunnerSuccess(t *testing.T) { assert.NotContains(t, recorder.Body.String(), "secret-value") } +func TestRemoteQueryExecuteHandlerRunnerSuccessWithFixtureTableQuery(t *testing.T) { + runner := &fakeRunnerCheck{ + fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-value\n"}, + response: `{"status":"SUCCEEDED","rows":[{"city":"Beautiful city of lights","country":"France"},{"city":"New York","country":"USA"}]}`, + } + handler := &remoteQueryExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{fakeWrappedCheck{Check: runner}}}} + + recorder := callExecuteHandler(handler, `{"integration":"postgres","target":{"host":"LOCALHOST.","port":5432,"dbname":"postgres"},"query":"SELECT city, country FROM cities ORDER BY city"}`) + + assert.Equal(t, http.StatusOK, recorder.Code) + assert.JSONEq(t, `{"status":"SUCCEEDED","rows":[{"city":"Beautiful city of lights","country":"France"},{"city":"New York","country":"USA"}]}`, recorder.Body.String()) + assert.JSONEq(t, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT city, country FROM cities ORDER BY city"}`, runner.seenRequest()) + assert.NotContains(t, runner.seenRequest(), "integration") + assert.NotContains(t, recorder.Body.String(), "secret-value") +} + func TestRemoteQueryExecuteHandlerRejectsInvalidIntegration(t *testing.T) { handler := &remoteQueryExecuteHandler{enabled: true, collector: fakeCollector{}} diff --git a/pkg/privateactionrunner/bundles/remotequeries/execute_test.go b/pkg/privateactionrunner/bundles/remotequeries/execute_test.go index 9c2e27711084..a6c6db9e327b 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/execute_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/execute_test.go @@ -22,9 +22,11 @@ import ( ) func TestExecuteActionUsesCredentialFreeAgentSecureRequestShape(t *testing.T) { - row, err := structpb.NewStruct(map[string]interface{}{"value": 1}) + franceRow, err := structpb.NewStruct(map[string]interface{}{"city": "Beautiful city of lights", "country": "France"}) require.NoError(t, err) - client := &captureBridgeClient{response: &pb.RemoteQueryExecuteResponse{Status: "SUCCEEDED", Rows: []*structpb.Struct{row}}} + usaRow, err := structpb.NewStruct(map[string]interface{}{"city": "New York", "country": "USA"}) + require.NoError(t, err) + client := &captureBridgeClient{response: &pb.RemoteQueryExecuteResponse{Status: "SUCCEEDED", Rows: []*structpb.Struct{franceRow, usaRow}}} action := NewExecuteAction(func() (BridgeClient, error) { return client, nil }) @@ -36,9 +38,9 @@ func TestExecuteActionUsesCredentialFreeAgentSecureRequestShape(t *testing.T) { "port": 5432, "dbname": "postgres", }, - "query": "SELECT 1 AS value", + "query": "SELECT city, country FROM cities ORDER BY city", "limits": map[string]interface{}{ - "maxRows": 1, + "maxRows": 2, "maxBytes": 1024, "timeoutMs": 1000, }, @@ -50,15 +52,16 @@ func TestExecuteActionUsesCredentialFreeAgentSecureRequestShape(t *testing.T) { assert.Equal(t, "localhost", client.request.GetTarget().GetHost()) assert.Equal(t, int32(5432), client.request.GetTarget().GetPort()) assert.Equal(t, "postgres", client.request.GetTarget().GetDbname()) - assert.Equal(t, "SELECT 1 AS value", client.request.GetQuery()) - assert.Equal(t, int32(1), client.request.GetLimits().GetMaxRows()) + assert.Equal(t, "SELECT city, country FROM cities ORDER BY city", client.request.GetQuery()) + assert.Equal(t, int32(2), client.request.GetLimits().GetMaxRows()) requestEvidence, err := json.Marshal(client.request) require.NoError(t, err) assert.NotContains(t, string(requestEvidence), "secret-value") assert.Equal(t, map[string]interface{}{ "status": "SUCCEEDED", "rows": []interface{}{ - map[string]interface{}{"value": float64(1)}, + map[string]interface{}{"city": "Beautiful city of lights", "country": "France"}, + map[string]interface{}{"city": "New York", "country": "USA"}, }, }, output) } diff --git a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go index b0af7c3350ce..5484ef2115c8 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go @@ -28,7 +28,12 @@ import ( "github.com/stretchr/testify/require" ) -const fusedLocalProofEnv = "RQ_FUSED_PROOF" +const ( + fusedLocalProofEnv = "RQ_FUSED_PROOF" + remoteQueriesSeedProofQuery = "SELECT 1 AS value" + remoteQueriesFixtureTableProofQuery = "SELECT city, country FROM cities ORDER BY city" + remoteQueriesProofQueryOverrideEnv = "RQ_REMOTE_QUERY" +) func TestRemoteQueriesActionRunsThroughLivePARLoopWithRealAgentIPC(t *testing.T) { if os.Getenv(fusedLocalProofEnv) != "1" { @@ -73,15 +78,12 @@ func TestRemoteQueriesActionRunsThroughLivePARLoopWithRealAgentIPC(t *testing.T) }() taskID := fmt.Sprintf("remotequeries-fused-local-proof-%d", time.Now().UnixNano()) + proofQuery := remoteQueriesProofQueryFromEnv() inputs := map[string]interface{}{ "integration": "postgres", "target": remoteQueriesPostgresTargetFromEnv(t), - "query": "SELECT 1 AS value", - "limits": map[string]interface{}{ - "maxRows": 1, - "maxBytes": 1024, - "timeoutMs": 1000, - }, + "query": proofQuery, + "limits": remoteQueriesProofLimits(proofQuery), } requestEvidence, err := json.Marshal(inputs) require.NoError(t, err) @@ -106,10 +108,7 @@ func TestRemoteQueriesActionRunsThroughLivePARLoopWithRealAgentIPC(t *testing.T) rows, ok := result.Outputs["rows"].([]interface{}) require.True(t, ok) - require.Len(t, rows, 1) - firstRow, ok := rows[0].(map[string]interface{}) - require.True(t, ok) - assert.Equal(t, float64(1), firstRow["value"]) + assertRemoteQueriesProofRows(t, proofQuery, rows) resultEvidence, err := json.Marshal(result.Outputs) require.NoError(t, err) @@ -136,6 +135,49 @@ func getenvOptional(name string) string { return os.Getenv(name) } +func remoteQueriesProofQueryFromEnv() string { + if query := os.Getenv(remoteQueriesProofQueryOverrideEnv); query != "" { + return query + } + return remoteQueriesFixtureTableProofQuery +} + +func remoteQueriesProofLimits(query string) map[string]interface{} { + maxRows := 1 + if query == remoteQueriesFixtureTableProofQuery { + maxRows = 2 + } + return map[string]interface{}{ + "maxRows": maxRows, + "maxBytes": 1024, + "timeoutMs": 1000, + } +} + +func assertRemoteQueriesProofRows(t *testing.T, query string, rows []interface{}) { + t.Helper() + + switch query { + case remoteQueriesFixtureTableProofQuery: + require.Len(t, rows, 2) + firstRow, ok := rows[0].(map[string]interface{}) + require.True(t, ok) + assert.Equal(t, "Beautiful city of lights", firstRow["city"]) + assert.Equal(t, "France", firstRow["country"]) + secondRow, ok := rows[1].(map[string]interface{}) + require.True(t, ok) + assert.Equal(t, "New York", secondRow["city"]) + assert.Equal(t, "USA", secondRow["country"]) + case remoteQueriesSeedProofQuery: + require.Len(t, rows, 1) + firstRow, ok := rows[0].(map[string]interface{}) + require.True(t, ok) + assert.Equal(t, float64(1), firstRow["value"]) + default: + require.FailNowf(t, "unsupported proof query", "%s=%q must use a bridge-allowlisted proof query", remoteQueriesProofQueryOverrideEnv, query) + } +} + func remoteQueriesPostgresTargetFromEnv(t *testing.T) map[string]interface{} { t.Helper() diff --git a/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go b/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go index 927af14899f5..f6eeac770835 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go @@ -78,15 +78,12 @@ func TestRemoteQueriesActionRunsThroughStandalonePARProcessWithRealAgentIPC(t *t waitForStandalonePARPolling(t, fakeintakeClient, cmd, parLog, &stdout, &stderr) taskID := fmt.Sprintf("remotequeries-standalone-par-proof-%d", time.Now().UnixNano()) + proofQuery := remoteQueriesProofQueryFromEnv() inputs := map[string]interface{}{ "integration": "postgres", "target": remoteQueriesPostgresTargetFromEnv(t), - "query": "SELECT 1 AS value", - "limits": map[string]interface{}{ - "maxRows": 1, - "maxBytes": 1024, - "timeoutMs": 1000, - }, + "query": proofQuery, + "limits": remoteQueriesProofLimits(proofQuery), } requestEvidence, err := json.Marshal(inputs) require.NoError(t, err) @@ -110,10 +107,7 @@ func TestRemoteQueriesActionRunsThroughStandalonePARProcessWithRealAgentIPC(t *t rows, ok := result.Outputs["rows"].([]interface{}) require.True(t, ok) - require.Len(t, rows, 1) - firstRow, ok := rows[0].(map[string]interface{}) - require.True(t, ok) - assert.Equal(t, float64(1), firstRow["value"]) + assertRemoteQueriesProofRows(t, proofQuery, rows) resultEvidence, err := json.Marshal(result.Outputs) require.NoError(t, err) diff --git a/test/remotequeries/fused-local-par-agent-postgres-proof.sh b/test/remotequeries/fused-local-par-agent-postgres-proof.sh index 1fb72a2a170e..e593ccae1ea7 100755 --- a/test/remotequeries/fused-local-par-agent-postgres-proof.sh +++ b/test/remotequeries/fused-local-par-agent-postgres-proof.sh @@ -2,13 +2,13 @@ # Runs the fused local-only Remote Queries proof: # fakeintake -> live WorkflowRunner PAR loop -> com.datadoghq.remotequeries.execute # -> real local AgentSecure gRPC RemoteQueryExecute over Agent IPC TLS/auth -# -> loaded Postgres check -> SELECT 1 AS value -> fakeintake publish. +# -> loaded Postgres check -> fixture-table proof query -> fakeintake publish. # The HTTP execute endpoint remains as a dev preflight for local evidence only. # # Defaults assume the remote-queries-poc worktree layout and reuse the # integrations-core Postgres integration test compose fixture. Override # AGENT_REPO, INTEGRATIONS_CORE, TMP_ROOT, CMD_PORT, POSTGRES_COMPOSE_FILE, -# POSTGRES_COMPOSE_PROJECT, POSTGRES_IMAGE, RQ_POSTGRES_*, or +# POSTGRES_COMPOSE_PROJECT, POSTGRES_IMAGE, RQ_REMOTE_QUERY, RQ_POSTGRES_*, or # AGENT_PYTHON_VERSION / AGENT_PYTHON_ABI if needed. The proof is local-only and intentionally sets # DD_INTERNAL_PAR_SKIP_TASK_VERIFICATION=true inside the Go proof test. @@ -22,11 +22,12 @@ POSTGRES_COMPOSE_FILE=${POSTGRES_COMPOSE_FILE:-$INTEGRATIONS_CORE/postgres/tests POSTGRES_COMPOSE_PROJECT=${POSTGRES_COMPOSE_PROJECT:-rq-fused-local-par-agent-postgres-$$} POSTGRES_IMAGE=${POSTGRES_IMAGE:-13-alpine} POSTGRES_LOCALE=${POSTGRES_LOCALE:-UTF8} +RQ_REMOTE_QUERY=${RQ_REMOTE_QUERY:-SELECT city, country FROM cities ORDER BY city} RQ_POSTGRES_HOST=${RQ_POSTGRES_HOST:-localhost} RQ_POSTGRES_PORT=${RQ_POSTGRES_PORT:-5432} RQ_POSTGRES_DBNAME=${RQ_POSTGRES_DBNAME:-datadog_test} -RQ_POSTGRES_USERNAME=${RQ_POSTGRES_USERNAME:-datadog} -RQ_POSTGRES_PASSWORD=${RQ_POSTGRES_PASSWORD:-datadog} +RQ_POSTGRES_USERNAME=${RQ_POSTGRES_USERNAME:-bob} +RQ_POSTGRES_PASSWORD=${RQ_POSTGRES_PASSWORD:-bob} PIP_PLATFORM=${PIP_PLATFORM:-manylinux2014_x86_64} AGENT_PYTHON_VERSION=${AGENT_PYTHON_VERSION:-} AGENT_PYTHON_ABI=${AGENT_PYTHON_ABI:-} @@ -205,8 +206,8 @@ postgres_target_is_default_fixture() { [[ "$RQ_POSTGRES_HOST" == "localhost" && \ "$RQ_POSTGRES_PORT" == "5432" && \ "$RQ_POSTGRES_DBNAME" == "datadog_test" && \ - "$RQ_POSTGRES_USERNAME" == "datadog" && \ - "$RQ_POSTGRES_PASSWORD" == "datadog" ]] + "$RQ_POSTGRES_USERNAME" == "bob" && \ + "$RQ_POSTGRES_PASSWORD" == "bob" ]] } start_postgres_fixture() { @@ -289,7 +290,7 @@ start_agent_and_wait_for_postgres_check() { call_agent_execute_preflight() { local payload - payload=$(RQ_POSTGRES_HOST="$RQ_POSTGRES_HOST" RQ_POSTGRES_PORT="$RQ_POSTGRES_PORT" RQ_POSTGRES_DBNAME="$RQ_POSTGRES_DBNAME" python3 - <<'PY' + payload=$(RQ_POSTGRES_HOST="$RQ_POSTGRES_HOST" RQ_POSTGRES_PORT="$RQ_POSTGRES_PORT" RQ_POSTGRES_DBNAME="$RQ_POSTGRES_DBNAME" RQ_REMOTE_QUERY="$RQ_REMOTE_QUERY" python3 - <<'PY' import json import os print(json.dumps({ @@ -299,8 +300,8 @@ print(json.dumps({ "port": int(os.environ["RQ_POSTGRES_PORT"]), "dbname": os.environ["RQ_POSTGRES_DBNAME"], }, - "query": "SELECT 1 AS value", - "limits": {"maxRows": 1, "maxBytes": 1024, "timeoutMs": 1000}, + "query": os.environ["RQ_REMOTE_QUERY"], + "limits": {"maxRows": 2 if os.environ["RQ_REMOTE_QUERY"] == "SELECT city, country FROM cities ORDER BY city" else 1, "maxBytes": 1024, "timeoutMs": 1000}, })) PY ) @@ -322,10 +323,29 @@ PY echo "FAIL: expected Agent execute preflight HTTP 200, got $status" >&2 exit 1 fi - if ! grep -Eq '"status"[[:space:]]*:[[:space:]]*"SUCCEEDED".*"value"[[:space:]]*:[[:space:]]*1|"value"[[:space:]]*:[[:space:]]*1.*"status"[[:space:]]*:[[:space:]]*"SUCCEEDED"' "$TMP_ROOT/results/agent-execute-preflight.body"; then - echo "FAIL: Agent execute preflight response did not contain SUCCEEDED row value=1" >&2 - exit 1 - fi + RQ_REMOTE_QUERY="$RQ_REMOTE_QUERY" python3 - "$TMP_ROOT/results/agent-execute-preflight.body" <<'PY' +import json +import os +import sys + +with open(sys.argv[1], encoding="utf-8") as f: + body = json.load(f) + +if body.get("status") != "SUCCEEDED": + raise SystemExit(f"Agent execute preflight response status was not SUCCEEDED: {body}") + +rows = body.get("rows") +query = os.environ["RQ_REMOTE_QUERY"] +if query == "SELECT city, country FROM cities ORDER BY city": + expected = [ + {"city": "Beautiful city of lights", "country": "France"}, + {"city": "New York", "country": "USA"}, + ] +else: + expected = [{"value": 1}] +if rows != expected: + raise SystemExit(f"Agent execute preflight rows did not match expected fixture data: rows={rows!r} expected={expected!r}") +PY if grep -Eq 'password|token|secret' "$TMP_ROOT/results/agent-execute-preflight.body"; then echo "FAIL: Agent execute preflight response contained credential-shaped text" >&2 exit 1 @@ -344,6 +364,7 @@ run_fused_go_proof() { RQ_POSTGRES_HOST="$RQ_POSTGRES_HOST" \ RQ_POSTGRES_PORT="$RQ_POSTGRES_PORT" \ RQ_POSTGRES_DBNAME="$RQ_POSTGRES_DBNAME" \ + RQ_REMOTE_QUERY="$RQ_REMOTE_QUERY" \ dda inv test --targets=./pkg/privateactionrunner/bundles/remotequeries \ --extra-args='-run TestRemoteQueriesActionRunsThroughLivePARLoopWithRealAgentIPC -count=1 -v' ) | tee "$TMP_ROOT/results/fused-proof-test.log" diff --git a/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh b/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh index 12a5c32e2a97..c954c9d28e2b 100755 --- a/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh +++ b/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh @@ -2,13 +2,13 @@ # Runs the standalone local-only Remote Queries proof: # fakeintake -> standalone OS private-action-runner process -> com.datadoghq.remotequeries.execute # -> real local AgentSecure gRPC RemoteQueryExecute over Agent IPC TLS/auth -# -> loaded Postgres check -> SELECT 1 AS value -> fakeintake publish. +# -> loaded Postgres check -> fixture-table proof query -> fakeintake publish. # The HTTP execute endpoint remains as a dev preflight for local evidence only. # # Defaults assume the remote-queries-poc worktree layout and reuse the # integrations-core Postgres integration test compose fixture. Override # AGENT_REPO, INTEGRATIONS_CORE, TMP_ROOT, CMD_PORT, POSTGRES_COMPOSE_FILE, -# POSTGRES_COMPOSE_PROJECT, POSTGRES_IMAGE, RQ_POSTGRES_*, or +# POSTGRES_COMPOSE_PROJECT, POSTGRES_IMAGE, RQ_REMOTE_QUERY, RQ_POSTGRES_*, or # AGENT_PYTHON_VERSION / AGENT_PYTHON_ABI if needed. This proof intentionally follows the repository's # fakeintake/OPMS precedent and sets DD_INTERNAL_PAR_SKIP_TASK_VERIFICATION=true # for the standalone-process tracer bullet. Signed task verification is postponed @@ -24,11 +24,12 @@ POSTGRES_COMPOSE_FILE=${POSTGRES_COMPOSE_FILE:-$INTEGRATIONS_CORE/postgres/tests POSTGRES_COMPOSE_PROJECT=${POSTGRES_COMPOSE_PROJECT:-rq-standalone-par-agent-postgres-$$} POSTGRES_IMAGE=${POSTGRES_IMAGE:-13-alpine} POSTGRES_LOCALE=${POSTGRES_LOCALE:-UTF8} +RQ_REMOTE_QUERY=${RQ_REMOTE_QUERY:-SELECT city, country FROM cities ORDER BY city} RQ_POSTGRES_HOST=${RQ_POSTGRES_HOST:-localhost} RQ_POSTGRES_PORT=${RQ_POSTGRES_PORT:-5432} RQ_POSTGRES_DBNAME=${RQ_POSTGRES_DBNAME:-datadog_test} -RQ_POSTGRES_USERNAME=${RQ_POSTGRES_USERNAME:-datadog} -RQ_POSTGRES_PASSWORD=${RQ_POSTGRES_PASSWORD:-datadog} +RQ_POSTGRES_USERNAME=${RQ_POSTGRES_USERNAME:-bob} +RQ_POSTGRES_PASSWORD=${RQ_POSTGRES_PASSWORD:-bob} PIP_PLATFORM=${PIP_PLATFORM:-manylinux2014_x86_64} AGENT_PYTHON_VERSION=${AGENT_PYTHON_VERSION:-} AGENT_PYTHON_ABI=${AGENT_PYTHON_ABI:-} @@ -207,8 +208,8 @@ postgres_target_is_default_fixture() { [[ "$RQ_POSTGRES_HOST" == "localhost" && \ "$RQ_POSTGRES_PORT" == "5432" && \ "$RQ_POSTGRES_DBNAME" == "datadog_test" && \ - "$RQ_POSTGRES_USERNAME" == "datadog" && \ - "$RQ_POSTGRES_PASSWORD" == "datadog" ]] + "$RQ_POSTGRES_USERNAME" == "bob" && \ + "$RQ_POSTGRES_PASSWORD" == "bob" ]] } start_postgres_fixture() { @@ -291,7 +292,7 @@ start_agent_and_wait_for_postgres_check() { call_agent_execute_preflight() { local payload - payload=$(RQ_POSTGRES_HOST="$RQ_POSTGRES_HOST" RQ_POSTGRES_PORT="$RQ_POSTGRES_PORT" RQ_POSTGRES_DBNAME="$RQ_POSTGRES_DBNAME" python3 - <<'PY' + payload=$(RQ_POSTGRES_HOST="$RQ_POSTGRES_HOST" RQ_POSTGRES_PORT="$RQ_POSTGRES_PORT" RQ_POSTGRES_DBNAME="$RQ_POSTGRES_DBNAME" RQ_REMOTE_QUERY="$RQ_REMOTE_QUERY" python3 - <<'PY' import json import os print(json.dumps({ @@ -301,8 +302,8 @@ print(json.dumps({ "port": int(os.environ["RQ_POSTGRES_PORT"]), "dbname": os.environ["RQ_POSTGRES_DBNAME"], }, - "query": "SELECT 1 AS value", - "limits": {"maxRows": 1, "maxBytes": 1024, "timeoutMs": 1000}, + "query": os.environ["RQ_REMOTE_QUERY"], + "limits": {"maxRows": 2 if os.environ["RQ_REMOTE_QUERY"] == "SELECT city, country FROM cities ORDER BY city" else 1, "maxBytes": 1024, "timeoutMs": 1000}, })) PY ) @@ -324,10 +325,29 @@ PY echo "FAIL: expected Agent execute preflight HTTP 200, got $status" >&2 exit 1 fi - if ! grep -Eq '"status"[[:space:]]*:[[:space:]]*"SUCCEEDED".*"value"[[:space:]]*:[[:space:]]*1|"value"[[:space:]]*:[[:space:]]*1.*"status"[[:space:]]*:[[:space:]]*"SUCCEEDED"' "$TMP_ROOT/results/agent-execute-preflight.body"; then - echo "FAIL: Agent execute preflight response did not contain SUCCEEDED row value=1" >&2 - exit 1 - fi + RQ_REMOTE_QUERY="$RQ_REMOTE_QUERY" python3 - "$TMP_ROOT/results/agent-execute-preflight.body" <<'PY' +import json +import os +import sys + +with open(sys.argv[1], encoding="utf-8") as f: + body = json.load(f) + +if body.get("status") != "SUCCEEDED": + raise SystemExit(f"Agent execute preflight response status was not SUCCEEDED: {body}") + +rows = body.get("rows") +query = os.environ["RQ_REMOTE_QUERY"] +if query == "SELECT city, country FROM cities ORDER BY city": + expected = [ + {"city": "Beautiful city of lights", "country": "France"}, + {"city": "New York", "country": "USA"}, + ] +else: + expected = [{"value": 1}] +if rows != expected: + raise SystemExit(f"Agent execute preflight rows did not match expected fixture data: rows={rows!r} expected={expected!r}") +PY if grep -Eq 'password|token|secret' "$TMP_ROOT/results/agent-execute-preflight.body"; then echo "FAIL: Agent execute preflight response contained credential-shaped text" >&2 exit 1 @@ -348,6 +368,7 @@ run_standalone_go_proof() { RQ_POSTGRES_HOST="$RQ_POSTGRES_HOST" \ RQ_POSTGRES_PORT="$RQ_POSTGRES_PORT" \ RQ_POSTGRES_DBNAME="$RQ_POSTGRES_DBNAME" \ + RQ_REMOTE_QUERY="$RQ_REMOTE_QUERY" \ dda inv test --targets=./pkg/privateactionrunner/bundles/remotequeries \ --extra-args='-run TestRemoteQueriesActionRunsThroughStandalonePARProcessWithRealAgentIPC -count=1 -v' ) | tee "$TMP_ROOT/results/standalone-proof-test.log" From 686602a9d446dcf29e1520953dbc13e2e0fbf06b Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 18:55:49 +0000 Subject: [PATCH 20/33] Add large Remote Queries payload proof cases --- .../impl/remote_query_execute.go | 12 +- comp/remotequeries/impl/remote_query_test.go | 11 ++ .../live_agent_ipc_par_loop_test.go | 75 +++++++- .../standalone_par_process_proof_test.go | 14 +- ...andalone-par-agentsecure-postgres-proof.sh | 164 +++++++++++++++--- 5 files changed, 245 insertions(+), 31 deletions(-) diff --git a/comp/remotequeries/impl/remote_query_execute.go b/comp/remotequeries/impl/remote_query_execute.go index b98db6a7f42c..79aac1b7e8dd 100644 --- a/comp/remotequeries/impl/remote_query_execute.go +++ b/comp/remotequeries/impl/remote_query_execute.go @@ -30,6 +30,15 @@ const ( statusExecutorUnavailable = "executor_unavailable" ) +var remoteQueryLargePayloadProofQueries = map[string]int{ + "SELECT repeat('x', 1048576) AS payload": 1 << 20, + "SELECT repeat('x', 2097152) AS payload": 2 << 20, + "SELECT repeat('x', 4194304) AS payload": 4 << 20, + "SELECT repeat('x', 8388608) AS payload": 8 << 20, + "SELECT repeat('x', 16777216) AS payload": 16 << 20, + "SELECT repeat('x', 33554432) AS payload": 32 << 20, +} + type remoteQueryRunner interface { RunRemoteQueryJSON(integration string, requestJSON string) (string, error) } @@ -39,7 +48,8 @@ func isRemoteQueryAllowedProofQuery(query string) bool { case remoteQueryProofSeedQuery, remoteQueryFixtureTableProofQuery: return true default: - return false + _, ok := remoteQueryLargePayloadProofQueries[query] + return ok } } diff --git a/comp/remotequeries/impl/remote_query_test.go b/comp/remotequeries/impl/remote_query_test.go index 19982d213a98..a4d780916465 100644 --- a/comp/remotequeries/impl/remote_query_test.go +++ b/comp/remotequeries/impl/remote_query_test.go @@ -6,6 +6,7 @@ package remotequeriesimpl import ( + "fmt" "net/http" "net/http/httptest" "strings" @@ -386,6 +387,16 @@ func TestNewRemoteQueryExecuteRequestAllowsFixtureTableProofQuery(t *testing.T) assert.Equal(t, remoteQueryFixtureTableProofQuery, req.Query) } +func TestNewRemoteQueryExecuteRequestAllowsLargePayloadProofQueries(t *testing.T) { + for query, payloadBytes := range remoteQueryLargePayloadProofQueries { + t.Run(fmt.Sprintf("%d", payloadBytes), func(t *testing.T) { + req, err := NewRemoteQueryExecuteRequest("postgres", RemoteQueryExecuteTarget{Host: "localhost", Port: 5432, DBName: "postgres"}, query, &RemoteQueryExecuteLimits{MaxRows: 1, MaxBytes: payloadBytes + (1 << 20), TimeoutMs: 60_000}) + require.NoError(t, err) + assert.Equal(t, query, req.Query) + }) + } +} + func TestRemoteQueryExecuteHandlerDisabled(t *testing.T) { handler := &remoteQueryExecuteHandler{enabled: false, collector: fakeCollector{}} diff --git a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go index 5484ef2115c8..a626e098ffa0 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go @@ -35,6 +35,15 @@ const ( remoteQueriesProofQueryOverrideEnv = "RQ_REMOTE_QUERY" ) +var remoteQueriesLargePayloadProofQueries = map[string]int{ + "SELECT repeat('x', 1048576) AS payload": 1 << 20, + "SELECT repeat('x', 2097152) AS payload": 2 << 20, + "SELECT repeat('x', 4194304) AS payload": 4 << 20, + "SELECT repeat('x', 8388608) AS payload": 8 << 20, + "SELECT repeat('x', 16777216) AS payload": 16 << 20, + "SELECT repeat('x', 33554432) AS payload": 32 << 20, +} + func TestRemoteQueriesActionRunsThroughLivePARLoopWithRealAgentIPC(t *testing.T) { if os.Getenv(fusedLocalProofEnv) != "1" { t.Skipf("set %s=1 and start a local Agent with a loaded Postgres check to run the fused local proof", fusedLocalProofEnv) @@ -96,10 +105,16 @@ func TestRemoteQueriesActionRunsThroughLivePARLoopWithRealAgentIPC(t *testing.T) t.Logf("real AgentSecure IPC configured: 127.0.0.1:%d RemoteQueryExecute", cmdPortInt) require.NoError(t, fakeintakeClient.EnqueuePARTask(taskID, fqn, inputs)) - result, err := fakeintakeClient.GetPARTaskResult(taskID, 20*time.Second) + result, err := fakeintakeClient.GetPARTaskResult(taskID, remoteQueriesProofResultTimeout(proofQuery)) require.NoError(t, err) if !result.Success { - t.Logf("failed PAR task result: %+v", result) + t.Logf("failed PAR task result: %+v", summarizeRemoteQueriesProofPayload(map[string]interface{}{ + "task_id": result.TaskID, + "success": result.Success, + "outputs": result.Outputs, + "error_code": result.ErrorCode, + "error_details": result.ErrorDetails, + })) } require.True(t, result.Success) require.Equal(t, taskID, result.TaskID) @@ -110,7 +125,7 @@ func TestRemoteQueriesActionRunsThroughLivePARLoopWithRealAgentIPC(t *testing.T) require.True(t, ok) assertRemoteQueriesProofRows(t, proofQuery, rows) - resultEvidence, err := json.Marshal(result.Outputs) + resultEvidence, err := json.Marshal(summarizeRemoteQueriesProofPayload(result.Outputs)) require.NoError(t, err) require.NotContains(t, string(resultEvidence), "password") require.NotContains(t, string(resultEvidence), "token") @@ -144,14 +159,32 @@ func remoteQueriesProofQueryFromEnv() string { func remoteQueriesProofLimits(query string) map[string]interface{} { maxRows := 1 + maxBytes := 4 << 10 + timeoutMs := 5_000 if query == remoteQueriesFixtureTableProofQuery { maxRows = 2 } + if payloadBytes, ok := remoteQueriesLargePayloadBytes(query); ok { + maxBytes = payloadBytes + (1 << 20) + timeoutMs = 60_000 + } return map[string]interface{}{ "maxRows": maxRows, - "maxBytes": 1024, - "timeoutMs": 1000, + "maxBytes": maxBytes, + "timeoutMs": timeoutMs, + } +} + +func remoteQueriesProofResultTimeout(query string) time.Duration { + if _, ok := remoteQueriesLargePayloadBytes(query); ok { + return 2 * time.Minute } + return 30 * time.Second +} + +func remoteQueriesLargePayloadBytes(query string) (int, bool) { + payloadBytes, ok := remoteQueriesLargePayloadProofQueries[query] + return payloadBytes, ok } func assertRemoteQueriesProofRows(t *testing.T, query string, rows []interface{}) { @@ -174,10 +207,42 @@ func assertRemoteQueriesProofRows(t *testing.T, query string, rows []interface{} require.True(t, ok) assert.Equal(t, float64(1), firstRow["value"]) default: + expectedPayloadBytes, ok := remoteQueriesLargePayloadBytes(query) + if ok { + require.Len(t, rows, 1) + firstRow, ok := rows[0].(map[string]interface{}) + require.True(t, ok) + payload, ok := firstRow["payload"].(string) + require.True(t, ok, "payload field must be a string") + assert.Len(t, payload, expectedPayloadBytes) + return + } require.FailNowf(t, "unsupported proof query", "%s=%q must use a bridge-allowlisted proof query", remoteQueriesProofQueryOverrideEnv, query) } } +func summarizeRemoteQueriesProofPayload(value interface{}) interface{} { + switch typed := value.(type) { + case map[string]interface{}: + copy := make(map[string]interface{}, len(typed)) + for key, nested := range typed { + copy[key] = summarizeRemoteQueriesProofPayload(nested) + } + return copy + case []interface{}: + copy := make([]interface{}, 0, len(typed)) + for _, nested := range typed { + copy = append(copy, summarizeRemoteQueriesProofPayload(nested)) + } + return copy + case string: + if len(typed) > 4096 { + return fmt.Sprintf("<%d bytes>", len(typed)) + } + } + return value +} + func remoteQueriesPostgresTargetFromEnv(t *testing.T) map[string]interface{} { t.Helper() diff --git a/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go b/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go index f6eeac770835..4a8437aa7c1d 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go @@ -94,10 +94,16 @@ func TestRemoteQueriesActionRunsThroughStandalonePARProcessWithRealAgentIPC(t *t t.Logf("real AgentSecure IPC configured for standalone PAR: 127.0.0.1:%d RemoteQueryExecute", cmdPortInt) require.NoError(t, fakeintakeClient.EnqueuePARTask(taskID, fqn, inputs)) - result, err := fakeintakeClient.GetPARTaskResult(taskID, 30*time.Second) + result, err := fakeintakeClient.GetPARTaskResult(taskID, remoteQueriesProofResultTimeout(proofQuery)) require.NoError(t, err) if !result.Success { - t.Logf("failed PAR task result: %+v", result) + t.Logf("failed PAR task result: %+v", summarizeRemoteQueriesProofPayload(map[string]interface{}{ + "task_id": result.TaskID, + "success": result.Success, + "outputs": result.Outputs, + "error_code": result.ErrorCode, + "error_details": result.ErrorDetails, + })) t.Logf("PAR log tail:\n%s", readTail(parLog, 120)) } require.True(t, result.Success) @@ -109,7 +115,7 @@ func TestRemoteQueriesActionRunsThroughStandalonePARProcessWithRealAgentIPC(t *t require.True(t, ok) assertRemoteQueriesProofRows(t, proofQuery, rows) - resultEvidence, err := json.Marshal(result.Outputs) + resultEvidence, err := json.Marshal(summarizeRemoteQueriesProofPayload(result.Outputs)) require.NoError(t, err) requireNoCredentialShape(t, resultEvidence) t.Logf("fakeintake captured successful PAR task result: %s", resultEvidence) @@ -162,7 +168,7 @@ private_action_runner: actions_allowlist: - "com.datadoghq.remotequeries.execute" task_concurrency: 1 - task_timeout_seconds: 10 + task_timeout_seconds: 120 `, fakeintakeURL, cmdPort, authTokenFile, ipcCertFile, privateKeyB64, logFile) require.NoError(t, os.WriteFile(filepath.Join(dir, "datadog.yaml"), []byte(cfg), 0o600)) } diff --git a/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh b/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh index c954c9d28e2b..38cab41906ed 100755 --- a/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh +++ b/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh @@ -24,7 +24,11 @@ POSTGRES_COMPOSE_FILE=${POSTGRES_COMPOSE_FILE:-$INTEGRATIONS_CORE/postgres/tests POSTGRES_COMPOSE_PROJECT=${POSTGRES_COMPOSE_PROJECT:-rq-standalone-par-agent-postgres-$$} POSTGRES_IMAGE=${POSTGRES_IMAGE:-13-alpine} POSTGRES_LOCALE=${POSTGRES_LOCALE:-UTF8} -RQ_REMOTE_QUERY=${RQ_REMOTE_QUERY:-SELECT city, country FROM cities ORDER BY city} +RQ_REMOTE_QUERY_WAS_SET=0 +if [[ -n "${RQ_REMOTE_QUERY+x}" ]]; then + RQ_REMOTE_QUERY_WAS_SET=1 +fi +RQ_REMOTE_QUERY=${RQ_REMOTE_QUERY:-} RQ_POSTGRES_HOST=${RQ_POSTGRES_HOST:-localhost} RQ_POSTGRES_PORT=${RQ_POSTGRES_PORT:-5432} RQ_POSTGRES_DBNAME=${RQ_POSTGRES_DBNAME:-datadog_test} @@ -36,6 +40,30 @@ AGENT_PYTHON_ABI=${AGENT_PYTHON_ABI:-} AGENT_PID="" POSTGRES_COMPOSE_STARTED=0 +PROOF_CASE_NAME="" +CASE_RESULTS_DIR="" + +PROOF_CASE_NAMES=( + "seed" + "fixture-city" + "payload-1mib" + "payload-2mib" + "payload-4mib" + "payload-8mib" + "payload-16mib" + "payload-32mib" +) + +PROOF_CASE_QUERIES=( + "SELECT 1 AS value" + "SELECT city, country FROM cities ORDER BY city" + "SELECT repeat('x', 1048576) AS payload" + "SELECT repeat('x', 2097152) AS payload" + "SELECT repeat('x', 4194304) AS payload" + "SELECT repeat('x', 8388608) AS payload" + "SELECT repeat('x', 16777216) AS payload" + "SELECT repeat('x', 33554432) AS payload" +) log() { printf '\n[%s] %s\n' "$(date -u +%H:%M:%S)" "$*" @@ -264,6 +292,52 @@ instances: YAML } +remote_query_limits_json() { + RQ_REMOTE_QUERY="$RQ_REMOTE_QUERY" python3 - <<'PY' +import json +import os +import re + +query = os.environ["RQ_REMOTE_QUERY"] +max_rows = 2 if query == "SELECT city, country FROM cities ORDER BY city" else 1 +max_bytes = 4 * 1024 +timeout_ms = 5000 +match = re.fullmatch(r"SELECT repeat\('x', ([0-9]+)\) AS payload", query) +if match: + payload_bytes = int(match.group(1)) + max_bytes = payload_bytes + 1024 * 1024 + timeout_ms = 60000 +print(json.dumps({"maxRows": max_rows, "maxBytes": max_bytes, "timeoutMs": timeout_ms})) +PY +} + +summarize_json_file() { + local input=$1 + local output=$2 + python3 - "$input" "$output" <<'PY' +import json +import sys + +with open(sys.argv[1], encoding="utf-8") as f: + body = json.load(f) + +def summarize(value): + if isinstance(value, dict): + return {key: summarize(nested) for key, nested in value.items()} + if isinstance(value, list): + return [summarize(nested) for nested in value] + if isinstance(value, str) and len(value) > 4096: + return f"<{len(value)} bytes>" + return value + +summary = summarize(body) +with open(sys.argv[2], "w", encoding="utf-8") as f: + json.dump(summary, f, indent=2, sort_keys=True) + f.write("\n") +print(json.dumps(summary, sort_keys=True)) +PY +} + start_agent_and_wait_for_postgres_check() { : > "$TMP_ROOT/agent.log" PYTHONPATH="$TMP_ROOT/checks.d:$TMP_ROOT/pydeps" \ @@ -292,7 +366,9 @@ start_agent_and_wait_for_postgres_check() { call_agent_execute_preflight() { local payload - payload=$(RQ_POSTGRES_HOST="$RQ_POSTGRES_HOST" RQ_POSTGRES_PORT="$RQ_POSTGRES_PORT" RQ_POSTGRES_DBNAME="$RQ_POSTGRES_DBNAME" RQ_REMOTE_QUERY="$RQ_REMOTE_QUERY" python3 - <<'PY' + local limits + limits=$(remote_query_limits_json) + payload=$(RQ_POSTGRES_HOST="$RQ_POSTGRES_HOST" RQ_POSTGRES_PORT="$RQ_POSTGRES_PORT" RQ_POSTGRES_DBNAME="$RQ_POSTGRES_DBNAME" RQ_REMOTE_QUERY="$RQ_REMOTE_QUERY" RQ_LIMITS_JSON="$limits" python3 - <<'PY' import json import os print(json.dumps({ @@ -303,31 +379,37 @@ print(json.dumps({ "dbname": os.environ["RQ_POSTGRES_DBNAME"], }, "query": os.environ["RQ_REMOTE_QUERY"], - "limits": {"maxRows": 2 if os.environ["RQ_REMOTE_QUERY"] == "SELECT city, country FROM cities ORDER BY city" else 1, "maxBytes": 1024, "timeoutMs": 1000}, + "limits": json.loads(os.environ["RQ_LIMITS_JSON"]), })) PY ) local token token=$(cat "$TMP_ROOT/run/auth_token") - log "Preflight real Agent IPC HTTP execute endpoint (dev evidence only)" + log "[$PROOF_CASE_NAME] Preflight real Agent IPC HTTP execute endpoint (dev evidence only)" + local body_file="$CASE_RESULTS_DIR/agent-execute-preflight.raw-body" + local status_file="$CASE_RESULTS_DIR/agent-execute-preflight.status" + local summary_file="$CASE_RESULTS_DIR/agent-execute-preflight.summary.json" + local sanitized_body_file="$CASE_RESULTS_DIR/agent-execute-preflight.body" local status - status=$(curl -sS -k -o "$TMP_ROOT/results/agent-execute-preflight.body" -w '%{http_code}' \ + status=$(curl -sS -k -o "$body_file" -w '%{http_code}' \ -H "Authorization: Bearer ${token}" \ -H 'Content-Type: application/json' \ --data "$payload" \ "https://127.0.0.1:${CMD_PORT}/agent/remote-queries/execute") - printf 'agent_execute_http_status=%s\n' "$status" | tee "$TMP_ROOT/results/agent-execute-preflight.status" - cat "$TMP_ROOT/results/agent-execute-preflight.body" - printf '\n' + printf 'agent_execute_http_status=%s\n' "$status" | tee "$status_file" + if [[ -s "$body_file" ]]; then + summarize_json_file "$body_file" "$summary_file" + fi if [[ "$status" != "200" ]]; then echo "FAIL: expected Agent execute preflight HTTP 200, got $status" >&2 exit 1 fi - RQ_REMOTE_QUERY="$RQ_REMOTE_QUERY" python3 - "$TMP_ROOT/results/agent-execute-preflight.body" <<'PY' + RQ_REMOTE_QUERY="$RQ_REMOTE_QUERY" python3 - "$body_file" <<'PY' import json import os +import re import sys with open(sys.argv[1], encoding="utf-8") as f: @@ -343,19 +425,35 @@ if query == "SELECT city, country FROM cities ORDER BY city": {"city": "Beautiful city of lights", "country": "France"}, {"city": "New York", "country": "USA"}, ] -else: + if rows != expected: + raise SystemExit(f"Agent execute preflight rows did not match expected fixture data: rows={rows!r} expected={expected!r}") +elif query == "SELECT 1 AS value": expected = [{"value": 1}] -if rows != expected: - raise SystemExit(f"Agent execute preflight rows did not match expected fixture data: rows={rows!r} expected={expected!r}") + if rows != expected: + raise SystemExit(f"Agent execute preflight rows did not match expected seed data: rows={rows!r} expected={expected!r}") +else: + match = re.fullmatch(r"SELECT repeat\('x', ([0-9]+)\) AS payload", query) + if not match: + raise SystemExit(f"Unsupported proof query: {query!r}") + expected_len = int(match.group(1)) + if not isinstance(rows, list) or len(rows) != 1: + raise SystemExit(f"Agent execute preflight expected one payload row, got rows={rows!r}") + payload = rows[0].get("payload") if isinstance(rows[0], dict) else None + if not isinstance(payload, str): + raise SystemExit("Agent execute preflight payload field was not a string") + if len(payload) != expected_len: + raise SystemExit(f"Agent execute preflight payload length mismatch: got {len(payload)} expected {expected_len}") PY - if grep -Eq 'password|token|secret' "$TMP_ROOT/results/agent-execute-preflight.body"; then + if grep -Eq 'password|token|secret' "$body_file"; then echo "FAIL: Agent execute preflight response contained credential-shaped text" >&2 exit 1 fi + cp "$summary_file" "$sanitized_body_file" + rm -f "$body_file" } run_standalone_go_proof() { - log "Running standalone PAR process -> real AgentSecure gRPC IPC -> Postgres -> fakeintake proof test" + log "[$PROOF_CASE_NAME] Running standalone PAR process -> real AgentSecure gRPC IPC -> Postgres -> fakeintake proof test" ( cd "$AGENT_REPO" RQ_STANDALONE_PROOF=1 \ @@ -364,14 +462,39 @@ run_standalone_go_proof() { RQ_STANDALONE_AGENT_CMD_PORT="$CMD_PORT" \ RQ_STANDALONE_AGENT_AUTH_TOKEN_FILE="$TMP_ROOT/run/auth_token" \ RQ_STANDALONE_AGENT_IPC_CERT_FILE="$TMP_ROOT/run/ipc_cert.pem" \ - RQ_STANDALONE_EVIDENCE_FILE="$TMP_ROOT/results/standalone-proof-evidence.txt" \ + RQ_STANDALONE_EVIDENCE_FILE="$CASE_RESULTS_DIR/standalone-proof-evidence.txt" \ RQ_POSTGRES_HOST="$RQ_POSTGRES_HOST" \ RQ_POSTGRES_PORT="$RQ_POSTGRES_PORT" \ RQ_POSTGRES_DBNAME="$RQ_POSTGRES_DBNAME" \ RQ_REMOTE_QUERY="$RQ_REMOTE_QUERY" \ dda inv test --targets=./pkg/privateactionrunner/bundles/remotequeries \ --extra-args='-run TestRemoteQueriesActionRunsThroughStandalonePARProcessWithRealAgentIPC -count=1 -v' - ) | tee "$TMP_ROOT/results/standalone-proof-test.log" + ) | tee "$CASE_RESULTS_DIR/standalone-proof-test.log" +} + +run_proof_case() { + PROOF_CASE_NAME=$1 + RQ_REMOTE_QUERY=$2 + CASE_RESULTS_DIR="$TMP_ROOT/results/$PROOF_CASE_NAME" + mkdir -p "$CASE_RESULTS_DIR" + printf '%s\n' "$RQ_REMOTE_QUERY" > "$CASE_RESULTS_DIR/query.sql" + + log "[$PROOF_CASE_NAME] Starting proof case: $RQ_REMOTE_QUERY" + call_agent_execute_preflight + run_standalone_go_proof + printf 'case=%s status=passed\n' "$PROOF_CASE_NAME" | tee "$CASE_RESULTS_DIR/status.txt" +} + +run_proof_cases() { + if [[ "$RQ_REMOTE_QUERY_WAS_SET" == "1" ]]; then + run_proof_case "single" "$RQ_REMOTE_QUERY" + return + fi + + local idx + for idx in "${!PROOF_CASE_NAMES[@]}"; do + run_proof_case "${PROOF_CASE_NAMES[$idx]}" "${PROOF_CASE_QUERIES[$idx]}" + done } main() { @@ -401,14 +524,13 @@ main() { start_postgres_fixture write_postgres_config start_agent_and_wait_for_postgres_check - call_agent_execute_preflight - run_standalone_go_proof + run_proof_cases log "Sanitized standalone proof evidence" - cat "$TMP_ROOT/results/standalone-proof-evidence.txt" + find "$TMP_ROOT/results" -name 'standalone-proof-evidence.txt' -print -exec cat {} \; - log "Done. Sanitized artifacts left in $TMP_ROOT" - log "Key evidence: fakeintake enqueue/dequeue/publish, standalone PAR PID, and real AgentSecure IPC evidence are in $TMP_ROOT/results/standalone-proof-evidence.txt" + log "Done. Sanitized artifacts left in $TMP_ROOT/results" + log "Key evidence: per-case preflight summaries, fakeintake enqueue/dequeue/publish, standalone PAR PID, and real AgentSecure IPC evidence are under $TMP_ROOT/results//" } main "$@" From 6d5bc7ffce600db04235a036d80e731ee5977f50 Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 20:51:10 +0000 Subject: [PATCH 21/33] Stream Remote Queries proof payloads to PAR --- .../impl-agent/remote_query_execute_test.go | 37 +++ comp/api/grpcserver/impl-agent/server.go | 74 +++++- .../bundles/remotequeries/execute.go | 84 +++++- .../bundles/remotequeries/execute_test.go | 42 +++ .../live_agent_ipc_par_loop_test.go | 4 +- .../remotequeries/live_par_loop_test.go | 24 ++ .../standalone_par_process_proof_test.go | 4 +- pkg/proto/datadog/api/v1/api.proto | 9 + pkg/proto/pbgo/core/api.pb.go | 241 +++++++++++------- pkg/proto/pbgo/core/api_grpc.pb.go | 45 +++- pkg/proto/pbgo/mocks/core/api_mockgen.pb.go | 34 +++ ...andalone-par-agentsecure-postgres-proof.sh | 4 +- 12 files changed, 496 insertions(+), 106 deletions(-) diff --git a/comp/api/grpcserver/impl-agent/remote_query_execute_test.go b/comp/api/grpcserver/impl-agent/remote_query_execute_test.go index 7bf21f7ec4d9..bed1c89bf38c 100644 --- a/comp/api/grpcserver/impl-agent/remote_query_execute_test.go +++ b/comp/api/grpcserver/impl-agent/remote_query_execute_test.go @@ -9,6 +9,8 @@ import ( "context" "testing" + "google.golang.org/grpc" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -36,3 +38,38 @@ func TestRemoteQueryExecuteReturnsSanitizedUnavailableWhenServiceMissing(t *test require.NotNil(t, resp.GetError()) assert.Equal(t, "executor_unavailable", resp.GetError().GetCode()) } + +func TestRemoteQueryExecuteStreamReturnsSanitizedUnavailableWhenServiceMissing(t *testing.T) { + stream := &captureRemoteQueryExecuteStreamServer{} + err := (&serverSecure{}).RemoteQueryExecuteStream(&pb.RemoteQueryExecuteRequest{}, stream) + + require.NoError(t, err) + require.Len(t, stream.chunks, 1) + assert.True(t, stream.chunks[0].GetFinal()) + assert.JSONEq(t, `{"status":"executor_unavailable","error":{"code":"executor_unavailable","message":"remote query executor is unavailable"}}`, string(stream.chunks[0].GetResponseJsonChunk())) +} + +func TestRemoteQueryExecuteStreamJSONChunksResponse(t *testing.T) { + stream := &captureRemoteQueryExecuteStreamServer{} + responseJSON := `{"status":"SUCCEEDED","rows":[{"payload":"` + string(make([]byte, remoteQueryExecuteStreamChunkSize+1)) + `"}]}` + + err := remoteQueryExecuteStreamJSON(responseJSON, stream) + + require.NoError(t, err) + require.Len(t, stream.chunks, 2) + assert.Equal(t, int32(0), stream.chunks[0].GetChunkIndex()) + assert.False(t, stream.chunks[0].GetFinal()) + assert.Equal(t, int32(1), stream.chunks[1].GetChunkIndex()) + assert.True(t, stream.chunks[1].GetFinal()) + assert.Equal(t, len(responseJSON), len(stream.chunks[0].GetResponseJsonChunk())+len(stream.chunks[1].GetResponseJsonChunk())) +} + +type captureRemoteQueryExecuteStreamServer struct { + grpc.ServerStream + chunks []*pb.RemoteQueryExecuteChunk +} + +func (s *captureRemoteQueryExecuteStreamServer) Send(chunk *pb.RemoteQueryExecuteChunk) error { + s.chunks = append(s.chunks, chunk) + return nil +} diff --git a/comp/api/grpcserver/impl-agent/server.go b/comp/api/grpcserver/impl-agent/server.go index 67a57ebe99b3..3d44096a8fdc 100644 --- a/comp/api/grpcserver/impl-agent/server.go +++ b/comp/api/grpcserver/impl-agent/server.go @@ -281,13 +281,47 @@ func (s *serverSecure) WorkloadFilterEvaluate(ctx context.Context, req *pb.Workl return s.workloadfilterServer.WorkloadFilterEvaluate(ctx, req) } +const remoteQueryExecuteStreamChunkSize = 1 << 20 + func (s *serverSecure) RemoteQueryExecute(_ context.Context, req *pb.RemoteQueryExecuteRequest) (*pb.RemoteQueryExecuteResponse, error) { if s.remoteQueries == nil { return remoteQueryExecuteErrorResponse(remotequeriesimpl.RemoteQueryStatusExecutorUnavailable, "remote query executor is unavailable"), nil } + execReq, err := remoteQueryExecuteRequestFromProto(req) + if err != nil { + return remoteQueryExecuteErrorResponse(remotequeriesimpl.RemoteQueryStatusInvalidRequest, err.Error()), nil + } + + result := s.remoteQueries.Execute(execReq) + if result.Error != nil { + return remoteQueryExecuteErrorResponse(result.Error.Code, result.Error.Message), nil + } + + return remoteQueryExecuteResponseFromJSON(result.ResponseJSON) +} + +func (s *serverSecure) RemoteQueryExecuteStream(req *pb.RemoteQueryExecuteRequest, stream pb.AgentSecure_RemoteQueryExecuteStreamServer) error { + if s.remoteQueries == nil { + return remoteQueryExecuteStreamJSON(remoteQueryExecuteErrorJSON(remotequeriesimpl.RemoteQueryStatusExecutorUnavailable, "remote query executor is unavailable"), stream) + } + + execReq, err := remoteQueryExecuteRequestFromProto(req) + if err != nil { + return remoteQueryExecuteStreamJSON(remoteQueryExecuteErrorJSON(remotequeriesimpl.RemoteQueryStatusInvalidRequest, err.Error()), stream) + } + + result := s.remoteQueries.Execute(execReq) + if result.Error != nil { + return remoteQueryExecuteStreamJSON(remoteQueryExecuteErrorJSON(result.Error.Code, result.Error.Message), stream) + } + + return remoteQueryExecuteStreamJSON(result.ResponseJSON, stream) +} + +func remoteQueryExecuteRequestFromProto(req *pb.RemoteQueryExecuteRequest) (remotequeriesimpl.RemoteQueryExecuteRequest, error) { limits := remoteQueryLimitsFromProto(req.GetLimits()) - execReq, err := remotequeriesimpl.NewRemoteQueryExecuteRequest( + return remotequeriesimpl.NewRemoteQueryExecuteRequest( req.GetIntegration(), remotequeriesimpl.RemoteQueryExecuteTarget{ Host: req.GetTarget().GetHost(), @@ -297,16 +331,27 @@ func (s *serverSecure) RemoteQueryExecute(_ context.Context, req *pb.RemoteQuery req.GetQuery(), limits, ) - if err != nil { - return remoteQueryExecuteErrorResponse(remotequeriesimpl.RemoteQueryStatusInvalidRequest, err.Error()), nil - } +} - result := s.remoteQueries.Execute(execReq) - if result.Error != nil { - return remoteQueryExecuteErrorResponse(result.Error.Code, result.Error.Message), nil +func remoteQueryExecuteStreamJSON(responseJSON string, stream pb.AgentSecure_RemoteQueryExecuteStreamServer) error { + responseBytes := []byte(responseJSON) + if len(responseBytes) == 0 { + return stream.Send(&pb.RemoteQueryExecuteChunk{Final: true}) } - - return remoteQueryExecuteResponseFromJSON(result.ResponseJSON) + for offset, chunkIndex := 0, int32(0); offset < len(responseBytes); offset, chunkIndex = offset+remoteQueryExecuteStreamChunkSize, chunkIndex+1 { + end := offset + remoteQueryExecuteStreamChunkSize + if end > len(responseBytes) { + end = len(responseBytes) + } + if err := stream.Send(&pb.RemoteQueryExecuteChunk{ + ResponseJsonChunk: responseBytes[offset:end], + ChunkIndex: chunkIndex, + Final: end == len(responseBytes), + }); err != nil { + return err + } + } + return nil } func remoteQueryLimitsFromProto(limits *pb.RemoteQueryExecuteLimits) *remotequeriesimpl.RemoteQueryExecuteLimits { @@ -327,6 +372,17 @@ func remoteQueryExecuteErrorResponse(code string, message string) *pb.RemoteQuer } } +func remoteQueryExecuteErrorJSON(code string, message string) string { + payload, err := json.Marshal(remoteQueryExecuteJSONResponse{ + Status: code, + Error: &remoteQueryExecuteError{Code: code, Message: message}, + }) + if err != nil { + return `{"status":"executor_unavailable","error":{"code":"executor_unavailable","message":"remote query executor is unavailable"}}` + } + return string(payload) +} + type remoteQueryExecuteJSONResponse struct { Status string `json:"status"` Error *remoteQueryExecuteError `json:"error,omitempty"` diff --git a/pkg/privateactionrunner/bundles/remotequeries/execute.go b/pkg/privateactionrunner/bundles/remotequeries/execute.go index 59b347b52030..4a591513b608 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/execute.go +++ b/pkg/privateactionrunner/bundles/remotequeries/execute.go @@ -6,8 +6,11 @@ package com_datadoghq_remotequeries import ( + "bytes" "context" + "encoding/json" "fmt" + "io" "google.golang.org/grpc" @@ -20,6 +23,7 @@ import ( // BridgeClient is the narrow AgentSecure gRPC client surface required by this action. type BridgeClient interface { RemoteQueryExecute(ctx context.Context, in *pb.RemoteQueryExecuteRequest, opts ...grpc.CallOption) (*pb.RemoteQueryExecuteResponse, error) + RemoteQueryExecuteStream(ctx context.Context, in *pb.RemoteQueryExecuteRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[pb.RemoteQueryExecuteChunk], error) } // BridgeClientFactory returns an authenticated AgentSecure client over the local Agent IPC channel. @@ -76,13 +80,13 @@ func (a *ExecuteAction) Run( return nil, util.DefaultActionError(fmt.Errorf("remote query action requires an AgentSecure client")) } - resp, err := client.RemoteQueryExecute(ctx, remoteQueryExecuteRequestFromInputs(inputs)) + stream, err := client.RemoteQueryExecuteStream(ctx, remoteQueryExecuteRequestFromInputs(inputs)) if err != nil { - return nil, util.DefaultActionErrorWithDisplayError(err, "remote query AgentSecure RPC failed") + return nil, util.DefaultActionErrorWithDisplayError(err, "remote query AgentSecure streaming RPC failed") } - output, err := remoteQueryExecuteOutputFromProto(resp) + output, err := remoteQueryExecuteOutputFromStream(stream) if err != nil { - return nil, util.DefaultActionErrorWithDisplayError(err, "remote query AgentSecure RPC response was invalid") + return nil, util.DefaultActionErrorWithDisplayError(err, "remote query AgentSecure streaming RPC response was invalid") } return output, nil } @@ -107,6 +111,54 @@ func remoteQueryExecuteRequestFromInputs(inputs ExecuteInputs) *pb.RemoteQueryEx return req } +func remoteQueryExecuteOutputFromStream(stream grpc.ServerStreamingClient[pb.RemoteQueryExecuteChunk]) (map[string]interface{}, error) { + if stream == nil { + return nil, fmt.Errorf("remote query response stream missing") + } + + var assembled bytes.Buffer + expectedChunkIndex := int32(0) + seenFinal := false + for { + chunk, err := stream.Recv() + if err == io.EOF { + break + } + if err != nil { + return nil, err + } + if chunk == nil { + return nil, fmt.Errorf("remote query response stream returned nil chunk") + } + if chunk.GetChunkIndex() != expectedChunkIndex { + return nil, fmt.Errorf("remote query response stream chunk index mismatch") + } + if seenFinal { + return nil, fmt.Errorf("remote query response stream sent chunk after final") + } + _, _ = assembled.Write(chunk.GetResponseJsonChunk()) + seenFinal = chunk.GetFinal() + expectedChunkIndex++ + } + if !seenFinal { + return nil, fmt.Errorf("remote query response stream missing final chunk") + } + return remoteQueryExecuteOutputFromJSON(assembled.Bytes()) +} + +func remoteQueryExecuteOutputFromJSON(responseJSON []byte) (map[string]interface{}, error) { + var output map[string]interface{} + decoder := json.NewDecoder(bytes.NewReader(responseJSON)) + if err := decoder.Decode(&output); err != nil { + return nil, err + } + status, ok := output["status"].(string) + if !ok || status == "" { + return nil, fmt.Errorf("remote query response missing status") + } + return normalizeRemoteQueryOutput(output).(map[string]interface{}), nil +} + func remoteQueryExecuteOutputFromProto(resp *pb.RemoteQueryExecuteResponse) (map[string]interface{}, error) { if resp == nil || resp.GetStatus() == "" { return nil, fmt.Errorf("remote query response missing status") @@ -141,3 +193,27 @@ func remoteQueryExecuteOutputFromProto(resp *pb.RemoteQueryExecuteResponse) (map } return output, nil } + +func normalizeRemoteQueryOutput(value interface{}) interface{} { + switch v := value.(type) { + case map[string]interface{}: + out := make(map[string]interface{}, len(v)) + for key, item := range v { + out[key] = normalizeRemoteQueryOutput(item) + } + return out + case []interface{}: + out := make([]interface{}, len(v)) + for i, item := range v { + out[i] = normalizeRemoteQueryOutput(item) + } + return out + case json.Number: + if f, err := v.Float64(); err == nil { + return f + } + return v.String() + default: + return v + } +} diff --git a/pkg/privateactionrunner/bundles/remotequeries/execute_test.go b/pkg/privateactionrunner/bundles/remotequeries/execute_test.go index a6c6db9e327b..2c583da91c31 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/execute_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/execute_test.go @@ -8,6 +8,7 @@ package com_datadoghq_remotequeries import ( "context" "encoding/json" + "io" "testing" "google.golang.org/grpc" @@ -132,3 +133,44 @@ func (c *captureBridgeClient) RemoteQueryExecute(_ context.Context, req *pb.Remo c.request = req return c.response, c.err } + +func (c *captureBridgeClient) RemoteQueryExecuteStream(_ context.Context, req *pb.RemoteQueryExecuteRequest, _ ...grpc.CallOption) (grpc.ServerStreamingClient[pb.RemoteQueryExecuteChunk], error) { + c.request = req + if c.err != nil { + return nil, c.err + } + responseJSON, err := json.Marshal(remoteQueryExecuteOutputFromProtoForTest(c.response)) + if err != nil { + return nil, err + } + return &captureRemoteQueryExecuteStream{chunks: []*pb.RemoteQueryExecuteChunk{{ResponseJsonChunk: responseJSON, Final: true}}}, nil +} + +func remoteQueryExecuteOutputFromProtoForTest(resp *pb.RemoteQueryExecuteResponse) map[string]interface{} { + output := map[string]interface{}{"status": resp.GetStatus()} + if resp.GetError() != nil { + output["error"] = map[string]interface{}{"code": resp.GetError().GetCode(), "message": resp.GetError().GetMessage()} + } + if len(resp.GetRows()) > 0 { + rows := make([]interface{}, 0, len(resp.GetRows())) + for _, row := range resp.GetRows() { + rows = append(rows, row.AsMap()) + } + output["rows"] = rows + } + return output +} + +type captureRemoteQueryExecuteStream struct { + grpc.ClientStream + chunks []*pb.RemoteQueryExecuteChunk +} + +func (s *captureRemoteQueryExecuteStream) Recv() (*pb.RemoteQueryExecuteChunk, error) { + if len(s.chunks) == 0 { + return nil, io.EOF + } + chunk := s.chunks[0] + s.chunks = s.chunks[1:] + return chunk, nil +} diff --git a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go index a626e098ffa0..1281ffd596a9 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go @@ -102,7 +102,7 @@ func TestRemoteQueriesActionRunsThroughLivePARLoopWithRealAgentIPC(t *testing.T) fqn := com_datadoghq_remotequeries.BundleID + "." + com_datadoghq_remotequeries.ExecuteActionName t.Logf("fakeintake task enqueued: task_id=%s action_fqn=%s inputs=%s", taskID, fqn, requestEvidence) - t.Logf("real AgentSecure IPC configured: 127.0.0.1:%d RemoteQueryExecute", cmdPortInt) + t.Logf("real AgentSecure IPC configured: 127.0.0.1:%d RemoteQueryExecuteStream", cmdPortInt) require.NoError(t, fakeintakeClient.EnqueuePARTask(taskID, fqn, inputs)) result, err := fakeintakeClient.GetPARTaskResult(taskID, remoteQueriesProofResultTimeout(proofQuery)) @@ -139,7 +139,7 @@ func TestRemoteQueriesActionRunsThroughLivePARLoopWithRealAgentIPC(t *testing.T) writeFusedEvidence(t, getenvOptional("RQ_FUSED_EVIDENCE_FILE"), []string{ fmt.Sprintf("fakeintake task enqueued: task_id=%s action_fqn=%s inputs=%s", taskID, fqn, requestEvidence), "live PAR loop dequeued the fakeintake OPMS task and invoked the registered action", - fmt.Sprintf("real AgentSecure IPC called via NewDefaultBridgeClient: 127.0.0.1:%d RemoteQueryExecute", cmdPortInt), + fmt.Sprintf("real AgentSecure IPC called via NewDefaultBridgeClient: 127.0.0.1:%d RemoteQueryExecuteStream", cmdPortInt), fmt.Sprintf("fakeintake captured successful PAR task result: %s", resultEvidence), fmt.Sprintf("dequeue_calls=%d", dequeueCalls), "task verification skipped locally with DD_INTERNAL_PAR_SKIP_TASK_VERIFICATION=true", diff --git a/pkg/privateactionrunner/bundles/remotequeries/live_par_loop_test.go b/pkg/privateactionrunner/bundles/remotequeries/live_par_loop_test.go index 6ccb94d4f669..fb15dccf49b0 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/live_par_loop_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/live_par_loop_test.go @@ -13,6 +13,7 @@ import ( "crypto/elliptic" "crypto/rand" "encoding/json" + "io" "testing" "time" @@ -153,3 +154,26 @@ func (c *captureAgentSecureClient) RemoteQueryExecute(_ context.Context, req *pb c.requests <- req return c.response, nil } + +func (c *captureAgentSecureClient) RemoteQueryExecuteStream(_ context.Context, req *pb.RemoteQueryExecuteRequest, _ ...grpc.CallOption) (grpc.ServerStreamingClient[pb.RemoteQueryExecuteChunk], error) { + c.requests <- req + responseJSON, err := json.Marshal(map[string]interface{}{"status": c.response.GetStatus(), "rows": []interface{}{map[string]interface{}{"value": 1}}}) + if err != nil { + return nil, err + } + return &captureRemoteQueryExecuteStream{chunks: []*pb.RemoteQueryExecuteChunk{{ResponseJsonChunk: responseJSON, Final: true}}}, nil +} + +type captureRemoteQueryExecuteStream struct { + grpc.ClientStream + chunks []*pb.RemoteQueryExecuteChunk +} + +func (s *captureRemoteQueryExecuteStream) Recv() (*pb.RemoteQueryExecuteChunk, error) { + if len(s.chunks) == 0 { + return nil, io.EOF + } + chunk := s.chunks[0] + s.chunks = s.chunks[1:] + return chunk, nil +} diff --git a/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go b/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go index 4a8437aa7c1d..be3882feed75 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go @@ -91,7 +91,7 @@ func TestRemoteQueriesActionRunsThroughStandalonePARProcessWithRealAgentIPC(t *t fqn := com_datadoghq_remotequeries.BundleID + "." + com_datadoghq_remotequeries.ExecuteActionName t.Logf("fakeintake task enqueued: task_id=%s action_fqn=%s inputs=%s", taskID, fqn, requestEvidence) - t.Logf("real AgentSecure IPC configured for standalone PAR: 127.0.0.1:%d RemoteQueryExecute", cmdPortInt) + t.Logf("real AgentSecure IPC configured for standalone PAR: 127.0.0.1:%d RemoteQueryExecuteStream", cmdPortInt) require.NoError(t, fakeintakeClient.EnqueuePARTask(taskID, fqn, inputs)) result, err := fakeintakeClient.GetPARTaskResult(taskID, remoteQueriesProofResultTimeout(proofQuery)) @@ -130,7 +130,7 @@ func TestRemoteQueriesActionRunsThroughStandalonePARProcessWithRealAgentIPC(t *t fmt.Sprintf("separate Agent process pid=%s", agentPID), fmt.Sprintf("fakeintake task enqueued: task_id=%s action_fqn=%s inputs=%s", taskID, fqn, requestEvidence), "standalone PAR process dequeued the fakeintake OPMS task and invoked the registered action", - fmt.Sprintf("real AgentSecure IPC called by standalone PAR: 127.0.0.1:%d RemoteQueryExecute", cmdPortInt), + fmt.Sprintf("real AgentSecure IPC called by standalone PAR: 127.0.0.1:%d RemoteQueryExecuteStream", cmdPortInt), fmt.Sprintf("fakeintake captured successful PAR task result: %s", resultEvidence), fmt.Sprintf("dequeue_calls=%d", dequeueCalls), "task verification skipped for this standalone tracer bullet with DD_INTERNAL_PAR_SKIP_TASK_VERIFICATION=true", diff --git a/pkg/proto/datadog/api/v1/api.proto b/pkg/proto/datadog/api/v1/api.proto index bc0856677cd8..81e5b6ce78ce 100644 --- a/pkg/proto/datadog/api/v1/api.proto +++ b/pkg/proto/datadog/api/v1/api.proto @@ -54,6 +54,12 @@ message RemoteQueryExecuteResponse { google.protobuf.Struct stats = 6; } +message RemoteQueryExecuteChunk { + bytes response_json_chunk = 1; + int32 chunk_index = 2; + bool final = 3; +} + service AgentSecure { // subscribes to added, removed, or changed entities in the Tagger // and streams them to clients as events. @@ -108,6 +114,9 @@ service AgentSecure { // Executes an Agent-local Remote Queries request through a matched integration check. rpc RemoteQueryExecute(RemoteQueryExecuteRequest) returns (RemoteQueryExecuteResponse); + // Executes an Agent-local Remote Queries request and streams the JSON response in chunks. + rpc RemoteQueryExecuteStream(RemoteQueryExecuteRequest) returns (stream RemoteQueryExecuteChunk); + // Streams pod-to-service metadata for a specific node. rpc StreamKubeMetadata(datadog.kubemetadata.KubeMetadataStreamRequest) returns (stream datadog.kubemetadata.KubeMetadataStreamResponse); } diff --git a/pkg/proto/pbgo/core/api.pb.go b/pkg/proto/pbgo/core/api.pb.go index a438caa58b51..d6b495227ede 100644 --- a/pkg/proto/pbgo/core/api.pb.go +++ b/pkg/proto/pbgo/core/api.pb.go @@ -347,6 +347,66 @@ func (x *RemoteQueryExecuteResponse) GetStats() *structpb.Struct { return nil } +type RemoteQueryExecuteChunk struct { + state protoimpl.MessageState `protogen:"open.v1"` + ResponseJsonChunk []byte `protobuf:"bytes,1,opt,name=response_json_chunk,json=responseJsonChunk,proto3" json:"response_json_chunk,omitempty"` + ChunkIndex int32 `protobuf:"varint,2,opt,name=chunk_index,json=chunkIndex,proto3" json:"chunk_index,omitempty"` + Final bool `protobuf:"varint,3,opt,name=final,proto3" json:"final,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RemoteQueryExecuteChunk) Reset() { + *x = RemoteQueryExecuteChunk{} + mi := &file_datadog_api_v1_api_proto_msgTypes[5] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RemoteQueryExecuteChunk) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RemoteQueryExecuteChunk) ProtoMessage() {} + +func (x *RemoteQueryExecuteChunk) ProtoReflect() protoreflect.Message { + mi := &file_datadog_api_v1_api_proto_msgTypes[5] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RemoteQueryExecuteChunk.ProtoReflect.Descriptor instead. +func (*RemoteQueryExecuteChunk) Descriptor() ([]byte, []int) { + return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{5} +} + +func (x *RemoteQueryExecuteChunk) GetResponseJsonChunk() []byte { + if x != nil { + return x.ResponseJsonChunk + } + return nil +} + +func (x *RemoteQueryExecuteChunk) GetChunkIndex() int32 { + if x != nil { + return x.ChunkIndex + } + return 0 +} + +func (x *RemoteQueryExecuteChunk) GetFinal() bool { + if x != nil { + return x.Final + } + return false +} + var File_datadog_api_v1_api_proto protoreflect.FileDescriptor const file_datadog_api_v1_api_proto_rawDesc = "" + @@ -375,9 +435,14 @@ const file_datadog_api_v1_api_proto_rawDesc = "" + "\acolumns\x18\x03 \x03(\v2\x17.google.protobuf.StructR\acolumns\x12+\n" + "\x04rows\x18\x04 \x03(\v2\x17.google.protobuf.StructR\x04rows\x12\x1c\n" + "\ttruncated\x18\x05 \x01(\bR\ttruncated\x12-\n" + - "\x05stats\x18\x06 \x01(\v2\x17.google.protobuf.StructR\x05stats2Z\n" + + "\x05stats\x18\x06 \x01(\v2\x17.google.protobuf.StructR\x05stats\"\x80\x01\n" + + "\x17RemoteQueryExecuteChunk\x12.\n" + + "\x13response_json_chunk\x18\x01 \x01(\fR\x11responseJsonChunk\x12\x1f\n" + + "\vchunk_index\x18\x02 \x01(\x05R\n" + + "chunkIndex\x12\x14\n" + + "\x05final\x18\x03 \x01(\bR\x05final2Z\n" + "\x05Agent\x12Q\n" + - "\vGetHostname\x12!.datadog.model.v1.HostnameRequest\x1a\x1f.datadog.model.v1.HostnameReply2\x98\x11\n" + + "\vGetHostname\x12!.datadog.model.v1.HostnameRequest\x1a\x1f.datadog.model.v1.HostnameReply2\x8a\x12\n" + "\vAgentSecure\x12c\n" + "\x14TaggerStreamEntities\x12#.datadog.model.v1.StreamTagsRequest\x1a$.datadog.model.v1.StreamTagsResponse0\x01\x12\xa2\x01\n" + "'TaggerGenerateContainerIDFromOriginInfo\x12:.datadog.model.v1.GenerateContainerIDFromOriginInfoRequest\x1a;.datadog.model.v1.GenerateContainerIDFromOriginInfoResponse\x12`\n" + @@ -397,7 +462,8 @@ const file_datadog_api_v1_api_proto_rawDesc = "" + "\vGetHostTags\x12 .datadog.model.v1.HostTagRequest\x1a\x1e.datadog.model.v1.HostTagReply\x12\\\n" + "\x12StreamConfigEvents\x12%.datadog.model.v1.ConfigStreamRequest\x1a\x1d.datadog.model.v1.ConfigEvent0\x01\x12\x87\x01\n" + "\x16WorkloadFilterEvaluate\x125.datadog.workloadfilter.WorkloadFilterEvaluateRequest\x1a6.datadog.workloadfilter.WorkloadFilterEvaluateResponse\x12k\n" + - "\x12RemoteQueryExecute\x12).datadog.api.v1.RemoteQueryExecuteRequest\x1a*.datadog.api.v1.RemoteQueryExecuteResponse\x12y\n" + + "\x12RemoteQueryExecute\x12).datadog.api.v1.RemoteQueryExecuteRequest\x1a*.datadog.api.v1.RemoteQueryExecuteResponse\x12p\n" + + "\x18RemoteQueryExecuteStream\x12).datadog.api.v1.RemoteQueryExecuteRequest\x1a'.datadog.api.v1.RemoteQueryExecuteChunk0\x01\x12y\n" + "\x12StreamKubeMetadata\x12/.datadog.kubemetadata.KubeMetadataStreamRequest\x1a0.datadog.kubemetadata.KubeMetadataStreamResponse0\x01B\x15Z\x13pkg/proto/pbgo/coreb\x06proto3" var ( @@ -412,100 +478,103 @@ func file_datadog_api_v1_api_proto_rawDescGZIP() []byte { return file_datadog_api_v1_api_proto_rawDescData } -var file_datadog_api_v1_api_proto_msgTypes = make([]protoimpl.MessageInfo, 5) +var file_datadog_api_v1_api_proto_msgTypes = make([]protoimpl.MessageInfo, 6) var file_datadog_api_v1_api_proto_goTypes = []any{ (*RemoteQueryTarget)(nil), // 0: datadog.api.v1.RemoteQueryTarget (*RemoteQueryExecuteLimits)(nil), // 1: datadog.api.v1.RemoteQueryExecuteLimits (*RemoteQueryExecuteRequest)(nil), // 2: datadog.api.v1.RemoteQueryExecuteRequest (*RemoteQueryExecuteError)(nil), // 3: datadog.api.v1.RemoteQueryExecuteError (*RemoteQueryExecuteResponse)(nil), // 4: datadog.api.v1.RemoteQueryExecuteResponse - (*structpb.Struct)(nil), // 5: google.protobuf.Struct - (*HostnameRequest)(nil), // 6: datadog.model.v1.HostnameRequest - (*StreamTagsRequest)(nil), // 7: datadog.model.v1.StreamTagsRequest - (*GenerateContainerIDFromOriginInfoRequest)(nil), // 8: datadog.model.v1.GenerateContainerIDFromOriginInfoRequest - (*FetchEntityRequest)(nil), // 9: datadog.model.v1.FetchEntityRequest - (*CaptureTriggerRequest)(nil), // 10: datadog.model.v1.CaptureTriggerRequest - (*TaggerState)(nil), // 11: datadog.model.v1.TaggerState - (*ClientGetConfigsRequest)(nil), // 12: datadog.config.ClientGetConfigsRequest - (*emptypb.Empty)(nil), // 13: google.protobuf.Empty - (*ConfigSubscriptionRequest)(nil), // 14: datadog.config.ConfigSubscriptionRequest - (*WorkloadmetaStreamRequest)(nil), // 15: datadog.workloadmeta.WorkloadmetaStreamRequest - (*RegisterRemoteAgentRequest)(nil), // 16: datadog.remoteagent.v1.RegisterRemoteAgentRequest - (*RefreshRemoteAgentRequest)(nil), // 17: datadog.remoteagent.v1.RefreshRemoteAgentRequest - (*HostTagRequest)(nil), // 18: datadog.model.v1.HostTagRequest - (*ConfigStreamRequest)(nil), // 19: datadog.model.v1.ConfigStreamRequest - (*WorkloadFilterEvaluateRequest)(nil), // 20: datadog.workloadfilter.WorkloadFilterEvaluateRequest - (*KubeMetadataStreamRequest)(nil), // 21: datadog.kubemetadata.KubeMetadataStreamRequest - (*HostnameReply)(nil), // 22: datadog.model.v1.HostnameReply - (*StreamTagsResponse)(nil), // 23: datadog.model.v1.StreamTagsResponse - (*GenerateContainerIDFromOriginInfoResponse)(nil), // 24: datadog.model.v1.GenerateContainerIDFromOriginInfoResponse - (*FetchEntityResponse)(nil), // 25: datadog.model.v1.FetchEntityResponse - (*CaptureTriggerResponse)(nil), // 26: datadog.model.v1.CaptureTriggerResponse - (*TaggerStateResponse)(nil), // 27: datadog.model.v1.TaggerStateResponse - (*ClientGetConfigsResponse)(nil), // 28: datadog.config.ClientGetConfigsResponse - (*GetStateConfigResponse)(nil), // 29: datadog.config.GetStateConfigResponse - (*ConfigSubscriptionResponse)(nil), // 30: datadog.config.ConfigSubscriptionResponse - (*ResetStateConfigResponse)(nil), // 31: datadog.config.ResetStateConfigResponse - (*WorkloadmetaStreamResponse)(nil), // 32: datadog.workloadmeta.WorkloadmetaStreamResponse - (*RegisterRemoteAgentResponse)(nil), // 33: datadog.remoteagent.v1.RegisterRemoteAgentResponse - (*RefreshRemoteAgentResponse)(nil), // 34: datadog.remoteagent.v1.RefreshRemoteAgentResponse - (*AutodiscoveryStreamResponse)(nil), // 35: datadog.autodiscovery.AutodiscoveryStreamResponse - (*HostTagReply)(nil), // 36: datadog.model.v1.HostTagReply - (*ConfigEvent)(nil), // 37: datadog.model.v1.ConfigEvent - (*WorkloadFilterEvaluateResponse)(nil), // 38: datadog.workloadfilter.WorkloadFilterEvaluateResponse - (*KubeMetadataStreamResponse)(nil), // 39: datadog.kubemetadata.KubeMetadataStreamResponse + (*RemoteQueryExecuteChunk)(nil), // 5: datadog.api.v1.RemoteQueryExecuteChunk + (*structpb.Struct)(nil), // 6: google.protobuf.Struct + (*HostnameRequest)(nil), // 7: datadog.model.v1.HostnameRequest + (*StreamTagsRequest)(nil), // 8: datadog.model.v1.StreamTagsRequest + (*GenerateContainerIDFromOriginInfoRequest)(nil), // 9: datadog.model.v1.GenerateContainerIDFromOriginInfoRequest + (*FetchEntityRequest)(nil), // 10: datadog.model.v1.FetchEntityRequest + (*CaptureTriggerRequest)(nil), // 11: datadog.model.v1.CaptureTriggerRequest + (*TaggerState)(nil), // 12: datadog.model.v1.TaggerState + (*ClientGetConfigsRequest)(nil), // 13: datadog.config.ClientGetConfigsRequest + (*emptypb.Empty)(nil), // 14: google.protobuf.Empty + (*ConfigSubscriptionRequest)(nil), // 15: datadog.config.ConfigSubscriptionRequest + (*WorkloadmetaStreamRequest)(nil), // 16: datadog.workloadmeta.WorkloadmetaStreamRequest + (*RegisterRemoteAgentRequest)(nil), // 17: datadog.remoteagent.v1.RegisterRemoteAgentRequest + (*RefreshRemoteAgentRequest)(nil), // 18: datadog.remoteagent.v1.RefreshRemoteAgentRequest + (*HostTagRequest)(nil), // 19: datadog.model.v1.HostTagRequest + (*ConfigStreamRequest)(nil), // 20: datadog.model.v1.ConfigStreamRequest + (*WorkloadFilterEvaluateRequest)(nil), // 21: datadog.workloadfilter.WorkloadFilterEvaluateRequest + (*KubeMetadataStreamRequest)(nil), // 22: datadog.kubemetadata.KubeMetadataStreamRequest + (*HostnameReply)(nil), // 23: datadog.model.v1.HostnameReply + (*StreamTagsResponse)(nil), // 24: datadog.model.v1.StreamTagsResponse + (*GenerateContainerIDFromOriginInfoResponse)(nil), // 25: datadog.model.v1.GenerateContainerIDFromOriginInfoResponse + (*FetchEntityResponse)(nil), // 26: datadog.model.v1.FetchEntityResponse + (*CaptureTriggerResponse)(nil), // 27: datadog.model.v1.CaptureTriggerResponse + (*TaggerStateResponse)(nil), // 28: datadog.model.v1.TaggerStateResponse + (*ClientGetConfigsResponse)(nil), // 29: datadog.config.ClientGetConfigsResponse + (*GetStateConfigResponse)(nil), // 30: datadog.config.GetStateConfigResponse + (*ConfigSubscriptionResponse)(nil), // 31: datadog.config.ConfigSubscriptionResponse + (*ResetStateConfigResponse)(nil), // 32: datadog.config.ResetStateConfigResponse + (*WorkloadmetaStreamResponse)(nil), // 33: datadog.workloadmeta.WorkloadmetaStreamResponse + (*RegisterRemoteAgentResponse)(nil), // 34: datadog.remoteagent.v1.RegisterRemoteAgentResponse + (*RefreshRemoteAgentResponse)(nil), // 35: datadog.remoteagent.v1.RefreshRemoteAgentResponse + (*AutodiscoveryStreamResponse)(nil), // 36: datadog.autodiscovery.AutodiscoveryStreamResponse + (*HostTagReply)(nil), // 37: datadog.model.v1.HostTagReply + (*ConfigEvent)(nil), // 38: datadog.model.v1.ConfigEvent + (*WorkloadFilterEvaluateResponse)(nil), // 39: datadog.workloadfilter.WorkloadFilterEvaluateResponse + (*KubeMetadataStreamResponse)(nil), // 40: datadog.kubemetadata.KubeMetadataStreamResponse } var file_datadog_api_v1_api_proto_depIdxs = []int32{ 0, // 0: datadog.api.v1.RemoteQueryExecuteRequest.target:type_name -> datadog.api.v1.RemoteQueryTarget 1, // 1: datadog.api.v1.RemoteQueryExecuteRequest.limits:type_name -> datadog.api.v1.RemoteQueryExecuteLimits 3, // 2: datadog.api.v1.RemoteQueryExecuteResponse.error:type_name -> datadog.api.v1.RemoteQueryExecuteError - 5, // 3: datadog.api.v1.RemoteQueryExecuteResponse.columns:type_name -> google.protobuf.Struct - 5, // 4: datadog.api.v1.RemoteQueryExecuteResponse.rows:type_name -> google.protobuf.Struct - 5, // 5: datadog.api.v1.RemoteQueryExecuteResponse.stats:type_name -> google.protobuf.Struct - 6, // 6: datadog.api.v1.Agent.GetHostname:input_type -> datadog.model.v1.HostnameRequest - 7, // 7: datadog.api.v1.AgentSecure.TaggerStreamEntities:input_type -> datadog.model.v1.StreamTagsRequest - 8, // 8: datadog.api.v1.AgentSecure.TaggerGenerateContainerIDFromOriginInfo:input_type -> datadog.model.v1.GenerateContainerIDFromOriginInfoRequest - 9, // 9: datadog.api.v1.AgentSecure.TaggerFetchEntity:input_type -> datadog.model.v1.FetchEntityRequest - 10, // 10: datadog.api.v1.AgentSecure.DogstatsdCaptureTrigger:input_type -> datadog.model.v1.CaptureTriggerRequest - 11, // 11: datadog.api.v1.AgentSecure.DogstatsdSetTaggerState:input_type -> datadog.model.v1.TaggerState - 12, // 12: datadog.api.v1.AgentSecure.ClientGetConfigs:input_type -> datadog.config.ClientGetConfigsRequest - 13, // 13: datadog.api.v1.AgentSecure.GetConfigState:input_type -> google.protobuf.Empty - 12, // 14: datadog.api.v1.AgentSecure.ClientGetConfigsHA:input_type -> datadog.config.ClientGetConfigsRequest - 13, // 15: datadog.api.v1.AgentSecure.GetConfigStateHA:input_type -> google.protobuf.Empty - 14, // 16: datadog.api.v1.AgentSecure.CreateConfigSubscription:input_type -> datadog.config.ConfigSubscriptionRequest - 13, // 17: datadog.api.v1.AgentSecure.ResetConfigState:input_type -> google.protobuf.Empty - 15, // 18: datadog.api.v1.AgentSecure.WorkloadmetaStreamEntities:input_type -> datadog.workloadmeta.WorkloadmetaStreamRequest - 16, // 19: datadog.api.v1.AgentSecure.RegisterRemoteAgent:input_type -> datadog.remoteagent.v1.RegisterRemoteAgentRequest - 17, // 20: datadog.api.v1.AgentSecure.RefreshRemoteAgent:input_type -> datadog.remoteagent.v1.RefreshRemoteAgentRequest - 13, // 21: datadog.api.v1.AgentSecure.AutodiscoveryStreamConfig:input_type -> google.protobuf.Empty - 18, // 22: datadog.api.v1.AgentSecure.GetHostTags:input_type -> datadog.model.v1.HostTagRequest - 19, // 23: datadog.api.v1.AgentSecure.StreamConfigEvents:input_type -> datadog.model.v1.ConfigStreamRequest - 20, // 24: datadog.api.v1.AgentSecure.WorkloadFilterEvaluate:input_type -> datadog.workloadfilter.WorkloadFilterEvaluateRequest + 6, // 3: datadog.api.v1.RemoteQueryExecuteResponse.columns:type_name -> google.protobuf.Struct + 6, // 4: datadog.api.v1.RemoteQueryExecuteResponse.rows:type_name -> google.protobuf.Struct + 6, // 5: datadog.api.v1.RemoteQueryExecuteResponse.stats:type_name -> google.protobuf.Struct + 7, // 6: datadog.api.v1.Agent.GetHostname:input_type -> datadog.model.v1.HostnameRequest + 8, // 7: datadog.api.v1.AgentSecure.TaggerStreamEntities:input_type -> datadog.model.v1.StreamTagsRequest + 9, // 8: datadog.api.v1.AgentSecure.TaggerGenerateContainerIDFromOriginInfo:input_type -> datadog.model.v1.GenerateContainerIDFromOriginInfoRequest + 10, // 9: datadog.api.v1.AgentSecure.TaggerFetchEntity:input_type -> datadog.model.v1.FetchEntityRequest + 11, // 10: datadog.api.v1.AgentSecure.DogstatsdCaptureTrigger:input_type -> datadog.model.v1.CaptureTriggerRequest + 12, // 11: datadog.api.v1.AgentSecure.DogstatsdSetTaggerState:input_type -> datadog.model.v1.TaggerState + 13, // 12: datadog.api.v1.AgentSecure.ClientGetConfigs:input_type -> datadog.config.ClientGetConfigsRequest + 14, // 13: datadog.api.v1.AgentSecure.GetConfigState:input_type -> google.protobuf.Empty + 13, // 14: datadog.api.v1.AgentSecure.ClientGetConfigsHA:input_type -> datadog.config.ClientGetConfigsRequest + 14, // 15: datadog.api.v1.AgentSecure.GetConfigStateHA:input_type -> google.protobuf.Empty + 15, // 16: datadog.api.v1.AgentSecure.CreateConfigSubscription:input_type -> datadog.config.ConfigSubscriptionRequest + 14, // 17: datadog.api.v1.AgentSecure.ResetConfigState:input_type -> google.protobuf.Empty + 16, // 18: datadog.api.v1.AgentSecure.WorkloadmetaStreamEntities:input_type -> datadog.workloadmeta.WorkloadmetaStreamRequest + 17, // 19: datadog.api.v1.AgentSecure.RegisterRemoteAgent:input_type -> datadog.remoteagent.v1.RegisterRemoteAgentRequest + 18, // 20: datadog.api.v1.AgentSecure.RefreshRemoteAgent:input_type -> datadog.remoteagent.v1.RefreshRemoteAgentRequest + 14, // 21: datadog.api.v1.AgentSecure.AutodiscoveryStreamConfig:input_type -> google.protobuf.Empty + 19, // 22: datadog.api.v1.AgentSecure.GetHostTags:input_type -> datadog.model.v1.HostTagRequest + 20, // 23: datadog.api.v1.AgentSecure.StreamConfigEvents:input_type -> datadog.model.v1.ConfigStreamRequest + 21, // 24: datadog.api.v1.AgentSecure.WorkloadFilterEvaluate:input_type -> datadog.workloadfilter.WorkloadFilterEvaluateRequest 2, // 25: datadog.api.v1.AgentSecure.RemoteQueryExecute:input_type -> datadog.api.v1.RemoteQueryExecuteRequest - 21, // 26: datadog.api.v1.AgentSecure.StreamKubeMetadata:input_type -> datadog.kubemetadata.KubeMetadataStreamRequest - 22, // 27: datadog.api.v1.Agent.GetHostname:output_type -> datadog.model.v1.HostnameReply - 23, // 28: datadog.api.v1.AgentSecure.TaggerStreamEntities:output_type -> datadog.model.v1.StreamTagsResponse - 24, // 29: datadog.api.v1.AgentSecure.TaggerGenerateContainerIDFromOriginInfo:output_type -> datadog.model.v1.GenerateContainerIDFromOriginInfoResponse - 25, // 30: datadog.api.v1.AgentSecure.TaggerFetchEntity:output_type -> datadog.model.v1.FetchEntityResponse - 26, // 31: datadog.api.v1.AgentSecure.DogstatsdCaptureTrigger:output_type -> datadog.model.v1.CaptureTriggerResponse - 27, // 32: datadog.api.v1.AgentSecure.DogstatsdSetTaggerState:output_type -> datadog.model.v1.TaggerStateResponse - 28, // 33: datadog.api.v1.AgentSecure.ClientGetConfigs:output_type -> datadog.config.ClientGetConfigsResponse - 29, // 34: datadog.api.v1.AgentSecure.GetConfigState:output_type -> datadog.config.GetStateConfigResponse - 28, // 35: datadog.api.v1.AgentSecure.ClientGetConfigsHA:output_type -> datadog.config.ClientGetConfigsResponse - 29, // 36: datadog.api.v1.AgentSecure.GetConfigStateHA:output_type -> datadog.config.GetStateConfigResponse - 30, // 37: datadog.api.v1.AgentSecure.CreateConfigSubscription:output_type -> datadog.config.ConfigSubscriptionResponse - 31, // 38: datadog.api.v1.AgentSecure.ResetConfigState:output_type -> datadog.config.ResetStateConfigResponse - 32, // 39: datadog.api.v1.AgentSecure.WorkloadmetaStreamEntities:output_type -> datadog.workloadmeta.WorkloadmetaStreamResponse - 33, // 40: datadog.api.v1.AgentSecure.RegisterRemoteAgent:output_type -> datadog.remoteagent.v1.RegisterRemoteAgentResponse - 34, // 41: datadog.api.v1.AgentSecure.RefreshRemoteAgent:output_type -> datadog.remoteagent.v1.RefreshRemoteAgentResponse - 35, // 42: datadog.api.v1.AgentSecure.AutodiscoveryStreamConfig:output_type -> datadog.autodiscovery.AutodiscoveryStreamResponse - 36, // 43: datadog.api.v1.AgentSecure.GetHostTags:output_type -> datadog.model.v1.HostTagReply - 37, // 44: datadog.api.v1.AgentSecure.StreamConfigEvents:output_type -> datadog.model.v1.ConfigEvent - 38, // 45: datadog.api.v1.AgentSecure.WorkloadFilterEvaluate:output_type -> datadog.workloadfilter.WorkloadFilterEvaluateResponse - 4, // 46: datadog.api.v1.AgentSecure.RemoteQueryExecute:output_type -> datadog.api.v1.RemoteQueryExecuteResponse - 39, // 47: datadog.api.v1.AgentSecure.StreamKubeMetadata:output_type -> datadog.kubemetadata.KubeMetadataStreamResponse - 27, // [27:48] is the sub-list for method output_type - 6, // [6:27] is the sub-list for method input_type + 2, // 26: datadog.api.v1.AgentSecure.RemoteQueryExecuteStream:input_type -> datadog.api.v1.RemoteQueryExecuteRequest + 22, // 27: datadog.api.v1.AgentSecure.StreamKubeMetadata:input_type -> datadog.kubemetadata.KubeMetadataStreamRequest + 23, // 28: datadog.api.v1.Agent.GetHostname:output_type -> datadog.model.v1.HostnameReply + 24, // 29: datadog.api.v1.AgentSecure.TaggerStreamEntities:output_type -> datadog.model.v1.StreamTagsResponse + 25, // 30: datadog.api.v1.AgentSecure.TaggerGenerateContainerIDFromOriginInfo:output_type -> datadog.model.v1.GenerateContainerIDFromOriginInfoResponse + 26, // 31: datadog.api.v1.AgentSecure.TaggerFetchEntity:output_type -> datadog.model.v1.FetchEntityResponse + 27, // 32: datadog.api.v1.AgentSecure.DogstatsdCaptureTrigger:output_type -> datadog.model.v1.CaptureTriggerResponse + 28, // 33: datadog.api.v1.AgentSecure.DogstatsdSetTaggerState:output_type -> datadog.model.v1.TaggerStateResponse + 29, // 34: datadog.api.v1.AgentSecure.ClientGetConfigs:output_type -> datadog.config.ClientGetConfigsResponse + 30, // 35: datadog.api.v1.AgentSecure.GetConfigState:output_type -> datadog.config.GetStateConfigResponse + 29, // 36: datadog.api.v1.AgentSecure.ClientGetConfigsHA:output_type -> datadog.config.ClientGetConfigsResponse + 30, // 37: datadog.api.v1.AgentSecure.GetConfigStateHA:output_type -> datadog.config.GetStateConfigResponse + 31, // 38: datadog.api.v1.AgentSecure.CreateConfigSubscription:output_type -> datadog.config.ConfigSubscriptionResponse + 32, // 39: datadog.api.v1.AgentSecure.ResetConfigState:output_type -> datadog.config.ResetStateConfigResponse + 33, // 40: datadog.api.v1.AgentSecure.WorkloadmetaStreamEntities:output_type -> datadog.workloadmeta.WorkloadmetaStreamResponse + 34, // 41: datadog.api.v1.AgentSecure.RegisterRemoteAgent:output_type -> datadog.remoteagent.v1.RegisterRemoteAgentResponse + 35, // 42: datadog.api.v1.AgentSecure.RefreshRemoteAgent:output_type -> datadog.remoteagent.v1.RefreshRemoteAgentResponse + 36, // 43: datadog.api.v1.AgentSecure.AutodiscoveryStreamConfig:output_type -> datadog.autodiscovery.AutodiscoveryStreamResponse + 37, // 44: datadog.api.v1.AgentSecure.GetHostTags:output_type -> datadog.model.v1.HostTagReply + 38, // 45: datadog.api.v1.AgentSecure.StreamConfigEvents:output_type -> datadog.model.v1.ConfigEvent + 39, // 46: datadog.api.v1.AgentSecure.WorkloadFilterEvaluate:output_type -> datadog.workloadfilter.WorkloadFilterEvaluateResponse + 4, // 47: datadog.api.v1.AgentSecure.RemoteQueryExecute:output_type -> datadog.api.v1.RemoteQueryExecuteResponse + 5, // 48: datadog.api.v1.AgentSecure.RemoteQueryExecuteStream:output_type -> datadog.api.v1.RemoteQueryExecuteChunk + 40, // 49: datadog.api.v1.AgentSecure.StreamKubeMetadata:output_type -> datadog.kubemetadata.KubeMetadataStreamResponse + 28, // [28:50] is the sub-list for method output_type + 6, // [6:28] is the sub-list for method input_type 6, // [6:6] is the sub-list for extension type_name 6, // [6:6] is the sub-list for extension extendee 0, // [0:6] is the sub-list for field type_name @@ -529,7 +598,7 @@ func file_datadog_api_v1_api_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_datadog_api_v1_api_proto_rawDesc), len(file_datadog_api_v1_api_proto_rawDesc)), NumEnums: 0, - NumMessages: 5, + NumMessages: 6, NumExtensions: 0, NumServices: 2, }, diff --git a/pkg/proto/pbgo/core/api_grpc.pb.go b/pkg/proto/pbgo/core/api_grpc.pb.go index 891713c46a26..0490e0d2dd2b 100644 --- a/pkg/proto/pbgo/core/api_grpc.pb.go +++ b/pkg/proto/pbgo/core/api_grpc.pb.go @@ -147,6 +147,7 @@ const ( AgentSecure_StreamConfigEvents_FullMethodName = "/datadog.api.v1.AgentSecure/StreamConfigEvents" AgentSecure_WorkloadFilterEvaluate_FullMethodName = "/datadog.api.v1.AgentSecure/WorkloadFilterEvaluate" AgentSecure_RemoteQueryExecute_FullMethodName = "/datadog.api.v1.AgentSecure/RemoteQueryExecute" + AgentSecure_RemoteQueryExecuteStream_FullMethodName = "/datadog.api.v1.AgentSecure/RemoteQueryExecuteStream" AgentSecure_StreamKubeMetadata_FullMethodName = "/datadog.api.v1.AgentSecure/StreamKubeMetadata" ) @@ -188,6 +189,8 @@ type AgentSecureClient interface { WorkloadFilterEvaluate(ctx context.Context, in *WorkloadFilterEvaluateRequest, opts ...grpc.CallOption) (*WorkloadFilterEvaluateResponse, error) // Executes an Agent-local Remote Queries request through a matched integration check. RemoteQueryExecute(ctx context.Context, in *RemoteQueryExecuteRequest, opts ...grpc.CallOption) (*RemoteQueryExecuteResponse, error) + // Executes an Agent-local Remote Queries request and streams the JSON response in chunks. + RemoteQueryExecuteStream(ctx context.Context, in *RemoteQueryExecuteRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[RemoteQueryExecuteChunk], error) // Streams pod-to-service metadata for a specific node. StreamKubeMetadata(ctx context.Context, in *KubeMetadataStreamRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[KubeMetadataStreamResponse], error) } @@ -429,9 +432,28 @@ func (c *agentSecureClient) RemoteQueryExecute(ctx context.Context, in *RemoteQu return out, nil } +func (c *agentSecureClient) RemoteQueryExecuteStream(ctx context.Context, in *RemoteQueryExecuteRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[RemoteQueryExecuteChunk], error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + stream, err := c.cc.NewStream(ctx, &AgentSecure_ServiceDesc.Streams[5], AgentSecure_RemoteQueryExecuteStream_FullMethodName, cOpts...) + if err != nil { + return nil, err + } + x := &grpc.GenericClientStream[RemoteQueryExecuteRequest, RemoteQueryExecuteChunk]{ClientStream: stream} + if err := x.ClientStream.SendMsg(in); err != nil { + return nil, err + } + if err := x.ClientStream.CloseSend(); err != nil { + return nil, err + } + return x, nil +} + +// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. +type AgentSecure_RemoteQueryExecuteStreamClient = grpc.ServerStreamingClient[RemoteQueryExecuteChunk] + func (c *agentSecureClient) StreamKubeMetadata(ctx context.Context, in *KubeMetadataStreamRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[KubeMetadataStreamResponse], error) { cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) - stream, err := c.cc.NewStream(ctx, &AgentSecure_ServiceDesc.Streams[5], AgentSecure_StreamKubeMetadata_FullMethodName, cOpts...) + stream, err := c.cc.NewStream(ctx, &AgentSecure_ServiceDesc.Streams[6], AgentSecure_StreamKubeMetadata_FullMethodName, cOpts...) if err != nil { return nil, err } @@ -486,6 +508,8 @@ type AgentSecureServer interface { WorkloadFilterEvaluate(context.Context, *WorkloadFilterEvaluateRequest) (*WorkloadFilterEvaluateResponse, error) // Executes an Agent-local Remote Queries request through a matched integration check. RemoteQueryExecute(context.Context, *RemoteQueryExecuteRequest) (*RemoteQueryExecuteResponse, error) + // Executes an Agent-local Remote Queries request and streams the JSON response in chunks. + RemoteQueryExecuteStream(*RemoteQueryExecuteRequest, grpc.ServerStreamingServer[RemoteQueryExecuteChunk]) error // Streams pod-to-service metadata for a specific node. StreamKubeMetadata(*KubeMetadataStreamRequest, grpc.ServerStreamingServer[KubeMetadataStreamResponse]) error mustEmbedUnimplementedAgentSecureServer() @@ -555,6 +579,9 @@ func (UnimplementedAgentSecureServer) WorkloadFilterEvaluate(context.Context, *W func (UnimplementedAgentSecureServer) RemoteQueryExecute(context.Context, *RemoteQueryExecuteRequest) (*RemoteQueryExecuteResponse, error) { return nil, status.Errorf(codes.Unimplemented, "method RemoteQueryExecute not implemented") } +func (UnimplementedAgentSecureServer) RemoteQueryExecuteStream(*RemoteQueryExecuteRequest, grpc.ServerStreamingServer[RemoteQueryExecuteChunk]) error { + return status.Errorf(codes.Unimplemented, "method RemoteQueryExecuteStream not implemented") +} func (UnimplementedAgentSecureServer) StreamKubeMetadata(*KubeMetadataStreamRequest, grpc.ServerStreamingServer[KubeMetadataStreamResponse]) error { return status.Errorf(codes.Unimplemented, "method StreamKubeMetadata not implemented") } @@ -882,6 +909,17 @@ func _AgentSecure_RemoteQueryExecute_Handler(srv interface{}, ctx context.Contex return interceptor(ctx, in, info, handler) } +func _AgentSecure_RemoteQueryExecuteStream_Handler(srv interface{}, stream grpc.ServerStream) error { + m := new(RemoteQueryExecuteRequest) + if err := stream.RecvMsg(m); err != nil { + return err + } + return srv.(AgentSecureServer).RemoteQueryExecuteStream(m, &grpc.GenericServerStream[RemoteQueryExecuteRequest, RemoteQueryExecuteChunk]{ServerStream: stream}) +} + +// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. +type AgentSecure_RemoteQueryExecuteStreamServer = grpc.ServerStreamingServer[RemoteQueryExecuteChunk] + func _AgentSecure_StreamKubeMetadata_Handler(srv interface{}, stream grpc.ServerStream) error { m := new(KubeMetadataStreamRequest) if err := stream.RecvMsg(m); err != nil { @@ -984,6 +1022,11 @@ var AgentSecure_ServiceDesc = grpc.ServiceDesc{ Handler: _AgentSecure_StreamConfigEvents_Handler, ServerStreams: true, }, + { + StreamName: "RemoteQueryExecuteStream", + Handler: _AgentSecure_RemoteQueryExecuteStream_Handler, + ServerStreams: true, + }, { StreamName: "StreamKubeMetadata", Handler: _AgentSecure_StreamKubeMetadata_Handler, diff --git a/pkg/proto/pbgo/mocks/core/api_mockgen.pb.go b/pkg/proto/pbgo/mocks/core/api_mockgen.pb.go index defe22225aa0..3715563dc831 100644 --- a/pkg/proto/pbgo/mocks/core/api_mockgen.pb.go +++ b/pkg/proto/pbgo/mocks/core/api_mockgen.pb.go @@ -405,6 +405,26 @@ func (mr *MockAgentSecureClientMockRecorder) RemoteQueryExecute(ctx, in interfac return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RemoteQueryExecute", reflect.TypeOf((*MockAgentSecureClient)(nil).RemoteQueryExecute), varargs...) } +// RemoteQueryExecuteStream mocks base method. +func (m *MockAgentSecureClient) RemoteQueryExecuteStream(ctx context.Context, in *core.RemoteQueryExecuteRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[core.RemoteQueryExecuteChunk], error) { + m.ctrl.T.Helper() + varargs := []interface{}{ctx, in} + for _, a := range opts { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "RemoteQueryExecuteStream", varargs...) + ret0, _ := ret[0].(grpc.ServerStreamingClient[core.RemoteQueryExecuteChunk]) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// RemoteQueryExecuteStream indicates an expected call of RemoteQueryExecuteStream. +func (mr *MockAgentSecureClientMockRecorder) RemoteQueryExecuteStream(ctx, in interface{}, opts ...interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]interface{}{ctx, in}, opts...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RemoteQueryExecuteStream", reflect.TypeOf((*MockAgentSecureClient)(nil).RemoteQueryExecuteStream), varargs...) +} + // ResetConfigState mocks base method. func (m *MockAgentSecureClient) ResetConfigState(ctx context.Context, in *emptypb.Empty, opts ...grpc.CallOption) (*core.ResetStateConfigResponse, error) { m.ctrl.T.Helper() @@ -766,6 +786,20 @@ func (mr *MockAgentSecureServerMockRecorder) RemoteQueryExecute(arg0, arg1 inter return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RemoteQueryExecute", reflect.TypeOf((*MockAgentSecureServer)(nil).RemoteQueryExecute), arg0, arg1) } +// RemoteQueryExecuteStream mocks base method. +func (m *MockAgentSecureServer) RemoteQueryExecuteStream(arg0 *core.RemoteQueryExecuteRequest, arg1 grpc.ServerStreamingServer[core.RemoteQueryExecuteChunk]) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "RemoteQueryExecuteStream", arg0, arg1) + ret0, _ := ret[0].(error) + return ret0 +} + +// RemoteQueryExecuteStream indicates an expected call of RemoteQueryExecuteStream. +func (mr *MockAgentSecureServerMockRecorder) RemoteQueryExecuteStream(arg0, arg1 interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RemoteQueryExecuteStream", reflect.TypeOf((*MockAgentSecureServer)(nil).RemoteQueryExecuteStream), arg0, arg1) +} + // ResetConfigState mocks base method. func (m *MockAgentSecureServer) ResetConfigState(arg0 context.Context, arg1 *emptypb.Empty) (*core.ResetStateConfigResponse, error) { m.ctrl.T.Helper() diff --git a/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh b/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh index 38cab41906ed..432a5927a28e 100755 --- a/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh +++ b/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh @@ -1,7 +1,7 @@ #!/usr/bin/env bash # Runs the standalone local-only Remote Queries proof: # fakeintake -> standalone OS private-action-runner process -> com.datadoghq.remotequeries.execute -# -> real local AgentSecure gRPC RemoteQueryExecute over Agent IPC TLS/auth +# -> real local AgentSecure gRPC RemoteQueryExecuteStream over Agent IPC TLS/auth # -> loaded Postgres check -> fixture-table proof query -> fakeintake publish. # The HTTP execute endpoint remains as a dev preflight for local evidence only. # @@ -453,7 +453,7 @@ PY } run_standalone_go_proof() { - log "[$PROOF_CASE_NAME] Running standalone PAR process -> real AgentSecure gRPC IPC -> Postgres -> fakeintake proof test" + log "[$PROOF_CASE_NAME] Running standalone PAR process -> real AgentSecure gRPC streaming IPC -> Postgres -> fakeintake proof test" ( cd "$AGENT_REPO" RQ_STANDALONE_PROOF=1 \ From 66e2a1285ef1f08003d610c3844c627dbcdd583b Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 11:57:07 +0000 Subject: [PATCH 22/33] Stream Remote Queries COPY results through Agent --- .../impl-agent/remote_query_execute_test.go | 19 + comp/api/grpcserver/impl-agent/server.go | 44 ++- .../impl/remote_query_execute.go | 257 +++++++++++++- comp/remotequeries/impl/remote_query_test.go | 52 ++- pkg/collector/python/check.go | 44 +++ pkg/collector/python/check_test.go | 12 + pkg/collector/python/remote_query_stream.go | 28 ++ pkg/collector/python/test_check.go | 79 +++++ .../bundles/remotequeries/execute.go | 83 ++++- .../bundles/remotequeries/execute_test.go | 31 ++ .../live_agent_ipc_par_loop_test.go | 28 ++ .../standalone_par_process_proof_test.go | 25 +- pkg/proto/datadog/api/v1/api.proto | 10 + pkg/proto/pbgo/core/api.pb.go | 331 ++++++++++++------ rtloader/include/datadog_agent_rtloader.h | 15 + rtloader/include/rtloader.h | 13 + rtloader/include/rtloader_types.h | 1 + rtloader/rtloader/api.cpp | 9 + rtloader/three/three.cpp | 141 ++++++++ rtloader/three/three.h | 2 + ...andalone-par-agentsecure-postgres-proof.sh | 14 + 21 files changed, 1087 insertions(+), 151 deletions(-) create mode 100644 pkg/collector/python/remote_query_stream.go diff --git a/comp/api/grpcserver/impl-agent/remote_query_execute_test.go b/comp/api/grpcserver/impl-agent/remote_query_execute_test.go index bed1c89bf38c..82cbbabb8052 100644 --- a/comp/api/grpcserver/impl-agent/remote_query_execute_test.go +++ b/comp/api/grpcserver/impl-agent/remote_query_execute_test.go @@ -49,6 +49,25 @@ func TestRemoteQueryExecuteStreamReturnsSanitizedUnavailableWhenServiceMissing(t assert.JSONEq(t, `{"status":"executor_unavailable","error":{"code":"executor_unavailable","message":"remote query executor is unavailable"}}`, string(stream.chunks[0].GetResponseJsonChunk())) } +func TestRemoteQueryExecuteRequestFromProtoPreservesCopyStream(t *testing.T) { + req, err := remoteQueryExecuteRequestFromProto(&pb.RemoteQueryExecuteRequest{ + Integration: "postgres", + Operation: "copy_stream", + Format: "csv", + Target: &pb.RemoteQueryTarget{Host: "LOCALHOST.", Port: 5432, Dbname: "postgres"}, + Query: "SELECT city, country FROM cities ORDER BY city", + CopyLimits: &pb.RemoteQueryExecuteCopyLimits{ChunkBytes: 32, MaxBytes: 1024, MaxRowBytes: 1024, TimeoutMs: 1000}, + }) + + require.NoError(t, err) + assert.Equal(t, "postgres", req.Integration) + assert.Equal(t, "copy_stream", req.Operation) + assert.Equal(t, "csv", req.Format) + require.NotNil(t, req.CopyLimits) + assert.Equal(t, 32, req.CopyLimits.ChunkBytes) + assert.Nil(t, req.Limits) +} + func TestRemoteQueryExecuteStreamJSONChunksResponse(t *testing.T) { stream := &captureRemoteQueryExecuteStreamServer{} responseJSON := `{"status":"SUCCEEDED","rows":[{"payload":"` + string(make([]byte, remoteQueryExecuteStreamChunkSize+1)) + `"}]}` diff --git a/comp/api/grpcserver/impl-agent/server.go b/comp/api/grpcserver/impl-agent/server.go index 3d44096a8fdc..3938ddae20f1 100644 --- a/comp/api/grpcserver/impl-agent/server.go +++ b/comp/api/grpcserver/impl-agent/server.go @@ -311,6 +311,19 @@ func (s *serverSecure) RemoteQueryExecuteStream(req *pb.RemoteQueryExecuteReques return remoteQueryExecuteStreamJSON(remoteQueryExecuteErrorJSON(remotequeriesimpl.RemoteQueryStatusInvalidRequest, err.Error()), stream) } + if execReq.Operation == "copy_stream" { + chunkIndex := int32(0) + result := s.remoteQueries.ExecuteStream(execReq, func(eventJSON string) error { + err := stream.Send(&pb.RemoteQueryExecuteChunk{ResponseJsonChunk: []byte(eventJSON), ChunkIndex: chunkIndex}) + chunkIndex++ + return err + }) + if result.Error != nil { + return remoteQueryExecuteStreamJSON(remoteQueryExecuteErrorJSON(result.Error.Code, result.Error.Message), stream) + } + return stream.Send(&pb.RemoteQueryExecuteChunk{ChunkIndex: chunkIndex, Final: true}) + } + result := s.remoteQueries.Execute(execReq) if result.Error != nil { return remoteQueryExecuteStreamJSON(remoteQueryExecuteErrorJSON(result.Error.Code, result.Error.Message), stream) @@ -320,17 +333,16 @@ func (s *serverSecure) RemoteQueryExecuteStream(req *pb.RemoteQueryExecuteReques } func remoteQueryExecuteRequestFromProto(req *pb.RemoteQueryExecuteRequest) (remotequeriesimpl.RemoteQueryExecuteRequest, error) { + target := remotequeriesimpl.RemoteQueryExecuteTarget{ + Host: req.GetTarget().GetHost(), + Port: int(req.GetTarget().GetPort()), + DBName: req.GetTarget().GetDbname(), + } + if req.GetOperation() == "copy_stream" { + return remotequeriesimpl.NewRemoteQueryCopyStreamExecuteRequest(req.GetIntegration(), target, req.GetQuery(), req.GetFormat(), remoteQueryCopyLimitsFromProto(req.GetCopyLimits())) + } limits := remoteQueryLimitsFromProto(req.GetLimits()) - return remotequeriesimpl.NewRemoteQueryExecuteRequest( - req.GetIntegration(), - remotequeriesimpl.RemoteQueryExecuteTarget{ - Host: req.GetTarget().GetHost(), - Port: int(req.GetTarget().GetPort()), - DBName: req.GetTarget().GetDbname(), - }, - req.GetQuery(), - limits, - ) + return remotequeriesimpl.NewRemoteQueryExecuteRequest(req.GetIntegration(), target, req.GetQuery(), limits) } func remoteQueryExecuteStreamJSON(responseJSON string, stream pb.AgentSecure_RemoteQueryExecuteStreamServer) error { @@ -365,6 +377,18 @@ func remoteQueryLimitsFromProto(limits *pb.RemoteQueryExecuteLimits) *remotequer } } +func remoteQueryCopyLimitsFromProto(limits *pb.RemoteQueryExecuteCopyLimits) *remotequeriesimpl.RemoteQueryExecuteCopyLimits { + if limits == nil { + return nil + } + return &remotequeriesimpl.RemoteQueryExecuteCopyLimits{ + ChunkBytes: int(limits.GetChunkBytes()), + MaxBytes: int(limits.GetMaxBytes()), + MaxRowBytes: int(limits.GetMaxRowBytes()), + TimeoutMs: int(limits.GetTimeoutMs()), + } +} + func remoteQueryExecuteErrorResponse(code string, message string) *pb.RemoteQueryExecuteResponse { return &pb.RemoteQueryExecuteResponse{ Status: code, diff --git a/comp/remotequeries/impl/remote_query_execute.go b/comp/remotequeries/impl/remote_query_execute.go index 79aac1b7e8dd..09ee5e67aa44 100644 --- a/comp/remotequeries/impl/remote_query_execute.go +++ b/comp/remotequeries/impl/remote_query_execute.go @@ -43,6 +43,10 @@ type remoteQueryRunner interface { RunRemoteQueryJSON(integration string, requestJSON string) (string, error) } +type remoteQueryStreamRunner interface { + RunRemoteQueryStream(integration string, requestJSON string, emit func(string) error) error +} + func isRemoteQueryAllowedProofQuery(query string) bool { switch query { case remoteQueryProofSeedQuery, remoteQueryFixtureTableProofQuery: @@ -75,6 +79,24 @@ func remoteQueryRunnerFor(chk check.Check) (remoteQueryRunner, bool) { return nil, false } +func remoteQueryStreamRunnerFor(chk check.Check) (remoteQueryStreamRunner, bool) { + for chk != nil { + if runner, ok := chk.(remoteQueryStreamRunner); ok { + return runner, true + } + unwrapper, ok := chk.(remoteQueryCheckUnwrapper) + if !ok { + break + } + unwrapped := unwrapper.Unwrap() + if unwrapped == chk { + break + } + chk = unwrapped + } + return nil, false +} + // NewRemoteQueryExecuteEndpointProvider registers the remote query execute endpoint on the internal Agent API. func NewRemoteQueryExecuteEndpointProvider(reqs Requires) api.AgentEndpointProvider { h := &remoteQueryExecuteHandler{ @@ -114,12 +136,60 @@ type RemoteQueryExecuteLimits struct { TimeoutMs int } +// RemoteQueryExecuteCopyLimits contains COPY stream execution limits. +type RemoteQueryExecuteCopyLimits struct { + ChunkBytes int + MaxBytes int + MaxRowBytes int + TimeoutMs int +} + // RemoteQueryExecuteRequest is the typed internal request shape shared by HTTP and gRPC callers. type RemoteQueryExecuteRequest struct { Integration string + Operation string Target RemoteQueryExecuteTarget Query string + Format string Limits *RemoteQueryExecuteLimits + CopyLimits *RemoteQueryExecuteCopyLimits +} + +// NewRemoteQueryCopyStreamExecuteRequest validates and normalizes a typed COPY stream request. +func NewRemoteQueryCopyStreamExecuteRequest(integration string, target RemoteQueryExecuteTarget, query string, format string, limits *RemoteQueryExecuteCopyLimits) (RemoteQueryExecuteRequest, error) { + parsedIntegration, err := parseIntegration(integration) + if err != nil { + return RemoteQueryExecuteRequest{}, err + } + parsedTarget, err := parseTarget(&remoteQueryTargetRequestJSON{Host: target.Host, Port: &target.Port, DBName: target.DBName}) + if err != nil { + return RemoteQueryExecuteRequest{}, err + } + if query == "" { + return RemoteQueryExecuteRequest{}, fmt.Errorf("query is required") + } + if !isRemoteQueryAllowedProofQuery(query) { + return RemoteQueryExecuteRequest{}, fmt.Errorf("query is not allowed") + } + if format == "" { + format = "csv" + } + if format != "csv" { + return RemoteQueryExecuteRequest{}, fmt.Errorf("format must be csv") + } + var parsedLimits *remoteQueryExecuteCopyLimits + if limits != nil { + parsedLimits, err = parseExecuteCopyLimits(&remoteQueryExecuteCopyLimitsRequestJSON{ + ChunkBytes: &limits.ChunkBytes, + MaxBytes: &limits.MaxBytes, + MaxRowBytes: &limits.MaxRowBytes, + TimeoutMs: &limits.TimeoutMs, + }) + if err != nil { + return RemoteQueryExecuteRequest{}, err + } + } + return remoteQueryExecuteRequestFromInternal(remoteQueryExecuteRequest{Integration: parsedIntegration, Operation: "copy_stream", Target: parsedTarget, Query: query, Format: format, CopyLimits: parsedLimits}), nil } // RemoteQueryExecuteError is a sanitized remote query bridge error. @@ -184,16 +254,22 @@ func NewRemoteQueryExecuteRequest(integration string, target RemoteQueryExecuteT type remoteQueryExecuteRequest struct { Integration string + Operation string Target remoteQueryTarget Query string + Format string Limits *remoteQueryExecuteLimits + CopyLimits *remoteQueryExecuteCopyLimits } type remoteQueryExecuteRequestJSON struct { - Integration string `json:"integration"` - Target *remoteQueryTargetRequestJSON `json:"target"` - Query string `json:"query"` - Limits *remoteQueryExecuteLimitsRequestJSON `json:"limits,omitempty"` + Integration string `json:"integration"` + Operation string `json:"operation,omitempty"` + Target *remoteQueryTargetRequestJSON `json:"target"` + Query string `json:"query"` + Format string `json:"format,omitempty"` + Limits *remoteQueryExecuteLimitsRequestJSON `json:"limits,omitempty"` + CopyLimits *remoteQueryExecuteCopyLimitsRequestJSON `json:"copyLimits,omitempty"` } type remoteQueryExecuteLimitsRequestJSON struct { @@ -208,12 +284,41 @@ type remoteQueryExecuteLimits struct { TimeoutMs int } +type remoteQueryExecuteCopyLimitsRequestJSON struct { + ChunkBytes *int `json:"chunkBytes"` + MaxBytes *int `json:"maxBytes"` + MaxRowBytes *int `json:"maxRowBytes"` + TimeoutMs *int `json:"timeoutMs"` +} + +type remoteQueryExecuteCopyLimits struct { + ChunkBytes int + MaxBytes int + MaxRowBytes int + TimeoutMs int +} + type remoteQueryExecutorRequestJSON struct { Target remoteQueryTargetJSON `json:"target"` Query string `json:"query"` Limits *remoteQueryExecuteLimitsJSON `json:"limits,omitempty"` } +type remoteQueryCopyExecutorRequestJSON struct { + Operation string `json:"operation"` + Target remoteQueryTargetJSON `json:"target"` + Query string `json:"query"` + Format string `json:"format"` + Limits *remoteQueryExecuteCopyLimitsJSON `json:"limits,omitempty"` +} + +type remoteQueryExecuteCopyLimitsJSON struct { + ChunkBytes int `json:"chunkBytes"` + MaxBytes int `json:"maxBytes"` + MaxRowBytes int `json:"maxRowBytes"` + TimeoutMs int `json:"timeoutMs"` +} + type remoteQueryTargetJSON struct { Host string `json:"host"` Port int `json:"port"` @@ -286,8 +391,12 @@ func parseExecuteRequest(r *http.Request) (remoteQueryExecuteRequest, string, er if err != nil { return remoteQueryExecuteRequest{}, "", err } + copyLimits, err := parseExecuteCopyLimits(wireReq.CopyLimits) + if err != nil { + return remoteQueryExecuteRequest{}, "", err + } - req := remoteQueryExecuteRequest{Integration: integration, Target: target, Query: wireReq.Query, Limits: limits} + req := remoteQueryExecuteRequest{Integration: integration, Operation: wireReq.Operation, Target: target, Query: wireReq.Query, Format: wireReq.Format, Limits: limits, CopyLimits: copyLimits} requestJSON, err := marshalExecuteRequest(req) if err != nil { return remoteQueryExecuteRequest{}, "", fmt.Errorf("malformed JSON request") @@ -300,6 +409,23 @@ var ( errLimitsMustBeObject = errors.New("limits must be an object") ) +func (l *remoteQueryExecuteCopyLimitsRequestJSON) UnmarshalJSON(data []byte) error { + if !isJSONObject(data) { + return errLimitsMustBeObject + } + + type limitsAlias remoteQueryExecuteCopyLimitsRequestJSON + var limits limitsAlias + if err := decodeStrictJSON(bytes.NewReader(data), &limits); err != nil { + if isUnknownJSONFieldError(err) { + return errLimitsUnknownField + } + return err + } + *l = remoteQueryExecuteCopyLimitsRequestJSON(limits) + return nil +} + func (l *remoteQueryExecuteLimitsRequestJSON) UnmarshalJSON(data []byte) error { if !isJSONObject(data) { return errLimitsMustBeObject @@ -317,6 +443,29 @@ func (l *remoteQueryExecuteLimitsRequestJSON) UnmarshalJSON(data []byte) error { return nil } +func parseExecuteCopyLimits(limits *remoteQueryExecuteCopyLimitsRequestJSON) (*remoteQueryExecuteCopyLimits, error) { + if limits == nil { + return nil, nil + } + chunkBytes, err := parseRequiredPositiveInt(limits.ChunkBytes, "copyLimits.chunkBytes") + if err != nil { + return nil, err + } + maxBytes, err := parseRequiredPositiveInt(limits.MaxBytes, "copyLimits.maxBytes") + if err != nil { + return nil, err + } + maxRowBytes, err := parseRequiredPositiveInt(limits.MaxRowBytes, "copyLimits.maxRowBytes") + if err != nil { + return nil, err + } + timeoutMs, err := parseRequiredPositiveInt(limits.TimeoutMs, "copyLimits.timeoutMs") + if err != nil { + return nil, err + } + return &remoteQueryExecuteCopyLimits{ChunkBytes: chunkBytes, MaxBytes: maxBytes, MaxRowBytes: maxRowBytes, TimeoutMs: timeoutMs}, nil +} + func parseExecuteLimits(limits *remoteQueryExecuteLimitsRequestJSON) (*remoteQueryExecuteLimits, error) { if limits == nil { return nil, nil @@ -349,6 +498,9 @@ func parseRequiredPositiveInt(value *int, name string) (int, error) { } func (s *RemoteQueryExecuteService) Execute(req RemoteQueryExecuteRequest) RemoteQueryExecuteResult { + if req.Operation == "copy_stream" { + return remoteQueryExecuteErrorResult(http.StatusBadRequest, statusInvalidRequest, "copy_stream requires the streaming executor") + } if s == nil || !s.enabled { return remoteQueryExecuteErrorResult(http.StatusServiceUnavailable, statusBridgeDisabled, "remote queries bridge is disabled") } @@ -357,17 +509,12 @@ func (s *RemoteQueryExecuteService) Execute(req RemoteQueryExecuteRequest) Remot } internal := req.internal() - matches := findIntegrationMatches(s.collector, internal.Integration, internal.Target) - switch len(matches) { - case 0: - return remoteQueryExecuteErrorResult(http.StatusNotFound, statusTargetNotFound, "no matching integration check found") - case 1: - // continue below - default: - return remoteQueryExecuteErrorResult(http.StatusConflict, statusAmbiguous, "multiple matching integration checks found") + match, result := s.matchExecutor(internal) + if result.Error != nil { + return result } - runner, ok := remoteQueryRunnerFor(matches[0].check) + runner, ok := remoteQueryRunnerFor(match.check) if !ok { return remoteQueryExecuteErrorResult(http.StatusFailedDependency, statusExecutorUnavailable, "matched integration check does not support remote query execution") } @@ -385,6 +532,52 @@ func (s *RemoteQueryExecuteService) Execute(req RemoteQueryExecuteRequest) Remot return RemoteQueryExecuteResult{HTTPStatus: http.StatusOK, ResponseJSON: responseJSON} } +// ExecuteStream executes a COPY streaming request and emits serialized stream events without materializing the full result. +func (s *RemoteQueryExecuteService) ExecuteStream(req RemoteQueryExecuteRequest, emit func(string) error) RemoteQueryExecuteResult { + if req.Operation != "copy_stream" { + return s.Execute(req) + } + if emit == nil { + return remoteQueryExecuteErrorResult(http.StatusFailedDependency, statusExecutorUnavailable, "remote query stream emitter is unavailable") + } + if s == nil || !s.enabled { + return remoteQueryExecuteErrorResult(http.StatusServiceUnavailable, statusBridgeDisabled, "remote queries bridge is disabled") + } + if s.collector == nil { + return remoteQueryExecuteErrorResult(http.StatusFailedDependency, statusExecutorUnavailable, "remote query executor is unavailable") + } + + internal := req.internal() + match, result := s.matchExecutor(internal) + if result.Error != nil { + return result + } + runner, ok := remoteQueryStreamRunnerFor(match.check) + if !ok { + return remoteQueryExecuteErrorResult(http.StatusFailedDependency, statusExecutorUnavailable, "matched integration check does not support remote query streaming") + } + requestJSON, err := marshalExecuteRequest(internal) + if err != nil { + return remoteQueryExecuteErrorResult(http.StatusBadRequest, statusInvalidRequest, "malformed JSON request") + } + if err := runner.RunRemoteQueryStream(internal.Integration, requestJSON, emit); err != nil { + return remoteQueryExecuteErrorResult(http.StatusBadGateway, statusExecutorUnavailable, "remote query stream executor failed") + } + return RemoteQueryExecuteResult{HTTPStatus: http.StatusOK, Status: "SUCCEEDED"} +} + +func (s *RemoteQueryExecuteService) matchExecutor(internal remoteQueryExecuteRequest) (integrationCheckMatch, RemoteQueryExecuteResult) { + matches := findIntegrationMatches(s.collector, internal.Integration, internal.Target) + switch len(matches) { + case 0: + return integrationCheckMatch{}, remoteQueryExecuteErrorResult(http.StatusNotFound, statusTargetNotFound, "no matching integration check found") + case 1: + return matches[0], RemoteQueryExecuteResult{HTTPStatus: http.StatusOK} + default: + return integrationCheckMatch{}, remoteQueryExecuteErrorResult(http.StatusConflict, statusAmbiguous, "multiple matching integration checks found") + } +} + func remoteQueryExecuteErrorResult(httpStatus int, status string, message string) RemoteQueryExecuteResult { return RemoteQueryExecuteResult{ HTTPStatus: httpStatus, @@ -396,28 +589,64 @@ func remoteQueryExecuteErrorResult(httpStatus int, status string, message string func (r RemoteQueryExecuteRequest) internal() remoteQueryExecuteRequest { internal := remoteQueryExecuteRequest{ Integration: r.Integration, + Operation: r.Operation, Target: remoteQueryTarget{Host: r.Target.Host, Port: r.Target.Port, DBName: r.Target.DBName}, Query: r.Query, + Format: r.Format, } if r.Limits != nil { internal.Limits = &remoteQueryExecuteLimits{MaxRows: r.Limits.MaxRows, MaxBytes: r.Limits.MaxBytes, TimeoutMs: r.Limits.TimeoutMs} } + if r.CopyLimits != nil { + internal.CopyLimits = &remoteQueryExecuteCopyLimits{ChunkBytes: r.CopyLimits.ChunkBytes, MaxBytes: r.CopyLimits.MaxBytes, MaxRowBytes: r.CopyLimits.MaxRowBytes, TimeoutMs: r.CopyLimits.TimeoutMs} + } return internal } func remoteQueryExecuteRequestFromInternal(req remoteQueryExecuteRequest) RemoteQueryExecuteRequest { out := RemoteQueryExecuteRequest{ Integration: req.Integration, + Operation: req.Operation, Target: RemoteQueryExecuteTarget{Host: req.Target.Host, Port: req.Target.Port, DBName: req.Target.DBName}, Query: req.Query, + Format: req.Format, } if req.Limits != nil { out.Limits = &RemoteQueryExecuteLimits{MaxRows: req.Limits.MaxRows, MaxBytes: req.Limits.MaxBytes, TimeoutMs: req.Limits.TimeoutMs} } + if req.CopyLimits != nil { + out.CopyLimits = &RemoteQueryExecuteCopyLimits{ChunkBytes: req.CopyLimits.ChunkBytes, MaxBytes: req.CopyLimits.MaxBytes, MaxRowBytes: req.CopyLimits.MaxRowBytes, TimeoutMs: req.CopyLimits.TimeoutMs} + } return out } func marshalExecuteRequest(req remoteQueryExecuteRequest) (string, error) { + if req.Operation == "copy_stream" { + format := req.Format + if format == "" { + format = "csv" + } + wireReq := remoteQueryCopyExecutorRequestJSON{ + Operation: req.Operation, + Target: remoteQueryTargetJSON{Host: req.Target.Host, Port: req.Target.Port, DBName: req.Target.DBName}, + Query: req.Query, + Format: format, + } + if req.CopyLimits != nil { + wireReq.Limits = &remoteQueryExecuteCopyLimitsJSON{ + ChunkBytes: req.CopyLimits.ChunkBytes, + MaxBytes: req.CopyLimits.MaxBytes, + MaxRowBytes: req.CopyLimits.MaxRowBytes, + TimeoutMs: req.CopyLimits.TimeoutMs, + } + } + requestJSON, err := json.Marshal(wireReq) + if err != nil { + return "", err + } + return string(requestJSON), nil + } + wireReq := remoteQueryExecutorRequestJSON{ Target: remoteQueryTargetJSON{Host: req.Target.Host, Port: req.Target.Port, DBName: req.Target.DBName}, Query: req.Query, diff --git a/comp/remotequeries/impl/remote_query_test.go b/comp/remotequeries/impl/remote_query_test.go index a4d780916465..ac9a8f84b83f 100644 --- a/comp/remotequeries/impl/remote_query_test.go +++ b/comp/remotequeries/impl/remote_query_test.go @@ -438,6 +438,29 @@ func TestRemoteQueryExecuteHandlerRunnerSuccessWithFixtureTableQuery(t *testing. assert.NotContains(t, recorder.Body.String(), "secret-value") } +func TestRemoteQueryExecuteServiceCopyStreamDispatch(t *testing.T) { + runner := &fakeStreamRunnerCheck{ + fakeRunnerCheck: fakeRunnerCheck{fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-value\n"}}, + events: []string{`{"type":"metadata","status":"STARTED"}`, `{"type":"data","sequence":0,"data":"a,b\n","bytes":4}`, `{"type":"final","status":"SUCCEEDED"}`}, + } + service := NewRemoteQueryExecuteService(fakeCollector{checks: []check.Check{fakeWrappedCheck{Check: runner}}}, true) + req, err := NewRemoteQueryCopyStreamExecuteRequest("postgres", RemoteQueryExecuteTarget{Host: "LOCALHOST.", Port: 5432, DBName: "postgres"}, "SELECT city, country FROM cities ORDER BY city", "csv", &RemoteQueryExecuteCopyLimits{ChunkBytes: 4, MaxBytes: 1024, MaxRowBytes: 1024, TimeoutMs: 1000}) + require.NoError(t, err) + + var events []string + result := service.ExecuteStream(req, func(event string) error { + events = append(events, event) + return nil + }) + + require.Nil(t, result.Error) + assert.Equal(t, runner.events, events) + assert.Equal(t, 1, runner.streamCalls) + assert.Equal(t, 0, runner.jsonCalls) + assert.JSONEq(t, `{"operation":"copy_stream","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT city, country FROM cities ORDER BY city","format":"csv","limits":{"chunkBytes":4,"maxBytes":1024,"maxRowBytes":1024,"timeoutMs":1000}}`, runner.streamSeen) + assert.NotContains(t, runner.streamSeen, "integration") +} + func TestRemoteQueryExecuteHandlerRejectsInvalidIntegration(t *testing.T) { handler := &remoteQueryExecuteHandler{enabled: true, collector: fakeCollector{}} @@ -523,15 +546,17 @@ func (f fakeWrappedCheck) Unwrap() check.Check { type fakeRunnerCheck struct { fakeCheck - response string - err error - seen string + response string + err error + seen string + jsonCalls int } func (f *fakeRunnerCheck) RunRemoteQueryJSON(integration string, requestJSON string) (string, error) { if integration != "postgres" { return "", assert.AnError } + f.jsonCalls++ f.seen = requestJSON return f.response, f.err } @@ -539,3 +564,24 @@ func (f *fakeRunnerCheck) RunRemoteQueryJSON(integration string, requestJSON str func (f *fakeRunnerCheck) seenRequest() string { return f.seen } + +type fakeStreamRunnerCheck struct { + fakeRunnerCheck + events []string + streamSeen string + streamCalls int +} + +func (f *fakeStreamRunnerCheck) RunRemoteQueryStream(integration string, requestJSON string, emit func(string) error) error { + if integration != "postgres" { + return assert.AnError + } + f.streamCalls++ + f.streamSeen = requestJSON + for _, event := range f.events { + if err := emit(event); err != nil { + return err + } + } + return nil +} diff --git a/pkg/collector/python/check.go b/pkg/collector/python/check.go index 7f13184fc8aa..0966d4a794af 100644 --- a/pkg/collector/python/check.go +++ b/pkg/collector/python/check.go @@ -13,6 +13,7 @@ import ( "errors" "fmt" "runtime" + "runtime/cgo" "runtime/pprof" "strings" "time" @@ -40,6 +41,12 @@ import ( char *getStringAddr(char **array, unsigned int idx); char *run_remote_query(rtloader_t *, rtloader_pyobject_t *check, const char *integration, const char *request_json); +extern int remoteQueryStreamEmitBridge(const char *event_json, void *userdata); +int run_remote_query_stream(rtloader_t *, rtloader_pyobject_t *check, const char *integration, const char *request_json, int (*emit)(const char *, void *), void *userdata); + +static inline int call_run_remote_query_stream(rtloader_t *rtloader, rtloader_pyobject_t *check, const char *integration, const char *request_json, void *userdata) { + return run_remote_query_stream(rtloader, check, integration, request_json, remoteQueryStreamEmitBridge, userdata); +} static inline void call_free(void* ptr) { _free(ptr); @@ -192,6 +199,43 @@ func (c *PythonCheck) RunRemoteQueryJSON(integration string, requestJSON string) return C.GoString(cResult), nil } +// RunRemoteQueryStream runs a streaming remote query helper for this Python check. +func (c *PythonCheck) RunRemoteQueryStream(integration string, requestJSON string, emit func(string) error) error { + integration = strings.ToLower(strings.TrimSpace(integration)) + if integration == "" { + return fmt.Errorf("integration is required") + } + if emit == nil { + return fmt.Errorf("emit callback is required") + } + + gstate, err := newStickyLock() + if err != nil { + return err + } + defer gstate.unlock() + + if c.cancelled { + return fmt.Errorf("check %s is already cancelled", c.ModuleName) + } + + cIntegration := C.CString(integration) + defer C.free(unsafe.Pointer(cIntegration)) + cRequestJSON := C.CString(requestJSON) + defer C.free(unsafe.Pointer(cRequestJSON)) + + h := cgo.NewHandle(emit) + defer h.Delete() + ok := C.call_run_remote_query_stream(rtloader, c.instance, cIntegration, cRequestJSON, unsafe.Pointer(h)) + if ok == 0 { + if err := getRtLoaderError(); err != nil { + return err + } + return fmt.Errorf("an error occurred while running remote query stream") + } + return nil +} + // Stop does nothing func (c *PythonCheck) Stop() {} diff --git a/pkg/collector/python/check_test.go b/pkg/collector/python/check_test.go index 8a27336aa5e9..a1ba6a2bfb41 100644 --- a/pkg/collector/python/check_test.go +++ b/pkg/collector/python/check_test.go @@ -90,3 +90,15 @@ func TestRunRemoteQueryJSONWithRuntimeNotInitializedError(t *testing.T) { func TestRunRemoteQueryJSONAfterCancel(t *testing.T) { testRunRemoteQueryJSONAfterCancel(t) } + +func TestRunRemoteQueryStream(t *testing.T) { + testRunRemoteQueryStream(t) +} + +func TestRunRemoteQueryStreamEmitError(t *testing.T) { + testRunRemoteQueryStreamEmitError(t) +} + +func TestRunRemoteQueryStreamAfterCancel(t *testing.T) { + testRunRemoteQueryStreamAfterCancel(t) +} diff --git a/pkg/collector/python/remote_query_stream.go b/pkg/collector/python/remote_query_stream.go new file mode 100644 index 000000000000..5c7d1012ac53 --- /dev/null +++ b/pkg/collector/python/remote_query_stream.go @@ -0,0 +1,28 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2026-present Datadog, Inc. + +//go:build python + +package python + +/* +#include +*/ +import "C" + +import ( + "runtime/cgo" + "unsafe" +) + +//export remoteQueryStreamEmitBridge +func remoteQueryStreamEmitBridge(eventJSON *C.char, userdata unsafe.Pointer) C.int { + h := cgo.Handle(uintptr(userdata)) + emit := h.Value().(func(string) error) + if err := emit(C.GoString(eventJSON)); err != nil { + return 1 + } + return 0 +} diff --git a/pkg/collector/python/test_check.go b/pkg/collector/python/test_check.go index c704da3aa67d..39162d3d052f 100644 --- a/pkg/collector/python/test_check.go +++ b/pkg/collector/python/test_check.go @@ -117,6 +117,26 @@ char *run_remote_query(rtloader_t *s, rtloader_pyobject_t *check, const char *in return run_remote_query_return; } +int run_remote_query_stream_return = 1; +int run_remote_query_stream_calls = 0; +rtloader_pyobject_t *run_remote_query_stream_instance = NULL; +const char *run_remote_query_stream_integration = NULL; +const char *run_remote_query_stream_request_json = NULL; +const char *run_remote_query_stream_event_json = NULL; +int run_remote_query_stream(rtloader_t *s, rtloader_pyobject_t *check, const char *integration, const char *request_json, int (*emit)(const char *, void *), void *userdata) { + run_remote_query_stream_instance = check; + run_remote_query_stream_integration = strdup(integration); + run_remote_query_stream_request_json = strdup(request_json); + run_remote_query_stream_calls++; + run_remote_query_stream_event_json = "{\"type\":\"final\",\"status\":\"SUCCEEDED\"}"; + if (run_remote_query_stream_return && emit != NULL) { + if (emit(run_remote_query_stream_event_json, userdata) != 0) { + return 0; + } + } + return run_remote_query_stream_return; +} + // // get_check MOCK // @@ -221,6 +241,12 @@ void reset_check_mock() { run_remote_query_instance = NULL; run_remote_query_integration = NULL; run_remote_query_request_json = NULL; + run_remote_query_stream_return = 1; + run_remote_query_stream_calls = 0; + run_remote_query_stream_instance = NULL; + run_remote_query_stream_integration = NULL; + run_remote_query_stream_request_json = NULL; + run_remote_query_stream_event_json = NULL; } */ import "C" @@ -779,6 +805,59 @@ func testRunRemoteQueryJSONAfterCancel(t *testing.T) { assert.Equal(t, C.int(0), C.run_remote_query_calls) } +func testRunRemoteQueryStream(t *testing.T) { + mockRtloader(t) + + check, err := NewPythonFakeCheck(aggregator.NewNoOpSenderManager()) + require.NoError(t, err) + check.instance = newMockPyObjectPtr() + + C.reset_check_mock() + var events []string + err = check.RunRemoteQueryStream(" Postgres ", `{"operation":"copy_stream"}`, func(event string) error { + events = append(events, event) + return nil + }) + + require.NoError(t, err) + assert.Equal(t, []string{`{"type":"final","status":"SUCCEEDED"}`}, events) + assert.Equal(t, C.int(1), C.gil_locked_calls) + assert.Equal(t, C.int(1), C.gil_unlocked_calls) + assert.Equal(t, C.int(1), C.run_remote_query_stream_calls) + assert.Equal(t, check.instance, C.run_remote_query_stream_instance) + assert.Equal(t, "postgres", C.GoString(C.run_remote_query_stream_integration)) + assert.JSONEq(t, `{"operation":"copy_stream"}`, C.GoString(C.run_remote_query_stream_request_json)) +} + +func testRunRemoteQueryStreamEmitError(t *testing.T) { + mockRtloader(t) + + check, err := NewPythonFakeCheck(aggregator.NewNoOpSenderManager()) + require.NoError(t, err) + check.instance = newMockPyObjectPtr() + + C.reset_check_mock() + err = check.RunRemoteQueryStream("postgres", `{"operation":"copy_stream"}`, func(string) error { return assert.AnError }) + + require.Error(t, err) + assert.Equal(t, C.int(1), C.run_remote_query_stream_calls) +} + +func testRunRemoteQueryStreamAfterCancel(t *testing.T) { + mockRtloader(t) + + check, err := NewPythonFakeCheck(aggregator.NewNoOpSenderManager()) + require.NoError(t, err) + check.instance = newMockPyObjectPtr() + + C.reset_check_mock() + check.Cancel() + + err = check.RunRemoteQueryStream("postgres", `{"operation":"copy_stream"}`, func(string) error { return nil }) + assert.EqualError(t, err, "check fake_check is already cancelled") + assert.Equal(t, C.int(0), C.run_remote_query_stream_calls) +} + func testRunAfterCancel(t *testing.T) { mockRtloader(t) diff --git a/pkg/privateactionrunner/bundles/remotequeries/execute.go b/pkg/privateactionrunner/bundles/remotequeries/execute.go index 4a591513b608..d2e45e828caa 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/execute.go +++ b/pkg/privateactionrunner/bundles/remotequeries/execute.go @@ -38,10 +38,13 @@ func NewExecuteAction(newBridgeClient BridgeClientFactory) *ExecuteAction { } type ExecuteInputs struct { - Integration string `json:"integration"` - Target TargetInputs `json:"target"` - Query string `json:"query"` - Limits *LimitsInputs `json:"limits,omitempty"` + Integration string `json:"integration"` + Operation string `json:"operation,omitempty"` + Target TargetInputs `json:"target"` + Query string `json:"query"` + Format string `json:"format,omitempty"` + Limits *LimitsInputs `json:"limits,omitempty"` + CopyLimits *CopyLimitsInputs `json:"copyLimits,omitempty"` } type TargetInputs struct { @@ -56,6 +59,13 @@ type LimitsInputs struct { TimeoutMs int `json:"timeoutMs"` } +type CopyLimitsInputs struct { + ChunkBytes int `json:"chunkBytes"` + MaxBytes int `json:"maxBytes"` + MaxRowBytes int `json:"maxRowBytes"` + TimeoutMs int `json:"timeoutMs"` +} + func (a *ExecuteAction) Run( ctx context.Context, task *types.Task, @@ -94,6 +104,8 @@ func (a *ExecuteAction) Run( func remoteQueryExecuteRequestFromInputs(inputs ExecuteInputs) *pb.RemoteQueryExecuteRequest { req := &pb.RemoteQueryExecuteRequest{ Integration: inputs.Integration, + Operation: inputs.Operation, + Format: inputs.Format, Target: &pb.RemoteQueryTarget{ Host: inputs.Target.Host, Port: int32(inputs.Target.Port), @@ -108,6 +120,14 @@ func remoteQueryExecuteRequestFromInputs(inputs ExecuteInputs) *pb.RemoteQueryEx TimeoutMs: int32(inputs.Limits.TimeoutMs), } } + if inputs.CopyLimits != nil { + req.CopyLimits = &pb.RemoteQueryExecuteCopyLimits{ + ChunkBytes: int32(inputs.CopyLimits.ChunkBytes), + MaxBytes: int32(inputs.CopyLimits.MaxBytes), + MaxRowBytes: int32(inputs.CopyLimits.MaxRowBytes), + TimeoutMs: int32(inputs.CopyLimits.TimeoutMs), + } + } return req } @@ -117,8 +137,10 @@ func remoteQueryExecuteOutputFromStream(stream grpc.ServerStreamingClient[pb.Rem } var assembled bytes.Buffer + streamEvents := make([]json.RawMessage, 0) expectedChunkIndex := int32(0) seenFinal := false + finalChunkWasEmpty := false for { chunk, err := stream.Recv() if err == io.EOF { @@ -136,16 +158,67 @@ func remoteQueryExecuteOutputFromStream(stream grpc.ServerStreamingClient[pb.Rem if seenFinal { return nil, fmt.Errorf("remote query response stream sent chunk after final") } - _, _ = assembled.Write(chunk.GetResponseJsonChunk()) + payload := chunk.GetResponseJsonChunk() + if chunk.GetFinal() && len(payload) == 0 && len(streamEvents) > 0 { + finalChunkWasEmpty = true + } else { + _, _ = assembled.Write(payload) + if !chunk.GetFinal() { + streamEvents = append(streamEvents, append(json.RawMessage(nil), payload...)) + } + } seenFinal = chunk.GetFinal() expectedChunkIndex++ } if !seenFinal { return nil, fmt.Errorf("remote query response stream missing final chunk") } + if finalChunkWasEmpty { + return remoteQueryExecuteOutputFromEvents(streamEvents) + } return remoteQueryExecuteOutputFromJSON(assembled.Bytes()) } +func remoteQueryExecuteOutputFromEvents(rawEvents []json.RawMessage) (map[string]interface{}, error) { + events := make([]interface{}, 0, len(rawEvents)) + var finalEvent map[string]interface{} + var data bytes.Buffer + for _, raw := range rawEvents { + var event map[string]interface{} + if err := json.Unmarshal(raw, &event); err != nil { + return nil, err + } + events = append(events, normalizeRemoteQueryOutput(event)) + if event["type"] == "data" { + if chunk, ok := event["data"].(string); ok { + _, _ = data.WriteString(chunk) + } + } + if event["type"] == "final" { + finalEvent = event + } + } + if finalEvent == nil { + return nil, fmt.Errorf("remote query stream response missing final event") + } + status, _ := finalEvent["status"].(string) + if status == "" { + return nil, fmt.Errorf("remote query stream final event missing status") + } + output := map[string]interface{}{ + "status": status, + "events": events, + "data": data.String(), + } + if v, ok := finalEvent["error"]; ok { + output["error"] = normalizeRemoteQueryOutput(v) + } + if v, ok := finalEvent["stats"]; ok { + output["stats"] = normalizeRemoteQueryOutput(v) + } + return output, nil +} + func remoteQueryExecuteOutputFromJSON(responseJSON []byte) (map[string]interface{}, error) { var output map[string]interface{} decoder := json.NewDecoder(bytes.NewReader(responseJSON)) diff --git a/pkg/privateactionrunner/bundles/remotequeries/execute_test.go b/pkg/privateactionrunner/bundles/remotequeries/execute_test.go index 2c583da91c31..dd7f9a8d6700 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/execute_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/execute_test.go @@ -67,6 +67,33 @@ func TestExecuteActionUsesCredentialFreeAgentSecureRequestShape(t *testing.T) { }, output) } +func TestExecuteActionPreservesCopyStreamEvents(t *testing.T) { + client := &captureBridgeClient{chunks: []*pb.RemoteQueryExecuteChunk{ + {ResponseJsonChunk: []byte(`{"type":"metadata","status":"STARTED","operation":"copy_stream"}`), ChunkIndex: 0}, + {ResponseJsonChunk: []byte(`{"type":"data","sequence":0,"data":"Beautiful city of lights,France\n","bytes":32}`), ChunkIndex: 1}, + {ResponseJsonChunk: []byte(`{"type":"data","sequence":1,"data":"New York,USA\n","bytes":13}`), ChunkIndex: 2}, + {ResponseJsonChunk: []byte(`{"type":"final","status":"SUCCEEDED","stats":{"bytesEmitted":45}}`), ChunkIndex: 3}, + {ChunkIndex: 4, Final: true}, + }} + action := NewExecuteAction(func() (BridgeClient, error) { return client, nil }) + + output, err := action.Run(context.Background(), taskWithInputs(map[string]interface{}{ + "integration": "postgres", + "operation": "copy_stream", + "format": "csv", + "target": map[string]interface{}{"host": "localhost", "port": 5432, "dbname": "postgres"}, + "query": "SELECT city, country FROM cities ORDER BY city", + "copyLimits": map[string]interface{}{"chunkBytes": 32, "maxBytes": 1024, "maxRowBytes": 1024, "timeoutMs": 1000}, + }), nil) + + require.NoError(t, err) + assert.Equal(t, "copy_stream", client.request.GetOperation()) + assert.Equal(t, "csv", client.request.GetFormat()) + assert.Equal(t, int32(32), client.request.GetCopyLimits().GetChunkBytes()) + assert.Equal(t, "Beautiful city of lights,France\nNew York,USA\n", output.(map[string]interface{})["data"]) + assert.Equal(t, "SUCCEEDED", output.(map[string]interface{})["status"]) +} + func TestExecuteActionPreservesSanitizedBridgeErrorBody(t *testing.T) { client := &captureBridgeClient{ response: &pb.RemoteQueryExecuteResponse{Status: "target_not_found", Error: &pb.RemoteQueryExecuteError{Code: "target_not_found", Message: "no matching integration check found"}}, @@ -126,6 +153,7 @@ func taskWithInputs(inputs map[string]interface{}) *types.Task { type captureBridgeClient struct { request *pb.RemoteQueryExecuteRequest response *pb.RemoteQueryExecuteResponse + chunks []*pb.RemoteQueryExecuteChunk err error } @@ -139,6 +167,9 @@ func (c *captureBridgeClient) RemoteQueryExecuteStream(_ context.Context, req *p if c.err != nil { return nil, c.err } + if c.chunks != nil { + return &captureRemoteQueryExecuteStream{chunks: c.chunks}, nil + } responseJSON, err := json.Marshal(remoteQueryExecuteOutputFromProtoForTest(c.response)) if err != nil { return nil, err diff --git a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go index 1281ffd596a9..4276d1706a86 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go @@ -175,6 +175,21 @@ func remoteQueriesProofLimits(query string) map[string]interface{} { } } +func remoteQueriesProofCopyLimits(query string) map[string]interface{} { + maxBytes := 4 << 10 + timeoutMs := 5_000 + if payloadBytes, ok := remoteQueriesLargePayloadBytes(query); ok { + maxBytes = payloadBytes + (1 << 20) + timeoutMs = 60_000 + } + return map[string]interface{}{ + "chunkBytes": 32, + "maxBytes": maxBytes, + "maxRowBytes": maxBytes, + "timeoutMs": timeoutMs, + } +} + func remoteQueriesProofResultTimeout(query string) time.Duration { if _, ok := remoteQueriesLargePayloadBytes(query); ok { return 2 * time.Minute @@ -187,6 +202,19 @@ func remoteQueriesLargePayloadBytes(query string) (int, bool) { return payloadBytes, ok } +func assertRemoteQueriesProofCopyData(t *testing.T, query string, data string) { + t.Helper() + + switch query { + case remoteQueriesFixtureTableProofQuery: + assert.Equal(t, "Beautiful city of lights,France\nNew York,USA\n", data) + case remoteQueriesSeedProofQuery: + assert.Equal(t, "1\n", data) + default: + require.FailNowf(t, "unsupported COPY proof query", "%s=%q must use a COPY bridge-allowlisted proof query", remoteQueriesProofQueryOverrideEnv, query) + } +} + func assertRemoteQueriesProofRows(t *testing.T, query string, rows []interface{}) { t.Helper() diff --git a/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go b/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go index be3882feed75..2c57db6c4a07 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go @@ -83,7 +83,14 @@ func TestRemoteQueriesActionRunsThroughStandalonePARProcessWithRealAgentIPC(t *t "integration": "postgres", "target": remoteQueriesPostgresTargetFromEnv(t), "query": proofQuery, - "limits": remoteQueriesProofLimits(proofQuery), + } + copyStream := os.Getenv("RQ_REMOTE_OPERATION") == "copy_stream" + if copyStream { + inputs["operation"] = "copy_stream" + inputs["format"] = "csv" + inputs["copyLimits"] = remoteQueriesProofCopyLimits(proofQuery) + } else { + inputs["limits"] = remoteQueriesProofLimits(proofQuery) } requestEvidence, err := json.Marshal(inputs) require.NoError(t, err) @@ -109,11 +116,17 @@ func TestRemoteQueriesActionRunsThroughStandalonePARProcessWithRealAgentIPC(t *t require.True(t, result.Success) require.Equal(t, taskID, result.TaskID) assert.Equal(t, "SUCCEEDED", result.Outputs["status"]) - require.Contains(t, result.Outputs, "rows") - - rows, ok := result.Outputs["rows"].([]interface{}) - require.True(t, ok) - assertRemoteQueriesProofRows(t, proofQuery, rows) + if copyStream { + t.Logf("copy stream PAR outputs: %+v", summarizeRemoteQueriesProofPayload(result.Outputs)) + data, ok := result.Outputs["data"].(string) + require.True(t, ok) + assertRemoteQueriesProofCopyData(t, proofQuery, data) + } else { + require.Contains(t, result.Outputs, "rows") + rows, ok := result.Outputs["rows"].([]interface{}) + require.True(t, ok) + assertRemoteQueriesProofRows(t, proofQuery, rows) + } resultEvidence, err := json.Marshal(summarizeRemoteQueriesProofPayload(result.Outputs)) require.NoError(t, err) diff --git a/pkg/proto/datadog/api/v1/api.proto b/pkg/proto/datadog/api/v1/api.proto index 81e5b6ce78ce..5973d5c9b5be 100644 --- a/pkg/proto/datadog/api/v1/api.proto +++ b/pkg/proto/datadog/api/v1/api.proto @@ -33,11 +33,21 @@ message RemoteQueryExecuteLimits { int32 timeout_ms = 3; } +message RemoteQueryExecuteCopyLimits { + int32 chunk_bytes = 1; + int32 max_bytes = 2; + int32 max_row_bytes = 3; + int32 timeout_ms = 4; +} + message RemoteQueryExecuteRequest { string integration = 1; RemoteQueryTarget target = 2; string query = 3; RemoteQueryExecuteLimits limits = 4; + string operation = 5; + string format = 6; + RemoteQueryExecuteCopyLimits copy_limits = 7; } message RemoteQueryExecuteError { diff --git a/pkg/proto/pbgo/core/api.pb.go b/pkg/proto/pbgo/core/api.pb.go index d6b495227ede..b65269ec75ce 100644 --- a/pkg/proto/pbgo/core/api.pb.go +++ b/pkg/proto/pbgo/core/api.pb.go @@ -143,19 +143,90 @@ func (x *RemoteQueryExecuteLimits) GetTimeoutMs() int32 { return 0 } +type RemoteQueryExecuteCopyLimits struct { + state protoimpl.MessageState `protogen:"open.v1"` + ChunkBytes int32 `protobuf:"varint,1,opt,name=chunk_bytes,json=chunkBytes,proto3" json:"chunk_bytes,omitempty"` + MaxBytes int32 `protobuf:"varint,2,opt,name=max_bytes,json=maxBytes,proto3" json:"max_bytes,omitempty"` + MaxRowBytes int32 `protobuf:"varint,3,opt,name=max_row_bytes,json=maxRowBytes,proto3" json:"max_row_bytes,omitempty"` + TimeoutMs int32 `protobuf:"varint,4,opt,name=timeout_ms,json=timeoutMs,proto3" json:"timeout_ms,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RemoteQueryExecuteCopyLimits) Reset() { + *x = RemoteQueryExecuteCopyLimits{} + mi := &file_datadog_api_v1_api_proto_msgTypes[2] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RemoteQueryExecuteCopyLimits) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RemoteQueryExecuteCopyLimits) ProtoMessage() {} + +func (x *RemoteQueryExecuteCopyLimits) ProtoReflect() protoreflect.Message { + mi := &file_datadog_api_v1_api_proto_msgTypes[2] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RemoteQueryExecuteCopyLimits.ProtoReflect.Descriptor instead. +func (*RemoteQueryExecuteCopyLimits) Descriptor() ([]byte, []int) { + return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{2} +} + +func (x *RemoteQueryExecuteCopyLimits) GetChunkBytes() int32 { + if x != nil { + return x.ChunkBytes + } + return 0 +} + +func (x *RemoteQueryExecuteCopyLimits) GetMaxBytes() int32 { + if x != nil { + return x.MaxBytes + } + return 0 +} + +func (x *RemoteQueryExecuteCopyLimits) GetMaxRowBytes() int32 { + if x != nil { + return x.MaxRowBytes + } + return 0 +} + +func (x *RemoteQueryExecuteCopyLimits) GetTimeoutMs() int32 { + if x != nil { + return x.TimeoutMs + } + return 0 +} + type RemoteQueryExecuteRequest struct { - state protoimpl.MessageState `protogen:"open.v1"` - Integration string `protobuf:"bytes,1,opt,name=integration,proto3" json:"integration,omitempty"` - Target *RemoteQueryTarget `protobuf:"bytes,2,opt,name=target,proto3" json:"target,omitempty"` - Query string `protobuf:"bytes,3,opt,name=query,proto3" json:"query,omitempty"` - Limits *RemoteQueryExecuteLimits `protobuf:"bytes,4,opt,name=limits,proto3" json:"limits,omitempty"` + state protoimpl.MessageState `protogen:"open.v1"` + Integration string `protobuf:"bytes,1,opt,name=integration,proto3" json:"integration,omitempty"` + Target *RemoteQueryTarget `protobuf:"bytes,2,opt,name=target,proto3" json:"target,omitempty"` + Query string `protobuf:"bytes,3,opt,name=query,proto3" json:"query,omitempty"` + Limits *RemoteQueryExecuteLimits `protobuf:"bytes,4,opt,name=limits,proto3" json:"limits,omitempty"` + Operation string `protobuf:"bytes,5,opt,name=operation,proto3" json:"operation,omitempty"` + Format string `protobuf:"bytes,6,opt,name=format,proto3" json:"format,omitempty"` + CopyLimits *RemoteQueryExecuteCopyLimits `protobuf:"bytes,7,opt,name=copy_limits,json=copyLimits,proto3" json:"copy_limits,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } func (x *RemoteQueryExecuteRequest) Reset() { *x = RemoteQueryExecuteRequest{} - mi := &file_datadog_api_v1_api_proto_msgTypes[2] + mi := &file_datadog_api_v1_api_proto_msgTypes[3] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -167,7 +238,7 @@ func (x *RemoteQueryExecuteRequest) String() string { func (*RemoteQueryExecuteRequest) ProtoMessage() {} func (x *RemoteQueryExecuteRequest) ProtoReflect() protoreflect.Message { - mi := &file_datadog_api_v1_api_proto_msgTypes[2] + mi := &file_datadog_api_v1_api_proto_msgTypes[3] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -180,7 +251,7 @@ func (x *RemoteQueryExecuteRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use RemoteQueryExecuteRequest.ProtoReflect.Descriptor instead. func (*RemoteQueryExecuteRequest) Descriptor() ([]byte, []int) { - return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{2} + return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{3} } func (x *RemoteQueryExecuteRequest) GetIntegration() string { @@ -211,6 +282,27 @@ func (x *RemoteQueryExecuteRequest) GetLimits() *RemoteQueryExecuteLimits { return nil } +func (x *RemoteQueryExecuteRequest) GetOperation() string { + if x != nil { + return x.Operation + } + return "" +} + +func (x *RemoteQueryExecuteRequest) GetFormat() string { + if x != nil { + return x.Format + } + return "" +} + +func (x *RemoteQueryExecuteRequest) GetCopyLimits() *RemoteQueryExecuteCopyLimits { + if x != nil { + return x.CopyLimits + } + return nil +} + type RemoteQueryExecuteError struct { state protoimpl.MessageState `protogen:"open.v1"` Code string `protobuf:"bytes,1,opt,name=code,proto3" json:"code,omitempty"` @@ -221,7 +313,7 @@ type RemoteQueryExecuteError struct { func (x *RemoteQueryExecuteError) Reset() { *x = RemoteQueryExecuteError{} - mi := &file_datadog_api_v1_api_proto_msgTypes[3] + mi := &file_datadog_api_v1_api_proto_msgTypes[4] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -233,7 +325,7 @@ func (x *RemoteQueryExecuteError) String() string { func (*RemoteQueryExecuteError) ProtoMessage() {} func (x *RemoteQueryExecuteError) ProtoReflect() protoreflect.Message { - mi := &file_datadog_api_v1_api_proto_msgTypes[3] + mi := &file_datadog_api_v1_api_proto_msgTypes[4] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -246,7 +338,7 @@ func (x *RemoteQueryExecuteError) ProtoReflect() protoreflect.Message { // Deprecated: Use RemoteQueryExecuteError.ProtoReflect.Descriptor instead. func (*RemoteQueryExecuteError) Descriptor() ([]byte, []int) { - return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{3} + return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{4} } func (x *RemoteQueryExecuteError) GetCode() string { @@ -277,7 +369,7 @@ type RemoteQueryExecuteResponse struct { func (x *RemoteQueryExecuteResponse) Reset() { *x = RemoteQueryExecuteResponse{} - mi := &file_datadog_api_v1_api_proto_msgTypes[4] + mi := &file_datadog_api_v1_api_proto_msgTypes[5] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -289,7 +381,7 @@ func (x *RemoteQueryExecuteResponse) String() string { func (*RemoteQueryExecuteResponse) ProtoMessage() {} func (x *RemoteQueryExecuteResponse) ProtoReflect() protoreflect.Message { - mi := &file_datadog_api_v1_api_proto_msgTypes[4] + mi := &file_datadog_api_v1_api_proto_msgTypes[5] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -302,7 +394,7 @@ func (x *RemoteQueryExecuteResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use RemoteQueryExecuteResponse.ProtoReflect.Descriptor instead. func (*RemoteQueryExecuteResponse) Descriptor() ([]byte, []int) { - return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{4} + return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{5} } func (x *RemoteQueryExecuteResponse) GetStatus() string { @@ -358,7 +450,7 @@ type RemoteQueryExecuteChunk struct { func (x *RemoteQueryExecuteChunk) Reset() { *x = RemoteQueryExecuteChunk{} - mi := &file_datadog_api_v1_api_proto_msgTypes[5] + mi := &file_datadog_api_v1_api_proto_msgTypes[6] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -370,7 +462,7 @@ func (x *RemoteQueryExecuteChunk) String() string { func (*RemoteQueryExecuteChunk) ProtoMessage() {} func (x *RemoteQueryExecuteChunk) ProtoReflect() protoreflect.Message { - mi := &file_datadog_api_v1_api_proto_msgTypes[5] + mi := &file_datadog_api_v1_api_proto_msgTypes[6] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -383,7 +475,7 @@ func (x *RemoteQueryExecuteChunk) ProtoReflect() protoreflect.Message { // Deprecated: Use RemoteQueryExecuteChunk.ProtoReflect.Descriptor instead. func (*RemoteQueryExecuteChunk) Descriptor() ([]byte, []int) { - return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{5} + return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{6} } func (x *RemoteQueryExecuteChunk) GetResponseJsonChunk() []byte { @@ -420,12 +512,23 @@ const file_datadog_api_v1_api_proto_rawDesc = "" + "\bmax_rows\x18\x01 \x01(\x05R\amaxRows\x12\x1b\n" + "\tmax_bytes\x18\x02 \x01(\x05R\bmaxBytes\x12\x1d\n" + "\n" + - "timeout_ms\x18\x03 \x01(\x05R\ttimeoutMs\"\xd0\x01\n" + + "timeout_ms\x18\x03 \x01(\x05R\ttimeoutMs\"\x9f\x01\n" + + "\x1cRemoteQueryExecuteCopyLimits\x12\x1f\n" + + "\vchunk_bytes\x18\x01 \x01(\x05R\n" + + "chunkBytes\x12\x1b\n" + + "\tmax_bytes\x18\x02 \x01(\x05R\bmaxBytes\x12\"\n" + + "\rmax_row_bytes\x18\x03 \x01(\x05R\vmaxRowBytes\x12\x1d\n" + + "\n" + + "timeout_ms\x18\x04 \x01(\x05R\ttimeoutMs\"\xd5\x02\n" + "\x19RemoteQueryExecuteRequest\x12 \n" + "\vintegration\x18\x01 \x01(\tR\vintegration\x129\n" + "\x06target\x18\x02 \x01(\v2!.datadog.api.v1.RemoteQueryTargetR\x06target\x12\x14\n" + "\x05query\x18\x03 \x01(\tR\x05query\x12@\n" + - "\x06limits\x18\x04 \x01(\v2(.datadog.api.v1.RemoteQueryExecuteLimitsR\x06limits\"G\n" + + "\x06limits\x18\x04 \x01(\v2(.datadog.api.v1.RemoteQueryExecuteLimitsR\x06limits\x12\x1c\n" + + "\toperation\x18\x05 \x01(\tR\toperation\x12\x16\n" + + "\x06format\x18\x06 \x01(\tR\x06format\x12M\n" + + "\vcopy_limits\x18\a \x01(\v2,.datadog.api.v1.RemoteQueryExecuteCopyLimitsR\n" + + "copyLimits\"G\n" + "\x17RemoteQueryExecuteError\x12\x12\n" + "\x04code\x18\x01 \x01(\tR\x04code\x12\x18\n" + "\amessage\x18\x02 \x01(\tR\amessage\"\xa0\x02\n" + @@ -478,106 +581,108 @@ func file_datadog_api_v1_api_proto_rawDescGZIP() []byte { return file_datadog_api_v1_api_proto_rawDescData } -var file_datadog_api_v1_api_proto_msgTypes = make([]protoimpl.MessageInfo, 6) +var file_datadog_api_v1_api_proto_msgTypes = make([]protoimpl.MessageInfo, 7) var file_datadog_api_v1_api_proto_goTypes = []any{ (*RemoteQueryTarget)(nil), // 0: datadog.api.v1.RemoteQueryTarget (*RemoteQueryExecuteLimits)(nil), // 1: datadog.api.v1.RemoteQueryExecuteLimits - (*RemoteQueryExecuteRequest)(nil), // 2: datadog.api.v1.RemoteQueryExecuteRequest - (*RemoteQueryExecuteError)(nil), // 3: datadog.api.v1.RemoteQueryExecuteError - (*RemoteQueryExecuteResponse)(nil), // 4: datadog.api.v1.RemoteQueryExecuteResponse - (*RemoteQueryExecuteChunk)(nil), // 5: datadog.api.v1.RemoteQueryExecuteChunk - (*structpb.Struct)(nil), // 6: google.protobuf.Struct - (*HostnameRequest)(nil), // 7: datadog.model.v1.HostnameRequest - (*StreamTagsRequest)(nil), // 8: datadog.model.v1.StreamTagsRequest - (*GenerateContainerIDFromOriginInfoRequest)(nil), // 9: datadog.model.v1.GenerateContainerIDFromOriginInfoRequest - (*FetchEntityRequest)(nil), // 10: datadog.model.v1.FetchEntityRequest - (*CaptureTriggerRequest)(nil), // 11: datadog.model.v1.CaptureTriggerRequest - (*TaggerState)(nil), // 12: datadog.model.v1.TaggerState - (*ClientGetConfigsRequest)(nil), // 13: datadog.config.ClientGetConfigsRequest - (*emptypb.Empty)(nil), // 14: google.protobuf.Empty - (*ConfigSubscriptionRequest)(nil), // 15: datadog.config.ConfigSubscriptionRequest - (*WorkloadmetaStreamRequest)(nil), // 16: datadog.workloadmeta.WorkloadmetaStreamRequest - (*RegisterRemoteAgentRequest)(nil), // 17: datadog.remoteagent.v1.RegisterRemoteAgentRequest - (*RefreshRemoteAgentRequest)(nil), // 18: datadog.remoteagent.v1.RefreshRemoteAgentRequest - (*HostTagRequest)(nil), // 19: datadog.model.v1.HostTagRequest - (*ConfigStreamRequest)(nil), // 20: datadog.model.v1.ConfigStreamRequest - (*WorkloadFilterEvaluateRequest)(nil), // 21: datadog.workloadfilter.WorkloadFilterEvaluateRequest - (*KubeMetadataStreamRequest)(nil), // 22: datadog.kubemetadata.KubeMetadataStreamRequest - (*HostnameReply)(nil), // 23: datadog.model.v1.HostnameReply - (*StreamTagsResponse)(nil), // 24: datadog.model.v1.StreamTagsResponse - (*GenerateContainerIDFromOriginInfoResponse)(nil), // 25: datadog.model.v1.GenerateContainerIDFromOriginInfoResponse - (*FetchEntityResponse)(nil), // 26: datadog.model.v1.FetchEntityResponse - (*CaptureTriggerResponse)(nil), // 27: datadog.model.v1.CaptureTriggerResponse - (*TaggerStateResponse)(nil), // 28: datadog.model.v1.TaggerStateResponse - (*ClientGetConfigsResponse)(nil), // 29: datadog.config.ClientGetConfigsResponse - (*GetStateConfigResponse)(nil), // 30: datadog.config.GetStateConfigResponse - (*ConfigSubscriptionResponse)(nil), // 31: datadog.config.ConfigSubscriptionResponse - (*ResetStateConfigResponse)(nil), // 32: datadog.config.ResetStateConfigResponse - (*WorkloadmetaStreamResponse)(nil), // 33: datadog.workloadmeta.WorkloadmetaStreamResponse - (*RegisterRemoteAgentResponse)(nil), // 34: datadog.remoteagent.v1.RegisterRemoteAgentResponse - (*RefreshRemoteAgentResponse)(nil), // 35: datadog.remoteagent.v1.RefreshRemoteAgentResponse - (*AutodiscoveryStreamResponse)(nil), // 36: datadog.autodiscovery.AutodiscoveryStreamResponse - (*HostTagReply)(nil), // 37: datadog.model.v1.HostTagReply - (*ConfigEvent)(nil), // 38: datadog.model.v1.ConfigEvent - (*WorkloadFilterEvaluateResponse)(nil), // 39: datadog.workloadfilter.WorkloadFilterEvaluateResponse - (*KubeMetadataStreamResponse)(nil), // 40: datadog.kubemetadata.KubeMetadataStreamResponse + (*RemoteQueryExecuteCopyLimits)(nil), // 2: datadog.api.v1.RemoteQueryExecuteCopyLimits + (*RemoteQueryExecuteRequest)(nil), // 3: datadog.api.v1.RemoteQueryExecuteRequest + (*RemoteQueryExecuteError)(nil), // 4: datadog.api.v1.RemoteQueryExecuteError + (*RemoteQueryExecuteResponse)(nil), // 5: datadog.api.v1.RemoteQueryExecuteResponse + (*RemoteQueryExecuteChunk)(nil), // 6: datadog.api.v1.RemoteQueryExecuteChunk + (*structpb.Struct)(nil), // 7: google.protobuf.Struct + (*HostnameRequest)(nil), // 8: datadog.model.v1.HostnameRequest + (*StreamTagsRequest)(nil), // 9: datadog.model.v1.StreamTagsRequest + (*GenerateContainerIDFromOriginInfoRequest)(nil), // 10: datadog.model.v1.GenerateContainerIDFromOriginInfoRequest + (*FetchEntityRequest)(nil), // 11: datadog.model.v1.FetchEntityRequest + (*CaptureTriggerRequest)(nil), // 12: datadog.model.v1.CaptureTriggerRequest + (*TaggerState)(nil), // 13: datadog.model.v1.TaggerState + (*ClientGetConfigsRequest)(nil), // 14: datadog.config.ClientGetConfigsRequest + (*emptypb.Empty)(nil), // 15: google.protobuf.Empty + (*ConfigSubscriptionRequest)(nil), // 16: datadog.config.ConfigSubscriptionRequest + (*WorkloadmetaStreamRequest)(nil), // 17: datadog.workloadmeta.WorkloadmetaStreamRequest + (*RegisterRemoteAgentRequest)(nil), // 18: datadog.remoteagent.v1.RegisterRemoteAgentRequest + (*RefreshRemoteAgentRequest)(nil), // 19: datadog.remoteagent.v1.RefreshRemoteAgentRequest + (*HostTagRequest)(nil), // 20: datadog.model.v1.HostTagRequest + (*ConfigStreamRequest)(nil), // 21: datadog.model.v1.ConfigStreamRequest + (*WorkloadFilterEvaluateRequest)(nil), // 22: datadog.workloadfilter.WorkloadFilterEvaluateRequest + (*KubeMetadataStreamRequest)(nil), // 23: datadog.kubemetadata.KubeMetadataStreamRequest + (*HostnameReply)(nil), // 24: datadog.model.v1.HostnameReply + (*StreamTagsResponse)(nil), // 25: datadog.model.v1.StreamTagsResponse + (*GenerateContainerIDFromOriginInfoResponse)(nil), // 26: datadog.model.v1.GenerateContainerIDFromOriginInfoResponse + (*FetchEntityResponse)(nil), // 27: datadog.model.v1.FetchEntityResponse + (*CaptureTriggerResponse)(nil), // 28: datadog.model.v1.CaptureTriggerResponse + (*TaggerStateResponse)(nil), // 29: datadog.model.v1.TaggerStateResponse + (*ClientGetConfigsResponse)(nil), // 30: datadog.config.ClientGetConfigsResponse + (*GetStateConfigResponse)(nil), // 31: datadog.config.GetStateConfigResponse + (*ConfigSubscriptionResponse)(nil), // 32: datadog.config.ConfigSubscriptionResponse + (*ResetStateConfigResponse)(nil), // 33: datadog.config.ResetStateConfigResponse + (*WorkloadmetaStreamResponse)(nil), // 34: datadog.workloadmeta.WorkloadmetaStreamResponse + (*RegisterRemoteAgentResponse)(nil), // 35: datadog.remoteagent.v1.RegisterRemoteAgentResponse + (*RefreshRemoteAgentResponse)(nil), // 36: datadog.remoteagent.v1.RefreshRemoteAgentResponse + (*AutodiscoveryStreamResponse)(nil), // 37: datadog.autodiscovery.AutodiscoveryStreamResponse + (*HostTagReply)(nil), // 38: datadog.model.v1.HostTagReply + (*ConfigEvent)(nil), // 39: datadog.model.v1.ConfigEvent + (*WorkloadFilterEvaluateResponse)(nil), // 40: datadog.workloadfilter.WorkloadFilterEvaluateResponse + (*KubeMetadataStreamResponse)(nil), // 41: datadog.kubemetadata.KubeMetadataStreamResponse } var file_datadog_api_v1_api_proto_depIdxs = []int32{ 0, // 0: datadog.api.v1.RemoteQueryExecuteRequest.target:type_name -> datadog.api.v1.RemoteQueryTarget 1, // 1: datadog.api.v1.RemoteQueryExecuteRequest.limits:type_name -> datadog.api.v1.RemoteQueryExecuteLimits - 3, // 2: datadog.api.v1.RemoteQueryExecuteResponse.error:type_name -> datadog.api.v1.RemoteQueryExecuteError - 6, // 3: datadog.api.v1.RemoteQueryExecuteResponse.columns:type_name -> google.protobuf.Struct - 6, // 4: datadog.api.v1.RemoteQueryExecuteResponse.rows:type_name -> google.protobuf.Struct - 6, // 5: datadog.api.v1.RemoteQueryExecuteResponse.stats:type_name -> google.protobuf.Struct - 7, // 6: datadog.api.v1.Agent.GetHostname:input_type -> datadog.model.v1.HostnameRequest - 8, // 7: datadog.api.v1.AgentSecure.TaggerStreamEntities:input_type -> datadog.model.v1.StreamTagsRequest - 9, // 8: datadog.api.v1.AgentSecure.TaggerGenerateContainerIDFromOriginInfo:input_type -> datadog.model.v1.GenerateContainerIDFromOriginInfoRequest - 10, // 9: datadog.api.v1.AgentSecure.TaggerFetchEntity:input_type -> datadog.model.v1.FetchEntityRequest - 11, // 10: datadog.api.v1.AgentSecure.DogstatsdCaptureTrigger:input_type -> datadog.model.v1.CaptureTriggerRequest - 12, // 11: datadog.api.v1.AgentSecure.DogstatsdSetTaggerState:input_type -> datadog.model.v1.TaggerState - 13, // 12: datadog.api.v1.AgentSecure.ClientGetConfigs:input_type -> datadog.config.ClientGetConfigsRequest - 14, // 13: datadog.api.v1.AgentSecure.GetConfigState:input_type -> google.protobuf.Empty - 13, // 14: datadog.api.v1.AgentSecure.ClientGetConfigsHA:input_type -> datadog.config.ClientGetConfigsRequest - 14, // 15: datadog.api.v1.AgentSecure.GetConfigStateHA:input_type -> google.protobuf.Empty - 15, // 16: datadog.api.v1.AgentSecure.CreateConfigSubscription:input_type -> datadog.config.ConfigSubscriptionRequest - 14, // 17: datadog.api.v1.AgentSecure.ResetConfigState:input_type -> google.protobuf.Empty - 16, // 18: datadog.api.v1.AgentSecure.WorkloadmetaStreamEntities:input_type -> datadog.workloadmeta.WorkloadmetaStreamRequest - 17, // 19: datadog.api.v1.AgentSecure.RegisterRemoteAgent:input_type -> datadog.remoteagent.v1.RegisterRemoteAgentRequest - 18, // 20: datadog.api.v1.AgentSecure.RefreshRemoteAgent:input_type -> datadog.remoteagent.v1.RefreshRemoteAgentRequest - 14, // 21: datadog.api.v1.AgentSecure.AutodiscoveryStreamConfig:input_type -> google.protobuf.Empty - 19, // 22: datadog.api.v1.AgentSecure.GetHostTags:input_type -> datadog.model.v1.HostTagRequest - 20, // 23: datadog.api.v1.AgentSecure.StreamConfigEvents:input_type -> datadog.model.v1.ConfigStreamRequest - 21, // 24: datadog.api.v1.AgentSecure.WorkloadFilterEvaluate:input_type -> datadog.workloadfilter.WorkloadFilterEvaluateRequest - 2, // 25: datadog.api.v1.AgentSecure.RemoteQueryExecute:input_type -> datadog.api.v1.RemoteQueryExecuteRequest - 2, // 26: datadog.api.v1.AgentSecure.RemoteQueryExecuteStream:input_type -> datadog.api.v1.RemoteQueryExecuteRequest - 22, // 27: datadog.api.v1.AgentSecure.StreamKubeMetadata:input_type -> datadog.kubemetadata.KubeMetadataStreamRequest - 23, // 28: datadog.api.v1.Agent.GetHostname:output_type -> datadog.model.v1.HostnameReply - 24, // 29: datadog.api.v1.AgentSecure.TaggerStreamEntities:output_type -> datadog.model.v1.StreamTagsResponse - 25, // 30: datadog.api.v1.AgentSecure.TaggerGenerateContainerIDFromOriginInfo:output_type -> datadog.model.v1.GenerateContainerIDFromOriginInfoResponse - 26, // 31: datadog.api.v1.AgentSecure.TaggerFetchEntity:output_type -> datadog.model.v1.FetchEntityResponse - 27, // 32: datadog.api.v1.AgentSecure.DogstatsdCaptureTrigger:output_type -> datadog.model.v1.CaptureTriggerResponse - 28, // 33: datadog.api.v1.AgentSecure.DogstatsdSetTaggerState:output_type -> datadog.model.v1.TaggerStateResponse - 29, // 34: datadog.api.v1.AgentSecure.ClientGetConfigs:output_type -> datadog.config.ClientGetConfigsResponse - 30, // 35: datadog.api.v1.AgentSecure.GetConfigState:output_type -> datadog.config.GetStateConfigResponse - 29, // 36: datadog.api.v1.AgentSecure.ClientGetConfigsHA:output_type -> datadog.config.ClientGetConfigsResponse - 30, // 37: datadog.api.v1.AgentSecure.GetConfigStateHA:output_type -> datadog.config.GetStateConfigResponse - 31, // 38: datadog.api.v1.AgentSecure.CreateConfigSubscription:output_type -> datadog.config.ConfigSubscriptionResponse - 32, // 39: datadog.api.v1.AgentSecure.ResetConfigState:output_type -> datadog.config.ResetStateConfigResponse - 33, // 40: datadog.api.v1.AgentSecure.WorkloadmetaStreamEntities:output_type -> datadog.workloadmeta.WorkloadmetaStreamResponse - 34, // 41: datadog.api.v1.AgentSecure.RegisterRemoteAgent:output_type -> datadog.remoteagent.v1.RegisterRemoteAgentResponse - 35, // 42: datadog.api.v1.AgentSecure.RefreshRemoteAgent:output_type -> datadog.remoteagent.v1.RefreshRemoteAgentResponse - 36, // 43: datadog.api.v1.AgentSecure.AutodiscoveryStreamConfig:output_type -> datadog.autodiscovery.AutodiscoveryStreamResponse - 37, // 44: datadog.api.v1.AgentSecure.GetHostTags:output_type -> datadog.model.v1.HostTagReply - 38, // 45: datadog.api.v1.AgentSecure.StreamConfigEvents:output_type -> datadog.model.v1.ConfigEvent - 39, // 46: datadog.api.v1.AgentSecure.WorkloadFilterEvaluate:output_type -> datadog.workloadfilter.WorkloadFilterEvaluateResponse - 4, // 47: datadog.api.v1.AgentSecure.RemoteQueryExecute:output_type -> datadog.api.v1.RemoteQueryExecuteResponse - 5, // 48: datadog.api.v1.AgentSecure.RemoteQueryExecuteStream:output_type -> datadog.api.v1.RemoteQueryExecuteChunk - 40, // 49: datadog.api.v1.AgentSecure.StreamKubeMetadata:output_type -> datadog.kubemetadata.KubeMetadataStreamResponse - 28, // [28:50] is the sub-list for method output_type - 6, // [6:28] is the sub-list for method input_type - 6, // [6:6] is the sub-list for extension type_name - 6, // [6:6] is the sub-list for extension extendee - 0, // [0:6] is the sub-list for field type_name + 2, // 2: datadog.api.v1.RemoteQueryExecuteRequest.copy_limits:type_name -> datadog.api.v1.RemoteQueryExecuteCopyLimits + 4, // 3: datadog.api.v1.RemoteQueryExecuteResponse.error:type_name -> datadog.api.v1.RemoteQueryExecuteError + 7, // 4: datadog.api.v1.RemoteQueryExecuteResponse.columns:type_name -> google.protobuf.Struct + 7, // 5: datadog.api.v1.RemoteQueryExecuteResponse.rows:type_name -> google.protobuf.Struct + 7, // 6: datadog.api.v1.RemoteQueryExecuteResponse.stats:type_name -> google.protobuf.Struct + 8, // 7: datadog.api.v1.Agent.GetHostname:input_type -> datadog.model.v1.HostnameRequest + 9, // 8: datadog.api.v1.AgentSecure.TaggerStreamEntities:input_type -> datadog.model.v1.StreamTagsRequest + 10, // 9: datadog.api.v1.AgentSecure.TaggerGenerateContainerIDFromOriginInfo:input_type -> datadog.model.v1.GenerateContainerIDFromOriginInfoRequest + 11, // 10: datadog.api.v1.AgentSecure.TaggerFetchEntity:input_type -> datadog.model.v1.FetchEntityRequest + 12, // 11: datadog.api.v1.AgentSecure.DogstatsdCaptureTrigger:input_type -> datadog.model.v1.CaptureTriggerRequest + 13, // 12: datadog.api.v1.AgentSecure.DogstatsdSetTaggerState:input_type -> datadog.model.v1.TaggerState + 14, // 13: datadog.api.v1.AgentSecure.ClientGetConfigs:input_type -> datadog.config.ClientGetConfigsRequest + 15, // 14: datadog.api.v1.AgentSecure.GetConfigState:input_type -> google.protobuf.Empty + 14, // 15: datadog.api.v1.AgentSecure.ClientGetConfigsHA:input_type -> datadog.config.ClientGetConfigsRequest + 15, // 16: datadog.api.v1.AgentSecure.GetConfigStateHA:input_type -> google.protobuf.Empty + 16, // 17: datadog.api.v1.AgentSecure.CreateConfigSubscription:input_type -> datadog.config.ConfigSubscriptionRequest + 15, // 18: datadog.api.v1.AgentSecure.ResetConfigState:input_type -> google.protobuf.Empty + 17, // 19: datadog.api.v1.AgentSecure.WorkloadmetaStreamEntities:input_type -> datadog.workloadmeta.WorkloadmetaStreamRequest + 18, // 20: datadog.api.v1.AgentSecure.RegisterRemoteAgent:input_type -> datadog.remoteagent.v1.RegisterRemoteAgentRequest + 19, // 21: datadog.api.v1.AgentSecure.RefreshRemoteAgent:input_type -> datadog.remoteagent.v1.RefreshRemoteAgentRequest + 15, // 22: datadog.api.v1.AgentSecure.AutodiscoveryStreamConfig:input_type -> google.protobuf.Empty + 20, // 23: datadog.api.v1.AgentSecure.GetHostTags:input_type -> datadog.model.v1.HostTagRequest + 21, // 24: datadog.api.v1.AgentSecure.StreamConfigEvents:input_type -> datadog.model.v1.ConfigStreamRequest + 22, // 25: datadog.api.v1.AgentSecure.WorkloadFilterEvaluate:input_type -> datadog.workloadfilter.WorkloadFilterEvaluateRequest + 3, // 26: datadog.api.v1.AgentSecure.RemoteQueryExecute:input_type -> datadog.api.v1.RemoteQueryExecuteRequest + 3, // 27: datadog.api.v1.AgentSecure.RemoteQueryExecuteStream:input_type -> datadog.api.v1.RemoteQueryExecuteRequest + 23, // 28: datadog.api.v1.AgentSecure.StreamKubeMetadata:input_type -> datadog.kubemetadata.KubeMetadataStreamRequest + 24, // 29: datadog.api.v1.Agent.GetHostname:output_type -> datadog.model.v1.HostnameReply + 25, // 30: datadog.api.v1.AgentSecure.TaggerStreamEntities:output_type -> datadog.model.v1.StreamTagsResponse + 26, // 31: datadog.api.v1.AgentSecure.TaggerGenerateContainerIDFromOriginInfo:output_type -> datadog.model.v1.GenerateContainerIDFromOriginInfoResponse + 27, // 32: datadog.api.v1.AgentSecure.TaggerFetchEntity:output_type -> datadog.model.v1.FetchEntityResponse + 28, // 33: datadog.api.v1.AgentSecure.DogstatsdCaptureTrigger:output_type -> datadog.model.v1.CaptureTriggerResponse + 29, // 34: datadog.api.v1.AgentSecure.DogstatsdSetTaggerState:output_type -> datadog.model.v1.TaggerStateResponse + 30, // 35: datadog.api.v1.AgentSecure.ClientGetConfigs:output_type -> datadog.config.ClientGetConfigsResponse + 31, // 36: datadog.api.v1.AgentSecure.GetConfigState:output_type -> datadog.config.GetStateConfigResponse + 30, // 37: datadog.api.v1.AgentSecure.ClientGetConfigsHA:output_type -> datadog.config.ClientGetConfigsResponse + 31, // 38: datadog.api.v1.AgentSecure.GetConfigStateHA:output_type -> datadog.config.GetStateConfigResponse + 32, // 39: datadog.api.v1.AgentSecure.CreateConfigSubscription:output_type -> datadog.config.ConfigSubscriptionResponse + 33, // 40: datadog.api.v1.AgentSecure.ResetConfigState:output_type -> datadog.config.ResetStateConfigResponse + 34, // 41: datadog.api.v1.AgentSecure.WorkloadmetaStreamEntities:output_type -> datadog.workloadmeta.WorkloadmetaStreamResponse + 35, // 42: datadog.api.v1.AgentSecure.RegisterRemoteAgent:output_type -> datadog.remoteagent.v1.RegisterRemoteAgentResponse + 36, // 43: datadog.api.v1.AgentSecure.RefreshRemoteAgent:output_type -> datadog.remoteagent.v1.RefreshRemoteAgentResponse + 37, // 44: datadog.api.v1.AgentSecure.AutodiscoveryStreamConfig:output_type -> datadog.autodiscovery.AutodiscoveryStreamResponse + 38, // 45: datadog.api.v1.AgentSecure.GetHostTags:output_type -> datadog.model.v1.HostTagReply + 39, // 46: datadog.api.v1.AgentSecure.StreamConfigEvents:output_type -> datadog.model.v1.ConfigEvent + 40, // 47: datadog.api.v1.AgentSecure.WorkloadFilterEvaluate:output_type -> datadog.workloadfilter.WorkloadFilterEvaluateResponse + 5, // 48: datadog.api.v1.AgentSecure.RemoteQueryExecute:output_type -> datadog.api.v1.RemoteQueryExecuteResponse + 6, // 49: datadog.api.v1.AgentSecure.RemoteQueryExecuteStream:output_type -> datadog.api.v1.RemoteQueryExecuteChunk + 41, // 50: datadog.api.v1.AgentSecure.StreamKubeMetadata:output_type -> datadog.kubemetadata.KubeMetadataStreamResponse + 29, // [29:51] is the sub-list for method output_type + 7, // [7:29] is the sub-list for method input_type + 7, // [7:7] is the sub-list for extension type_name + 7, // [7:7] is the sub-list for extension extendee + 0, // [0:7] is the sub-list for field type_name } func init() { file_datadog_api_v1_api_proto_init() } @@ -598,7 +703,7 @@ func file_datadog_api_v1_api_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_datadog_api_v1_api_proto_rawDesc), len(file_datadog_api_v1_api_proto_rawDesc)), NumEnums: 0, - NumMessages: 6, + NumMessages: 7, NumExtensions: 0, NumServices: 2, }, diff --git a/rtloader/include/datadog_agent_rtloader.h b/rtloader/include/datadog_agent_rtloader.h index 2efbf61f641e..373713f7ca79 100644 --- a/rtloader/include/datadog_agent_rtloader.h +++ b/rtloader/include/datadog_agent_rtloader.h @@ -202,6 +202,21 @@ DATADOG_AGENT_RTLOADER_API char *run_check(rtloader_t *, rtloader_pyobject_t *ch DATADOG_AGENT_RTLOADER_API char *run_remote_query(rtloader_t *, rtloader_pyobject_t *check, const char *integration, const char *request_json); +/*! \fn int run_remote_query_stream(rtloader_t *, rtloader_pyobject_t *check, const char *integration, const char *request_json, remote_query_stream_emit_cb emit, void *userdata) + \brief Runs the integration remote query streaming helper for a check instance. + \param rtloader_t A rtloader_t * pointer to the RtLoader instance. + \param check A rtloader_pyobject_t * pointer to the check instance we wish to use. + \param integration The integration helper module name. + \param request_json A credential-free JSON request string. + \param emit Callback invoked once per serialized stream event. + \param userdata Opaque pointer passed to emit. + \return 1 on success, 0 on failure. + \sa rtloader_pyobject_t, rtloader_t +*/ +DATADOG_AGENT_RTLOADER_API int run_remote_query_stream(rtloader_t *, rtloader_pyobject_t *check, + const char *integration, const char *request_json, + remote_query_stream_emit_cb emit, void *userdata); + /*! \fn char *cancel_check(rtloader_t *, rtloader_pyobject_t *check) \brief Cancels a check instance. This allow check to be notified when they're unscheduled and can free any remaining resources. diff --git a/rtloader/include/rtloader.h b/rtloader/include/rtloader.h index 37c10cca3e9d..24bf498edf35 100644 --- a/rtloader/include/rtloader.h +++ b/rtloader/include/rtloader.h @@ -130,6 +130,19 @@ class RtLoader */ virtual char *runRemoteQuery(RtLoaderPyObject *check, const char *integration, const char *request_json) = 0; + //! Pure virtual runRemoteQueryStream member. + /*! + \param check The python object pointer to the check we wish to use. + \param integration The integration helper module name. + \param request_json A credential-free JSON request string. + \param emit Callback invoked once per serialized stream event. + \param userdata Opaque pointer passed to emit. + \return A boolean indicating the success or not of the operation. + */ + virtual bool runRemoteQueryStream(RtLoaderPyObject *check, const char *integration, const char *request_json, + remote_query_stream_emit_cb emit, void *userdata) + = 0; + //! Pure virtual cancelCheck member. /*! \param check The python object pointer to the check we wish to cancel. diff --git a/rtloader/include/rtloader_types.h b/rtloader/include/rtloader_types.h index fa5f9d50e943..b247260e51f3 100644 --- a/rtloader/include/rtloader_types.h +++ b/rtloader/include/rtloader_types.h @@ -37,6 +37,7 @@ typedef enum rtloader_gilstate_e { typedef void *(*rtloader_malloc_t)(size_t); typedef void (*rtloader_free_t)(void *); +typedef int (*remote_query_stream_emit_cb)(const char *event_json, void *userdata); typedef enum { DATADOG_AGENT_RTLOADER_GAUGE = 0, diff --git a/rtloader/rtloader/api.cpp b/rtloader/rtloader/api.cpp index 6aa8fd0f909b..5b5998d6584a 100644 --- a/rtloader/rtloader/api.cpp +++ b/rtloader/rtloader/api.cpp @@ -283,6 +283,15 @@ char *run_remote_query(rtloader_t *rtloader, rtloader_pyobject_t *check, const c return AS_TYPE(RtLoader, rtloader)->runRemoteQuery(AS_TYPE(RtLoaderPyObject, check), integration, request_json); } +int run_remote_query_stream(rtloader_t *rtloader, rtloader_pyobject_t *check, const char *integration, + const char *request_json, remote_query_stream_emit_cb emit, void *userdata) +{ + return AS_TYPE(RtLoader, rtloader) + ->runRemoteQueryStream(AS_TYPE(RtLoaderPyObject, check), integration, request_json, emit, userdata) + ? 1 + : 0; +} + void cancel_check(rtloader_t *rtloader, rtloader_pyobject_t *check) { AS_TYPE(RtLoader, rtloader)->cancelCheck(AS_TYPE(RtLoaderPyObject, check)); diff --git a/rtloader/three/three.cpp b/rtloader/three/three.cpp index 90667d7b4e28..0e908e2b2345 100644 --- a/rtloader/three/three.cpp +++ b/rtloader/three/three.cpp @@ -531,6 +531,77 @@ bool isValidRemoteQueryIntegration(const std::string &integration) return std::islower(ch) || std::isdigit(ch) || ch == '_'; }); } + +struct RemoteQueryStreamEmitContext { + remote_query_stream_emit_cb emit; + void *userdata; +}; + +PyObject *remoteQueryStreamEmit(PyObject *self, PyObject *event) +{ + RemoteQueryStreamEmitContext *ctx = static_cast(PyCapsule_GetPointer(self, "remote_query_stream_emit")); + if (ctx == NULL || ctx->emit == NULL) { + PyErr_SetString(PyExc_RuntimeError, "remote query stream emit callback is unavailable"); + return NULL; + } + + PyObject *normalized_event = event; + PyObject *event_copy = NULL; + if (PyDict_Check(event)) { + PyObject *data = PyDict_GetItemString(event, "data"); + if (data != NULL && PyBytes_Check(data)) { + event_copy = PyDict_Copy(event); + if (event_copy == NULL) { + return NULL; + } + char *bytes = PyBytes_AsString(data); + Py_ssize_t size = PyBytes_Size(data); + PyObject *data_str = PyUnicode_FromStringAndSize(bytes, size); + if (data_str == NULL) { + Py_DECREF(event_copy); + return NULL; + } + if (PyDict_SetItemString(event_copy, "data", data_str) != 0) { + Py_DECREF(data_str); + Py_DECREF(event_copy); + return NULL; + } + Py_DECREF(data_str); + normalized_event = event_copy; + } + } + + PyObject *json_module = PyImport_ImportModule("json"); + if (json_module == NULL) { + Py_XDECREF(event_copy); + return NULL; + } + PyObject *event_json = PyObject_CallMethod(json_module, const_cast("dumps"), const_cast("O"), normalized_event); + Py_DECREF(json_module); + Py_XDECREF(event_copy); + if (event_json == NULL || !PyUnicode_Check(event_json)) { + Py_XDECREF(event_json); + return NULL; + } + + PyObject *utf8 = PyUnicode_AsUTF8String(event_json); + Py_DECREF(event_json); + if (utf8 == NULL) { + return NULL; + } + const char *event_json_c = PyBytes_AsString(utf8); + int emit_result = ctx->emit(event_json_c, ctx->userdata); + Py_DECREF(utf8); + if (emit_result != 0) { + PyErr_SetString(PyExc_RuntimeError, "remote query stream emit callback failed"); + return NULL; + } + + Py_RETURN_NONE; +} + +PyMethodDef remoteQueryStreamEmitMethod = {"remote_query_stream_emit", remoteQueryStreamEmit, METH_O, + "Emit a serialized remote query stream event."}; } // namespace char *Three::runRemoteQuery(RtLoaderPyObject *check, const char *integration, const char *request_json) @@ -592,6 +663,76 @@ char *Three::runRemoteQuery(RtLoaderPyObject *check, const char *integration, co return ret; } +bool Three::runRemoteQueryStream(RtLoaderPyObject *check, const char *integration, const char *request_json, + remote_query_stream_emit_cb emit, void *userdata) +{ + if (check == NULL || request_json == NULL || emit == NULL) { + return false; + } + + std::string normalized_integration = normalizeRemoteQueryIntegration(integration); + if (!isValidRemoteQueryIntegration(normalized_integration)) { + setError("invalid remote query integration name"); + return false; + } + + PyObject *py_check = reinterpret_cast(check); + PyObject *remote_query_module = NULL; + PyObject *execute_func = NULL; + PyObject *py_request_json = NULL; + PyObject *capsule = NULL; + PyObject *emit_func = NULL; + PyObject *result = NULL; + std::string module_name = "datadog_checks." + normalized_integration + ".remote_query"; + RemoteQueryStreamEmitContext ctx{emit, userdata}; + bool ok = false; + + remote_query_module = PyImport_ImportModule(module_name.c_str()); + if (remote_query_module == NULL) { + setError("error importing remote query helper: " + _fetchPythonError()); + goto done; + } + + execute_func = PyObject_GetAttrString(remote_query_module, "execute_agent_rpc_stream_copy"); + if (execute_func == NULL || !PyCallable_Check(execute_func)) { + setError("error loading remote query stream helper: " + _fetchPythonError()); + goto done; + } + + py_request_json = PyUnicode_FromString(request_json); + if (py_request_json == NULL) { + setError("error converting remote query stream request to Python string: " + _fetchPythonError()); + goto done; + } + + capsule = PyCapsule_New(&ctx, "remote_query_stream_emit", NULL); + if (capsule == NULL) { + setError("error creating remote query stream emit context: " + _fetchPythonError()); + goto done; + } + emit_func = PyCFunction_NewEx(&remoteQueryStreamEmitMethod, capsule, NULL); + if (emit_func == NULL) { + setError("error creating remote query stream emit callback: " + _fetchPythonError()); + goto done; + } + + result = PyObject_CallFunctionObjArgs(execute_func, py_request_json, py_check, emit_func, NULL); + if (result == NULL) { + setError("error invoking remote query stream helper: " + _fetchPythonError()); + goto done; + } + ok = true; + + done: + Py_XDECREF(result); + Py_XDECREF(emit_func); + Py_XDECREF(capsule); + Py_XDECREF(py_request_json); + Py_XDECREF(execute_func); + Py_XDECREF(remote_query_module); + return ok; +} + void Three::cancelCheck(RtLoaderPyObject *check) { if (check == NULL) { diff --git a/rtloader/three/three.h b/rtloader/three/three.h index f946886ef16a..0c4a33769105 100644 --- a/rtloader/three/three.h +++ b/rtloader/three/three.h @@ -66,6 +66,8 @@ class Three : public RtLoader char *runCheck(RtLoaderPyObject *check); char *runRemoteQuery(RtLoaderPyObject *check, const char *integration, const char *request_json); + bool runRemoteQueryStream(RtLoaderPyObject *check, const char *integration, const char *request_json, + remote_query_stream_emit_cb emit, void *userdata); void cancelCheck(RtLoaderPyObject *check); char **getCheckWarnings(RtLoaderPyObject *check); char *getCheckDiagnoses(RtLoaderPyObject *check); diff --git a/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh b/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh index 432a5927a28e..cc20b57dfce7 100755 --- a/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh +++ b/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh @@ -29,6 +29,11 @@ if [[ -n "${RQ_REMOTE_QUERY+x}" ]]; then RQ_REMOTE_QUERY_WAS_SET=1 fi RQ_REMOTE_QUERY=${RQ_REMOTE_QUERY:-} +RQ_REMOTE_OPERATION_WAS_SET=0 +if [[ -n "${RQ_REMOTE_OPERATION+x}" ]]; then + RQ_REMOTE_OPERATION_WAS_SET=1 +fi +RQ_REMOTE_OPERATION=${RQ_REMOTE_OPERATION:-} RQ_POSTGRES_HOST=${RQ_POSTGRES_HOST:-localhost} RQ_POSTGRES_PORT=${RQ_POSTGRES_PORT:-5432} RQ_POSTGRES_DBNAME=${RQ_POSTGRES_DBNAME:-datadog_test} @@ -46,6 +51,7 @@ CASE_RESULTS_DIR="" PROOF_CASE_NAMES=( "seed" "fixture-city" + "copy-fixture-city" "payload-1mib" "payload-2mib" "payload-4mib" @@ -57,6 +63,7 @@ PROOF_CASE_NAMES=( PROOF_CASE_QUERIES=( "SELECT 1 AS value" "SELECT city, country FROM cities ORDER BY city" + "SELECT city, country FROM cities ORDER BY city" "SELECT repeat('x', 1048576) AS payload" "SELECT repeat('x', 2097152) AS payload" "SELECT repeat('x', 4194304) AS payload" @@ -467,6 +474,7 @@ run_standalone_go_proof() { RQ_POSTGRES_PORT="$RQ_POSTGRES_PORT" \ RQ_POSTGRES_DBNAME="$RQ_POSTGRES_DBNAME" \ RQ_REMOTE_QUERY="$RQ_REMOTE_QUERY" \ + RQ_REMOTE_OPERATION="${RQ_REMOTE_OPERATION:-}" \ dda inv test --targets=./pkg/privateactionrunner/bundles/remotequeries \ --extra-args='-run TestRemoteQueriesActionRunsThroughStandalonePARProcessWithRealAgentIPC -count=1 -v' ) | tee "$CASE_RESULTS_DIR/standalone-proof-test.log" @@ -475,6 +483,12 @@ run_standalone_go_proof() { run_proof_case() { PROOF_CASE_NAME=$1 RQ_REMOTE_QUERY=$2 + if [[ "$RQ_REMOTE_OPERATION_WAS_SET" != "1" ]]; then + RQ_REMOTE_OPERATION="" + if [[ "$PROOF_CASE_NAME" == copy-* ]]; then + RQ_REMOTE_OPERATION=copy_stream + fi + fi CASE_RESULTS_DIR="$TMP_ROOT/results/$PROOF_CASE_NAME" mkdir -p "$CASE_RESULTS_DIR" printf '%s\n' "$RQ_REMOTE_QUERY" > "$CASE_RESULTS_DIR/query.sql" From 9f6f700b13398d04a2a5c437b1e3b381620fceae Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 14:11:54 +0000 Subject: [PATCH 23/33] Use binary-safe Remote Queries stream events --- .../impl-agent/remote_query_execute_test.go | 16 + comp/api/grpcserver/impl-agent/server.go | 113 ++- .../impl/remote_query_execute.go | 14 +- comp/remotequeries/impl/remote_query_test.go | 14 +- pkg/collector/check/remote_query_stream.go | 13 + pkg/collector/python/check.go | 7 +- pkg/collector/python/remote_query_stream.go | 18 +- pkg/collector/python/test_check.go | 27 +- .../bundles/remotequeries/execute.go | 103 ++- .../bundles/remotequeries/execute_test.go | 33 +- .../live_agent_ipc_par_loop_test.go | 26 +- .../standalone_par_process_proof_test.go | 18 +- pkg/proto/datadog/api/v1/api.proto | 38 + pkg/proto/pbgo/core/api.pb.go | 656 +++++++++++++++--- rtloader/include/rtloader_types.h | 3 +- rtloader/three/three.cpp | 56 +- ...andalone-par-agentsecure-postgres-proof.sh | 19 + 17 files changed, 984 insertions(+), 190 deletions(-) create mode 100644 pkg/collector/check/remote_query_stream.go diff --git a/comp/api/grpcserver/impl-agent/remote_query_execute_test.go b/comp/api/grpcserver/impl-agent/remote_query_execute_test.go index 82cbbabb8052..493d7e7edbfa 100644 --- a/comp/api/grpcserver/impl-agent/remote_query_execute_test.go +++ b/comp/api/grpcserver/impl-agent/remote_query_execute_test.go @@ -14,6 +14,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/DataDog/datadog-agent/pkg/collector/check" pb "github.com/DataDog/datadog-agent/pkg/proto/pbgo/core" ) @@ -68,6 +69,21 @@ func TestRemoteQueryExecuteRequestFromProtoPreservesCopyStream(t *testing.T) { assert.Nil(t, req.Limits) } +func TestRemoteQueryStreamEventFromCheckEventPreservesBinaryPayload(t *testing.T) { + event, err := remoteQueryStreamEventFromCheckEvent(check.RemoteQueryStreamEvent{ + Type: "data", + MetadataJSON: `{"sequence":7,"offset":11,"bytes":3}`, + Payload: []byte{0x00, 0xff, 0x80}, + }) + + require.NoError(t, err) + assert.Equal(t, uint64(7), event.GetSequence()) + require.NotNil(t, event.GetData()) + assert.Equal(t, []byte{0x00, 0xff, 0x80}, event.GetData().GetPayload()) + assert.Equal(t, uint64(11), event.GetData().GetOffset()) + assert.Equal(t, uint64(3), event.GetData().GetBytes()) +} + func TestRemoteQueryExecuteStreamJSONChunksResponse(t *testing.T) { stream := &captureRemoteQueryExecuteStreamServer{} responseJSON := `{"status":"SUCCEEDED","rows":[{"payload":"` + string(make([]byte, remoteQueryExecuteStreamChunkSize+1)) + `"}]}` diff --git a/comp/api/grpcserver/impl-agent/server.go b/comp/api/grpcserver/impl-agent/server.go index 3938ddae20f1..c920c58d8353 100644 --- a/comp/api/grpcserver/impl-agent/server.go +++ b/comp/api/grpcserver/impl-agent/server.go @@ -38,6 +38,7 @@ import ( rcservice "github.com/DataDog/datadog-agent/comp/remote-config/rcservice/def" rcservicemrf "github.com/DataDog/datadog-agent/comp/remote-config/rcservicemrf/def" remotequeriesimpl "github.com/DataDog/datadog-agent/comp/remotequeries/impl" + "github.com/DataDog/datadog-agent/pkg/collector/check" pb "github.com/DataDog/datadog-agent/pkg/proto/pbgo/core" "github.com/DataDog/datadog-agent/pkg/util/grpc" "github.com/DataDog/datadog-agent/pkg/util/log" @@ -313,8 +314,12 @@ func (s *serverSecure) RemoteQueryExecuteStream(req *pb.RemoteQueryExecuteReques if execReq.Operation == "copy_stream" { chunkIndex := int32(0) - result := s.remoteQueries.ExecuteStream(execReq, func(eventJSON string) error { - err := stream.Send(&pb.RemoteQueryExecuteChunk{ResponseJsonChunk: []byte(eventJSON), ChunkIndex: chunkIndex}) + result := s.remoteQueries.ExecuteStream(execReq, func(event check.RemoteQueryStreamEvent) error { + protoEvent, err := remoteQueryStreamEventFromCheckEvent(event) + if err != nil { + return err + } + err = stream.Send(&pb.RemoteQueryExecuteChunk{Event: protoEvent, ChunkIndex: chunkIndex}) chunkIndex++ return err }) @@ -345,6 +350,110 @@ func remoteQueryExecuteRequestFromProto(req *pb.RemoteQueryExecuteRequest) (remo return remotequeriesimpl.NewRemoteQueryExecuteRequest(req.GetIntegration(), target, req.GetQuery(), limits) } +func remoteQueryStreamEventFromCheckEvent(event check.RemoteQueryStreamEvent) (*pb.RemoteQueryExecuteStreamEvent, error) { + metadata := map[string]interface{}{} + if strings.TrimSpace(event.MetadataJSON) != "" { + if err := json.Unmarshal([]byte(event.MetadataJSON), &metadata); err != nil { + return nil, err + } + } + sequence := uint64FromMetadata(metadata, "sequence") + out := &pb.RemoteQueryExecuteStreamEvent{Sequence: sequence} + switch event.Type { + case "metadata": + attrs := stringAttributes(metadata, "operation", "integration", "format", "sequence") + out.Event = &pb.RemoteQueryExecuteStreamEvent_Metadata{Metadata: &pb.RemoteQueryStreamMetadata{ + Operation: stringFromMetadata(metadata, "operation"), + Integration: stringFromMetadata(metadata, "integration"), + Format: stringFromMetadata(metadata, "format"), + Attributes: attrs, + }} + case "data": + out.Event = &pb.RemoteQueryExecuteStreamEvent_Data{Data: &pb.RemoteQueryStreamData{ + Payload: append([]byte(nil), event.Payload...), + Offset: uint64FromMetadata(metadata, "offset"), + Bytes: uint64FromMetadata(metadata, "bytes"), + }} + case "final": + out.Event = &pb.RemoteQueryExecuteStreamEvent_Final{Final: &pb.RemoteQueryStreamFinal{ + Status: stringFromMetadata(metadata, "status"), + BytesEmitted: uint64FromMetadata(metadata, "bytes_emitted", "bytesEmitted", "bytes"), + ChunksEmitted: uint64FromMetadata(metadata, "chunks_emitted", "chunksEmitted", "chunks"), + Attributes: stringAttributes(metadata, "status", "sequence", "bytes_emitted", "bytesEmitted", "chunks_emitted", "chunksEmitted"), + }} + case "error": + out.Event = &pb.RemoteQueryExecuteStreamEvent_Error{Error: &pb.RemoteQueryStreamError{ + Code: stringFromMetadata(metadata, "code"), + Message: stringFromMetadata(metadata, "message"), + Retryable: boolFromMetadata(metadata, "retryable"), + Attributes: stringAttributes(metadata, "code", "message", "retryable", "sequence"), + }} + default: + return nil, errors.New("unknown remote query stream event type") + } + return out, nil +} + +func stringFromMetadata(metadata map[string]interface{}, key string) string { + if v, ok := metadata[key].(string); ok { + return v + } + return "" +} + +func boolFromMetadata(metadata map[string]interface{}, key string) bool { + if v, ok := metadata[key].(bool); ok { + return v + } + return false +} + +func uint64FromMetadata(metadata map[string]interface{}, keys ...string) uint64 { + for _, key := range keys { + switch v := metadata[key].(type) { + case float64: + if v > 0 { + return uint64(v) + } + case int: + if v > 0 { + return uint64(v) + } + case json.Number: + if n, err := strconv.ParseUint(string(v), 10, 64); err == nil { + return n + } + case string: + if n, err := strconv.ParseUint(v, 10, 64); err == nil { + return n + } + } + } + return 0 +} + +func stringAttributes(metadata map[string]interface{}, exclude ...string) map[string]string { + excluded := make(map[string]struct{}, len(exclude)) + for _, key := range exclude { + excluded[key] = struct{}{} + } + attrs := make(map[string]string) + for key, value := range metadata { + if _, ok := excluded[key]; ok { + continue + } + switch v := value.(type) { + case string: + attrs[key] = v + case float64: + attrs[key] = strconv.FormatFloat(v, 'f', -1, 64) + case bool: + attrs[key] = strconv.FormatBool(v) + } + } + return attrs +} + func remoteQueryExecuteStreamJSON(responseJSON string, stream pb.AgentSecure_RemoteQueryExecuteStreamServer) error { responseBytes := []byte(responseJSON) if len(responseBytes) == 0 { diff --git a/comp/remotequeries/impl/remote_query_execute.go b/comp/remotequeries/impl/remote_query_execute.go index 09ee5e67aa44..2c0ddd7b186e 100644 --- a/comp/remotequeries/impl/remote_query_execute.go +++ b/comp/remotequeries/impl/remote_query_execute.go @@ -30,6 +30,8 @@ const ( statusExecutorUnavailable = "executor_unavailable" ) +const remoteQueryBinaryPayloadProofQuery = "SELECT decode('00ff80', 'hex') AS payload" + var remoteQueryLargePayloadProofQueries = map[string]int{ "SELECT repeat('x', 1048576) AS payload": 1 << 20, "SELECT repeat('x', 2097152) AS payload": 2 << 20, @@ -44,12 +46,12 @@ type remoteQueryRunner interface { } type remoteQueryStreamRunner interface { - RunRemoteQueryStream(integration string, requestJSON string, emit func(string) error) error + RunRemoteQueryStream(integration string, requestJSON string, emit func(check.RemoteQueryStreamEvent) error) error } func isRemoteQueryAllowedProofQuery(query string) bool { switch query { - case remoteQueryProofSeedQuery, remoteQueryFixtureTableProofQuery: + case remoteQueryProofSeedQuery, remoteQueryFixtureTableProofQuery, remoteQueryBinaryPayloadProofQuery: return true default: _, ok := remoteQueryLargePayloadProofQueries[query] @@ -174,8 +176,8 @@ func NewRemoteQueryCopyStreamExecuteRequest(integration string, target RemoteQue if format == "" { format = "csv" } - if format != "csv" { - return RemoteQueryExecuteRequest{}, fmt.Errorf("format must be csv") + if format != "csv" && format != "binary" { + return RemoteQueryExecuteRequest{}, fmt.Errorf("format must be csv or binary") } var parsedLimits *remoteQueryExecuteCopyLimits if limits != nil { @@ -532,8 +534,8 @@ func (s *RemoteQueryExecuteService) Execute(req RemoteQueryExecuteRequest) Remot return RemoteQueryExecuteResult{HTTPStatus: http.StatusOK, ResponseJSON: responseJSON} } -// ExecuteStream executes a COPY streaming request and emits serialized stream events without materializing the full result. -func (s *RemoteQueryExecuteService) ExecuteStream(req RemoteQueryExecuteRequest, emit func(string) error) RemoteQueryExecuteResult { +// ExecuteStream executes a COPY streaming request and emits binary-safe stream events without materializing the full result. +func (s *RemoteQueryExecuteService) ExecuteStream(req RemoteQueryExecuteRequest, emit func(check.RemoteQueryStreamEvent) error) RemoteQueryExecuteResult { if req.Operation != "copy_stream" { return s.Execute(req) } diff --git a/comp/remotequeries/impl/remote_query_test.go b/comp/remotequeries/impl/remote_query_test.go index ac9a8f84b83f..ac990f13a25c 100644 --- a/comp/remotequeries/impl/remote_query_test.go +++ b/comp/remotequeries/impl/remote_query_test.go @@ -441,14 +441,18 @@ func TestRemoteQueryExecuteHandlerRunnerSuccessWithFixtureTableQuery(t *testing. func TestRemoteQueryExecuteServiceCopyStreamDispatch(t *testing.T) { runner := &fakeStreamRunnerCheck{ fakeRunnerCheck: fakeRunnerCheck{fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-value\n"}}, - events: []string{`{"type":"metadata","status":"STARTED"}`, `{"type":"data","sequence":0,"data":"a,b\n","bytes":4}`, `{"type":"final","status":"SUCCEEDED"}`}, + events: []check.RemoteQueryStreamEvent{ + {Type: "metadata", MetadataJSON: `{"status":"STARTED"}`}, + {Type: "data", MetadataJSON: `{"sequence":0,"offset":0,"bytes":3}`, Payload: []byte{0x00, 0xff, 0x80}}, + {Type: "final", MetadataJSON: `{"status":"SUCCEEDED"}`}, + }, } service := NewRemoteQueryExecuteService(fakeCollector{checks: []check.Check{fakeWrappedCheck{Check: runner}}}, true) req, err := NewRemoteQueryCopyStreamExecuteRequest("postgres", RemoteQueryExecuteTarget{Host: "LOCALHOST.", Port: 5432, DBName: "postgres"}, "SELECT city, country FROM cities ORDER BY city", "csv", &RemoteQueryExecuteCopyLimits{ChunkBytes: 4, MaxBytes: 1024, MaxRowBytes: 1024, TimeoutMs: 1000}) require.NoError(t, err) - var events []string - result := service.ExecuteStream(req, func(event string) error { + var events []check.RemoteQueryStreamEvent + result := service.ExecuteStream(req, func(event check.RemoteQueryStreamEvent) error { events = append(events, event) return nil }) @@ -567,12 +571,12 @@ func (f *fakeRunnerCheck) seenRequest() string { type fakeStreamRunnerCheck struct { fakeRunnerCheck - events []string + events []check.RemoteQueryStreamEvent streamSeen string streamCalls int } -func (f *fakeStreamRunnerCheck) RunRemoteQueryStream(integration string, requestJSON string, emit func(string) error) error { +func (f *fakeStreamRunnerCheck) RunRemoteQueryStream(integration string, requestJSON string, emit func(check.RemoteQueryStreamEvent) error) error { if integration != "postgres" { return assert.AnError } diff --git a/pkg/collector/check/remote_query_stream.go b/pkg/collector/check/remote_query_stream.go new file mode 100644 index 000000000000..b0b8de488a03 --- /dev/null +++ b/pkg/collector/check/remote_query_stream.go @@ -0,0 +1,13 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2026-present Datadog, Inc. + +package check + +// RemoteQueryStreamEvent is the binary-safe event emitted by Remote Queries COPY stream helpers. +type RemoteQueryStreamEvent struct { + Type string + MetadataJSON string + Payload []byte +} diff --git a/pkg/collector/python/check.go b/pkg/collector/python/check.go index 0966d4a794af..1e26582238ec 100644 --- a/pkg/collector/python/check.go +++ b/pkg/collector/python/check.go @@ -34,6 +34,7 @@ import ( ) /* +#include #include #include "datadog_agent_rtloader.h" @@ -41,8 +42,8 @@ import ( char *getStringAddr(char **array, unsigned int idx); char *run_remote_query(rtloader_t *, rtloader_pyobject_t *check, const char *integration, const char *request_json); -extern int remoteQueryStreamEmitBridge(const char *event_json, void *userdata); -int run_remote_query_stream(rtloader_t *, rtloader_pyobject_t *check, const char *integration, const char *request_json, int (*emit)(const char *, void *), void *userdata); +extern int remoteQueryStreamEmitBridge(const char *event_type, const char *metadata_json, const uint8_t *payload, size_t payload_len, void *userdata); +int run_remote_query_stream(rtloader_t *, rtloader_pyobject_t *check, const char *integration, const char *request_json, int (*emit)(const char *, const char *, const uint8_t *, size_t, void *), void *userdata); static inline int call_run_remote_query_stream(rtloader_t *rtloader, rtloader_pyobject_t *check, const char *integration, const char *request_json, void *userdata) { return run_remote_query_stream(rtloader, check, integration, request_json, remoteQueryStreamEmitBridge, userdata); @@ -200,7 +201,7 @@ func (c *PythonCheck) RunRemoteQueryJSON(integration string, requestJSON string) } // RunRemoteQueryStream runs a streaming remote query helper for this Python check. -func (c *PythonCheck) RunRemoteQueryStream(integration string, requestJSON string, emit func(string) error) error { +func (c *PythonCheck) RunRemoteQueryStream(integration string, requestJSON string, emit func(checkbase.RemoteQueryStreamEvent) error) error { integration = strings.ToLower(strings.TrimSpace(integration)) if integration == "" { return fmt.Errorf("integration is required") diff --git a/pkg/collector/python/remote_query_stream.go b/pkg/collector/python/remote_query_stream.go index 5c7d1012ac53..1a4bd0e3a600 100644 --- a/pkg/collector/python/remote_query_stream.go +++ b/pkg/collector/python/remote_query_stream.go @@ -9,19 +9,31 @@ package python /* #include +#include */ import "C" import ( "runtime/cgo" "unsafe" + + checkbase "github.com/DataDog/datadog-agent/pkg/collector/check" ) //export remoteQueryStreamEmitBridge -func remoteQueryStreamEmitBridge(eventJSON *C.char, userdata unsafe.Pointer) C.int { +func remoteQueryStreamEmitBridge(eventType *C.char, metadataJSON *C.char, payload *C.uint8_t, payloadLen C.size_t, userdata unsafe.Pointer) C.int { h := cgo.Handle(uintptr(userdata)) - emit := h.Value().(func(string) error) - if err := emit(C.GoString(eventJSON)); err != nil { + emit := h.Value().(func(checkbase.RemoteQueryStreamEvent) error) + event := checkbase.RemoteQueryStreamEvent{ + Type: C.GoString(eventType), + MetadataJSON: C.GoString(metadataJSON), + } + if payloadLen > 0 { + event.Payload = C.GoBytes(unsafe.Pointer(payload), C.int(payloadLen)) + } else { + event.Payload = []byte{} + } + if err := emit(event); err != nil { return 1 } return 0 diff --git a/pkg/collector/python/test_check.go b/pkg/collector/python/test_check.go index 39162d3d052f..7d03e7b46527 100644 --- a/pkg/collector/python/test_check.go +++ b/pkg/collector/python/test_check.go @@ -22,6 +22,7 @@ import ( "github.com/DataDog/datadog-agent/pkg/aggregator" "github.com/DataDog/datadog-agent/pkg/aggregator/mocksender" "github.com/DataDog/datadog-agent/pkg/aggregator/sender" + checkbase "github.com/DataDog/datadog-agent/pkg/collector/check" checkid "github.com/DataDog/datadog-agent/pkg/collector/check/id" ) @@ -123,14 +124,18 @@ rtloader_pyobject_t *run_remote_query_stream_instance = NULL; const char *run_remote_query_stream_integration = NULL; const char *run_remote_query_stream_request_json = NULL; const char *run_remote_query_stream_event_json = NULL; -int run_remote_query_stream(rtloader_t *s, rtloader_pyobject_t *check, const char *integration, const char *request_json, int (*emit)(const char *, void *), void *userdata) { +int run_remote_query_stream(rtloader_t *s, rtloader_pyobject_t *check, const char *integration, const char *request_json, int (*emit)(const char *, const char *, const uint8_t *, size_t, void *), void *userdata) { run_remote_query_stream_instance = check; run_remote_query_stream_integration = strdup(integration); run_remote_query_stream_request_json = strdup(request_json); run_remote_query_stream_calls++; - run_remote_query_stream_event_json = "{\"type\":\"final\",\"status\":\"SUCCEEDED\"}"; + run_remote_query_stream_event_json = "{\"status\":\"SUCCEEDED\"}"; + uint8_t payload[] = {0x00, 0xff, 0x80}; if (run_remote_query_stream_return && emit != NULL) { - if (emit(run_remote_query_stream_event_json, userdata) != 0) { + if (emit("data", "{\"sequence\":0,\"offset\":0,\"bytes\":3}", payload, sizeof(payload), userdata) != 0) { + return 0; + } + if (emit("final", run_remote_query_stream_event_json, NULL, 0, userdata) != 0) { return 0; } } @@ -813,14 +818,20 @@ func testRunRemoteQueryStream(t *testing.T) { check.instance = newMockPyObjectPtr() C.reset_check_mock() - var events []string - err = check.RunRemoteQueryStream(" Postgres ", `{"operation":"copy_stream"}`, func(event string) error { + var events []checkbase.RemoteQueryStreamEvent + err = check.RunRemoteQueryStream(" Postgres ", `{"operation":"copy_stream"}`, func(event checkbase.RemoteQueryStreamEvent) error { events = append(events, event) return nil }) require.NoError(t, err) - assert.Equal(t, []string{`{"type":"final","status":"SUCCEEDED"}`}, events) + require.Len(t, events, 2) + assert.Equal(t, "data", events[0].Type) + assert.JSONEq(t, `{"sequence":0,"offset":0,"bytes":3}`, events[0].MetadataJSON) + assert.Equal(t, []byte{0x00, 0xff, 0x80}, events[0].Payload) + assert.Equal(t, "final", events[1].Type) + assert.JSONEq(t, `{"status":"SUCCEEDED"}`, events[1].MetadataJSON) + assert.Empty(t, events[1].Payload) assert.Equal(t, C.int(1), C.gil_locked_calls) assert.Equal(t, C.int(1), C.gil_unlocked_calls) assert.Equal(t, C.int(1), C.run_remote_query_stream_calls) @@ -837,7 +848,7 @@ func testRunRemoteQueryStreamEmitError(t *testing.T) { check.instance = newMockPyObjectPtr() C.reset_check_mock() - err = check.RunRemoteQueryStream("postgres", `{"operation":"copy_stream"}`, func(string) error { return assert.AnError }) + err = check.RunRemoteQueryStream("postgres", `{"operation":"copy_stream"}`, func(checkbase.RemoteQueryStreamEvent) error { return assert.AnError }) require.Error(t, err) assert.Equal(t, C.int(1), C.run_remote_query_stream_calls) @@ -853,7 +864,7 @@ func testRunRemoteQueryStreamAfterCancel(t *testing.T) { C.reset_check_mock() check.Cancel() - err = check.RunRemoteQueryStream("postgres", `{"operation":"copy_stream"}`, func(string) error { return nil }) + err = check.RunRemoteQueryStream("postgres", `{"operation":"copy_stream"}`, func(checkbase.RemoteQueryStreamEvent) error { return nil }) assert.EqualError(t, err, "check fake_check is already cancelled") assert.Equal(t, C.int(0), C.run_remote_query_stream_calls) } diff --git a/pkg/privateactionrunner/bundles/remotequeries/execute.go b/pkg/privateactionrunner/bundles/remotequeries/execute.go index d2e45e828caa..95c6eddf3fbc 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/execute.go +++ b/pkg/privateactionrunner/bundles/remotequeries/execute.go @@ -11,6 +11,7 @@ import ( "encoding/json" "fmt" "io" + "unicode/utf8" "google.golang.org/grpc" @@ -137,7 +138,8 @@ func remoteQueryExecuteOutputFromStream(stream grpc.ServerStreamingClient[pb.Rem } var assembled bytes.Buffer - streamEvents := make([]json.RawMessage, 0) + legacyStreamEvents := make([]json.RawMessage, 0) + typedStreamEvents := make([]map[string]interface{}, 0) expectedChunkIndex := int32(0) seenFinal := false finalChunkWasEmpty := false @@ -158,13 +160,21 @@ func remoteQueryExecuteOutputFromStream(stream grpc.ServerStreamingClient[pb.Rem if seenFinal { return nil, fmt.Errorf("remote query response stream sent chunk after final") } - payload := chunk.GetResponseJsonChunk() - if chunk.GetFinal() && len(payload) == 0 && len(streamEvents) > 0 { - finalChunkWasEmpty = true + if event := chunk.GetEvent(); event != nil { + streamEvent, err := remoteQueryStreamEventFromProto(event) + if err != nil { + return nil, err + } + typedStreamEvents = append(typedStreamEvents, streamEvent) } else { - _, _ = assembled.Write(payload) - if !chunk.GetFinal() { - streamEvents = append(streamEvents, append(json.RawMessage(nil), payload...)) + payload := chunk.GetResponseJsonChunk() + if chunk.GetFinal() && len(payload) == 0 && len(legacyStreamEvents) > 0 { + finalChunkWasEmpty = true + } else { + _, _ = assembled.Write(payload) + if !chunk.GetFinal() { + legacyStreamEvents = append(legacyStreamEvents, append(json.RawMessage(nil), payload...)) + } } } seenFinal = chunk.GetFinal() @@ -173,12 +183,89 @@ func remoteQueryExecuteOutputFromStream(stream grpc.ServerStreamingClient[pb.Rem if !seenFinal { return nil, fmt.Errorf("remote query response stream missing final chunk") } + if len(typedStreamEvents) > 0 { + return remoteQueryExecuteOutputFromTypedEvents(typedStreamEvents) + } if finalChunkWasEmpty { - return remoteQueryExecuteOutputFromEvents(streamEvents) + return remoteQueryExecuteOutputFromEvents(legacyStreamEvents) } return remoteQueryExecuteOutputFromJSON(assembled.Bytes()) } +func remoteQueryStreamEventFromProto(event *pb.RemoteQueryExecuteStreamEvent) (map[string]interface{}, error) { + out := map[string]interface{}{"sequence": event.GetSequence()} + switch e := event.GetEvent().(type) { + case *pb.RemoteQueryExecuteStreamEvent_Metadata: + out["type"] = "metadata" + out["operation"] = e.Metadata.GetOperation() + out["integration"] = e.Metadata.GetIntegration() + out["format"] = e.Metadata.GetFormat() + if len(e.Metadata.GetAttributes()) > 0 { + out["attributes"] = e.Metadata.GetAttributes() + } + case *pb.RemoteQueryExecuteStreamEvent_Data: + out["type"] = "data" + payload := append([]byte(nil), e.Data.GetPayload()...) + out["payload"] = payload + out["offset"] = e.Data.GetOffset() + out["bytes"] = e.Data.GetBytes() + if utf8.Valid(payload) { + out["data"] = string(payload) + } + case *pb.RemoteQueryExecuteStreamEvent_Final: + out["type"] = "final" + out["status"] = e.Final.GetStatus() + out["bytes_emitted"] = e.Final.GetBytesEmitted() + out["chunks_emitted"] = e.Final.GetChunksEmitted() + if len(e.Final.GetAttributes()) > 0 { + out["attributes"] = e.Final.GetAttributes() + } + case *pb.RemoteQueryExecuteStreamEvent_Error: + out["type"] = "error" + out["code"] = e.Error.GetCode() + out["message"] = e.Error.GetMessage() + out["retryable"] = e.Error.GetRetryable() + if len(e.Error.GetAttributes()) > 0 { + out["attributes"] = e.Error.GetAttributes() + } + default: + return nil, fmt.Errorf("remote query stream response contained unknown event") + } + return out, nil +} + +func remoteQueryExecuteOutputFromTypedEvents(events []map[string]interface{}) (map[string]interface{}, error) { + var finalEvent map[string]interface{} + var data bytes.Buffer + for _, event := range events { + if event["type"] == "data" { + if payload, ok := event["payload"].([]byte); ok { + _, _ = data.Write(payload) + } + } + if event["type"] == "final" { + finalEvent = event + } + } + if finalEvent == nil { + return nil, fmt.Errorf("remote query stream response missing final event") + } + status, _ := finalEvent["status"].(string) + if status == "" { + return nil, fmt.Errorf("remote query stream final event missing status") + } + dataBytes := data.Bytes() + output := map[string]interface{}{ + "status": status, + "events": normalizeRemoteQueryOutput(events), + "data_bytes": append([]byte(nil), dataBytes...), + } + if utf8.Valid(dataBytes) { + output["data"] = string(dataBytes) + } + return output, nil +} + func remoteQueryExecuteOutputFromEvents(rawEvents []json.RawMessage) (map[string]interface{}, error) { events := make([]interface{}, 0, len(rawEvents)) var finalEvent map[string]interface{} diff --git a/pkg/privateactionrunner/bundles/remotequeries/execute_test.go b/pkg/privateactionrunner/bundles/remotequeries/execute_test.go index dd7f9a8d6700..65202b887d26 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/execute_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/execute_test.go @@ -69,10 +69,10 @@ func TestExecuteActionUsesCredentialFreeAgentSecureRequestShape(t *testing.T) { func TestExecuteActionPreservesCopyStreamEvents(t *testing.T) { client := &captureBridgeClient{chunks: []*pb.RemoteQueryExecuteChunk{ - {ResponseJsonChunk: []byte(`{"type":"metadata","status":"STARTED","operation":"copy_stream"}`), ChunkIndex: 0}, - {ResponseJsonChunk: []byte(`{"type":"data","sequence":0,"data":"Beautiful city of lights,France\n","bytes":32}`), ChunkIndex: 1}, - {ResponseJsonChunk: []byte(`{"type":"data","sequence":1,"data":"New York,USA\n","bytes":13}`), ChunkIndex: 2}, - {ResponseJsonChunk: []byte(`{"type":"final","status":"SUCCEEDED","stats":{"bytesEmitted":45}}`), ChunkIndex: 3}, + {Event: &pb.RemoteQueryExecuteStreamEvent{Sequence: 0, Event: &pb.RemoteQueryExecuteStreamEvent_Metadata{Metadata: &pb.RemoteQueryStreamMetadata{Operation: "copy_stream", Format: "csv"}}}, ChunkIndex: 0}, + {Event: &pb.RemoteQueryExecuteStreamEvent{Sequence: 1, Event: &pb.RemoteQueryExecuteStreamEvent_Data{Data: &pb.RemoteQueryStreamData{Payload: []byte("Beautiful city of lights,France\n"), Offset: 0, Bytes: 32}}}, ChunkIndex: 1}, + {Event: &pb.RemoteQueryExecuteStreamEvent{Sequence: 2, Event: &pb.RemoteQueryExecuteStreamEvent_Data{Data: &pb.RemoteQueryStreamData{Payload: []byte("New York,USA\n"), Offset: 32, Bytes: 13}}}, ChunkIndex: 2}, + {Event: &pb.RemoteQueryExecuteStreamEvent{Sequence: 3, Event: &pb.RemoteQueryExecuteStreamEvent_Final{Final: &pb.RemoteQueryStreamFinal{Status: "SUCCEEDED", BytesEmitted: 45}}}, ChunkIndex: 3}, {ChunkIndex: 4, Final: true}, }} action := NewExecuteAction(func() (BridgeClient, error) { return client, nil }) @@ -91,9 +91,34 @@ func TestExecuteActionPreservesCopyStreamEvents(t *testing.T) { assert.Equal(t, "csv", client.request.GetFormat()) assert.Equal(t, int32(32), client.request.GetCopyLimits().GetChunkBytes()) assert.Equal(t, "Beautiful city of lights,France\nNew York,USA\n", output.(map[string]interface{})["data"]) + assert.Equal(t, []byte("Beautiful city of lights,France\nNew York,USA\n"), output.(map[string]interface{})["data_bytes"]) assert.Equal(t, "SUCCEEDED", output.(map[string]interface{})["status"]) } +func TestExecuteActionPreservesBinaryCopyStreamPayload(t *testing.T) { + client := &captureBridgeClient{chunks: []*pb.RemoteQueryExecuteChunk{ + {Event: &pb.RemoteQueryExecuteStreamEvent{Sequence: 0, Event: &pb.RemoteQueryExecuteStreamEvent_Data{Data: &pb.RemoteQueryStreamData{Payload: []byte{0x00, 0xff, 0x80}, Offset: 0, Bytes: 3}}}, ChunkIndex: 0}, + {Event: &pb.RemoteQueryExecuteStreamEvent{Sequence: 1, Event: &pb.RemoteQueryExecuteStreamEvent_Final{Final: &pb.RemoteQueryStreamFinal{Status: "SUCCEEDED", BytesEmitted: 3, ChunksEmitted: 1}}}, ChunkIndex: 1}, + {ChunkIndex: 2, Final: true}, + }} + action := NewExecuteAction(func() (BridgeClient, error) { return client, nil }) + + output, err := action.Run(context.Background(), taskWithInputs(map[string]interface{}{ + "integration": "postgres", + "operation": "copy_stream", + "format": "binary", + "target": map[string]interface{}{"host": "localhost", "port": 5432, "dbname": "postgres"}, + "query": "SELECT decode('00ff80', 'hex') AS payload", + "copyLimits": map[string]interface{}{"chunkBytes": 32, "maxBytes": 1024, "maxRowBytes": 1024, "timeoutMs": 1000}, + }), nil) + + require.NoError(t, err) + out := output.(map[string]interface{}) + assert.Equal(t, []byte{0x00, 0xff, 0x80}, out["data_bytes"]) + assert.NotContains(t, out, "data") + assert.Equal(t, "SUCCEEDED", out["status"]) +} + func TestExecuteActionPreservesSanitizedBridgeErrorBody(t *testing.T) { client := &captureBridgeClient{ response: &pb.RemoteQueryExecuteResponse{Status: "target_not_found", Error: &pb.RemoteQueryExecuteError{Code: "target_not_found", Message: "no matching integration check found"}}, diff --git a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go index 4276d1706a86..eb9161684996 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go @@ -8,7 +8,9 @@ package com_datadoghq_remotequeries_test import ( + "bytes" "context" + "encoding/base64" "encoding/json" "fmt" "os" @@ -29,10 +31,11 @@ import ( ) const ( - fusedLocalProofEnv = "RQ_FUSED_PROOF" - remoteQueriesSeedProofQuery = "SELECT 1 AS value" - remoteQueriesFixtureTableProofQuery = "SELECT city, country FROM cities ORDER BY city" - remoteQueriesProofQueryOverrideEnv = "RQ_REMOTE_QUERY" + fusedLocalProofEnv = "RQ_FUSED_PROOF" + remoteQueriesSeedProofQuery = "SELECT 1 AS value" + remoteQueriesFixtureTableProofQuery = "SELECT city, country FROM cities ORDER BY city" + remoteQueriesBinaryPayloadProofQuery = "SELECT decode('00ff80', 'hex') AS payload" + remoteQueriesProofQueryOverrideEnv = "RQ_REMOTE_QUERY" ) var remoteQueriesLargePayloadProofQueries = map[string]int{ @@ -202,6 +205,21 @@ func remoteQueriesLargePayloadBytes(query string) (int, bool) { return payloadBytes, ok } +func assertRemoteQueriesProofBinaryCopyData(t *testing.T, query string, dataBase64 string) { + t.Helper() + + data, err := base64.StdEncoding.DecodeString(dataBase64) + require.NoError(t, err) + switch query { + case remoteQueriesBinaryPayloadProofQuery: + assert.True(t, bytes.HasPrefix(data, []byte("PGCOPY\n\xff\r\n\x00")), "binary COPY payload should keep the PostgreSQL binary header") + assert.Contains(t, data, byte(0x00)) + assert.True(t, bytes.Contains(data, []byte{0x00, 0xff, 0x80}), "binary COPY payload should contain the row bytes from decode('00ff80','hex')") + default: + require.FailNowf(t, "unsupported binary COPY proof query", "%s=%q must use a binary COPY bridge-allowlisted proof query", remoteQueriesProofQueryOverrideEnv, query) + } +} + func assertRemoteQueriesProofCopyData(t *testing.T, query string, data string) { t.Helper() diff --git a/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go b/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go index 2c57db6c4a07..1b02ea95f1a6 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go @@ -87,7 +87,11 @@ func TestRemoteQueriesActionRunsThroughStandalonePARProcessWithRealAgentIPC(t *t copyStream := os.Getenv("RQ_REMOTE_OPERATION") == "copy_stream" if copyStream { inputs["operation"] = "copy_stream" - inputs["format"] = "csv" + format := os.Getenv("RQ_REMOTE_FORMAT") + if format == "" { + format = "csv" + } + inputs["format"] = format inputs["copyLimits"] = remoteQueriesProofCopyLimits(proofQuery) } else { inputs["limits"] = remoteQueriesProofLimits(proofQuery) @@ -118,9 +122,15 @@ func TestRemoteQueriesActionRunsThroughStandalonePARProcessWithRealAgentIPC(t *t assert.Equal(t, "SUCCEEDED", result.Outputs["status"]) if copyStream { t.Logf("copy stream PAR outputs: %+v", summarizeRemoteQueriesProofPayload(result.Outputs)) - data, ok := result.Outputs["data"].(string) - require.True(t, ok) - assertRemoteQueriesProofCopyData(t, proofQuery, data) + if os.Getenv("RQ_REMOTE_FORMAT") == "binary" { + dataBytes, ok := result.Outputs["data_bytes"].(string) + require.True(t, ok) + assertRemoteQueriesProofBinaryCopyData(t, proofQuery, dataBytes) + } else { + data, ok := result.Outputs["data"].(string) + require.True(t, ok) + assertRemoteQueriesProofCopyData(t, proofQuery, data) + } } else { require.Contains(t, result.Outputs, "rows") rows, ok := result.Outputs["rows"].([]interface{}) diff --git a/pkg/proto/datadog/api/v1/api.proto b/pkg/proto/datadog/api/v1/api.proto index 5973d5c9b5be..8d4086825d8e 100644 --- a/pkg/proto/datadog/api/v1/api.proto +++ b/pkg/proto/datadog/api/v1/api.proto @@ -64,10 +64,48 @@ message RemoteQueryExecuteResponse { google.protobuf.Struct stats = 6; } +message RemoteQueryStreamMetadata { + string operation = 1; + string integration = 2; + string format = 3; + map attributes = 4; +} + +message RemoteQueryStreamData { + bytes payload = 1; + uint64 offset = 2; + uint64 bytes = 3; +} + +message RemoteQueryStreamFinal { + string status = 1; + uint64 bytes_emitted = 2; + uint64 chunks_emitted = 3; + map attributes = 4; +} + +message RemoteQueryStreamError { + string code = 1; + string message = 2; + bool retryable = 3; + map attributes = 4; +} + +message RemoteQueryExecuteStreamEvent { + uint64 sequence = 1; + oneof event { + RemoteQueryStreamMetadata metadata = 2; + RemoteQueryStreamData data = 3; + RemoteQueryStreamFinal final = 4; + RemoteQueryStreamError error = 5; + } +} + message RemoteQueryExecuteChunk { bytes response_json_chunk = 1; int32 chunk_index = 2; bool final = 3; + RemoteQueryExecuteStreamEvent event = 4; } service AgentSecure { diff --git a/pkg/proto/pbgo/core/api.pb.go b/pkg/proto/pbgo/core/api.pb.go index b65269ec75ce..b14c28de1a31 100644 --- a/pkg/proto/pbgo/core/api.pb.go +++ b/pkg/proto/pbgo/core/api.pb.go @@ -439,18 +439,405 @@ func (x *RemoteQueryExecuteResponse) GetStats() *structpb.Struct { return nil } +type RemoteQueryStreamMetadata struct { + state protoimpl.MessageState `protogen:"open.v1"` + Operation string `protobuf:"bytes,1,opt,name=operation,proto3" json:"operation,omitempty"` + Integration string `protobuf:"bytes,2,opt,name=integration,proto3" json:"integration,omitempty"` + Format string `protobuf:"bytes,3,opt,name=format,proto3" json:"format,omitempty"` + Attributes map[string]string `protobuf:"bytes,4,rep,name=attributes,proto3" json:"attributes,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RemoteQueryStreamMetadata) Reset() { + *x = RemoteQueryStreamMetadata{} + mi := &file_datadog_api_v1_api_proto_msgTypes[6] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RemoteQueryStreamMetadata) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RemoteQueryStreamMetadata) ProtoMessage() {} + +func (x *RemoteQueryStreamMetadata) ProtoReflect() protoreflect.Message { + mi := &file_datadog_api_v1_api_proto_msgTypes[6] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RemoteQueryStreamMetadata.ProtoReflect.Descriptor instead. +func (*RemoteQueryStreamMetadata) Descriptor() ([]byte, []int) { + return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{6} +} + +func (x *RemoteQueryStreamMetadata) GetOperation() string { + if x != nil { + return x.Operation + } + return "" +} + +func (x *RemoteQueryStreamMetadata) GetIntegration() string { + if x != nil { + return x.Integration + } + return "" +} + +func (x *RemoteQueryStreamMetadata) GetFormat() string { + if x != nil { + return x.Format + } + return "" +} + +func (x *RemoteQueryStreamMetadata) GetAttributes() map[string]string { + if x != nil { + return x.Attributes + } + return nil +} + +type RemoteQueryStreamData struct { + state protoimpl.MessageState `protogen:"open.v1"` + Payload []byte `protobuf:"bytes,1,opt,name=payload,proto3" json:"payload,omitempty"` + Offset uint64 `protobuf:"varint,2,opt,name=offset,proto3" json:"offset,omitempty"` + Bytes uint64 `protobuf:"varint,3,opt,name=bytes,proto3" json:"bytes,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RemoteQueryStreamData) Reset() { + *x = RemoteQueryStreamData{} + mi := &file_datadog_api_v1_api_proto_msgTypes[7] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RemoteQueryStreamData) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RemoteQueryStreamData) ProtoMessage() {} + +func (x *RemoteQueryStreamData) ProtoReflect() protoreflect.Message { + mi := &file_datadog_api_v1_api_proto_msgTypes[7] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RemoteQueryStreamData.ProtoReflect.Descriptor instead. +func (*RemoteQueryStreamData) Descriptor() ([]byte, []int) { + return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{7} +} + +func (x *RemoteQueryStreamData) GetPayload() []byte { + if x != nil { + return x.Payload + } + return nil +} + +func (x *RemoteQueryStreamData) GetOffset() uint64 { + if x != nil { + return x.Offset + } + return 0 +} + +func (x *RemoteQueryStreamData) GetBytes() uint64 { + if x != nil { + return x.Bytes + } + return 0 +} + +type RemoteQueryStreamFinal struct { + state protoimpl.MessageState `protogen:"open.v1"` + Status string `protobuf:"bytes,1,opt,name=status,proto3" json:"status,omitempty"` + BytesEmitted uint64 `protobuf:"varint,2,opt,name=bytes_emitted,json=bytesEmitted,proto3" json:"bytes_emitted,omitempty"` + ChunksEmitted uint64 `protobuf:"varint,3,opt,name=chunks_emitted,json=chunksEmitted,proto3" json:"chunks_emitted,omitempty"` + Attributes map[string]string `protobuf:"bytes,4,rep,name=attributes,proto3" json:"attributes,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RemoteQueryStreamFinal) Reset() { + *x = RemoteQueryStreamFinal{} + mi := &file_datadog_api_v1_api_proto_msgTypes[8] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RemoteQueryStreamFinal) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RemoteQueryStreamFinal) ProtoMessage() {} + +func (x *RemoteQueryStreamFinal) ProtoReflect() protoreflect.Message { + mi := &file_datadog_api_v1_api_proto_msgTypes[8] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RemoteQueryStreamFinal.ProtoReflect.Descriptor instead. +func (*RemoteQueryStreamFinal) Descriptor() ([]byte, []int) { + return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{8} +} + +func (x *RemoteQueryStreamFinal) GetStatus() string { + if x != nil { + return x.Status + } + return "" +} + +func (x *RemoteQueryStreamFinal) GetBytesEmitted() uint64 { + if x != nil { + return x.BytesEmitted + } + return 0 +} + +func (x *RemoteQueryStreamFinal) GetChunksEmitted() uint64 { + if x != nil { + return x.ChunksEmitted + } + return 0 +} + +func (x *RemoteQueryStreamFinal) GetAttributes() map[string]string { + if x != nil { + return x.Attributes + } + return nil +} + +type RemoteQueryStreamError struct { + state protoimpl.MessageState `protogen:"open.v1"` + Code string `protobuf:"bytes,1,opt,name=code,proto3" json:"code,omitempty"` + Message string `protobuf:"bytes,2,opt,name=message,proto3" json:"message,omitempty"` + Retryable bool `protobuf:"varint,3,opt,name=retryable,proto3" json:"retryable,omitempty"` + Attributes map[string]string `protobuf:"bytes,4,rep,name=attributes,proto3" json:"attributes,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RemoteQueryStreamError) Reset() { + *x = RemoteQueryStreamError{} + mi := &file_datadog_api_v1_api_proto_msgTypes[9] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RemoteQueryStreamError) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RemoteQueryStreamError) ProtoMessage() {} + +func (x *RemoteQueryStreamError) ProtoReflect() protoreflect.Message { + mi := &file_datadog_api_v1_api_proto_msgTypes[9] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RemoteQueryStreamError.ProtoReflect.Descriptor instead. +func (*RemoteQueryStreamError) Descriptor() ([]byte, []int) { + return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{9} +} + +func (x *RemoteQueryStreamError) GetCode() string { + if x != nil { + return x.Code + } + return "" +} + +func (x *RemoteQueryStreamError) GetMessage() string { + if x != nil { + return x.Message + } + return "" +} + +func (x *RemoteQueryStreamError) GetRetryable() bool { + if x != nil { + return x.Retryable + } + return false +} + +func (x *RemoteQueryStreamError) GetAttributes() map[string]string { + if x != nil { + return x.Attributes + } + return nil +} + +type RemoteQueryExecuteStreamEvent struct { + state protoimpl.MessageState `protogen:"open.v1"` + Sequence uint64 `protobuf:"varint,1,opt,name=sequence,proto3" json:"sequence,omitempty"` + // Types that are valid to be assigned to Event: + // + // *RemoteQueryExecuteStreamEvent_Metadata + // *RemoteQueryExecuteStreamEvent_Data + // *RemoteQueryExecuteStreamEvent_Final + // *RemoteQueryExecuteStreamEvent_Error + Event isRemoteQueryExecuteStreamEvent_Event `protobuf_oneof:"event"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RemoteQueryExecuteStreamEvent) Reset() { + *x = RemoteQueryExecuteStreamEvent{} + mi := &file_datadog_api_v1_api_proto_msgTypes[10] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RemoteQueryExecuteStreamEvent) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RemoteQueryExecuteStreamEvent) ProtoMessage() {} + +func (x *RemoteQueryExecuteStreamEvent) ProtoReflect() protoreflect.Message { + mi := &file_datadog_api_v1_api_proto_msgTypes[10] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RemoteQueryExecuteStreamEvent.ProtoReflect.Descriptor instead. +func (*RemoteQueryExecuteStreamEvent) Descriptor() ([]byte, []int) { + return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{10} +} + +func (x *RemoteQueryExecuteStreamEvent) GetSequence() uint64 { + if x != nil { + return x.Sequence + } + return 0 +} + +func (x *RemoteQueryExecuteStreamEvent) GetEvent() isRemoteQueryExecuteStreamEvent_Event { + if x != nil { + return x.Event + } + return nil +} + +func (x *RemoteQueryExecuteStreamEvent) GetMetadata() *RemoteQueryStreamMetadata { + if x != nil { + if x, ok := x.Event.(*RemoteQueryExecuteStreamEvent_Metadata); ok { + return x.Metadata + } + } + return nil +} + +func (x *RemoteQueryExecuteStreamEvent) GetData() *RemoteQueryStreamData { + if x != nil { + if x, ok := x.Event.(*RemoteQueryExecuteStreamEvent_Data); ok { + return x.Data + } + } + return nil +} + +func (x *RemoteQueryExecuteStreamEvent) GetFinal() *RemoteQueryStreamFinal { + if x != nil { + if x, ok := x.Event.(*RemoteQueryExecuteStreamEvent_Final); ok { + return x.Final + } + } + return nil +} + +func (x *RemoteQueryExecuteStreamEvent) GetError() *RemoteQueryStreamError { + if x != nil { + if x, ok := x.Event.(*RemoteQueryExecuteStreamEvent_Error); ok { + return x.Error + } + } + return nil +} + +type isRemoteQueryExecuteStreamEvent_Event interface { + isRemoteQueryExecuteStreamEvent_Event() +} + +type RemoteQueryExecuteStreamEvent_Metadata struct { + Metadata *RemoteQueryStreamMetadata `protobuf:"bytes,2,opt,name=metadata,proto3,oneof"` +} + +type RemoteQueryExecuteStreamEvent_Data struct { + Data *RemoteQueryStreamData `protobuf:"bytes,3,opt,name=data,proto3,oneof"` +} + +type RemoteQueryExecuteStreamEvent_Final struct { + Final *RemoteQueryStreamFinal `protobuf:"bytes,4,opt,name=final,proto3,oneof"` +} + +type RemoteQueryExecuteStreamEvent_Error struct { + Error *RemoteQueryStreamError `protobuf:"bytes,5,opt,name=error,proto3,oneof"` +} + +func (*RemoteQueryExecuteStreamEvent_Metadata) isRemoteQueryExecuteStreamEvent_Event() {} + +func (*RemoteQueryExecuteStreamEvent_Data) isRemoteQueryExecuteStreamEvent_Event() {} + +func (*RemoteQueryExecuteStreamEvent_Final) isRemoteQueryExecuteStreamEvent_Event() {} + +func (*RemoteQueryExecuteStreamEvent_Error) isRemoteQueryExecuteStreamEvent_Event() {} + type RemoteQueryExecuteChunk struct { - state protoimpl.MessageState `protogen:"open.v1"` - ResponseJsonChunk []byte `protobuf:"bytes,1,opt,name=response_json_chunk,json=responseJsonChunk,proto3" json:"response_json_chunk,omitempty"` - ChunkIndex int32 `protobuf:"varint,2,opt,name=chunk_index,json=chunkIndex,proto3" json:"chunk_index,omitempty"` - Final bool `protobuf:"varint,3,opt,name=final,proto3" json:"final,omitempty"` + state protoimpl.MessageState `protogen:"open.v1"` + ResponseJsonChunk []byte `protobuf:"bytes,1,opt,name=response_json_chunk,json=responseJsonChunk,proto3" json:"response_json_chunk,omitempty"` + ChunkIndex int32 `protobuf:"varint,2,opt,name=chunk_index,json=chunkIndex,proto3" json:"chunk_index,omitempty"` + Final bool `protobuf:"varint,3,opt,name=final,proto3" json:"final,omitempty"` + Event *RemoteQueryExecuteStreamEvent `protobuf:"bytes,4,opt,name=event,proto3" json:"event,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } func (x *RemoteQueryExecuteChunk) Reset() { *x = RemoteQueryExecuteChunk{} - mi := &file_datadog_api_v1_api_proto_msgTypes[6] + mi := &file_datadog_api_v1_api_proto_msgTypes[11] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -462,7 +849,7 @@ func (x *RemoteQueryExecuteChunk) String() string { func (*RemoteQueryExecuteChunk) ProtoMessage() {} func (x *RemoteQueryExecuteChunk) ProtoReflect() protoreflect.Message { - mi := &file_datadog_api_v1_api_proto_msgTypes[6] + mi := &file_datadog_api_v1_api_proto_msgTypes[11] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -475,7 +862,7 @@ func (x *RemoteQueryExecuteChunk) ProtoReflect() protoreflect.Message { // Deprecated: Use RemoteQueryExecuteChunk.ProtoReflect.Descriptor instead. func (*RemoteQueryExecuteChunk) Descriptor() ([]byte, []int) { - return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{6} + return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{11} } func (x *RemoteQueryExecuteChunk) GetResponseJsonChunk() []byte { @@ -499,6 +886,13 @@ func (x *RemoteQueryExecuteChunk) GetFinal() bool { return false } +func (x *RemoteQueryExecuteChunk) GetEvent() *RemoteQueryExecuteStreamEvent { + if x != nil { + return x.Event + } + return nil +} + var File_datadog_api_v1_api_proto protoreflect.FileDescriptor const file_datadog_api_v1_api_proto_rawDesc = "" + @@ -538,12 +932,54 @@ const file_datadog_api_v1_api_proto_rawDesc = "" + "\acolumns\x18\x03 \x03(\v2\x17.google.protobuf.StructR\acolumns\x12+\n" + "\x04rows\x18\x04 \x03(\v2\x17.google.protobuf.StructR\x04rows\x12\x1c\n" + "\ttruncated\x18\x05 \x01(\bR\ttruncated\x12-\n" + - "\x05stats\x18\x06 \x01(\v2\x17.google.protobuf.StructR\x05stats\"\x80\x01\n" + + "\x05stats\x18\x06 \x01(\v2\x17.google.protobuf.StructR\x05stats\"\x8d\x02\n" + + "\x19RemoteQueryStreamMetadata\x12\x1c\n" + + "\toperation\x18\x01 \x01(\tR\toperation\x12 \n" + + "\vintegration\x18\x02 \x01(\tR\vintegration\x12\x16\n" + + "\x06format\x18\x03 \x01(\tR\x06format\x12Y\n" + + "\n" + + "attributes\x18\x04 \x03(\v29.datadog.api.v1.RemoteQueryStreamMetadata.AttributesEntryR\n" + + "attributes\x1a=\n" + + "\x0fAttributesEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"_\n" + + "\x15RemoteQueryStreamData\x12\x18\n" + + "\apayload\x18\x01 \x01(\fR\apayload\x12\x16\n" + + "\x06offset\x18\x02 \x01(\x04R\x06offset\x12\x14\n" + + "\x05bytes\x18\x03 \x01(\x04R\x05bytes\"\x93\x02\n" + + "\x16RemoteQueryStreamFinal\x12\x16\n" + + "\x06status\x18\x01 \x01(\tR\x06status\x12#\n" + + "\rbytes_emitted\x18\x02 \x01(\x04R\fbytesEmitted\x12%\n" + + "\x0echunks_emitted\x18\x03 \x01(\x04R\rchunksEmitted\x12V\n" + + "\n" + + "attributes\x18\x04 \x03(\v26.datadog.api.v1.RemoteQueryStreamFinal.AttributesEntryR\n" + + "attributes\x1a=\n" + + "\x0fAttributesEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"\xfb\x01\n" + + "\x16RemoteQueryStreamError\x12\x12\n" + + "\x04code\x18\x01 \x01(\tR\x04code\x12\x18\n" + + "\amessage\x18\x02 \x01(\tR\amessage\x12\x1c\n" + + "\tretryable\x18\x03 \x01(\bR\tretryable\x12V\n" + + "\n" + + "attributes\x18\x04 \x03(\v26.datadog.api.v1.RemoteQueryStreamError.AttributesEntryR\n" + + "attributes\x1a=\n" + + "\x0fAttributesEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"\xca\x02\n" + + "\x1dRemoteQueryExecuteStreamEvent\x12\x1a\n" + + "\bsequence\x18\x01 \x01(\x04R\bsequence\x12G\n" + + "\bmetadata\x18\x02 \x01(\v2).datadog.api.v1.RemoteQueryStreamMetadataH\x00R\bmetadata\x12;\n" + + "\x04data\x18\x03 \x01(\v2%.datadog.api.v1.RemoteQueryStreamDataH\x00R\x04data\x12>\n" + + "\x05final\x18\x04 \x01(\v2&.datadog.api.v1.RemoteQueryStreamFinalH\x00R\x05final\x12>\n" + + "\x05error\x18\x05 \x01(\v2&.datadog.api.v1.RemoteQueryStreamErrorH\x00R\x05errorB\a\n" + + "\x05event\"\xc5\x01\n" + "\x17RemoteQueryExecuteChunk\x12.\n" + "\x13response_json_chunk\x18\x01 \x01(\fR\x11responseJsonChunk\x12\x1f\n" + "\vchunk_index\x18\x02 \x01(\x05R\n" + "chunkIndex\x12\x14\n" + - "\x05final\x18\x03 \x01(\bR\x05final2Z\n" + + "\x05final\x18\x03 \x01(\bR\x05final\x12C\n" + + "\x05event\x18\x04 \x01(\v2-.datadog.api.v1.RemoteQueryExecuteStreamEventR\x05event2Z\n" + "\x05Agent\x12Q\n" + "\vGetHostname\x12!.datadog.model.v1.HostnameRequest\x1a\x1f.datadog.model.v1.HostnameReply2\x8a\x12\n" + "\vAgentSecure\x12c\n" + @@ -581,7 +1017,7 @@ func file_datadog_api_v1_api_proto_rawDescGZIP() []byte { return file_datadog_api_v1_api_proto_rawDescData } -var file_datadog_api_v1_api_proto_msgTypes = make([]protoimpl.MessageInfo, 7) +var file_datadog_api_v1_api_proto_msgTypes = make([]protoimpl.MessageInfo, 15) var file_datadog_api_v1_api_proto_goTypes = []any{ (*RemoteQueryTarget)(nil), // 0: datadog.api.v1.RemoteQueryTarget (*RemoteQueryExecuteLimits)(nil), // 1: datadog.api.v1.RemoteQueryExecuteLimits @@ -589,100 +1025,116 @@ var file_datadog_api_v1_api_proto_goTypes = []any{ (*RemoteQueryExecuteRequest)(nil), // 3: datadog.api.v1.RemoteQueryExecuteRequest (*RemoteQueryExecuteError)(nil), // 4: datadog.api.v1.RemoteQueryExecuteError (*RemoteQueryExecuteResponse)(nil), // 5: datadog.api.v1.RemoteQueryExecuteResponse - (*RemoteQueryExecuteChunk)(nil), // 6: datadog.api.v1.RemoteQueryExecuteChunk - (*structpb.Struct)(nil), // 7: google.protobuf.Struct - (*HostnameRequest)(nil), // 8: datadog.model.v1.HostnameRequest - (*StreamTagsRequest)(nil), // 9: datadog.model.v1.StreamTagsRequest - (*GenerateContainerIDFromOriginInfoRequest)(nil), // 10: datadog.model.v1.GenerateContainerIDFromOriginInfoRequest - (*FetchEntityRequest)(nil), // 11: datadog.model.v1.FetchEntityRequest - (*CaptureTriggerRequest)(nil), // 12: datadog.model.v1.CaptureTriggerRequest - (*TaggerState)(nil), // 13: datadog.model.v1.TaggerState - (*ClientGetConfigsRequest)(nil), // 14: datadog.config.ClientGetConfigsRequest - (*emptypb.Empty)(nil), // 15: google.protobuf.Empty - (*ConfigSubscriptionRequest)(nil), // 16: datadog.config.ConfigSubscriptionRequest - (*WorkloadmetaStreamRequest)(nil), // 17: datadog.workloadmeta.WorkloadmetaStreamRequest - (*RegisterRemoteAgentRequest)(nil), // 18: datadog.remoteagent.v1.RegisterRemoteAgentRequest - (*RefreshRemoteAgentRequest)(nil), // 19: datadog.remoteagent.v1.RefreshRemoteAgentRequest - (*HostTagRequest)(nil), // 20: datadog.model.v1.HostTagRequest - (*ConfigStreamRequest)(nil), // 21: datadog.model.v1.ConfigStreamRequest - (*WorkloadFilterEvaluateRequest)(nil), // 22: datadog.workloadfilter.WorkloadFilterEvaluateRequest - (*KubeMetadataStreamRequest)(nil), // 23: datadog.kubemetadata.KubeMetadataStreamRequest - (*HostnameReply)(nil), // 24: datadog.model.v1.HostnameReply - (*StreamTagsResponse)(nil), // 25: datadog.model.v1.StreamTagsResponse - (*GenerateContainerIDFromOriginInfoResponse)(nil), // 26: datadog.model.v1.GenerateContainerIDFromOriginInfoResponse - (*FetchEntityResponse)(nil), // 27: datadog.model.v1.FetchEntityResponse - (*CaptureTriggerResponse)(nil), // 28: datadog.model.v1.CaptureTriggerResponse - (*TaggerStateResponse)(nil), // 29: datadog.model.v1.TaggerStateResponse - (*ClientGetConfigsResponse)(nil), // 30: datadog.config.ClientGetConfigsResponse - (*GetStateConfigResponse)(nil), // 31: datadog.config.GetStateConfigResponse - (*ConfigSubscriptionResponse)(nil), // 32: datadog.config.ConfigSubscriptionResponse - (*ResetStateConfigResponse)(nil), // 33: datadog.config.ResetStateConfigResponse - (*WorkloadmetaStreamResponse)(nil), // 34: datadog.workloadmeta.WorkloadmetaStreamResponse - (*RegisterRemoteAgentResponse)(nil), // 35: datadog.remoteagent.v1.RegisterRemoteAgentResponse - (*RefreshRemoteAgentResponse)(nil), // 36: datadog.remoteagent.v1.RefreshRemoteAgentResponse - (*AutodiscoveryStreamResponse)(nil), // 37: datadog.autodiscovery.AutodiscoveryStreamResponse - (*HostTagReply)(nil), // 38: datadog.model.v1.HostTagReply - (*ConfigEvent)(nil), // 39: datadog.model.v1.ConfigEvent - (*WorkloadFilterEvaluateResponse)(nil), // 40: datadog.workloadfilter.WorkloadFilterEvaluateResponse - (*KubeMetadataStreamResponse)(nil), // 41: datadog.kubemetadata.KubeMetadataStreamResponse + (*RemoteQueryStreamMetadata)(nil), // 6: datadog.api.v1.RemoteQueryStreamMetadata + (*RemoteQueryStreamData)(nil), // 7: datadog.api.v1.RemoteQueryStreamData + (*RemoteQueryStreamFinal)(nil), // 8: datadog.api.v1.RemoteQueryStreamFinal + (*RemoteQueryStreamError)(nil), // 9: datadog.api.v1.RemoteQueryStreamError + (*RemoteQueryExecuteStreamEvent)(nil), // 10: datadog.api.v1.RemoteQueryExecuteStreamEvent + (*RemoteQueryExecuteChunk)(nil), // 11: datadog.api.v1.RemoteQueryExecuteChunk + nil, // 12: datadog.api.v1.RemoteQueryStreamMetadata.AttributesEntry + nil, // 13: datadog.api.v1.RemoteQueryStreamFinal.AttributesEntry + nil, // 14: datadog.api.v1.RemoteQueryStreamError.AttributesEntry + (*structpb.Struct)(nil), // 15: google.protobuf.Struct + (*HostnameRequest)(nil), // 16: datadog.model.v1.HostnameRequest + (*StreamTagsRequest)(nil), // 17: datadog.model.v1.StreamTagsRequest + (*GenerateContainerIDFromOriginInfoRequest)(nil), // 18: datadog.model.v1.GenerateContainerIDFromOriginInfoRequest + (*FetchEntityRequest)(nil), // 19: datadog.model.v1.FetchEntityRequest + (*CaptureTriggerRequest)(nil), // 20: datadog.model.v1.CaptureTriggerRequest + (*TaggerState)(nil), // 21: datadog.model.v1.TaggerState + (*ClientGetConfigsRequest)(nil), // 22: datadog.config.ClientGetConfigsRequest + (*emptypb.Empty)(nil), // 23: google.protobuf.Empty + (*ConfigSubscriptionRequest)(nil), // 24: datadog.config.ConfigSubscriptionRequest + (*WorkloadmetaStreamRequest)(nil), // 25: datadog.workloadmeta.WorkloadmetaStreamRequest + (*RegisterRemoteAgentRequest)(nil), // 26: datadog.remoteagent.v1.RegisterRemoteAgentRequest + (*RefreshRemoteAgentRequest)(nil), // 27: datadog.remoteagent.v1.RefreshRemoteAgentRequest + (*HostTagRequest)(nil), // 28: datadog.model.v1.HostTagRequest + (*ConfigStreamRequest)(nil), // 29: datadog.model.v1.ConfigStreamRequest + (*WorkloadFilterEvaluateRequest)(nil), // 30: datadog.workloadfilter.WorkloadFilterEvaluateRequest + (*KubeMetadataStreamRequest)(nil), // 31: datadog.kubemetadata.KubeMetadataStreamRequest + (*HostnameReply)(nil), // 32: datadog.model.v1.HostnameReply + (*StreamTagsResponse)(nil), // 33: datadog.model.v1.StreamTagsResponse + (*GenerateContainerIDFromOriginInfoResponse)(nil), // 34: datadog.model.v1.GenerateContainerIDFromOriginInfoResponse + (*FetchEntityResponse)(nil), // 35: datadog.model.v1.FetchEntityResponse + (*CaptureTriggerResponse)(nil), // 36: datadog.model.v1.CaptureTriggerResponse + (*TaggerStateResponse)(nil), // 37: datadog.model.v1.TaggerStateResponse + (*ClientGetConfigsResponse)(nil), // 38: datadog.config.ClientGetConfigsResponse + (*GetStateConfigResponse)(nil), // 39: datadog.config.GetStateConfigResponse + (*ConfigSubscriptionResponse)(nil), // 40: datadog.config.ConfigSubscriptionResponse + (*ResetStateConfigResponse)(nil), // 41: datadog.config.ResetStateConfigResponse + (*WorkloadmetaStreamResponse)(nil), // 42: datadog.workloadmeta.WorkloadmetaStreamResponse + (*RegisterRemoteAgentResponse)(nil), // 43: datadog.remoteagent.v1.RegisterRemoteAgentResponse + (*RefreshRemoteAgentResponse)(nil), // 44: datadog.remoteagent.v1.RefreshRemoteAgentResponse + (*AutodiscoveryStreamResponse)(nil), // 45: datadog.autodiscovery.AutodiscoveryStreamResponse + (*HostTagReply)(nil), // 46: datadog.model.v1.HostTagReply + (*ConfigEvent)(nil), // 47: datadog.model.v1.ConfigEvent + (*WorkloadFilterEvaluateResponse)(nil), // 48: datadog.workloadfilter.WorkloadFilterEvaluateResponse + (*KubeMetadataStreamResponse)(nil), // 49: datadog.kubemetadata.KubeMetadataStreamResponse } var file_datadog_api_v1_api_proto_depIdxs = []int32{ 0, // 0: datadog.api.v1.RemoteQueryExecuteRequest.target:type_name -> datadog.api.v1.RemoteQueryTarget 1, // 1: datadog.api.v1.RemoteQueryExecuteRequest.limits:type_name -> datadog.api.v1.RemoteQueryExecuteLimits 2, // 2: datadog.api.v1.RemoteQueryExecuteRequest.copy_limits:type_name -> datadog.api.v1.RemoteQueryExecuteCopyLimits 4, // 3: datadog.api.v1.RemoteQueryExecuteResponse.error:type_name -> datadog.api.v1.RemoteQueryExecuteError - 7, // 4: datadog.api.v1.RemoteQueryExecuteResponse.columns:type_name -> google.protobuf.Struct - 7, // 5: datadog.api.v1.RemoteQueryExecuteResponse.rows:type_name -> google.protobuf.Struct - 7, // 6: datadog.api.v1.RemoteQueryExecuteResponse.stats:type_name -> google.protobuf.Struct - 8, // 7: datadog.api.v1.Agent.GetHostname:input_type -> datadog.model.v1.HostnameRequest - 9, // 8: datadog.api.v1.AgentSecure.TaggerStreamEntities:input_type -> datadog.model.v1.StreamTagsRequest - 10, // 9: datadog.api.v1.AgentSecure.TaggerGenerateContainerIDFromOriginInfo:input_type -> datadog.model.v1.GenerateContainerIDFromOriginInfoRequest - 11, // 10: datadog.api.v1.AgentSecure.TaggerFetchEntity:input_type -> datadog.model.v1.FetchEntityRequest - 12, // 11: datadog.api.v1.AgentSecure.DogstatsdCaptureTrigger:input_type -> datadog.model.v1.CaptureTriggerRequest - 13, // 12: datadog.api.v1.AgentSecure.DogstatsdSetTaggerState:input_type -> datadog.model.v1.TaggerState - 14, // 13: datadog.api.v1.AgentSecure.ClientGetConfigs:input_type -> datadog.config.ClientGetConfigsRequest - 15, // 14: datadog.api.v1.AgentSecure.GetConfigState:input_type -> google.protobuf.Empty - 14, // 15: datadog.api.v1.AgentSecure.ClientGetConfigsHA:input_type -> datadog.config.ClientGetConfigsRequest - 15, // 16: datadog.api.v1.AgentSecure.GetConfigStateHA:input_type -> google.protobuf.Empty - 16, // 17: datadog.api.v1.AgentSecure.CreateConfigSubscription:input_type -> datadog.config.ConfigSubscriptionRequest - 15, // 18: datadog.api.v1.AgentSecure.ResetConfigState:input_type -> google.protobuf.Empty - 17, // 19: datadog.api.v1.AgentSecure.WorkloadmetaStreamEntities:input_type -> datadog.workloadmeta.WorkloadmetaStreamRequest - 18, // 20: datadog.api.v1.AgentSecure.RegisterRemoteAgent:input_type -> datadog.remoteagent.v1.RegisterRemoteAgentRequest - 19, // 21: datadog.api.v1.AgentSecure.RefreshRemoteAgent:input_type -> datadog.remoteagent.v1.RefreshRemoteAgentRequest - 15, // 22: datadog.api.v1.AgentSecure.AutodiscoveryStreamConfig:input_type -> google.protobuf.Empty - 20, // 23: datadog.api.v1.AgentSecure.GetHostTags:input_type -> datadog.model.v1.HostTagRequest - 21, // 24: datadog.api.v1.AgentSecure.StreamConfigEvents:input_type -> datadog.model.v1.ConfigStreamRequest - 22, // 25: datadog.api.v1.AgentSecure.WorkloadFilterEvaluate:input_type -> datadog.workloadfilter.WorkloadFilterEvaluateRequest - 3, // 26: datadog.api.v1.AgentSecure.RemoteQueryExecute:input_type -> datadog.api.v1.RemoteQueryExecuteRequest - 3, // 27: datadog.api.v1.AgentSecure.RemoteQueryExecuteStream:input_type -> datadog.api.v1.RemoteQueryExecuteRequest - 23, // 28: datadog.api.v1.AgentSecure.StreamKubeMetadata:input_type -> datadog.kubemetadata.KubeMetadataStreamRequest - 24, // 29: datadog.api.v1.Agent.GetHostname:output_type -> datadog.model.v1.HostnameReply - 25, // 30: datadog.api.v1.AgentSecure.TaggerStreamEntities:output_type -> datadog.model.v1.StreamTagsResponse - 26, // 31: datadog.api.v1.AgentSecure.TaggerGenerateContainerIDFromOriginInfo:output_type -> datadog.model.v1.GenerateContainerIDFromOriginInfoResponse - 27, // 32: datadog.api.v1.AgentSecure.TaggerFetchEntity:output_type -> datadog.model.v1.FetchEntityResponse - 28, // 33: datadog.api.v1.AgentSecure.DogstatsdCaptureTrigger:output_type -> datadog.model.v1.CaptureTriggerResponse - 29, // 34: datadog.api.v1.AgentSecure.DogstatsdSetTaggerState:output_type -> datadog.model.v1.TaggerStateResponse - 30, // 35: datadog.api.v1.AgentSecure.ClientGetConfigs:output_type -> datadog.config.ClientGetConfigsResponse - 31, // 36: datadog.api.v1.AgentSecure.GetConfigState:output_type -> datadog.config.GetStateConfigResponse - 30, // 37: datadog.api.v1.AgentSecure.ClientGetConfigsHA:output_type -> datadog.config.ClientGetConfigsResponse - 31, // 38: datadog.api.v1.AgentSecure.GetConfigStateHA:output_type -> datadog.config.GetStateConfigResponse - 32, // 39: datadog.api.v1.AgentSecure.CreateConfigSubscription:output_type -> datadog.config.ConfigSubscriptionResponse - 33, // 40: datadog.api.v1.AgentSecure.ResetConfigState:output_type -> datadog.config.ResetStateConfigResponse - 34, // 41: datadog.api.v1.AgentSecure.WorkloadmetaStreamEntities:output_type -> datadog.workloadmeta.WorkloadmetaStreamResponse - 35, // 42: datadog.api.v1.AgentSecure.RegisterRemoteAgent:output_type -> datadog.remoteagent.v1.RegisterRemoteAgentResponse - 36, // 43: datadog.api.v1.AgentSecure.RefreshRemoteAgent:output_type -> datadog.remoteagent.v1.RefreshRemoteAgentResponse - 37, // 44: datadog.api.v1.AgentSecure.AutodiscoveryStreamConfig:output_type -> datadog.autodiscovery.AutodiscoveryStreamResponse - 38, // 45: datadog.api.v1.AgentSecure.GetHostTags:output_type -> datadog.model.v1.HostTagReply - 39, // 46: datadog.api.v1.AgentSecure.StreamConfigEvents:output_type -> datadog.model.v1.ConfigEvent - 40, // 47: datadog.api.v1.AgentSecure.WorkloadFilterEvaluate:output_type -> datadog.workloadfilter.WorkloadFilterEvaluateResponse - 5, // 48: datadog.api.v1.AgentSecure.RemoteQueryExecute:output_type -> datadog.api.v1.RemoteQueryExecuteResponse - 6, // 49: datadog.api.v1.AgentSecure.RemoteQueryExecuteStream:output_type -> datadog.api.v1.RemoteQueryExecuteChunk - 41, // 50: datadog.api.v1.AgentSecure.StreamKubeMetadata:output_type -> datadog.kubemetadata.KubeMetadataStreamResponse - 29, // [29:51] is the sub-list for method output_type - 7, // [7:29] is the sub-list for method input_type - 7, // [7:7] is the sub-list for extension type_name - 7, // [7:7] is the sub-list for extension extendee - 0, // [0:7] is the sub-list for field type_name + 15, // 4: datadog.api.v1.RemoteQueryExecuteResponse.columns:type_name -> google.protobuf.Struct + 15, // 5: datadog.api.v1.RemoteQueryExecuteResponse.rows:type_name -> google.protobuf.Struct + 15, // 6: datadog.api.v1.RemoteQueryExecuteResponse.stats:type_name -> google.protobuf.Struct + 12, // 7: datadog.api.v1.RemoteQueryStreamMetadata.attributes:type_name -> datadog.api.v1.RemoteQueryStreamMetadata.AttributesEntry + 13, // 8: datadog.api.v1.RemoteQueryStreamFinal.attributes:type_name -> datadog.api.v1.RemoteQueryStreamFinal.AttributesEntry + 14, // 9: datadog.api.v1.RemoteQueryStreamError.attributes:type_name -> datadog.api.v1.RemoteQueryStreamError.AttributesEntry + 6, // 10: datadog.api.v1.RemoteQueryExecuteStreamEvent.metadata:type_name -> datadog.api.v1.RemoteQueryStreamMetadata + 7, // 11: datadog.api.v1.RemoteQueryExecuteStreamEvent.data:type_name -> datadog.api.v1.RemoteQueryStreamData + 8, // 12: datadog.api.v1.RemoteQueryExecuteStreamEvent.final:type_name -> datadog.api.v1.RemoteQueryStreamFinal + 9, // 13: datadog.api.v1.RemoteQueryExecuteStreamEvent.error:type_name -> datadog.api.v1.RemoteQueryStreamError + 10, // 14: datadog.api.v1.RemoteQueryExecuteChunk.event:type_name -> datadog.api.v1.RemoteQueryExecuteStreamEvent + 16, // 15: datadog.api.v1.Agent.GetHostname:input_type -> datadog.model.v1.HostnameRequest + 17, // 16: datadog.api.v1.AgentSecure.TaggerStreamEntities:input_type -> datadog.model.v1.StreamTagsRequest + 18, // 17: datadog.api.v1.AgentSecure.TaggerGenerateContainerIDFromOriginInfo:input_type -> datadog.model.v1.GenerateContainerIDFromOriginInfoRequest + 19, // 18: datadog.api.v1.AgentSecure.TaggerFetchEntity:input_type -> datadog.model.v1.FetchEntityRequest + 20, // 19: datadog.api.v1.AgentSecure.DogstatsdCaptureTrigger:input_type -> datadog.model.v1.CaptureTriggerRequest + 21, // 20: datadog.api.v1.AgentSecure.DogstatsdSetTaggerState:input_type -> datadog.model.v1.TaggerState + 22, // 21: datadog.api.v1.AgentSecure.ClientGetConfigs:input_type -> datadog.config.ClientGetConfigsRequest + 23, // 22: datadog.api.v1.AgentSecure.GetConfigState:input_type -> google.protobuf.Empty + 22, // 23: datadog.api.v1.AgentSecure.ClientGetConfigsHA:input_type -> datadog.config.ClientGetConfigsRequest + 23, // 24: datadog.api.v1.AgentSecure.GetConfigStateHA:input_type -> google.protobuf.Empty + 24, // 25: datadog.api.v1.AgentSecure.CreateConfigSubscription:input_type -> datadog.config.ConfigSubscriptionRequest + 23, // 26: datadog.api.v1.AgentSecure.ResetConfigState:input_type -> google.protobuf.Empty + 25, // 27: datadog.api.v1.AgentSecure.WorkloadmetaStreamEntities:input_type -> datadog.workloadmeta.WorkloadmetaStreamRequest + 26, // 28: datadog.api.v1.AgentSecure.RegisterRemoteAgent:input_type -> datadog.remoteagent.v1.RegisterRemoteAgentRequest + 27, // 29: datadog.api.v1.AgentSecure.RefreshRemoteAgent:input_type -> datadog.remoteagent.v1.RefreshRemoteAgentRequest + 23, // 30: datadog.api.v1.AgentSecure.AutodiscoveryStreamConfig:input_type -> google.protobuf.Empty + 28, // 31: datadog.api.v1.AgentSecure.GetHostTags:input_type -> datadog.model.v1.HostTagRequest + 29, // 32: datadog.api.v1.AgentSecure.StreamConfigEvents:input_type -> datadog.model.v1.ConfigStreamRequest + 30, // 33: datadog.api.v1.AgentSecure.WorkloadFilterEvaluate:input_type -> datadog.workloadfilter.WorkloadFilterEvaluateRequest + 3, // 34: datadog.api.v1.AgentSecure.RemoteQueryExecute:input_type -> datadog.api.v1.RemoteQueryExecuteRequest + 3, // 35: datadog.api.v1.AgentSecure.RemoteQueryExecuteStream:input_type -> datadog.api.v1.RemoteQueryExecuteRequest + 31, // 36: datadog.api.v1.AgentSecure.StreamKubeMetadata:input_type -> datadog.kubemetadata.KubeMetadataStreamRequest + 32, // 37: datadog.api.v1.Agent.GetHostname:output_type -> datadog.model.v1.HostnameReply + 33, // 38: datadog.api.v1.AgentSecure.TaggerStreamEntities:output_type -> datadog.model.v1.StreamTagsResponse + 34, // 39: datadog.api.v1.AgentSecure.TaggerGenerateContainerIDFromOriginInfo:output_type -> datadog.model.v1.GenerateContainerIDFromOriginInfoResponse + 35, // 40: datadog.api.v1.AgentSecure.TaggerFetchEntity:output_type -> datadog.model.v1.FetchEntityResponse + 36, // 41: datadog.api.v1.AgentSecure.DogstatsdCaptureTrigger:output_type -> datadog.model.v1.CaptureTriggerResponse + 37, // 42: datadog.api.v1.AgentSecure.DogstatsdSetTaggerState:output_type -> datadog.model.v1.TaggerStateResponse + 38, // 43: datadog.api.v1.AgentSecure.ClientGetConfigs:output_type -> datadog.config.ClientGetConfigsResponse + 39, // 44: datadog.api.v1.AgentSecure.GetConfigState:output_type -> datadog.config.GetStateConfigResponse + 38, // 45: datadog.api.v1.AgentSecure.ClientGetConfigsHA:output_type -> datadog.config.ClientGetConfigsResponse + 39, // 46: datadog.api.v1.AgentSecure.GetConfigStateHA:output_type -> datadog.config.GetStateConfigResponse + 40, // 47: datadog.api.v1.AgentSecure.CreateConfigSubscription:output_type -> datadog.config.ConfigSubscriptionResponse + 41, // 48: datadog.api.v1.AgentSecure.ResetConfigState:output_type -> datadog.config.ResetStateConfigResponse + 42, // 49: datadog.api.v1.AgentSecure.WorkloadmetaStreamEntities:output_type -> datadog.workloadmeta.WorkloadmetaStreamResponse + 43, // 50: datadog.api.v1.AgentSecure.RegisterRemoteAgent:output_type -> datadog.remoteagent.v1.RegisterRemoteAgentResponse + 44, // 51: datadog.api.v1.AgentSecure.RefreshRemoteAgent:output_type -> datadog.remoteagent.v1.RefreshRemoteAgentResponse + 45, // 52: datadog.api.v1.AgentSecure.AutodiscoveryStreamConfig:output_type -> datadog.autodiscovery.AutodiscoveryStreamResponse + 46, // 53: datadog.api.v1.AgentSecure.GetHostTags:output_type -> datadog.model.v1.HostTagReply + 47, // 54: datadog.api.v1.AgentSecure.StreamConfigEvents:output_type -> datadog.model.v1.ConfigEvent + 48, // 55: datadog.api.v1.AgentSecure.WorkloadFilterEvaluate:output_type -> datadog.workloadfilter.WorkloadFilterEvaluateResponse + 5, // 56: datadog.api.v1.AgentSecure.RemoteQueryExecute:output_type -> datadog.api.v1.RemoteQueryExecuteResponse + 11, // 57: datadog.api.v1.AgentSecure.RemoteQueryExecuteStream:output_type -> datadog.api.v1.RemoteQueryExecuteChunk + 49, // 58: datadog.api.v1.AgentSecure.StreamKubeMetadata:output_type -> datadog.kubemetadata.KubeMetadataStreamResponse + 37, // [37:59] is the sub-list for method output_type + 15, // [15:37] is the sub-list for method input_type + 15, // [15:15] is the sub-list for extension type_name + 15, // [15:15] is the sub-list for extension extendee + 0, // [0:15] is the sub-list for field type_name } func init() { file_datadog_api_v1_api_proto_init() } @@ -697,13 +1149,19 @@ func file_datadog_api_v1_api_proto_init() { file_datadog_workloadfilter_workloadfilter_proto_init() file_datadog_autodiscovery_autodiscovery_proto_init() file_datadog_kubemetadata_kubemetadata_proto_init() + file_datadog_api_v1_api_proto_msgTypes[10].OneofWrappers = []any{ + (*RemoteQueryExecuteStreamEvent_Metadata)(nil), + (*RemoteQueryExecuteStreamEvent_Data)(nil), + (*RemoteQueryExecuteStreamEvent_Final)(nil), + (*RemoteQueryExecuteStreamEvent_Error)(nil), + } type x struct{} out := protoimpl.TypeBuilder{ File: protoimpl.DescBuilder{ GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_datadog_api_v1_api_proto_rawDesc), len(file_datadog_api_v1_api_proto_rawDesc)), NumEnums: 0, - NumMessages: 7, + NumMessages: 15, NumExtensions: 0, NumServices: 2, }, diff --git a/rtloader/include/rtloader_types.h b/rtloader/include/rtloader_types.h index b247260e51f3..641782b4a97c 100644 --- a/rtloader/include/rtloader_types.h +++ b/rtloader/include/rtloader_types.h @@ -6,6 +6,7 @@ #ifndef DATADOG_AGENT_RTLOADER_TYPES_H #define DATADOG_AGENT_RTLOADER_TYPES_H #include +#include #include #ifdef __cplusplus @@ -37,7 +38,7 @@ typedef enum rtloader_gilstate_e { typedef void *(*rtloader_malloc_t)(size_t); typedef void (*rtloader_free_t)(void *); -typedef int (*remote_query_stream_emit_cb)(const char *event_json, void *userdata); +typedef int (*remote_query_stream_emit_cb)(const char *event_type, const char *metadata_json, const uint8_t *payload, size_t payload_len, void *userdata); typedef enum { DATADOG_AGENT_RTLOADER_GAUGE = 0, diff --git a/rtloader/three/three.cpp b/rtloader/three/three.cpp index 0e908e2b2345..65dc04a51a42 100644 --- a/rtloader/three/three.cpp +++ b/rtloader/three/three.cpp @@ -537,7 +537,7 @@ struct RemoteQueryStreamEmitContext { void *userdata; }; -PyObject *remoteQueryStreamEmit(PyObject *self, PyObject *event) +PyObject *remoteQueryStreamEmit(PyObject *self, PyObject *args) { RemoteQueryStreamEmitContext *ctx = static_cast(PyCapsule_GetPointer(self, "remote_query_stream_emit")); if (ctx == NULL || ctx->emit == NULL) { @@ -545,53 +545,23 @@ PyObject *remoteQueryStreamEmit(PyObject *self, PyObject *event) return NULL; } - PyObject *normalized_event = event; - PyObject *event_copy = NULL; - if (PyDict_Check(event)) { - PyObject *data = PyDict_GetItemString(event, "data"); - if (data != NULL && PyBytes_Check(data)) { - event_copy = PyDict_Copy(event); - if (event_copy == NULL) { - return NULL; - } - char *bytes = PyBytes_AsString(data); - Py_ssize_t size = PyBytes_Size(data); - PyObject *data_str = PyUnicode_FromStringAndSize(bytes, size); - if (data_str == NULL) { - Py_DECREF(event_copy); - return NULL; - } - if (PyDict_SetItemString(event_copy, "data", data_str) != 0) { - Py_DECREF(data_str); - Py_DECREF(event_copy); - return NULL; - } - Py_DECREF(data_str); - normalized_event = event_copy; - } - } - - PyObject *json_module = PyImport_ImportModule("json"); - if (json_module == NULL) { - Py_XDECREF(event_copy); + const char *event_type = NULL; + const char *metadata_json = NULL; + PyObject *payload = NULL; + if (!PyArg_ParseTuple(args, "ssO:remote_query_stream_emit", &event_type, &metadata_json, &payload)) { return NULL; } - PyObject *event_json = PyObject_CallMethod(json_module, const_cast("dumps"), const_cast("O"), normalized_event); - Py_DECREF(json_module); - Py_XDECREF(event_copy); - if (event_json == NULL || !PyUnicode_Check(event_json)) { - Py_XDECREF(event_json); + if (!PyBytes_Check(payload)) { + PyErr_SetString(PyExc_TypeError, "remote query stream payload must be bytes"); return NULL; } - PyObject *utf8 = PyUnicode_AsUTF8String(event_json); - Py_DECREF(event_json); - if (utf8 == NULL) { + char *payload_bytes = NULL; + Py_ssize_t payload_len = 0; + if (PyBytes_AsStringAndSize(payload, &payload_bytes, &payload_len) != 0) { return NULL; } - const char *event_json_c = PyBytes_AsString(utf8); - int emit_result = ctx->emit(event_json_c, ctx->userdata); - Py_DECREF(utf8); + int emit_result = ctx->emit(event_type, metadata_json, reinterpret_cast(payload_bytes), static_cast(payload_len), ctx->userdata); if (emit_result != 0) { PyErr_SetString(PyExc_RuntimeError, "remote query stream emit callback failed"); return NULL; @@ -600,8 +570,8 @@ PyObject *remoteQueryStreamEmit(PyObject *self, PyObject *event) Py_RETURN_NONE; } -PyMethodDef remoteQueryStreamEmitMethod = {"remote_query_stream_emit", remoteQueryStreamEmit, METH_O, - "Emit a serialized remote query stream event."}; +PyMethodDef remoteQueryStreamEmitMethod = {"remote_query_stream_emit", remoteQueryStreamEmit, METH_VARARGS, + "Emit a remote query stream event."}; } // namespace char *Three::runRemoteQuery(RtLoaderPyObject *check, const char *integration, const char *request_json) diff --git a/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh b/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh index cc20b57dfce7..70e751105dde 100755 --- a/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh +++ b/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh @@ -34,6 +34,11 @@ if [[ -n "${RQ_REMOTE_OPERATION+x}" ]]; then RQ_REMOTE_OPERATION_WAS_SET=1 fi RQ_REMOTE_OPERATION=${RQ_REMOTE_OPERATION:-} +RQ_REMOTE_FORMAT_WAS_SET=0 +if [[ -n "${RQ_REMOTE_FORMAT+x}" ]]; then + RQ_REMOTE_FORMAT_WAS_SET=1 +fi +RQ_REMOTE_FORMAT=${RQ_REMOTE_FORMAT:-} RQ_POSTGRES_HOST=${RQ_POSTGRES_HOST:-localhost} RQ_POSTGRES_PORT=${RQ_POSTGRES_PORT:-5432} RQ_POSTGRES_DBNAME=${RQ_POSTGRES_DBNAME:-datadog_test} @@ -52,6 +57,7 @@ PROOF_CASE_NAMES=( "seed" "fixture-city" "copy-fixture-city" + "copy-binary-payload" "payload-1mib" "payload-2mib" "payload-4mib" @@ -64,6 +70,7 @@ PROOF_CASE_QUERIES=( "SELECT 1 AS value" "SELECT city, country FROM cities ORDER BY city" "SELECT city, country FROM cities ORDER BY city" + "SELECT decode('00ff80', 'hex') AS payload" "SELECT repeat('x', 1048576) AS payload" "SELECT repeat('x', 2097152) AS payload" "SELECT repeat('x', 4194304) AS payload" @@ -393,6 +400,11 @@ PY local token token=$(cat "$TMP_ROOT/run/auth_token") + if [[ "${RQ_REMOTE_OPERATION:-}" == "copy_stream" && "${RQ_REMOTE_FORMAT:-}" == "binary" ]]; then + log "[$PROOF_CASE_NAME] Skipping inline JSON preflight for binary COPY stream (streaming path only)" + return + fi + log "[$PROOF_CASE_NAME] Preflight real Agent IPC HTTP execute endpoint (dev evidence only)" local body_file="$CASE_RESULTS_DIR/agent-execute-preflight.raw-body" local status_file="$CASE_RESULTS_DIR/agent-execute-preflight.status" @@ -475,6 +487,7 @@ run_standalone_go_proof() { RQ_POSTGRES_DBNAME="$RQ_POSTGRES_DBNAME" \ RQ_REMOTE_QUERY="$RQ_REMOTE_QUERY" \ RQ_REMOTE_OPERATION="${RQ_REMOTE_OPERATION:-}" \ + RQ_REMOTE_FORMAT="${RQ_REMOTE_FORMAT:-}" \ dda inv test --targets=./pkg/privateactionrunner/bundles/remotequeries \ --extra-args='-run TestRemoteQueriesActionRunsThroughStandalonePARProcessWithRealAgentIPC -count=1 -v' ) | tee "$CASE_RESULTS_DIR/standalone-proof-test.log" @@ -489,6 +502,12 @@ run_proof_case() { RQ_REMOTE_OPERATION=copy_stream fi fi + if [[ "$RQ_REMOTE_FORMAT_WAS_SET" != "1" ]]; then + RQ_REMOTE_FORMAT="" + if [[ "$PROOF_CASE_NAME" == "copy-binary-payload" ]]; then + RQ_REMOTE_FORMAT=binary + fi + fi CASE_RESULTS_DIR="$TMP_ROOT/results/$PROOF_CASE_NAME" mkdir -p "$CASE_RESULTS_DIR" printf '%s\n' "$RQ_REMOTE_QUERY" > "$CASE_RESULTS_DIR/query.sql" From 5d5306c3cc99d2e85aaacfb2b9c4f90ae3e7e745 Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 15:18:23 +0000 Subject: [PATCH 24/33] Remove inline Remote Queries execution path --- .../impl-agent/remote_query_execute_test.go | 28 +-- comp/api/grpcserver/impl-agent/server.go | 109 +++-------- .../impl/remote_query_execute.go | 179 ++++-------------- .../impl/remote_query_par_poc.go | 10 +- .../impl/remote_query_par_poc_test.go | 31 +-- comp/remotequeries/impl/remote_query_test.go | 118 ++++-------- pkg/collector/python/check.go | 35 ---- pkg/collector/python/check_test.go | 20 -- pkg/collector/python/test_check.go | 111 ----------- .../bundles/remotequeries/execute.go | 89 ++------- .../bundles/remotequeries/execute_test.go | 81 +++----- .../live_agent_ipc_par_loop_test.go | 28 +-- .../remotequeries/live_par_loop_test.go | 28 +-- .../standalone_par_process_proof_test.go | 40 ++-- pkg/proto/datadog/api/v1/api.proto | 4 +- pkg/proto/pbgo/core/api.pb.go | 27 +-- pkg/proto/pbgo/core/api_grpc.pb.go | 4 +- rtloader/include/datadog_agent_rtloader.h | 12 -- rtloader/include/rtloader.h | 9 - rtloader/rtloader/api.cpp | 5 - rtloader/three/three.cpp | 59 ------ rtloader/three/three.h | 1 - ...andalone-par-agentsecure-postgres-proof.sh | 66 +------ 23 files changed, 233 insertions(+), 861 deletions(-) diff --git a/comp/api/grpcserver/impl-agent/remote_query_execute_test.go b/comp/api/grpcserver/impl-agent/remote_query_execute_test.go index 493d7e7edbfa..e8afe7a5222d 100644 --- a/comp/api/grpcserver/impl-agent/remote_query_execute_test.go +++ b/comp/api/grpcserver/impl-agent/remote_query_execute_test.go @@ -31,13 +31,14 @@ func TestRemoteQueryExecuteResponseFromJSONMapsStructuredRows(t *testing.T) { assert.Equal(t, float64(2), resp.GetStats().AsMap()["elapsed_ms"]) } -func TestRemoteQueryExecuteReturnsSanitizedUnavailableWhenServiceMissing(t *testing.T) { +func TestRemoteQueryExecuteRejectsUnaryInlineMode(t *testing.T) { resp, err := (&serverSecure{}).RemoteQueryExecute(context.Background(), &pb.RemoteQueryExecuteRequest{}) require.NoError(t, err) - assert.Equal(t, "executor_unavailable", resp.GetStatus()) + assert.Equal(t, "invalid_request", resp.GetStatus()) require.NotNil(t, resp.GetError()) - assert.Equal(t, "executor_unavailable", resp.GetError().GetCode()) + assert.Equal(t, "invalid_request", resp.GetError().GetCode()) + assert.Contains(t, resp.GetError().GetMessage(), "RemoteQueryExecuteStream") } func TestRemoteQueryExecuteStreamReturnsSanitizedUnavailableWhenServiceMissing(t *testing.T) { @@ -45,9 +46,9 @@ func TestRemoteQueryExecuteStreamReturnsSanitizedUnavailableWhenServiceMissing(t err := (&serverSecure{}).RemoteQueryExecuteStream(&pb.RemoteQueryExecuteRequest{}, stream) require.NoError(t, err) - require.Len(t, stream.chunks, 1) - assert.True(t, stream.chunks[0].GetFinal()) - assert.JSONEq(t, `{"status":"executor_unavailable","error":{"code":"executor_unavailable","message":"remote query executor is unavailable"}}`, string(stream.chunks[0].GetResponseJsonChunk())) + require.Len(t, stream.chunks, 2) + assert.Equal(t, "executor_unavailable", stream.chunks[0].GetEvent().GetError().GetCode()) + assert.True(t, stream.chunks[1].GetFinal()) } func TestRemoteQueryExecuteRequestFromProtoPreservesCopyStream(t *testing.T) { @@ -84,21 +85,6 @@ func TestRemoteQueryStreamEventFromCheckEventPreservesBinaryPayload(t *testing.T assert.Equal(t, uint64(3), event.GetData().GetBytes()) } -func TestRemoteQueryExecuteStreamJSONChunksResponse(t *testing.T) { - stream := &captureRemoteQueryExecuteStreamServer{} - responseJSON := `{"status":"SUCCEEDED","rows":[{"payload":"` + string(make([]byte, remoteQueryExecuteStreamChunkSize+1)) + `"}]}` - - err := remoteQueryExecuteStreamJSON(responseJSON, stream) - - require.NoError(t, err) - require.Len(t, stream.chunks, 2) - assert.Equal(t, int32(0), stream.chunks[0].GetChunkIndex()) - assert.False(t, stream.chunks[0].GetFinal()) - assert.Equal(t, int32(1), stream.chunks[1].GetChunkIndex()) - assert.True(t, stream.chunks[1].GetFinal()) - assert.Equal(t, len(responseJSON), len(stream.chunks[0].GetResponseJsonChunk())+len(stream.chunks[1].GetResponseJsonChunk())) -} - type captureRemoteQueryExecuteStreamServer struct { grpc.ServerStream chunks []*pb.RemoteQueryExecuteChunk diff --git a/comp/api/grpcserver/impl-agent/server.go b/comp/api/grpcserver/impl-agent/server.go index c920c58d8353..858c0a32fd9b 100644 --- a/comp/api/grpcserver/impl-agent/server.go +++ b/comp/api/grpcserver/impl-agent/server.go @@ -9,6 +9,7 @@ import ( "context" "encoding/json" "errors" + "fmt" "strconv" "strings" "time" @@ -282,59 +283,34 @@ func (s *serverSecure) WorkloadFilterEvaluate(ctx context.Context, req *pb.Workl return s.workloadfilterServer.WorkloadFilterEvaluate(ctx, req) } -const remoteQueryExecuteStreamChunkSize = 1 << 20 - func (s *serverSecure) RemoteQueryExecute(_ context.Context, req *pb.RemoteQueryExecuteRequest) (*pb.RemoteQueryExecuteResponse, error) { - if s.remoteQueries == nil { - return remoteQueryExecuteErrorResponse(remotequeriesimpl.RemoteQueryStatusExecutorUnavailable, "remote query executor is unavailable"), nil - } - - execReq, err := remoteQueryExecuteRequestFromProto(req) - if err != nil { - return remoteQueryExecuteErrorResponse(remotequeriesimpl.RemoteQueryStatusInvalidRequest, err.Error()), nil - } - - result := s.remoteQueries.Execute(execReq) - if result.Error != nil { - return remoteQueryExecuteErrorResponse(result.Error.Code, result.Error.Message), nil - } - - return remoteQueryExecuteResponseFromJSON(result.ResponseJSON) + return remoteQueryExecuteErrorResponse(remotequeriesimpl.RemoteQueryStatusInvalidRequest, "remote queries require RemoteQueryExecuteStream with operation copy_stream"), nil } func (s *serverSecure) RemoteQueryExecuteStream(req *pb.RemoteQueryExecuteRequest, stream pb.AgentSecure_RemoteQueryExecuteStreamServer) error { if s.remoteQueries == nil { - return remoteQueryExecuteStreamJSON(remoteQueryExecuteErrorJSON(remotequeriesimpl.RemoteQueryStatusExecutorUnavailable, "remote query executor is unavailable"), stream) + return remoteQueryExecuteStreamError(remotequeriesimpl.RemoteQueryStatusExecutorUnavailable, "remote query executor is unavailable", stream) } execReq, err := remoteQueryExecuteRequestFromProto(req) if err != nil { - return remoteQueryExecuteStreamJSON(remoteQueryExecuteErrorJSON(remotequeriesimpl.RemoteQueryStatusInvalidRequest, err.Error()), stream) + return remoteQueryExecuteStreamError(remotequeriesimpl.RemoteQueryStatusInvalidRequest, err.Error(), stream) } - if execReq.Operation == "copy_stream" { - chunkIndex := int32(0) - result := s.remoteQueries.ExecuteStream(execReq, func(event check.RemoteQueryStreamEvent) error { - protoEvent, err := remoteQueryStreamEventFromCheckEvent(event) - if err != nil { - return err - } - err = stream.Send(&pb.RemoteQueryExecuteChunk{Event: protoEvent, ChunkIndex: chunkIndex}) - chunkIndex++ + chunkIndex := int32(0) + result := s.remoteQueries.ExecuteStream(execReq, func(event check.RemoteQueryStreamEvent) error { + protoEvent, err := remoteQueryStreamEventFromCheckEvent(event) + if err != nil { return err - }) - if result.Error != nil { - return remoteQueryExecuteStreamJSON(remoteQueryExecuteErrorJSON(result.Error.Code, result.Error.Message), stream) } - return stream.Send(&pb.RemoteQueryExecuteChunk{ChunkIndex: chunkIndex, Final: true}) - } - - result := s.remoteQueries.Execute(execReq) + err = stream.Send(&pb.RemoteQueryExecuteChunk{Event: protoEvent, ChunkIndex: chunkIndex}) + chunkIndex++ + return err + }) if result.Error != nil { - return remoteQueryExecuteStreamJSON(remoteQueryExecuteErrorJSON(result.Error.Code, result.Error.Message), stream) + return remoteQueryExecuteStreamError(result.Error.Code, result.Error.Message, stream) } - - return remoteQueryExecuteStreamJSON(result.ResponseJSON, stream) + return stream.Send(&pb.RemoteQueryExecuteChunk{ChunkIndex: chunkIndex, Final: true}) } func remoteQueryExecuteRequestFromProto(req *pb.RemoteQueryExecuteRequest) (remotequeriesimpl.RemoteQueryExecuteRequest, error) { @@ -343,11 +319,10 @@ func remoteQueryExecuteRequestFromProto(req *pb.RemoteQueryExecuteRequest) (remo Port: int(req.GetTarget().GetPort()), DBName: req.GetTarget().GetDbname(), } - if req.GetOperation() == "copy_stream" { - return remotequeriesimpl.NewRemoteQueryCopyStreamExecuteRequest(req.GetIntegration(), target, req.GetQuery(), req.GetFormat(), remoteQueryCopyLimitsFromProto(req.GetCopyLimits())) + if req.GetOperation() != "copy_stream" { + return remotequeriesimpl.RemoteQueryExecuteRequest{}, fmt.Errorf("operation must be copy_stream") } - limits := remoteQueryLimitsFromProto(req.GetLimits()) - return remotequeriesimpl.NewRemoteQueryExecuteRequest(req.GetIntegration(), target, req.GetQuery(), limits) + return remotequeriesimpl.NewRemoteQueryCopyStreamExecuteRequest(req.GetIntegration(), target, req.GetQuery(), req.GetFormat(), remoteQueryCopyLimitsFromProto(req.GetCopyLimits())) } func remoteQueryStreamEventFromCheckEvent(event check.RemoteQueryStreamEvent) (*pb.RemoteQueryExecuteStreamEvent, error) { @@ -454,36 +429,17 @@ func stringAttributes(metadata map[string]interface{}, exclude ...string) map[st return attrs } -func remoteQueryExecuteStreamJSON(responseJSON string, stream pb.AgentSecure_RemoteQueryExecuteStreamServer) error { - responseBytes := []byte(responseJSON) - if len(responseBytes) == 0 { - return stream.Send(&pb.RemoteQueryExecuteChunk{Final: true}) - } - for offset, chunkIndex := 0, int32(0); offset < len(responseBytes); offset, chunkIndex = offset+remoteQueryExecuteStreamChunkSize, chunkIndex+1 { - end := offset + remoteQueryExecuteStreamChunkSize - if end > len(responseBytes) { - end = len(responseBytes) - } - if err := stream.Send(&pb.RemoteQueryExecuteChunk{ - ResponseJsonChunk: responseBytes[offset:end], - ChunkIndex: chunkIndex, - Final: end == len(responseBytes), - }); err != nil { - return err - } - } - return nil -} - -func remoteQueryLimitsFromProto(limits *pb.RemoteQueryExecuteLimits) *remotequeriesimpl.RemoteQueryExecuteLimits { - if limits == nil { - return nil - } - return &remotequeriesimpl.RemoteQueryExecuteLimits{ - MaxRows: int(limits.GetMaxRows()), - MaxBytes: int(limits.GetMaxBytes()), - TimeoutMs: int(limits.GetTimeoutMs()), +func remoteQueryExecuteStreamError(code string, message string, stream pb.AgentSecure_RemoteQueryExecuteStreamServer) error { + if err := stream.Send(&pb.RemoteQueryExecuteChunk{ + ChunkIndex: 0, + Event: &pb.RemoteQueryExecuteStreamEvent{Event: &pb.RemoteQueryExecuteStreamEvent_Error{Error: &pb.RemoteQueryStreamError{ + Code: code, + Message: message, + }}}, + }); err != nil { + return err } + return stream.Send(&pb.RemoteQueryExecuteChunk{ChunkIndex: 1, Final: true}) } func remoteQueryCopyLimitsFromProto(limits *pb.RemoteQueryExecuteCopyLimits) *remotequeriesimpl.RemoteQueryExecuteCopyLimits { @@ -505,17 +461,6 @@ func remoteQueryExecuteErrorResponse(code string, message string) *pb.RemoteQuer } } -func remoteQueryExecuteErrorJSON(code string, message string) string { - payload, err := json.Marshal(remoteQueryExecuteJSONResponse{ - Status: code, - Error: &remoteQueryExecuteError{Code: code, Message: message}, - }) - if err != nil { - return `{"status":"executor_unavailable","error":{"code":"executor_unavailable","message":"remote query executor is unavailable"}}` - } - return string(payload) -} - type remoteQueryExecuteJSONResponse struct { Status string `json:"status"` Error *remoteQueryExecuteError `json:"error,omitempty"` diff --git a/comp/remotequeries/impl/remote_query_execute.go b/comp/remotequeries/impl/remote_query_execute.go index 2c0ddd7b186e..049f75d479cf 100644 --- a/comp/remotequeries/impl/remote_query_execute.go +++ b/comp/remotequeries/impl/remote_query_execute.go @@ -41,10 +41,6 @@ var remoteQueryLargePayloadProofQueries = map[string]int{ "SELECT repeat('x', 33554432) AS payload": 32 << 20, } -type remoteQueryRunner interface { - RunRemoteQueryJSON(integration string, requestJSON string) (string, error) -} - type remoteQueryStreamRunner interface { RunRemoteQueryStream(integration string, requestJSON string, emit func(check.RemoteQueryStreamEvent) error) error } @@ -63,24 +59,6 @@ type remoteQueryCheckUnwrapper interface { Unwrap() check.Check } -func remoteQueryRunnerFor(chk check.Check) (remoteQueryRunner, bool) { - for chk != nil { - if runner, ok := chk.(remoteQueryRunner); ok { - return runner, true - } - unwrapper, ok := chk.(remoteQueryCheckUnwrapper) - if !ok { - break - } - unwrapped := unwrapper.Unwrap() - if unwrapped == chk { - break - } - chk = unwrapped - } - return nil, false -} - func remoteQueryStreamRunnerFor(chk check.Check) (remoteQueryStreamRunner, bool) { for chk != nil { if runner, ok := chk.(remoteQueryStreamRunner); ok { @@ -200,12 +178,11 @@ type RemoteQueryExecuteError struct { Message string } -// RemoteQueryExecuteResult is the service result. ResponseJSON is set only for successful executor responses. +// RemoteQueryExecuteResult is the service result. type RemoteQueryExecuteResult struct { - HTTPStatus int - Status string - Error *RemoteQueryExecuteError - ResponseJSON string + HTTPStatus int + Status string + Error *RemoteQueryExecuteError } const ( @@ -215,43 +192,9 @@ const ( RemoteQueryStatusExecutorUnavailable = statusExecutorUnavailable ) -// NewRemoteQueryExecuteRequest validates and normalizes a typed Remote Queries execute request. +// NewRemoteQueryExecuteRequest rejects legacy inline Remote Queries requests. func NewRemoteQueryExecuteRequest(integration string, target RemoteQueryExecuteTarget, query string, limits *RemoteQueryExecuteLimits) (RemoteQueryExecuteRequest, error) { - parsedIntegration, err := parseIntegration(integration) - if err != nil { - return RemoteQueryExecuteRequest{}, err - } - - parsedTarget, err := parseTarget(&remoteQueryTargetRequestJSON{Host: target.Host, Port: &target.Port, DBName: target.DBName}) - if err != nil { - return RemoteQueryExecuteRequest{}, err - } - - if query == "" { - return RemoteQueryExecuteRequest{}, fmt.Errorf("query is required") - } - if !isRemoteQueryAllowedProofQuery(query) { - return RemoteQueryExecuteRequest{}, fmt.Errorf("query is not allowed") - } - - var parsedLimits *remoteQueryExecuteLimits - if limits != nil { - parsedLimits, err = parseExecuteLimits(&remoteQueryExecuteLimitsRequestJSON{ - MaxRows: &limits.MaxRows, - MaxBytes: &limits.MaxBytes, - TimeoutMs: &limits.TimeoutMs, - }) - if err != nil { - return RemoteQueryExecuteRequest{}, err - } - } - - return remoteQueryExecuteRequestFromInternal(remoteQueryExecuteRequest{ - Integration: parsedIntegration, - Target: parsedTarget, - Query: query, - Limits: parsedLimits, - }), nil + return RemoteQueryExecuteRequest{}, fmt.Errorf("operation must be copy_stream") } type remoteQueryExecuteRequest struct { @@ -300,12 +243,6 @@ type remoteQueryExecuteCopyLimits struct { TimeoutMs int } -type remoteQueryExecutorRequestJSON struct { - Target remoteQueryTargetJSON `json:"target"` - Query string `json:"query"` - Limits *remoteQueryExecuteLimitsJSON `json:"limits,omitempty"` -} - type remoteQueryCopyExecutorRequestJSON struct { Operation string `json:"operation"` Target remoteQueryTargetJSON `json:"target"` @@ -327,12 +264,6 @@ type remoteQueryTargetJSON struct { DBName string `json:"dbname"` } -type remoteQueryExecuteLimitsJSON struct { - MaxRows int `json:"maxRows"` - MaxBytes int `json:"maxBytes"` - TimeoutMs int `json:"timeoutMs"` -} - func (h *remoteQueryExecuteHandler) handle(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") @@ -358,7 +289,7 @@ func (h *remoteQueryExecuteHandler) handle(w http.ResponseWriter, r *http.Reques } w.WriteHeader(http.StatusOK) - _, _ = io.WriteString(w, result.ResponseJSON) + _, _ = io.WriteString(w, `{"status":"SUCCEEDED"}`) } func parseExecuteRequest(r *http.Request) (remoteQueryExecuteRequest, string, error) { @@ -398,6 +329,16 @@ func parseExecuteRequest(r *http.Request) (remoteQueryExecuteRequest, string, er return remoteQueryExecuteRequest{}, "", err } + if wireReq.Operation != "copy_stream" { + return remoteQueryExecuteRequest{}, "", fmt.Errorf("operation must be copy_stream") + } + if wireReq.Format == "" { + wireReq.Format = "csv" + } + if wireReq.Format != "csv" && wireReq.Format != "binary" { + return remoteQueryExecuteRequest{}, "", fmt.Errorf("format must be csv or binary") + } + req := remoteQueryExecuteRequest{Integration: integration, Operation: wireReq.Operation, Target: target, Query: wireReq.Query, Format: wireReq.Format, Limits: limits, CopyLimits: copyLimits} requestJSON, err := marshalExecuteRequest(req) if err != nil { @@ -500,44 +441,13 @@ func parseRequiredPositiveInt(value *int, name string) (int, error) { } func (s *RemoteQueryExecuteService) Execute(req RemoteQueryExecuteRequest) RemoteQueryExecuteResult { - if req.Operation == "copy_stream" { - return remoteQueryExecuteErrorResult(http.StatusBadRequest, statusInvalidRequest, "copy_stream requires the streaming executor") - } - if s == nil || !s.enabled { - return remoteQueryExecuteErrorResult(http.StatusServiceUnavailable, statusBridgeDisabled, "remote queries bridge is disabled") - } - if s.collector == nil { - return remoteQueryExecuteErrorResult(http.StatusFailedDependency, statusExecutorUnavailable, "remote query executor is unavailable") - } - - internal := req.internal() - match, result := s.matchExecutor(internal) - if result.Error != nil { - return result - } - - runner, ok := remoteQueryRunnerFor(match.check) - if !ok { - return remoteQueryExecuteErrorResult(http.StatusFailedDependency, statusExecutorUnavailable, "matched integration check does not support remote query execution") - } - - requestJSON, err := marshalExecuteRequest(internal) - if err != nil { - return remoteQueryExecuteErrorResult(http.StatusBadRequest, statusInvalidRequest, "malformed JSON request") - } - - responseJSON, err := runner.RunRemoteQueryJSON(internal.Integration, requestJSON) - if err != nil { - return remoteQueryExecuteErrorResult(http.StatusBadGateway, statusExecutorUnavailable, "remote query executor failed") - } - - return RemoteQueryExecuteResult{HTTPStatus: http.StatusOK, ResponseJSON: responseJSON} + return remoteQueryExecuteErrorResult(http.StatusBadRequest, statusInvalidRequest, "remote queries require operation copy_stream and the streaming executor") } // ExecuteStream executes a COPY streaming request and emits binary-safe stream events without materializing the full result. func (s *RemoteQueryExecuteService) ExecuteStream(req RemoteQueryExecuteRequest, emit func(check.RemoteQueryStreamEvent) error) RemoteQueryExecuteResult { if req.Operation != "copy_stream" { - return s.Execute(req) + return remoteQueryExecuteErrorResult(http.StatusBadRequest, statusInvalidRequest, "operation must be copy_stream") } if emit == nil { return remoteQueryExecuteErrorResult(http.StatusFailedDependency, statusExecutorUnavailable, "remote query stream emitter is unavailable") @@ -623,44 +533,27 @@ func remoteQueryExecuteRequestFromInternal(req remoteQueryExecuteRequest) Remote } func marshalExecuteRequest(req remoteQueryExecuteRequest) (string, error) { - if req.Operation == "copy_stream" { - format := req.Format - if format == "" { - format = "csv" - } - wireReq := remoteQueryCopyExecutorRequestJSON{ - Operation: req.Operation, - Target: remoteQueryTargetJSON{Host: req.Target.Host, Port: req.Target.Port, DBName: req.Target.DBName}, - Query: req.Query, - Format: format, - } - if req.CopyLimits != nil { - wireReq.Limits = &remoteQueryExecuteCopyLimitsJSON{ - ChunkBytes: req.CopyLimits.ChunkBytes, - MaxBytes: req.CopyLimits.MaxBytes, - MaxRowBytes: req.CopyLimits.MaxRowBytes, - TimeoutMs: req.CopyLimits.TimeoutMs, - } - } - requestJSON, err := json.Marshal(wireReq) - if err != nil { - return "", err - } - return string(requestJSON), nil + if req.Operation != "copy_stream" { + return "", fmt.Errorf("operation must be copy_stream") } - - wireReq := remoteQueryExecutorRequestJSON{ - Target: remoteQueryTargetJSON{Host: req.Target.Host, Port: req.Target.Port, DBName: req.Target.DBName}, - Query: req.Query, + format := req.Format + if format == "" { + format = "csv" } - if req.Limits != nil { - wireReq.Limits = &remoteQueryExecuteLimitsJSON{ - MaxRows: req.Limits.MaxRows, - MaxBytes: req.Limits.MaxBytes, - TimeoutMs: req.Limits.TimeoutMs, + wireReq := remoteQueryCopyExecutorRequestJSON{ + Operation: req.Operation, + Target: remoteQueryTargetJSON{Host: req.Target.Host, Port: req.Target.Port, DBName: req.Target.DBName}, + Query: req.Query, + Format: format, + } + if req.CopyLimits != nil { + wireReq.Limits = &remoteQueryExecuteCopyLimitsJSON{ + ChunkBytes: req.CopyLimits.ChunkBytes, + MaxBytes: req.CopyLimits.MaxBytes, + MaxRowBytes: req.CopyLimits.MaxRowBytes, + TimeoutMs: req.CopyLimits.TimeoutMs, } } - requestJSON, err := json.Marshal(wireReq) if err != nil { return "", err diff --git a/comp/remotequeries/impl/remote_query_par_poc.go b/comp/remotequeries/impl/remote_query_par_poc.go index 1844a2aca334..f5993caf186b 100644 --- a/comp/remotequeries/impl/remote_query_par_poc.go +++ b/comp/remotequeries/impl/remote_query_par_poc.go @@ -37,10 +37,12 @@ type RemoteQueryPARHarness struct { // RemoteQueryPARInputs is the credential-free task input shape for the PAR-shaped POC harness. type RemoteQueryPARInputs struct { - Integration string `json:"integration"` - Target remoteQueryTargetJSON `json:"target"` - Query string `json:"query"` - Limits *remoteQueryExecuteLimitsJSON `json:"limits,omitempty"` + Integration string `json:"integration"` + Operation string `json:"operation"` + Format string `json:"format"` + Target remoteQueryTargetJSON `json:"target"` + Query string `json:"query"` + CopyLimits *remoteQueryExecuteCopyLimitsJSON `json:"copyLimits,omitempty"` } // RemoteQueryPARResult is the decoded execute bridge result or sanitized bridge error. diff --git a/comp/remotequeries/impl/remote_query_par_poc_test.go b/comp/remotequeries/impl/remote_query_par_poc_test.go index f4269b4f8d41..5ec718fd2268 100644 --- a/comp/remotequeries/impl/remote_query_par_poc_test.go +++ b/comp/remotequeries/impl/remote_query_par_poc_test.go @@ -25,27 +25,25 @@ func TestRemoteQueryPARHarnessUsesCredentialFreeIPCPostShape(t *testing.T) { result, err := harness.Execute(context.Background(), RemoteQueryPARInputs{ Integration: "postgres", + Operation: "copy_stream", + Format: "csv", Target: remoteQueryTargetJSON{Host: "localhost", Port: 5432, DBName: "postgres"}, Query: remoteQueryProofSeedQuery, - Limits: &remoteQueryExecuteLimitsJSON{MaxRows: 1, MaxBytes: 1024, TimeoutMs: 1000}, + CopyLimits: &remoteQueryExecuteCopyLimitsJSON{ChunkBytes: 1024, MaxBytes: 1024, MaxRowBytes: 1024, TimeoutMs: 1000}, }) require.NoError(t, err) assert.Equal(t, "https://localhost:5001"+AgentRemoteQueryExecuteEndpointPath, client.url) assert.Equal(t, "application/json", client.contentType) - assert.JSONEq(t, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":1,"maxBytes":1024,"timeoutMs":1000}}`, client.body) + assert.JSONEq(t, `{"integration":"postgres","operation":"copy_stream","format":"csv","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","copyLimits":{"chunkBytes":1024,"maxBytes":1024,"maxRowBytes":1024,"timeoutMs":1000}}`, client.body) assert.NotContains(t, client.body, "password") assert.NotContains(t, client.body, "secret") assert.Equal(t, "SUCCEEDED", result.Status) assert.JSONEq(t, `{"status":"SUCCEEDED","rows":[{"value":1}]}`, string(result.Raw)) } -func TestRemoteQueryPARHarnessWithRealAgentIPCClient(t *testing.T) { - runner := &fakeRunnerCheck{ - fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: datastore-secret\n"}, - response: `{"status":"SUCCEEDED","rows":[{"value":1}]}`, - } - handler := &remoteQueryExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{fakeWrappedCheck{Check: runner}}}} +func TestRemoteQueryPARHarnessWithRealAgentIPCClientRejectsHTTPExecution(t *testing.T) { + handler := &remoteQueryExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{fakeWrappedCheck{Check: &fakeRunnerCheck{fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: datastore-secret\n"}}}}}} ipc := ipcmock.New(t) mux := http.NewServeMux() mux.HandleFunc(AgentRemoteQueryExecuteEndpointPath, handler.handle) @@ -54,15 +52,16 @@ func TestRemoteQueryPARHarnessWithRealAgentIPCClient(t *testing.T) { result, err := harness.Execute(context.Background(), RemoteQueryPARInputs{ Integration: "postgres", + Operation: "copy_stream", + Format: "csv", Target: remoteQueryTargetJSON{Host: "LOCALHOST.", Port: 5432, DBName: "postgres"}, Query: remoteQueryProofSeedQuery, }) require.NoError(t, err) - assert.Equal(t, "SUCCEEDED", result.Status) - assert.JSONEq(t, `{"status":"SUCCEEDED","rows":[{"value":1}]}`, string(result.Raw)) - assert.JSONEq(t, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`, runner.seenRequest()) - assert.NotContains(t, runner.seenRequest(), "integration") + assert.Equal(t, statusInvalidRequest, result.Status) + require.NotNil(t, result.Error) + assert.Contains(t, result.Error.Message, "streaming executor") assert.NotContains(t, string(result.Raw), "datastore-secret") } @@ -86,16 +85,20 @@ func TestRemoteQueryPARHarnessPropagatesSanitizedBridgeErrors(t *testing.T) { name: "target not found", inputs: RemoteQueryPARInputs{ Integration: "postgres", + Operation: "copy_stream", + Format: "csv", Target: remoteQueryTargetJSON{Host: "localhost", Port: 5432, DBName: "other"}, Query: remoteQueryProofSeedQuery, }, - wantStatus: statusTargetNotFound, - wantCode: statusTargetNotFound, + wantStatus: statusInvalidRequest, + wantCode: statusInvalidRequest, }, { name: "invalid query", inputs: RemoteQueryPARInputs{ Integration: "postgres", + Operation: "copy_stream", + Format: "csv", Target: remoteQueryTargetJSON{Host: "localhost", Port: 5432, DBName: "postgres"}, Query: "SELECT 2 AS value", }, diff --git a/comp/remotequeries/impl/remote_query_test.go b/comp/remotequeries/impl/remote_query_test.go index ac990f13a25c..b50496af3051 100644 --- a/comp/remotequeries/impl/remote_query_test.go +++ b/comp/remotequeries/impl/remote_query_test.go @@ -6,7 +6,6 @@ package remotequeriesimpl import ( - "fmt" "net/http" "net/http/httptest" "strings" @@ -328,17 +327,19 @@ func TestParseExecuteRequestValidatesStrictShape(t *testing.T) { } } -func TestParseExecuteRequestNormalizesAndMarshalsExecutorJSON(t *testing.T) { +func TestParseExecuteRequestNormalizesAndMarshalsCopyStreamExecutorJSON(t *testing.T) { req := httptest.NewRequest(http.MethodPost, RemoteQueryExecuteEndpointPath, strings.NewReader( - `{"integration":"postgres","target":{"host":" LocalHost. ","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":10,"maxBytes":1048576,"timeoutMs":5000}}`, + `{"integration":"postgres","operation":"copy_stream","format":"csv","target":{"host":" LocalHost. ","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","copyLimits":{"chunkBytes":1024,"maxBytes":1048576,"maxRowBytes":1048576,"timeoutMs":5000}}`, )) req.Header.Set("Content-Type", "application/json; charset=utf-8") parsed, requestJSON, err := parseExecuteRequest(req) require.NoError(t, err) assert.Equal(t, "postgres", parsed.Integration) + assert.Equal(t, "copy_stream", parsed.Operation) + assert.Equal(t, "csv", parsed.Format) assert.Equal(t, remoteQueryTarget{Host: "localhost", Port: 5432, DBName: "postgres"}, parsed.Target) - assert.JSONEq(t, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","limits":{"maxRows":10,"maxBytes":1048576,"timeoutMs":5000}}`, requestJSON) + assert.JSONEq(t, `{"operation":"copy_stream","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","format":"csv","limits":{"chunkBytes":1024,"maxBytes":1048576,"maxRowBytes":1048576,"timeoutMs":5000}}`, requestJSON) assert.NotContains(t, requestJSON, "integration") } @@ -353,48 +354,34 @@ func TestParseExecuteRequestRejectsInvalidIntegration(t *testing.T) { assert.Equal(t, "integration contains invalid characters", err.Error()) } -func TestParseExecuteRequestAllowsOmittedLimits(t *testing.T) { +func TestParseExecuteRequestRejectsOmittedOperation(t *testing.T) { req := httptest.NewRequest(http.MethodPost, RemoteQueryExecuteEndpointPath, strings.NewReader( `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`, )) req.Header.Set("Content-Type", "application/json") - parsed, requestJSON, err := parseExecuteRequest(req) - require.NoError(t, err) - assert.Nil(t, parsed.Limits) - assert.JSONEq(t, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`, requestJSON) - assert.NotContains(t, requestJSON, "integration") + _, _, err := parseExecuteRequest(req) + require.Error(t, err) + assert.Equal(t, "operation must be copy_stream", err.Error()) } func TestParseExecuteRequestAllowsFixtureTableProofQuery(t *testing.T) { req := httptest.NewRequest(http.MethodPost, RemoteQueryExecuteEndpointPath, strings.NewReader( - `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT city, country FROM cities ORDER BY city","limits":{"maxRows":2,"maxBytes":1024,"timeoutMs":1000}}`, + `{"integration":"postgres","operation":"copy_stream","format":"csv","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT city, country FROM cities ORDER BY city"}`, )) req.Header.Set("Content-Type", "application/json") parsed, requestJSON, err := parseExecuteRequest(req) require.NoError(t, err) assert.Equal(t, remoteQueryFixtureTableProofQuery, parsed.Query) - assert.JSONEq(t, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT city, country FROM cities ORDER BY city","limits":{"maxRows":2,"maxBytes":1024,"timeoutMs":1000}}`, requestJSON) + assert.JSONEq(t, `{"operation":"copy_stream","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT city, country FROM cities ORDER BY city","format":"csv"}`, requestJSON) assert.NotContains(t, requestJSON, "integration") } -func TestNewRemoteQueryExecuteRequestAllowsFixtureTableProofQuery(t *testing.T) { - req, err := NewRemoteQueryExecuteRequest("postgres", RemoteQueryExecuteTarget{Host: " LocalHost. ", Port: 5432, DBName: "postgres"}, remoteQueryFixtureTableProofQuery, &RemoteQueryExecuteLimits{MaxRows: 2, MaxBytes: 1024, TimeoutMs: 1000}) - require.NoError(t, err) - assert.Equal(t, "postgres", req.Integration) - assert.Equal(t, RemoteQueryExecuteTarget{Host: "localhost", Port: 5432, DBName: "postgres"}, req.Target) - assert.Equal(t, remoteQueryFixtureTableProofQuery, req.Query) -} - -func TestNewRemoteQueryExecuteRequestAllowsLargePayloadProofQueries(t *testing.T) { - for query, payloadBytes := range remoteQueryLargePayloadProofQueries { - t.Run(fmt.Sprintf("%d", payloadBytes), func(t *testing.T) { - req, err := NewRemoteQueryExecuteRequest("postgres", RemoteQueryExecuteTarget{Host: "localhost", Port: 5432, DBName: "postgres"}, query, &RemoteQueryExecuteLimits{MaxRows: 1, MaxBytes: payloadBytes + (1 << 20), TimeoutMs: 60_000}) - require.NoError(t, err) - assert.Equal(t, query, req.Query) - }) - } +func TestNewRemoteQueryExecuteRequestRejectsInlineMode(t *testing.T) { + _, err := NewRemoteQueryExecuteRequest("postgres", RemoteQueryExecuteTarget{Host: " LocalHost. ", Port: 5432, DBName: "postgres"}, remoteQueryFixtureTableProofQuery, &RemoteQueryExecuteLimits{MaxRows: 2, MaxBytes: 1024, TimeoutMs: 1000}) + require.Error(t, err) + assert.EqualError(t, err, "operation must be copy_stream") } func TestRemoteQueryExecuteHandlerDisabled(t *testing.T) { @@ -406,35 +393,13 @@ func TestRemoteQueryExecuteHandlerDisabled(t *testing.T) { assert.Contains(t, recorder.Body.String(), `"status":"bridge_disabled"`) } -func TestRemoteQueryExecuteHandlerRunnerSuccess(t *testing.T) { - runner := &fakeRunnerCheck{ - fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-value\n"}, - response: `{"status":"SUCCEEDED","rows":[{"value":1}]}`, - } - handler := &remoteQueryExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{fakeWrappedCheck{Check: runner}}}} +func TestRemoteQueryExecuteHandlerRejectsInlineHTTPExecution(t *testing.T) { + handler := &remoteQueryExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{fakeWrappedCheck{Check: &fakeRunnerCheck{fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-value\n"}}}}}} - recorder := callExecuteHandler(handler, `{"integration":"postgres","target":{"host":"LOCALHOST.","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) + recorder := callExecuteHandler(handler, `{"integration":"postgres","operation":"copy_stream","format":"csv","target":{"host":"LOCALHOST.","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) - assert.Equal(t, http.StatusOK, recorder.Code) - assert.JSONEq(t, `{"status":"SUCCEEDED","rows":[{"value":1}]}`, recorder.Body.String()) - assert.JSONEq(t, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`, runner.seenRequest()) - assert.NotContains(t, runner.seenRequest(), "integration") - assert.NotContains(t, recorder.Body.String(), "secret-value") -} - -func TestRemoteQueryExecuteHandlerRunnerSuccessWithFixtureTableQuery(t *testing.T) { - runner := &fakeRunnerCheck{ - fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-value\n"}, - response: `{"status":"SUCCEEDED","rows":[{"city":"Beautiful city of lights","country":"France"},{"city":"New York","country":"USA"}]}`, - } - handler := &remoteQueryExecuteHandler{enabled: true, collector: fakeCollector{checks: []check.Check{fakeWrappedCheck{Check: runner}}}} - - recorder := callExecuteHandler(handler, `{"integration":"postgres","target":{"host":"LOCALHOST.","port":5432,"dbname":"postgres"},"query":"SELECT city, country FROM cities ORDER BY city"}`) - - assert.Equal(t, http.StatusOK, recorder.Code) - assert.JSONEq(t, `{"status":"SUCCEEDED","rows":[{"city":"Beautiful city of lights","country":"France"},{"city":"New York","country":"USA"}]}`, recorder.Body.String()) - assert.JSONEq(t, `{"target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT city, country FROM cities ORDER BY city"}`, runner.seenRequest()) - assert.NotContains(t, runner.seenRequest(), "integration") + assert.Equal(t, http.StatusBadRequest, recorder.Code) + assert.Contains(t, recorder.Body.String(), "streaming executor") assert.NotContains(t, recorder.Body.String(), "secret-value") } @@ -460,7 +425,6 @@ func TestRemoteQueryExecuteServiceCopyStreamDispatch(t *testing.T) { require.Nil(t, result.Error) assert.Equal(t, runner.events, events) assert.Equal(t, 1, runner.streamCalls) - assert.Equal(t, 0, runner.jsonCalls) assert.JSONEq(t, `{"operation":"copy_stream","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT city, country FROM cities ORDER BY city","format":"csv","limits":{"chunkBytes":4,"maxBytes":1024,"maxRowBytes":1024,"timeoutMs":1000}}`, runner.streamSeen) assert.NotContains(t, runner.streamSeen, "integration") } @@ -482,10 +446,10 @@ func TestRemoteQueryExecuteHandlerNoMatchAndAmbiguous(t *testing.T) { &fakeRunnerCheck{fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-value\n"}}, }}} - recorder := callExecuteHandler(handler, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"other"},"query":"SELECT 1 AS value"}`) + recorder := callExecuteHandler(handler, `{"integration":"postgres","operation":"copy_stream","format":"csv","target":{"host":"localhost","port":5432,"dbname":"other"},"query":"SELECT 1 AS value"}`) - assert.Equal(t, http.StatusNotFound, recorder.Code) - assert.Contains(t, recorder.Body.String(), `"status":"target_not_found"`) + assert.Equal(t, http.StatusBadRequest, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"status":"invalid_request"`) assert.NotContains(t, recorder.Body.String(), "secret-value") assert.NotContains(t, recorder.Body.String(), "other") }) @@ -496,10 +460,10 @@ func TestRemoteQueryExecuteHandlerNoMatchAndAmbiguous(t *testing.T) { &fakeRunnerCheck{fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-two\n"}}, }}} - recorder := callExecuteHandler(handler, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) + recorder := callExecuteHandler(handler, `{"integration":"postgres","operation":"copy_stream","format":"csv","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) - assert.Equal(t, http.StatusConflict, recorder.Code) - assert.Contains(t, recorder.Body.String(), `"status":"ambiguous_target"`) + assert.Equal(t, http.StatusBadRequest, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"status":"invalid_request"`) assert.NotContains(t, recorder.Body.String(), "secret-one") assert.NotContains(t, recorder.Body.String(), "secret-two") }) @@ -511,10 +475,10 @@ func TestRemoteQueryExecuteHandlerUnsupportedAndRunnerErrorAreSanitized(t *testi fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-value\n"}, }}} - recorder := callExecuteHandler(handler, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) + recorder := callExecuteHandler(handler, `{"integration":"postgres","operation":"copy_stream","format":"csv","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) - assert.Equal(t, http.StatusFailedDependency, recorder.Code) - assert.Contains(t, recorder.Body.String(), `"status":"executor_unavailable"`) + assert.Equal(t, http.StatusBadRequest, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"status":"invalid_request"`) assert.NotContains(t, recorder.Body.String(), "secret-value") }) @@ -523,10 +487,10 @@ func TestRemoteQueryExecuteHandlerUnsupportedAndRunnerErrorAreSanitized(t *testi &fakeRunnerCheck{fakeCheck: fakeCheck{name: "postgres", loader: "python", provider: "file", instance: "host: localhost\nport: 5432\ndbname: postgres\npassword: secret-value\n"}, err: assert.AnError}, }}} - recorder := callExecuteHandler(handler, `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) + recorder := callExecuteHandler(handler, `{"integration":"postgres","operation":"copy_stream","format":"csv","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) - assert.Equal(t, http.StatusBadGateway, recorder.Code) - assert.Contains(t, recorder.Body.String(), `"status":"executor_unavailable"`) + assert.Equal(t, http.StatusBadRequest, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"status":"invalid_request"`) assert.NotContains(t, recorder.Body.String(), "secret-value") assert.NotContains(t, recorder.Body.String(), assert.AnError.Error()) }) @@ -550,23 +514,7 @@ func (f fakeWrappedCheck) Unwrap() check.Check { type fakeRunnerCheck struct { fakeCheck - response string - err error - seen string - jsonCalls int -} - -func (f *fakeRunnerCheck) RunRemoteQueryJSON(integration string, requestJSON string) (string, error) { - if integration != "postgres" { - return "", assert.AnError - } - f.jsonCalls++ - f.seen = requestJSON - return f.response, f.err -} - -func (f *fakeRunnerCheck) seenRequest() string { - return f.seen + err error } type fakeStreamRunnerCheck struct { diff --git a/pkg/collector/python/check.go b/pkg/collector/python/check.go index 1e26582238ec..271c7c830730 100644 --- a/pkg/collector/python/check.go +++ b/pkg/collector/python/check.go @@ -41,7 +41,6 @@ import ( #include "rtloader_mem.h" char *getStringAddr(char **array, unsigned int idx); -char *run_remote_query(rtloader_t *, rtloader_pyobject_t *check, const char *integration, const char *request_json); extern int remoteQueryStreamEmitBridge(const char *event_type, const char *metadata_json, const uint8_t *payload, size_t payload_len, void *userdata); int run_remote_query_stream(rtloader_t *, rtloader_pyobject_t *check, const char *integration, const char *request_json, int (*emit)(const char *, const char *, const uint8_t *, size_t, void *), void *userdata); @@ -166,40 +165,6 @@ func (c *PythonCheck) RunSimple() error { return c.runCheck(false) } -// RunRemoteQueryJSON runs a remote query helper for this Python check. -func (c *PythonCheck) RunRemoteQueryJSON(integration string, requestJSON string) (string, error) { - integration = strings.ToLower(strings.TrimSpace(integration)) - if integration == "" { - return "", fmt.Errorf("integration is required") - } - - gstate, err := newStickyLock() - if err != nil { - return "", err - } - defer gstate.unlock() - - if c.cancelled { - return "", fmt.Errorf("check %s is already cancelled", c.ModuleName) - } - - cIntegration := C.CString(integration) - defer C.free(unsafe.Pointer(cIntegration)) - cRequestJSON := C.CString(requestJSON) - defer C.free(unsafe.Pointer(cRequestJSON)) - - cResult := C.run_remote_query(rtloader, c.instance, cIntegration, cRequestJSON) - if cResult == nil { - if err := getRtLoaderError(); err != nil { - return "", err - } - return "", fmt.Errorf("an error occurred while running remote query") - } - defer C.rtloader_free(rtloader, unsafe.Pointer(cResult)) - - return C.GoString(cResult), nil -} - // RunRemoteQueryStream runs a streaming remote query helper for this Python check. func (c *PythonCheck) RunRemoteQueryStream(integration string, requestJSON string, emit func(checkbase.RemoteQueryStreamEvent) error) error { integration = strings.ToLower(strings.TrimSpace(integration)) diff --git a/pkg/collector/python/check_test.go b/pkg/collector/python/check_test.go index a1ba6a2bfb41..5198a2fd425b 100644 --- a/pkg/collector/python/check_test.go +++ b/pkg/collector/python/check_test.go @@ -71,26 +71,6 @@ func TestRunAfterCancel(t *testing.T) { testRunAfterCancel(t) } -func TestRunRemoteQueryJSON(t *testing.T) { - testRunRemoteQueryJSON(t) -} - -func TestRunRemoteQueryJSONNormalizesIntegration(t *testing.T) { - testRunRemoteQueryJSONNormalizesIntegration(t) -} - -func TestRunRemoteQueryJSONError(t *testing.T) { - testRunRemoteQueryJSONError(t) -} - -func TestRunRemoteQueryJSONWithRuntimeNotInitializedError(t *testing.T) { - testRunRemoteQueryJSONWithRuntimeNotInitializedError(t) -} - -func TestRunRemoteQueryJSONAfterCancel(t *testing.T) { - testRunRemoteQueryJSONAfterCancel(t) -} - func TestRunRemoteQueryStream(t *testing.T) { testRunRemoteQueryStream(t) } diff --git a/pkg/collector/python/test_check.go b/pkg/collector/python/test_check.go index 7d03e7b46527..1fa6f167fc6e 100644 --- a/pkg/collector/python/test_check.go +++ b/pkg/collector/python/test_check.go @@ -105,19 +105,6 @@ char *get_check_diagnoses(rtloader_t *s, rtloader_pyobject_t *check) { return get_check_diagnoses_return; } -char *run_remote_query_return = NULL; -int run_remote_query_calls = 0; -rtloader_pyobject_t *run_remote_query_instance = NULL; -const char *run_remote_query_integration = NULL; -const char *run_remote_query_request_json = NULL; -char *run_remote_query(rtloader_t *s, rtloader_pyobject_t *check, const char *integration, const char *request_json) { - run_remote_query_instance = check; - run_remote_query_integration = strdup(integration); - run_remote_query_request_json = strdup(request_json); - run_remote_query_calls++; - return run_remote_query_return; -} - int run_remote_query_stream_return = 1; int run_remote_query_stream_calls = 0; rtloader_pyobject_t *run_remote_query_stream_instance = NULL; @@ -241,11 +228,6 @@ void reset_check_mock() { get_check_diagnoses_return = NULL; get_check_diagnoses_calls = 0; - run_remote_query_return = NULL; - run_remote_query_calls = 0; - run_remote_query_instance = NULL; - run_remote_query_integration = NULL; - run_remote_query_request_json = NULL; run_remote_query_stream_return = 1; run_remote_query_stream_calls = 0; run_remote_query_stream_instance = NULL; @@ -717,99 +699,6 @@ func testGetDiagnoses(t *testing.T) { assert.Zero(t, len(diagnoses[1].Remediation)) } -func testRunRemoteQueryJSON(t *testing.T) { - mockRtloader(t) - - check, err := NewPythonFakeCheck(aggregator.NewNoOpSenderManager()) - require.NoError(t, err) - check.instance = newMockPyObjectPtr() - - C.reset_check_mock() - C.run_remote_query_return = C.CString(`{"status":"SUCCEEDED"}`) - - result, err := check.RunRemoteQueryJSON("postgres", `{"integration":"postgres","query":"SELECT 1 AS value"}`) - - require.NoError(t, err) - assert.JSONEq(t, `{"status":"SUCCEEDED"}`, result) - assert.Equal(t, C.int(1), C.gil_locked_calls) - assert.Equal(t, C.int(1), C.gil_unlocked_calls) - assert.Equal(t, C.int(1), C.run_remote_query_calls) - assert.Equal(t, C.int(1), C.rtloader_free_calls) - assert.Equal(t, check.instance, C.run_remote_query_instance) - assert.Equal(t, "postgres", C.GoString(C.run_remote_query_integration)) - assert.JSONEq(t, `{"integration":"postgres","query":"SELECT 1 AS value"}`, C.GoString(C.run_remote_query_request_json)) -} - -func testRunRemoteQueryJSONNormalizesIntegration(t *testing.T) { - mockRtloader(t) - - check, err := NewPythonFakeCheck(aggregator.NewNoOpSenderManager()) - require.NoError(t, err) - check.instance = newMockPyObjectPtr() - - C.reset_check_mock() - C.run_remote_query_return = C.CString(`{"status":"SUCCEEDED"}`) - - result, err := check.RunRemoteQueryJSON(" MySQL ", `{"integration":"mysql","query":"SELECT 1"}`) - - require.NoError(t, err) - assert.JSONEq(t, `{"status":"SUCCEEDED"}`, result) - assert.Equal(t, C.int(1), C.run_remote_query_calls) - assert.Equal(t, "mysql", C.GoString(C.run_remote_query_integration)) -} - -func testRunRemoteQueryJSONError(t *testing.T) { - mockRtloader(t) - - check, err := NewPythonFakeCheck(aggregator.NewNoOpSenderManager()) - require.NoError(t, err) - check.instance = newMockPyObjectPtr() - - C.reset_check_mock() - C.run_remote_query_return = nil - C.has_error_return = 1 - C.get_error_return = C.CString("rtloader helper failed") - - result, err := check.RunRemoteQueryJSON("postgres", `{"query":"SELECT 1 AS value"}`) - - assert.Empty(t, result) - require.Error(t, err) - assert.EqualError(t, err, "rtloader helper failed") - assert.Equal(t, C.int(1), C.gil_locked_calls) - assert.Equal(t, C.int(1), C.gil_unlocked_calls) - assert.Equal(t, C.int(1), C.run_remote_query_calls) - assert.Equal(t, C.int(0), C.rtloader_free_calls) -} - -func testRunRemoteQueryJSONWithRuntimeNotInitializedError(t *testing.T) { - mockRtloader(t) - check, err := NewPythonFakeCheck(aggregator.NewNoOpSenderManager()) - require.NoError(t, err) - check.instance = newMockPyObjectPtr() - - C.reset_check_mock() - rtloader = nil - - _, err = check.RunRemoteQueryJSON("postgres", `{"query":"SELECT 1 AS value"}`) - assert.ErrorIs(t, err, ErrNotInitialized) - assert.Equal(t, C.int(0), C.run_remote_query_calls) -} - -func testRunRemoteQueryJSONAfterCancel(t *testing.T) { - mockRtloader(t) - - check, err := NewPythonFakeCheck(aggregator.NewNoOpSenderManager()) - require.NoError(t, err) - check.instance = newMockPyObjectPtr() - - C.reset_check_mock() - check.Cancel() - - _, err = check.RunRemoteQueryJSON("postgres", `{"query":"SELECT 1 AS value"}`) - assert.EqualError(t, err, "check fake_check is already cancelled") - assert.Equal(t, C.int(0), C.run_remote_query_calls) -} - func testRunRemoteQueryStream(t *testing.T) { mockRtloader(t) diff --git a/pkg/privateactionrunner/bundles/remotequeries/execute.go b/pkg/privateactionrunner/bundles/remotequeries/execute.go index 95c6eddf3fbc..16af35ba841a 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/execute.go +++ b/pkg/privateactionrunner/bundles/remotequeries/execute.go @@ -137,12 +137,9 @@ func remoteQueryExecuteOutputFromStream(stream grpc.ServerStreamingClient[pb.Rem return nil, fmt.Errorf("remote query response stream missing") } - var assembled bytes.Buffer - legacyStreamEvents := make([]json.RawMessage, 0) typedStreamEvents := make([]map[string]interface{}, 0) expectedChunkIndex := int32(0) seenFinal := false - finalChunkWasEmpty := false for { chunk, err := stream.Recv() if err == io.EOF { @@ -166,16 +163,8 @@ func remoteQueryExecuteOutputFromStream(stream grpc.ServerStreamingClient[pb.Rem return nil, err } typedStreamEvents = append(typedStreamEvents, streamEvent) - } else { - payload := chunk.GetResponseJsonChunk() - if chunk.GetFinal() && len(payload) == 0 && len(legacyStreamEvents) > 0 { - finalChunkWasEmpty = true - } else { - _, _ = assembled.Write(payload) - if !chunk.GetFinal() { - legacyStreamEvents = append(legacyStreamEvents, append(json.RawMessage(nil), payload...)) - } - } + } else if !chunk.GetFinal() { + return nil, fmt.Errorf("remote query response stream chunk missing typed event") } seenFinal = chunk.GetFinal() expectedChunkIndex++ @@ -183,13 +172,10 @@ func remoteQueryExecuteOutputFromStream(stream grpc.ServerStreamingClient[pb.Rem if !seenFinal { return nil, fmt.Errorf("remote query response stream missing final chunk") } - if len(typedStreamEvents) > 0 { - return remoteQueryExecuteOutputFromTypedEvents(typedStreamEvents) - } - if finalChunkWasEmpty { - return remoteQueryExecuteOutputFromEvents(legacyStreamEvents) + if len(typedStreamEvents) == 0 { + return nil, fmt.Errorf("remote query response stream missing typed events") } - return remoteQueryExecuteOutputFromJSON(assembled.Bytes()) + return remoteQueryExecuteOutputFromTypedEvents(typedStreamEvents) } func remoteQueryStreamEventFromProto(event *pb.RemoteQueryExecuteStreamEvent) (map[string]interface{}, error) { @@ -236,6 +222,7 @@ func remoteQueryStreamEventFromProto(event *pb.RemoteQueryExecuteStreamEvent) (m func remoteQueryExecuteOutputFromTypedEvents(events []map[string]interface{}) (map[string]interface{}, error) { var finalEvent map[string]interface{} + var errorEvent map[string]interface{} var data bytes.Buffer for _, event := range events { if event["type"] == "data" { @@ -246,8 +233,19 @@ func remoteQueryExecuteOutputFromTypedEvents(events []map[string]interface{}) (m if event["type"] == "final" { finalEvent = event } + if event["type"] == "error" { + errorEvent = event + } } if finalEvent == nil { + if errorEvent != nil { + code, _ := errorEvent["code"].(string) + message, _ := errorEvent["message"].(string) + return map[string]interface{}{ + "status": code, + "error": map[string]interface{}{"code": code, "message": message}, + }, nil + } return nil, fmt.Errorf("remote query stream response missing final event") } status, _ := finalEvent["status"].(string) @@ -266,59 +264,6 @@ func remoteQueryExecuteOutputFromTypedEvents(events []map[string]interface{}) (m return output, nil } -func remoteQueryExecuteOutputFromEvents(rawEvents []json.RawMessage) (map[string]interface{}, error) { - events := make([]interface{}, 0, len(rawEvents)) - var finalEvent map[string]interface{} - var data bytes.Buffer - for _, raw := range rawEvents { - var event map[string]interface{} - if err := json.Unmarshal(raw, &event); err != nil { - return nil, err - } - events = append(events, normalizeRemoteQueryOutput(event)) - if event["type"] == "data" { - if chunk, ok := event["data"].(string); ok { - _, _ = data.WriteString(chunk) - } - } - if event["type"] == "final" { - finalEvent = event - } - } - if finalEvent == nil { - return nil, fmt.Errorf("remote query stream response missing final event") - } - status, _ := finalEvent["status"].(string) - if status == "" { - return nil, fmt.Errorf("remote query stream final event missing status") - } - output := map[string]interface{}{ - "status": status, - "events": events, - "data": data.String(), - } - if v, ok := finalEvent["error"]; ok { - output["error"] = normalizeRemoteQueryOutput(v) - } - if v, ok := finalEvent["stats"]; ok { - output["stats"] = normalizeRemoteQueryOutput(v) - } - return output, nil -} - -func remoteQueryExecuteOutputFromJSON(responseJSON []byte) (map[string]interface{}, error) { - var output map[string]interface{} - decoder := json.NewDecoder(bytes.NewReader(responseJSON)) - if err := decoder.Decode(&output); err != nil { - return nil, err - } - status, ok := output["status"].(string) - if !ok || status == "" { - return nil, fmt.Errorf("remote query response missing status") - } - return normalizeRemoteQueryOutput(output).(map[string]interface{}), nil -} - func remoteQueryExecuteOutputFromProto(resp *pb.RemoteQueryExecuteResponse) (map[string]interface{}, error) { if resp == nil || resp.GetStatus() == "" { return nil, fmt.Errorf("remote query response missing status") diff --git a/pkg/privateactionrunner/bundles/remotequeries/execute_test.go b/pkg/privateactionrunner/bundles/remotequeries/execute_test.go index 65202b887d26..4e8d0ef77856 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/execute_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/execute_test.go @@ -11,60 +11,61 @@ import ( "io" "testing" - "google.golang.org/grpc" - "google.golang.org/protobuf/types/known/structpb" - "github.com/DataDog/datadog-agent/pkg/privateactionrunner/libs/privateconnection" "github.com/DataDog/datadog-agent/pkg/privateactionrunner/types" "github.com/DataDog/datadog-agent/pkg/privateactionrunner/util" pb "github.com/DataDog/datadog-agent/pkg/proto/pbgo/core" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "google.golang.org/grpc" ) func TestExecuteActionUsesCredentialFreeAgentSecureRequestShape(t *testing.T) { - franceRow, err := structpb.NewStruct(map[string]interface{}{"city": "Beautiful city of lights", "country": "France"}) - require.NoError(t, err) - usaRow, err := structpb.NewStruct(map[string]interface{}{"city": "New York", "country": "USA"}) - require.NoError(t, err) - client := &captureBridgeClient{response: &pb.RemoteQueryExecuteResponse{Status: "SUCCEEDED", Rows: []*structpb.Struct{franceRow, usaRow}}} + client := &captureBridgeClient{chunks: []*pb.RemoteQueryExecuteChunk{ + {Event: &pb.RemoteQueryExecuteStreamEvent{Sequence: 0, Event: &pb.RemoteQueryExecuteStreamEvent_Metadata{Metadata: &pb.RemoteQueryStreamMetadata{Operation: "copy_stream", Integration: "postgres", Format: "csv"}}}, ChunkIndex: 0}, + {Event: &pb.RemoteQueryExecuteStreamEvent{Sequence: 1, Event: &pb.RemoteQueryExecuteStreamEvent_Data{Data: &pb.RemoteQueryStreamData{Payload: []byte("Beautiful city of lights,France\nNew York,USA\n"), Offset: 0, Bytes: 42}}}, ChunkIndex: 1}, + {Event: &pb.RemoteQueryExecuteStreamEvent{Sequence: 2, Event: &pb.RemoteQueryExecuteStreamEvent_Final{Final: &pb.RemoteQueryStreamFinal{Status: "SUCCEEDED", BytesEmitted: 42, ChunksEmitted: 1}}}, ChunkIndex: 2}, + {ChunkIndex: 3, Final: true}, + }} action := NewExecuteAction(func() (BridgeClient, error) { return client, nil }) output, err := action.Run(context.Background(), taskWithInputs(map[string]interface{}{ "integration": "postgres", + "operation": "copy_stream", + "format": "csv", "target": map[string]interface{}{ "host": "localhost", "port": 5432, "dbname": "postgres", }, "query": "SELECT city, country FROM cities ORDER BY city", - "limits": map[string]interface{}{ - "maxRows": 2, - "maxBytes": 1024, - "timeoutMs": 1000, + "copyLimits": map[string]interface{}{ + "chunkBytes": 1024, + "maxBytes": 1024, + "maxRowBytes": 1024, + "timeoutMs": 1000, }, }), &privateconnection.PrivateCredentials{Tokens: []privateconnection.PrivateCredentialsToken{{Name: "password", Value: "secret-value"}}}) require.NoError(t, err) require.NotNil(t, client.request) assert.Equal(t, "postgres", client.request.GetIntegration()) + assert.Equal(t, "copy_stream", client.request.GetOperation()) + assert.Equal(t, "csv", client.request.GetFormat()) assert.Equal(t, "localhost", client.request.GetTarget().GetHost()) assert.Equal(t, int32(5432), client.request.GetTarget().GetPort()) assert.Equal(t, "postgres", client.request.GetTarget().GetDbname()) assert.Equal(t, "SELECT city, country FROM cities ORDER BY city", client.request.GetQuery()) - assert.Equal(t, int32(2), client.request.GetLimits().GetMaxRows()) + assert.Equal(t, int32(1024), client.request.GetCopyLimits().GetChunkBytes()) requestEvidence, err := json.Marshal(client.request) require.NoError(t, err) assert.NotContains(t, string(requestEvidence), "secret-value") - assert.Equal(t, map[string]interface{}{ - "status": "SUCCEEDED", - "rows": []interface{}{ - map[string]interface{}{"city": "Beautiful city of lights", "country": "France"}, - map[string]interface{}{"city": "New York", "country": "USA"}, - }, - }, output) + out, ok := output.(map[string]interface{}) + require.True(t, ok) + assert.Equal(t, "SUCCEEDED", out["status"]) + assert.Equal(t, "Beautiful city of lights,France\nNew York,USA\n", out["data"]) } func TestExecuteActionPreservesCopyStreamEvents(t *testing.T) { @@ -120,9 +121,10 @@ func TestExecuteActionPreservesBinaryCopyStreamPayload(t *testing.T) { } func TestExecuteActionPreservesSanitizedBridgeErrorBody(t *testing.T) { - client := &captureBridgeClient{ - response: &pb.RemoteQueryExecuteResponse{Status: "target_not_found", Error: &pb.RemoteQueryExecuteError{Code: "target_not_found", Message: "no matching integration check found"}}, - } + client := &captureBridgeClient{chunks: []*pb.RemoteQueryExecuteChunk{ + {Event: &pb.RemoteQueryExecuteStreamEvent{Sequence: 0, Event: &pb.RemoteQueryExecuteStreamEvent_Error{Error: &pb.RemoteQueryStreamError{Code: "target_not_found", Message: "no matching integration check found"}}}, ChunkIndex: 0}, + {ChunkIndex: 1, Final: true}, + }} action := NewExecuteAction(func() (BridgeClient, error) { return client, nil }) @@ -176,15 +178,14 @@ func taskWithInputs(inputs map[string]interface{}) *types.Task { } type captureBridgeClient struct { - request *pb.RemoteQueryExecuteRequest - response *pb.RemoteQueryExecuteResponse - chunks []*pb.RemoteQueryExecuteChunk - err error + request *pb.RemoteQueryExecuteRequest + chunks []*pb.RemoteQueryExecuteChunk + err error } func (c *captureBridgeClient) RemoteQueryExecute(_ context.Context, req *pb.RemoteQueryExecuteRequest, _ ...grpc.CallOption) (*pb.RemoteQueryExecuteResponse, error) { c.request = req - return c.response, c.err + return nil, c.err } func (c *captureBridgeClient) RemoteQueryExecuteStream(_ context.Context, req *pb.RemoteQueryExecuteRequest, _ ...grpc.CallOption) (grpc.ServerStreamingClient[pb.RemoteQueryExecuteChunk], error) { @@ -192,29 +193,7 @@ func (c *captureBridgeClient) RemoteQueryExecuteStream(_ context.Context, req *p if c.err != nil { return nil, c.err } - if c.chunks != nil { - return &captureRemoteQueryExecuteStream{chunks: c.chunks}, nil - } - responseJSON, err := json.Marshal(remoteQueryExecuteOutputFromProtoForTest(c.response)) - if err != nil { - return nil, err - } - return &captureRemoteQueryExecuteStream{chunks: []*pb.RemoteQueryExecuteChunk{{ResponseJsonChunk: responseJSON, Final: true}}}, nil -} - -func remoteQueryExecuteOutputFromProtoForTest(resp *pb.RemoteQueryExecuteResponse) map[string]interface{} { - output := map[string]interface{}{"status": resp.GetStatus()} - if resp.GetError() != nil { - output["error"] = map[string]interface{}{"code": resp.GetError().GetCode(), "message": resp.GetError().GetMessage()} - } - if len(resp.GetRows()) > 0 { - rows := make([]interface{}, 0, len(resp.GetRows())) - for _, row := range resp.GetRows() { - rows = append(rows, row.AsMap()) - } - output["rows"] = rows - } - return output + return &captureRemoteQueryExecuteStream{chunks: c.chunks}, nil } type captureRemoteQueryExecuteStream struct { diff --git a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go index eb9161684996..5e0170e83a55 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go @@ -93,9 +93,11 @@ func TestRemoteQueriesActionRunsThroughLivePARLoopWithRealAgentIPC(t *testing.T) proofQuery := remoteQueriesProofQueryFromEnv() inputs := map[string]interface{}{ "integration": "postgres", + "operation": "copy_stream", + "format": "csv", "target": remoteQueriesPostgresTargetFromEnv(t), "query": proofQuery, - "limits": remoteQueriesProofLimits(proofQuery), + "copyLimits": remoteQueriesProofCopyLimits(proofQuery), } requestEvidence, err := json.Marshal(inputs) require.NoError(t, err) @@ -160,24 +162,6 @@ func remoteQueriesProofQueryFromEnv() string { return remoteQueriesFixtureTableProofQuery } -func remoteQueriesProofLimits(query string) map[string]interface{} { - maxRows := 1 - maxBytes := 4 << 10 - timeoutMs := 5_000 - if query == remoteQueriesFixtureTableProofQuery { - maxRows = 2 - } - if payloadBytes, ok := remoteQueriesLargePayloadBytes(query); ok { - maxBytes = payloadBytes + (1 << 20) - timeoutMs = 60_000 - } - return map[string]interface{}{ - "maxRows": maxRows, - "maxBytes": maxBytes, - "timeoutMs": timeoutMs, - } -} - func remoteQueriesProofCopyLimits(query string) map[string]interface{} { maxBytes := 4 << 10 timeoutMs := 5_000 @@ -229,6 +213,12 @@ func assertRemoteQueriesProofCopyData(t *testing.T, query string, data string) { case remoteQueriesSeedProofQuery: assert.Equal(t, "1\n", data) default: + expectedPayloadBytes, ok := remoteQueriesLargePayloadBytes(query) + if ok { + assert.Len(t, data, expectedPayloadBytes+1) + assert.Equal(t, "\n", data[len(data)-1:]) + return + } require.FailNowf(t, "unsupported COPY proof query", "%s=%q must use a COPY bridge-allowlisted proof query", remoteQueriesProofQueryOverrideEnv, query) } } diff --git a/pkg/privateactionrunner/bundles/remotequeries/live_par_loop_test.go b/pkg/privateactionrunner/bundles/remotequeries/live_par_loop_test.go index fb15dccf49b0..35c1e6f454cb 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/live_par_loop_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/live_par_loop_test.go @@ -78,11 +78,14 @@ func TestRemoteQueriesActionRunsThroughLivePARLoopAndFakeintake(t *testing.T) { "port": 5432, "dbname": "postgres", }, - "query": "SELECT 1 AS value", - "limits": map[string]interface{}{ - "maxRows": 1, - "maxBytes": 1024, - "timeoutMs": 1000, + "operation": "copy_stream", + "format": "csv", + "query": "SELECT 1 AS value", + "copyLimits": map[string]interface{}{ + "chunkBytes": 1024, + "maxBytes": 1024, + "maxRowBytes": 1024, + "timeoutMs": 1000, }, })) @@ -91,7 +94,7 @@ func TestRemoteQueriesActionRunsThroughLivePARLoopAndFakeintake(t *testing.T) { require.True(t, result.Success) require.Equal(t, taskID, result.TaskID) assert.Equal(t, "SUCCEEDED", result.Outputs["status"]) - require.Contains(t, result.Outputs, "rows") + require.Contains(t, result.Outputs, "events") select { case req := <-bridgeRequests: @@ -101,6 +104,8 @@ func TestRemoteQueriesActionRunsThroughLivePARLoopAndFakeintake(t *testing.T) { require.NotContains(t, string(requestEvidence), "token") require.NotContains(t, string(requestEvidence), "secret") assert.Equal(t, "postgres", req.GetIntegration()) + assert.Equal(t, "copy_stream", req.GetOperation()) + assert.Equal(t, "csv", req.GetFormat()) assert.Equal(t, "SELECT 1 AS value", req.GetQuery()) assert.Equal(t, "localhost", req.GetTarget().GetHost()) assert.Equal(t, int32(5432), req.GetTarget().GetPort()) @@ -157,11 +162,12 @@ func (c *captureAgentSecureClient) RemoteQueryExecute(_ context.Context, req *pb func (c *captureAgentSecureClient) RemoteQueryExecuteStream(_ context.Context, req *pb.RemoteQueryExecuteRequest, _ ...grpc.CallOption) (grpc.ServerStreamingClient[pb.RemoteQueryExecuteChunk], error) { c.requests <- req - responseJSON, err := json.Marshal(map[string]interface{}{"status": c.response.GetStatus(), "rows": []interface{}{map[string]interface{}{"value": 1}}}) - if err != nil { - return nil, err - } - return &captureRemoteQueryExecuteStream{chunks: []*pb.RemoteQueryExecuteChunk{{ResponseJsonChunk: responseJSON, Final: true}}}, nil + return &captureRemoteQueryExecuteStream{chunks: []*pb.RemoteQueryExecuteChunk{ + {Event: &pb.RemoteQueryExecuteStreamEvent{Sequence: 0, Event: &pb.RemoteQueryExecuteStreamEvent_Metadata{Metadata: &pb.RemoteQueryStreamMetadata{Operation: "copy_stream", Integration: "postgres", Format: "csv"}}}, ChunkIndex: 0}, + {Event: &pb.RemoteQueryExecuteStreamEvent{Sequence: 1, Event: &pb.RemoteQueryExecuteStreamEvent_Data{Data: &pb.RemoteQueryStreamData{Payload: []byte("1\n"), Offset: 0, Bytes: 2}}}, ChunkIndex: 1}, + {Event: &pb.RemoteQueryExecuteStreamEvent{Sequence: 2, Event: &pb.RemoteQueryExecuteStreamEvent_Final{Final: &pb.RemoteQueryStreamFinal{Status: c.response.GetStatus(), BytesEmitted: 2, ChunksEmitted: 1}}}, ChunkIndex: 2}, + {ChunkIndex: 3, Final: true}, + }}, nil } type captureRemoteQueryExecuteStream struct { diff --git a/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go b/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go index 1b02ea95f1a6..0736e7336740 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go @@ -79,22 +79,17 @@ func TestRemoteQueriesActionRunsThroughStandalonePARProcessWithRealAgentIPC(t *t taskID := fmt.Sprintf("remotequeries-standalone-par-proof-%d", time.Now().UnixNano()) proofQuery := remoteQueriesProofQueryFromEnv() + format := os.Getenv("RQ_REMOTE_FORMAT") + if format == "" { + format = "csv" + } inputs := map[string]interface{}{ "integration": "postgres", + "operation": "copy_stream", + "format": format, "target": remoteQueriesPostgresTargetFromEnv(t), "query": proofQuery, - } - copyStream := os.Getenv("RQ_REMOTE_OPERATION") == "copy_stream" - if copyStream { - inputs["operation"] = "copy_stream" - format := os.Getenv("RQ_REMOTE_FORMAT") - if format == "" { - format = "csv" - } - inputs["format"] = format - inputs["copyLimits"] = remoteQueriesProofCopyLimits(proofQuery) - } else { - inputs["limits"] = remoteQueriesProofLimits(proofQuery) + "copyLimits": remoteQueriesProofCopyLimits(proofQuery), } requestEvidence, err := json.Marshal(inputs) require.NoError(t, err) @@ -120,22 +115,15 @@ func TestRemoteQueriesActionRunsThroughStandalonePARProcessWithRealAgentIPC(t *t require.True(t, result.Success) require.Equal(t, taskID, result.TaskID) assert.Equal(t, "SUCCEEDED", result.Outputs["status"]) - if copyStream { - t.Logf("copy stream PAR outputs: %+v", summarizeRemoteQueriesProofPayload(result.Outputs)) - if os.Getenv("RQ_REMOTE_FORMAT") == "binary" { - dataBytes, ok := result.Outputs["data_bytes"].(string) - require.True(t, ok) - assertRemoteQueriesProofBinaryCopyData(t, proofQuery, dataBytes) - } else { - data, ok := result.Outputs["data"].(string) - require.True(t, ok) - assertRemoteQueriesProofCopyData(t, proofQuery, data) - } + t.Logf("copy stream PAR outputs: %+v", summarizeRemoteQueriesProofPayload(result.Outputs)) + if os.Getenv("RQ_REMOTE_FORMAT") == "binary" { + dataBytes, ok := result.Outputs["data_bytes"].(string) + require.True(t, ok) + assertRemoteQueriesProofBinaryCopyData(t, proofQuery, dataBytes) } else { - require.Contains(t, result.Outputs, "rows") - rows, ok := result.Outputs["rows"].([]interface{}) + data, ok := result.Outputs["data"].(string) require.True(t, ok) - assertRemoteQueriesProofRows(t, proofQuery, rows) + assertRemoteQueriesProofCopyData(t, proofQuery, data) } resultEvidence, err := json.Marshal(summarizeRemoteQueriesProofPayload(result.Outputs)) diff --git a/pkg/proto/datadog/api/v1/api.proto b/pkg/proto/datadog/api/v1/api.proto index 8d4086825d8e..34d2cdad9718 100644 --- a/pkg/proto/datadog/api/v1/api.proto +++ b/pkg/proto/datadog/api/v1/api.proto @@ -102,7 +102,7 @@ message RemoteQueryExecuteStreamEvent { } message RemoteQueryExecuteChunk { - bytes response_json_chunk = 1; + reserved 1; int32 chunk_index = 2; bool final = 3; RemoteQueryExecuteStreamEvent event = 4; @@ -162,7 +162,7 @@ service AgentSecure { // Executes an Agent-local Remote Queries request through a matched integration check. rpc RemoteQueryExecute(RemoteQueryExecuteRequest) returns (RemoteQueryExecuteResponse); - // Executes an Agent-local Remote Queries request and streams the JSON response in chunks. + // Executes an Agent-local Remote Queries COPY request and streams typed binary-safe events. rpc RemoteQueryExecuteStream(RemoteQueryExecuteRequest) returns (stream RemoteQueryExecuteChunk); // Streams pod-to-service metadata for a specific node. diff --git a/pkg/proto/pbgo/core/api.pb.go b/pkg/proto/pbgo/core/api.pb.go index b14c28de1a31..f8e3c93a3187 100644 --- a/pkg/proto/pbgo/core/api.pb.go +++ b/pkg/proto/pbgo/core/api.pb.go @@ -826,13 +826,12 @@ func (*RemoteQueryExecuteStreamEvent_Final) isRemoteQueryExecuteStreamEvent_Even func (*RemoteQueryExecuteStreamEvent_Error) isRemoteQueryExecuteStreamEvent_Event() {} type RemoteQueryExecuteChunk struct { - state protoimpl.MessageState `protogen:"open.v1"` - ResponseJsonChunk []byte `protobuf:"bytes,1,opt,name=response_json_chunk,json=responseJsonChunk,proto3" json:"response_json_chunk,omitempty"` - ChunkIndex int32 `protobuf:"varint,2,opt,name=chunk_index,json=chunkIndex,proto3" json:"chunk_index,omitempty"` - Final bool `protobuf:"varint,3,opt,name=final,proto3" json:"final,omitempty"` - Event *RemoteQueryExecuteStreamEvent `protobuf:"bytes,4,opt,name=event,proto3" json:"event,omitempty"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + state protoimpl.MessageState `protogen:"open.v1"` + ChunkIndex int32 `protobuf:"varint,2,opt,name=chunk_index,json=chunkIndex,proto3" json:"chunk_index,omitempty"` + Final bool `protobuf:"varint,3,opt,name=final,proto3" json:"final,omitempty"` + Event *RemoteQueryExecuteStreamEvent `protobuf:"bytes,4,opt,name=event,proto3" json:"event,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *RemoteQueryExecuteChunk) Reset() { @@ -865,13 +864,6 @@ func (*RemoteQueryExecuteChunk) Descriptor() ([]byte, []int) { return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{11} } -func (x *RemoteQueryExecuteChunk) GetResponseJsonChunk() []byte { - if x != nil { - return x.ResponseJsonChunk - } - return nil -} - func (x *RemoteQueryExecuteChunk) GetChunkIndex() int32 { if x != nil { return x.ChunkIndex @@ -973,13 +965,12 @@ const file_datadog_api_v1_api_proto_rawDesc = "" + "\x04data\x18\x03 \x01(\v2%.datadog.api.v1.RemoteQueryStreamDataH\x00R\x04data\x12>\n" + "\x05final\x18\x04 \x01(\v2&.datadog.api.v1.RemoteQueryStreamFinalH\x00R\x05final\x12>\n" + "\x05error\x18\x05 \x01(\v2&.datadog.api.v1.RemoteQueryStreamErrorH\x00R\x05errorB\a\n" + - "\x05event\"\xc5\x01\n" + - "\x17RemoteQueryExecuteChunk\x12.\n" + - "\x13response_json_chunk\x18\x01 \x01(\fR\x11responseJsonChunk\x12\x1f\n" + + "\x05event\"\x9b\x01\n" + + "\x17RemoteQueryExecuteChunk\x12\x1f\n" + "\vchunk_index\x18\x02 \x01(\x05R\n" + "chunkIndex\x12\x14\n" + "\x05final\x18\x03 \x01(\bR\x05final\x12C\n" + - "\x05event\x18\x04 \x01(\v2-.datadog.api.v1.RemoteQueryExecuteStreamEventR\x05event2Z\n" + + "\x05event\x18\x04 \x01(\v2-.datadog.api.v1.RemoteQueryExecuteStreamEventR\x05eventJ\x04\b\x01\x10\x022Z\n" + "\x05Agent\x12Q\n" + "\vGetHostname\x12!.datadog.model.v1.HostnameRequest\x1a\x1f.datadog.model.v1.HostnameReply2\x8a\x12\n" + "\vAgentSecure\x12c\n" + diff --git a/pkg/proto/pbgo/core/api_grpc.pb.go b/pkg/proto/pbgo/core/api_grpc.pb.go index 0490e0d2dd2b..c8a02b47e5b9 100644 --- a/pkg/proto/pbgo/core/api_grpc.pb.go +++ b/pkg/proto/pbgo/core/api_grpc.pb.go @@ -189,7 +189,7 @@ type AgentSecureClient interface { WorkloadFilterEvaluate(ctx context.Context, in *WorkloadFilterEvaluateRequest, opts ...grpc.CallOption) (*WorkloadFilterEvaluateResponse, error) // Executes an Agent-local Remote Queries request through a matched integration check. RemoteQueryExecute(ctx context.Context, in *RemoteQueryExecuteRequest, opts ...grpc.CallOption) (*RemoteQueryExecuteResponse, error) - // Executes an Agent-local Remote Queries request and streams the JSON response in chunks. + // Executes an Agent-local Remote Queries COPY request and streams typed binary-safe events. RemoteQueryExecuteStream(ctx context.Context, in *RemoteQueryExecuteRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[RemoteQueryExecuteChunk], error) // Streams pod-to-service metadata for a specific node. StreamKubeMetadata(ctx context.Context, in *KubeMetadataStreamRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[KubeMetadataStreamResponse], error) @@ -508,7 +508,7 @@ type AgentSecureServer interface { WorkloadFilterEvaluate(context.Context, *WorkloadFilterEvaluateRequest) (*WorkloadFilterEvaluateResponse, error) // Executes an Agent-local Remote Queries request through a matched integration check. RemoteQueryExecute(context.Context, *RemoteQueryExecuteRequest) (*RemoteQueryExecuteResponse, error) - // Executes an Agent-local Remote Queries request and streams the JSON response in chunks. + // Executes an Agent-local Remote Queries COPY request and streams typed binary-safe events. RemoteQueryExecuteStream(*RemoteQueryExecuteRequest, grpc.ServerStreamingServer[RemoteQueryExecuteChunk]) error // Streams pod-to-service metadata for a specific node. StreamKubeMetadata(*KubeMetadataStreamRequest, grpc.ServerStreamingServer[KubeMetadataStreamResponse]) error diff --git a/rtloader/include/datadog_agent_rtloader.h b/rtloader/include/datadog_agent_rtloader.h index 373713f7ca79..a469e078c700 100644 --- a/rtloader/include/datadog_agent_rtloader.h +++ b/rtloader/include/datadog_agent_rtloader.h @@ -190,18 +190,6 @@ DATADOG_AGENT_RTLOADER_API int get_check_deprecated(rtloader_t *rtloader, rtload */ DATADOG_AGENT_RTLOADER_API char *run_check(rtloader_t *, rtloader_pyobject_t *check); -/*! \fn char *run_remote_query(rtloader_t *, rtloader_pyobject_t *check, const char *integration, const char *request_json) - \brief Runs the integration remote query helper for a check instance. - \param rtloader_t A rtloader_t * pointer to the RtLoader instance. - \param check A rtloader_pyobject_t * pointer to the check instance we wish to use. - \param integration The integration helper module name. - \param request_json A credential-free JSON request string. - \return A C-string with the JSON result. - \sa rtloader_pyobject_t, rtloader_t -*/ -DATADOG_AGENT_RTLOADER_API char *run_remote_query(rtloader_t *, rtloader_pyobject_t *check, const char *integration, - const char *request_json); - /*! \fn int run_remote_query_stream(rtloader_t *, rtloader_pyobject_t *check, const char *integration, const char *request_json, remote_query_stream_emit_cb emit, void *userdata) \brief Runs the integration remote query streaming helper for a check instance. \param rtloader_t A rtloader_t * pointer to the RtLoader instance. diff --git a/rtloader/include/rtloader.h b/rtloader/include/rtloader.h index 24bf498edf35..c7c7af7a4d8a 100644 --- a/rtloader/include/rtloader.h +++ b/rtloader/include/rtloader.h @@ -121,15 +121,6 @@ class RtLoader */ virtual char *runCheck(RtLoaderPyObject *check) = 0; - //! Pure virtual runRemoteQuery member. - /*! - \param check The python object pointer to the check we wish to use. - \param integration The integration helper module name. - \param request_json A credential-free JSON request string. - \return A C-string with the JSON result. - */ - virtual char *runRemoteQuery(RtLoaderPyObject *check, const char *integration, const char *request_json) = 0; - //! Pure virtual runRemoteQueryStream member. /*! \param check The python object pointer to the check we wish to use. diff --git a/rtloader/rtloader/api.cpp b/rtloader/rtloader/api.cpp index 5b5998d6584a..25b94b5ca919 100644 --- a/rtloader/rtloader/api.cpp +++ b/rtloader/rtloader/api.cpp @@ -278,11 +278,6 @@ char *run_check(rtloader_t *rtloader, rtloader_pyobject_t *check) return AS_TYPE(RtLoader, rtloader)->runCheck(AS_TYPE(RtLoaderPyObject, check)); } -char *run_remote_query(rtloader_t *rtloader, rtloader_pyobject_t *check, const char *integration, const char *request_json) -{ - return AS_TYPE(RtLoader, rtloader)->runRemoteQuery(AS_TYPE(RtLoaderPyObject, check), integration, request_json); -} - int run_remote_query_stream(rtloader_t *rtloader, rtloader_pyobject_t *check, const char *integration, const char *request_json, remote_query_stream_emit_cb emit, void *userdata) { diff --git a/rtloader/three/three.cpp b/rtloader/three/three.cpp index 65dc04a51a42..cec202ceac5a 100644 --- a/rtloader/three/three.cpp +++ b/rtloader/three/three.cpp @@ -574,65 +574,6 @@ PyMethodDef remoteQueryStreamEmitMethod = {"remote_query_stream_emit", remoteQue "Emit a remote query stream event."}; } // namespace -char *Three::runRemoteQuery(RtLoaderPyObject *check, const char *integration, const char *request_json) -{ - if (check == NULL || request_json == NULL) { - return NULL; - } - - std::string normalized_integration = normalizeRemoteQueryIntegration(integration); - if (!isValidRemoteQueryIntegration(normalized_integration)) { - setError("invalid remote query integration name"); - return NULL; - } - - PyObject *py_check = reinterpret_cast(check); - char *ret = NULL; - PyObject *remote_query_module = NULL; - PyObject *execute_func = NULL; - PyObject *py_request_json = NULL; - PyObject *result = NULL; - std::string module_name = "datadog_checks." + normalized_integration + ".remote_query"; - - remote_query_module = PyImport_ImportModule(module_name.c_str()); - if (remote_query_module == NULL) { - setError("error importing remote query helper: " + _fetchPythonError()); - goto done; - } - - execute_func = PyObject_GetAttrString(remote_query_module, "execute_agent_rpc_json"); - if (execute_func == NULL || !PyCallable_Check(execute_func)) { - setError("error loading remote query helper: " + _fetchPythonError()); - goto done; - } - - py_request_json = PyUnicode_FromString(request_json); - if (py_request_json == NULL) { - setError("error converting remote query request to Python string: " + _fetchPythonError()); - goto done; - } - - result = PyObject_CallFunctionObjArgs(execute_func, py_request_json, py_check, NULL); - if (result == NULL || !PyUnicode_Check(result)) { - setError("error invoking remote query helper: " + _fetchPythonError()); - goto done; - } - - ret = as_string(result); - if (ret == NULL) { - // as_string clears the error, so we can't fetch it here - setError("error converting remote query helper result to string"); - goto done; - } - - done: - Py_XDECREF(result); - Py_XDECREF(py_request_json); - Py_XDECREF(execute_func); - Py_XDECREF(remote_query_module); - return ret; -} - bool Three::runRemoteQueryStream(RtLoaderPyObject *check, const char *integration, const char *request_json, remote_query_stream_emit_cb emit, void *userdata) { diff --git a/rtloader/three/three.h b/rtloader/three/three.h index 0c4a33769105..6fc3cbdd1560 100644 --- a/rtloader/three/three.h +++ b/rtloader/three/three.h @@ -65,7 +65,6 @@ class Three : public RtLoader const char *provider_str, RtLoaderPyObject *&check); char *runCheck(RtLoaderPyObject *check); - char *runRemoteQuery(RtLoaderPyObject *check, const char *integration, const char *request_json); bool runRemoteQueryStream(RtLoaderPyObject *check, const char *integration, const char *request_json, remote_query_stream_emit_cb emit, void *userdata); void cancelCheck(RtLoaderPyObject *check); diff --git a/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh b/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh index 70e751105dde..cdbb6e0befa3 100755 --- a/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh +++ b/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh @@ -400,16 +400,10 @@ PY local token token=$(cat "$TMP_ROOT/run/auth_token") - if [[ "${RQ_REMOTE_OPERATION:-}" == "copy_stream" && "${RQ_REMOTE_FORMAT:-}" == "binary" ]]; then - log "[$PROOF_CASE_NAME] Skipping inline JSON preflight for binary COPY stream (streaming path only)" - return - fi - - log "[$PROOF_CASE_NAME] Preflight real Agent IPC HTTP execute endpoint (dev evidence only)" + log "[$PROOF_CASE_NAME] Preflight real Agent IPC HTTP execute endpoint is disabled; Remote Queries execution requires AgentSecure COPY streaming" local body_file="$CASE_RESULTS_DIR/agent-execute-preflight.raw-body" local status_file="$CASE_RESULTS_DIR/agent-execute-preflight.status" local summary_file="$CASE_RESULTS_DIR/agent-execute-preflight.summary.json" - local sanitized_body_file="$CASE_RESULTS_DIR/agent-execute-preflight.body" local status status=$(curl -sS -k -o "$body_file" -w '%{http_code}' \ -H "Authorization: Bearer ${token}" \ @@ -421,53 +415,10 @@ PY summarize_json_file "$body_file" "$summary_file" fi - if [[ "$status" != "200" ]]; then - echo "FAIL: expected Agent execute preflight HTTP 200, got $status" >&2 + if [[ "$status" != "400" ]]; then + echo "FAIL: expected Agent execute preflight HTTP 400 for disabled inline execution, got $status" >&2 exit 1 fi - RQ_REMOTE_QUERY="$RQ_REMOTE_QUERY" python3 - "$body_file" <<'PY' -import json -import os -import re -import sys - -with open(sys.argv[1], encoding="utf-8") as f: - body = json.load(f) - -if body.get("status") != "SUCCEEDED": - raise SystemExit(f"Agent execute preflight response status was not SUCCEEDED: {body}") - -rows = body.get("rows") -query = os.environ["RQ_REMOTE_QUERY"] -if query == "SELECT city, country FROM cities ORDER BY city": - expected = [ - {"city": "Beautiful city of lights", "country": "France"}, - {"city": "New York", "country": "USA"}, - ] - if rows != expected: - raise SystemExit(f"Agent execute preflight rows did not match expected fixture data: rows={rows!r} expected={expected!r}") -elif query == "SELECT 1 AS value": - expected = [{"value": 1}] - if rows != expected: - raise SystemExit(f"Agent execute preflight rows did not match expected seed data: rows={rows!r} expected={expected!r}") -else: - match = re.fullmatch(r"SELECT repeat\('x', ([0-9]+)\) AS payload", query) - if not match: - raise SystemExit(f"Unsupported proof query: {query!r}") - expected_len = int(match.group(1)) - if not isinstance(rows, list) or len(rows) != 1: - raise SystemExit(f"Agent execute preflight expected one payload row, got rows={rows!r}") - payload = rows[0].get("payload") if isinstance(rows[0], dict) else None - if not isinstance(payload, str): - raise SystemExit("Agent execute preflight payload field was not a string") - if len(payload) != expected_len: - raise SystemExit(f"Agent execute preflight payload length mismatch: got {len(payload)} expected {expected_len}") -PY - if grep -Eq 'password|token|secret' "$body_file"; then - echo "FAIL: Agent execute preflight response contained credential-shaped text" >&2 - exit 1 - fi - cp "$summary_file" "$sanitized_body_file" rm -f "$body_file" } @@ -486,8 +437,8 @@ run_standalone_go_proof() { RQ_POSTGRES_PORT="$RQ_POSTGRES_PORT" \ RQ_POSTGRES_DBNAME="$RQ_POSTGRES_DBNAME" \ RQ_REMOTE_QUERY="$RQ_REMOTE_QUERY" \ - RQ_REMOTE_OPERATION="${RQ_REMOTE_OPERATION:-}" \ - RQ_REMOTE_FORMAT="${RQ_REMOTE_FORMAT:-}" \ + RQ_REMOTE_OPERATION="copy_stream" \ + RQ_REMOTE_FORMAT="${RQ_REMOTE_FORMAT:-csv}" \ dda inv test --targets=./pkg/privateactionrunner/bundles/remotequeries \ --extra-args='-run TestRemoteQueriesActionRunsThroughStandalonePARProcessWithRealAgentIPC -count=1 -v' ) | tee "$CASE_RESULTS_DIR/standalone-proof-test.log" @@ -497,13 +448,10 @@ run_proof_case() { PROOF_CASE_NAME=$1 RQ_REMOTE_QUERY=$2 if [[ "$RQ_REMOTE_OPERATION_WAS_SET" != "1" ]]; then - RQ_REMOTE_OPERATION="" - if [[ "$PROOF_CASE_NAME" == copy-* ]]; then - RQ_REMOTE_OPERATION=copy_stream - fi + RQ_REMOTE_OPERATION=copy_stream fi if [[ "$RQ_REMOTE_FORMAT_WAS_SET" != "1" ]]; then - RQ_REMOTE_FORMAT="" + RQ_REMOTE_FORMAT=csv if [[ "$PROOF_CASE_NAME" == "copy-binary-payload" ]]; then RQ_REMOTE_FORMAT=binary fi From a2f9d8a2be6854b941e1dc7c4e09c65e5b8699f5 Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 15:51:26 +0000 Subject: [PATCH 25/33] Coalesce Remote Queries stream events for IPC --- .../impl-agent/remote_query_execute_test.go | 24 ++++ comp/api/grpcserver/impl-agent/server.go | 122 ++++++++++++++++-- .../bundles/remotequeries/execute.go | 61 ++++++++- 3 files changed, 194 insertions(+), 13 deletions(-) diff --git a/comp/api/grpcserver/impl-agent/remote_query_execute_test.go b/comp/api/grpcserver/impl-agent/remote_query_execute_test.go index e8afe7a5222d..d7d04e4133a9 100644 --- a/comp/api/grpcserver/impl-agent/remote_query_execute_test.go +++ b/comp/api/grpcserver/impl-agent/remote_query_execute_test.go @@ -70,6 +70,30 @@ func TestRemoteQueryExecuteRequestFromProtoPreservesCopyStream(t *testing.T) { assert.Nil(t, req.Limits) } +func TestRemoteQueryIPCStreamCoalescerFlushesDataAtFourMiB(t *testing.T) { + stream := &captureRemoteQueryExecuteStreamServer{} + coalescer := newRemoteQueryIPCStreamCoalescer(stream) + + require.NoError(t, coalescer.Send(check.RemoteQueryStreamEvent{Type: "metadata", MetadataJSON: `{"operation":"copy_stream","format":"csv"}`})) + require.NoError(t, coalescer.Send(check.RemoteQueryStreamEvent{Type: "data", MetadataJSON: `{"sequence":1,"offset":0,"bytes":3145728}`, Payload: make([]byte, 3<<20)})) + assert.Len(t, stream.chunks, 1, "data below 4MiB should be coalesced before secure IPC send") + require.NoError(t, coalescer.Send(check.RemoteQueryStreamEvent{Type: "data", MetadataJSON: `{"sequence":2,"offset":3145728,"bytes":2097152}`, Payload: make([]byte, 2<<20)})) + require.Len(t, stream.chunks, 2, "crossing 4MiB should flush one coalesced data event") + firstData := stream.chunks[1].GetEvent().GetData() + require.NotNil(t, firstData) + assert.Equal(t, uint64(0), firstData.GetOffset()) + assert.Equal(t, uint64(remoteQuerySecureIPCDataFlushBytes), firstData.GetBytes()) + assert.Len(t, firstData.GetPayload(), remoteQuerySecureIPCDataFlushBytes) + + require.NoError(t, coalescer.Flush()) + require.Len(t, stream.chunks, 3) + secondData := stream.chunks[2].GetEvent().GetData() + require.NotNil(t, secondData) + assert.Equal(t, uint64(remoteQuerySecureIPCDataFlushBytes), secondData.GetOffset()) + assert.Equal(t, uint64((5<<20)-remoteQuerySecureIPCDataFlushBytes), secondData.GetBytes()) + assert.Len(t, secondData.GetPayload(), (5<<20)-remoteQuerySecureIPCDataFlushBytes) +} + func TestRemoteQueryStreamEventFromCheckEventPreservesBinaryPayload(t *testing.T) { event, err := remoteQueryStreamEventFromCheckEvent(check.RemoteQueryStreamEvent{ Type: "data", diff --git a/comp/api/grpcserver/impl-agent/server.go b/comp/api/grpcserver/impl-agent/server.go index 858c0a32fd9b..b90656adc520 100644 --- a/comp/api/grpcserver/impl-agent/server.go +++ b/comp/api/grpcserver/impl-agent/server.go @@ -6,6 +6,7 @@ package agentimpl import ( + "bytes" "context" "encoding/json" "errors" @@ -297,20 +298,113 @@ func (s *serverSecure) RemoteQueryExecuteStream(req *pb.RemoteQueryExecuteReques return remoteQueryExecuteStreamError(remotequeriesimpl.RemoteQueryStatusInvalidRequest, err.Error(), stream) } - chunkIndex := int32(0) - result := s.remoteQueries.ExecuteStream(execReq, func(event check.RemoteQueryStreamEvent) error { - protoEvent, err := remoteQueryStreamEventFromCheckEvent(event) - if err != nil { + coalescer := newRemoteQueryIPCStreamCoalescer(stream) + result := s.remoteQueries.ExecuteStream(execReq, coalescer.Send) + if result.Error != nil { + if err := coalescer.Flush(); err != nil { return err } - err = stream.Send(&pb.RemoteQueryExecuteChunk{Event: protoEvent, ChunkIndex: chunkIndex}) - chunkIndex++ + return remoteQueryExecuteStreamErrorAt(result.Error.Code, result.Error.Message, stream, coalescer.NextChunkIndex()) + } + if err := coalescer.Flush(); err != nil { return err - }) - if result.Error != nil { - return remoteQueryExecuteStreamError(result.Error.Code, result.Error.Message, stream) } - return stream.Send(&pb.RemoteQueryExecuteChunk{ChunkIndex: chunkIndex, Final: true}) + return stream.Send(&pb.RemoteQueryExecuteChunk{ChunkIndex: coalescer.NextChunkIndex(), Final: true}) +} + +const remoteQuerySecureIPCDataFlushBytes = 4_000_000 + +type remoteQueryIPCStreamCoalescer struct { + stream pb.AgentSecure_RemoteQueryExecuteStreamServer + chunkIndex int32 + data bytes.Buffer + dataOffset uint64 + dataSeq uint64 + dataStarted bool + dataChunks uint64 +} + +func newRemoteQueryIPCStreamCoalescer(stream pb.AgentSecure_RemoteQueryExecuteStreamServer) *remoteQueryIPCStreamCoalescer { + return &remoteQueryIPCStreamCoalescer{stream: stream} +} + +func (c *remoteQueryIPCStreamCoalescer) NextChunkIndex() int32 { + return c.chunkIndex +} + +func (c *remoteQueryIPCStreamCoalescer) Send(event check.RemoteQueryStreamEvent) error { + protoEvent, err := remoteQueryStreamEventFromCheckEvent(event) + if err != nil { + return err + } + data := protoEvent.GetData() + if data == nil { + if err := c.Flush(); err != nil { + return err + } + return c.sendProtoEvent(protoEvent) + } + + if c.dataStarted && data.GetOffset() != c.dataOffset+uint64(c.data.Len()) { + if err := c.Flush(); err != nil { + return err + } + } + if !c.dataStarted { + c.dataStarted = true + c.dataOffset = data.GetOffset() + c.dataSeq = protoEvent.GetSequence() + } + if _, err := c.data.Write(data.GetPayload()); err != nil { + return err + } + for c.data.Len() >= remoteQuerySecureIPCDataFlushBytes { + if err := c.flushData(remoteQuerySecureIPCDataFlushBytes); err != nil { + return err + } + } + return nil +} + +func (c *remoteQueryIPCStreamCoalescer) Flush() error { + if !c.dataStarted || c.data.Len() == 0 { + c.data.Reset() + c.dataStarted = false + return nil + } + return c.flushData(c.data.Len()) +} + +func (c *remoteQueryIPCStreamCoalescer) flushData(size int) error { + payload := append([]byte(nil), c.data.Bytes()[:size]...) + protoEvent := &pb.RemoteQueryExecuteStreamEvent{ + Sequence: c.dataSeq + c.dataChunks, + Event: &pb.RemoteQueryExecuteStreamEvent_Data{Data: &pb.RemoteQueryStreamData{ + Payload: payload, + Offset: c.dataOffset, + Bytes: uint64(len(payload)), + }}, + } + if err := c.sendProtoEvent(protoEvent); err != nil { + return err + } + remaining := append([]byte(nil), c.data.Bytes()[size:]...) + c.data.Reset() + _, _ = c.data.Write(remaining) + c.dataOffset += uint64(len(payload)) + c.dataChunks++ + if c.data.Len() == 0 { + c.dataStarted = false + } + return nil +} + +func (c *remoteQueryIPCStreamCoalescer) sendProtoEvent(event *pb.RemoteQueryExecuteStreamEvent) error { + if err := c.stream.Send(&pb.RemoteQueryExecuteChunk{Event: event, ChunkIndex: c.chunkIndex}); err != nil { + return err + } + c.chunkIndex++ + return nil } func remoteQueryExecuteRequestFromProto(req *pb.RemoteQueryExecuteRequest) (remotequeriesimpl.RemoteQueryExecuteRequest, error) { @@ -430,8 +524,12 @@ func stringAttributes(metadata map[string]interface{}, exclude ...string) map[st } func remoteQueryExecuteStreamError(code string, message string, stream pb.AgentSecure_RemoteQueryExecuteStreamServer) error { + return remoteQueryExecuteStreamErrorAt(code, message, stream, 0) +} + +func remoteQueryExecuteStreamErrorAt(code string, message string, stream pb.AgentSecure_RemoteQueryExecuteStreamServer, chunkIndex int32) error { if err := stream.Send(&pb.RemoteQueryExecuteChunk{ - ChunkIndex: 0, + ChunkIndex: chunkIndex, Event: &pb.RemoteQueryExecuteStreamEvent{Event: &pb.RemoteQueryExecuteStreamEvent_Error{Error: &pb.RemoteQueryStreamError{ Code: code, Message: message, @@ -439,7 +537,7 @@ func remoteQueryExecuteStreamError(code string, message string, stream pb.AgentS }); err != nil { return err } - return stream.Send(&pb.RemoteQueryExecuteChunk{ChunkIndex: 1, Final: true}) + return stream.Send(&pb.RemoteQueryExecuteChunk{ChunkIndex: chunkIndex + 1, Final: true}) } func remoteQueryCopyLimitsFromProto(limits *pb.RemoteQueryExecuteCopyLimits) *remotequeriesimpl.RemoteQueryExecuteCopyLimits { diff --git a/pkg/privateactionrunner/bundles/remotequeries/execute.go b/pkg/privateactionrunner/bundles/remotequeries/execute.go index 16af35ba841a..98ad7a01b672 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/execute.go +++ b/pkg/privateactionrunner/bundles/remotequeries/execute.go @@ -11,6 +11,7 @@ import ( "encoding/json" "fmt" "io" + "time" "unicode/utf8" "google.golang.org/grpc" @@ -138,6 +139,11 @@ func remoteQueryExecuteOutputFromStream(stream grpc.ServerStreamingClient[pb.Rem } typedStreamEvents := make([]map[string]interface{}, 0) + streamStart := time.Now() + var firstChunkAt time.Time + var firstDataAt time.Time + var finalChunkAt time.Time + payloadBytes := 0 expectedChunkIndex := int32(0) seenFinal := false for { @@ -151,6 +157,9 @@ func remoteQueryExecuteOutputFromStream(stream grpc.ServerStreamingClient[pb.Rem if chunk == nil { return nil, fmt.Errorf("remote query response stream returned nil chunk") } + if firstChunkAt.IsZero() { + firstChunkAt = time.Now() + } if chunk.GetChunkIndex() != expectedChunkIndex { return nil, fmt.Errorf("remote query response stream chunk index mismatch") } @@ -162,11 +171,22 @@ func remoteQueryExecuteOutputFromStream(stream grpc.ServerStreamingClient[pb.Rem if err != nil { return nil, err } + if streamEvent["type"] == "data" { + if firstDataAt.IsZero() { + firstDataAt = time.Now() + } + if payload, ok := streamEvent["payload"].([]byte); ok { + payloadBytes += len(payload) + } + } typedStreamEvents = append(typedStreamEvents, streamEvent) } else if !chunk.GetFinal() { return nil, fmt.Errorf("remote query response stream chunk missing typed event") } seenFinal = chunk.GetFinal() + if seenFinal { + finalChunkAt = time.Now() + } expectedChunkIndex++ } if !seenFinal { @@ -175,7 +195,46 @@ func remoteQueryExecuteOutputFromStream(stream grpc.ServerStreamingClient[pb.Rem if len(typedStreamEvents) == 0 { return nil, fmt.Errorf("remote query response stream missing typed events") } - return remoteQueryExecuteOutputFromTypedEvents(typedStreamEvents) + output, err := remoteQueryExecuteOutputFromTypedEvents(typedStreamEvents) + if err != nil { + return nil, err + } + if payloadBytes > 0 { + output["stream_timing"] = remoteQueryStreamTiming(streamStart, firstChunkAt, firstDataAt, finalChunkAt, payloadBytes, int(expectedChunkIndex)) + } + return output, nil +} + +func remoteQueryStreamTiming(streamStart time.Time, firstChunkAt time.Time, firstDataAt time.Time, finalChunkAt time.Time, payloadBytes int, chunksReceived int) map[string]interface{} { + streamEnd := time.Now() + dataDuration := finalChunkAt.Sub(firstDataAt) + if firstDataAt.IsZero() || finalChunkAt.IsZero() { + dataDuration = 0 + } + return map[string]interface{}{ + "payload_bytes": payloadBytes, + "chunks_received": chunksReceived, + "first_chunk_latency_ms": durationMillis(firstChunkAt.Sub(streamStart)), + "first_data_latency_ms": durationMillis(firstDataAt.Sub(streamStart)), + "data_to_final_ms": durationMillis(dataDuration), + "stream_loop_total_ms": durationMillis(streamEnd.Sub(streamStart)), + "data_to_final_mib_per_second": mibPerSecond(payloadBytes, dataDuration), + "stream_loop_mib_per_second": mibPerSecond(payloadBytes, streamEnd.Sub(streamStart)), + } +} + +func durationMillis(duration time.Duration) float64 { + if duration <= 0 { + return 0 + } + return duration.Seconds() * 1000 +} + +func mibPerSecond(bytes int, duration time.Duration) float64 { + if bytes <= 0 || duration <= 0 { + return 0 + } + return (float64(bytes) / 1024 / 1024) / duration.Seconds() } func remoteQueryStreamEventFromProto(event *pb.RemoteQueryExecuteStreamEvent) (map[string]interface{}, error) { From 4e8c4b91796a24b9fcb3ee9378e7cfd244890fda Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 15:58:46 +0000 Subject: [PATCH 26/33] Add Remote Queries stream timing breakdown --- comp/api/grpcserver/impl-agent/server.go | 90 +++++++++++++++++-- .../bundles/remotequeries/execute.go | 9 ++ 2 files changed, 93 insertions(+), 6 deletions(-) diff --git a/comp/api/grpcserver/impl-agent/server.go b/comp/api/grpcserver/impl-agent/server.go index b90656adc520..60cf34a1a136 100644 --- a/comp/api/grpcserver/impl-agent/server.go +++ b/comp/api/grpcserver/impl-agent/server.go @@ -322,10 +322,23 @@ type remoteQueryIPCStreamCoalescer struct { dataSeq uint64 dataStarted bool dataChunks uint64 + + start time.Time + firstEventAt time.Time + firstDataAt time.Time + lastDataAt time.Time + upstreamDataEvents uint64 + upstreamDataBytes uint64 + coalescedDataEvents uint64 + sendCalls uint64 + sendDuration time.Duration + dataSendDuration time.Duration + maxSendDuration time.Duration + maxDataSendDuration time.Duration } func newRemoteQueryIPCStreamCoalescer(stream pb.AgentSecure_RemoteQueryExecuteStreamServer) *remoteQueryIPCStreamCoalescer { - return &remoteQueryIPCStreamCoalescer{stream: stream} + return &remoteQueryIPCStreamCoalescer{stream: stream, start: time.Now()} } func (c *remoteQueryIPCStreamCoalescer) NextChunkIndex() int32 { @@ -333,6 +346,9 @@ func (c *remoteQueryIPCStreamCoalescer) NextChunkIndex() int32 { } func (c *remoteQueryIPCStreamCoalescer) Send(event check.RemoteQueryStreamEvent) error { + if c.firstEventAt.IsZero() { + c.firstEventAt = time.Now() + } protoEvent, err := remoteQueryStreamEventFromCheckEvent(event) if err != nil { return err @@ -342,8 +358,18 @@ func (c *remoteQueryIPCStreamCoalescer) Send(event check.RemoteQueryStreamEvent) if err := c.Flush(); err != nil { return err } - return c.sendProtoEvent(protoEvent) + c.addTimingAttributes(protoEvent) + _, err := c.sendProtoEvent(protoEvent) + return err + } + + now := time.Now() + if c.firstDataAt.IsZero() { + c.firstDataAt = now } + c.lastDataAt = now + c.upstreamDataEvents++ + c.upstreamDataBytes += uint64(len(data.GetPayload())) if c.dataStarted && data.GetOffset() != c.dataOffset+uint64(c.data.Len()) { if err := c.Flush(); err != nil { @@ -385,9 +411,15 @@ func (c *remoteQueryIPCStreamCoalescer) flushData(size int) error { Bytes: uint64(len(payload)), }}, } - if err := c.sendProtoEvent(protoEvent); err != nil { + duration, err := c.sendProtoEvent(protoEvent) + if err != nil { return err } + c.coalescedDataEvents++ + c.dataSendDuration += duration + if duration > c.maxDataSendDuration { + c.maxDataSendDuration = duration + } remaining := append([]byte(nil), c.data.Bytes()[size:]...) c.data.Reset() _, _ = c.data.Write(remaining) @@ -399,12 +431,58 @@ func (c *remoteQueryIPCStreamCoalescer) flushData(size int) error { return nil } -func (c *remoteQueryIPCStreamCoalescer) sendProtoEvent(event *pb.RemoteQueryExecuteStreamEvent) error { +func (c *remoteQueryIPCStreamCoalescer) sendProtoEvent(event *pb.RemoteQueryExecuteStreamEvent) (time.Duration, error) { + start := time.Now() if err := c.stream.Send(&pb.RemoteQueryExecuteChunk{Event: event, ChunkIndex: c.chunkIndex}); err != nil { - return err + return 0, err + } + duration := time.Since(start) + c.sendCalls++ + c.sendDuration += duration + if duration > c.maxSendDuration { + c.maxSendDuration = duration } c.chunkIndex++ - return nil + return duration, nil +} + +func (c *remoteQueryIPCStreamCoalescer) addTimingAttributes(event *pb.RemoteQueryExecuteStreamEvent) { + final := event.GetFinal() + if final == nil { + return + } + if final.Attributes == nil { + final.Attributes = map[string]string{} + } + elapsed := time.Since(c.start) + final.Attributes["agent_coalesce_flush_bytes"] = strconv.Itoa(remoteQuerySecureIPCDataFlushBytes) + final.Attributes["agent_upstream_data_events"] = strconv.FormatUint(c.upstreamDataEvents, 10) + final.Attributes["agent_upstream_data_bytes"] = strconv.FormatUint(c.upstreamDataBytes, 10) + final.Attributes["agent_coalesced_data_events"] = strconv.FormatUint(c.coalescedDataEvents, 10) + final.Attributes["agent_ipc_send_calls"] = strconv.FormatUint(c.sendCalls, 10) + final.Attributes["agent_first_event_latency_ms"] = formatDurationMillis(c.firstEventAt.Sub(c.start)) + final.Attributes["agent_first_data_latency_ms"] = formatDurationMillis(c.firstDataAt.Sub(c.start)) + final.Attributes["agent_upstream_data_span_ms"] = formatDurationMillis(c.lastDataAt.Sub(c.firstDataAt)) + final.Attributes["agent_total_stream_ms"] = formatDurationMillis(elapsed) + final.Attributes["agent_total_stream_mib_per_second"] = formatMiBPerSecond(c.upstreamDataBytes, elapsed) + final.Attributes["agent_ipc_send_total_ms"] = formatDurationMillis(c.sendDuration) + final.Attributes["agent_ipc_send_max_ms"] = formatDurationMillis(c.maxSendDuration) + final.Attributes["agent_ipc_data_send_total_ms"] = formatDurationMillis(c.dataSendDuration) + final.Attributes["agent_ipc_data_send_max_ms"] = formatDurationMillis(c.maxDataSendDuration) +} + +func formatDurationMillis(duration time.Duration) string { + if duration <= 0 { + return "0" + } + return strconv.FormatFloat(duration.Seconds()*1000, 'f', 3, 64) +} + +func formatMiBPerSecond(bytes uint64, duration time.Duration) string { + if bytes == 0 || duration <= 0 { + return "0" + } + return strconv.FormatFloat((float64(bytes)/1024/1024)/duration.Seconds(), 'f', 3, 64) } func remoteQueryExecuteRequestFromProto(req *pb.RemoteQueryExecuteRequest) (remotequeriesimpl.RemoteQueryExecuteRequest, error) { diff --git a/pkg/privateactionrunner/bundles/remotequeries/execute.go b/pkg/privateactionrunner/bundles/remotequeries/execute.go index 98ad7a01b672..73e3d8c11529 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/execute.go +++ b/pkg/privateactionrunner/bundles/remotequeries/execute.go @@ -73,6 +73,7 @@ func (a *ExecuteAction) Run( task *types.Task, _ *privateconnection.PrivateCredentials, ) (interface{}, error) { + actionStart := time.Now() inputs, err := types.ExtractInputs[ExecuteInputs](task) if err != nil { return nil, util.DefaultActionErrorWithDisplayError( @@ -92,14 +93,22 @@ func (a *ExecuteAction) Run( return nil, util.DefaultActionError(fmt.Errorf("remote query action requires an AgentSecure client")) } + rpcStart := time.Now() stream, err := client.RemoteQueryExecuteStream(ctx, remoteQueryExecuteRequestFromInputs(inputs)) if err != nil { return nil, util.DefaultActionErrorWithDisplayError(err, "remote query AgentSecure streaming RPC failed") } + rpcCreatedAt := time.Now() output, err := remoteQueryExecuteOutputFromStream(stream) if err != nil { return nil, util.DefaultActionErrorWithDisplayError(err, "remote query AgentSecure streaming RPC response was invalid") } + if timing, ok := output["stream_timing"].(map[string]interface{}); ok { + now := time.Now() + timing["action_total_ms"] = durationMillis(now.Sub(actionStart)) + timing["rpc_create_ms"] = durationMillis(rpcCreatedAt.Sub(rpcStart)) + timing["rpc_receive_and_assemble_ms"] = durationMillis(now.Sub(rpcCreatedAt)) + } return output, nil } From 83f739b9a370421ce8d24e90e01a8a91a386a2b4 Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 16:47:58 +0000 Subject: [PATCH 27/33] Allow Remote Queries proof chunk size override --- .../remotequeries/live_agent_ipc_par_loop_test.go | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go index 5e0170e83a55..5f6a76c79c3d 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go @@ -169,8 +169,15 @@ func remoteQueriesProofCopyLimits(query string) map[string]interface{} { maxBytes = payloadBytes + (1 << 20) timeoutMs = 60_000 } + chunkBytes := 32 + if value := os.Getenv("RQ_REMOTE_CHUNK_BYTES"); value != "" { + parsed, err := strconv.Atoi(value) + if err == nil && parsed > 0 { + chunkBytes = parsed + } + } return map[string]interface{}{ - "chunkBytes": 32, + "chunkBytes": chunkBytes, "maxBytes": maxBytes, "maxRowBytes": maxBytes, "timeoutMs": timeoutMs, From f7713810bbcdd2fc46bceebaf90b168040e98010 Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 17:00:14 +0000 Subject: [PATCH 28/33] Use practical Remote Queries proof chunk size --- .../bundles/remotequeries/live_agent_ipc_par_loop_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go index 5f6a76c79c3d..3efa7961c8f1 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go @@ -169,7 +169,7 @@ func remoteQueriesProofCopyLimits(query string) map[string]interface{} { maxBytes = payloadBytes + (1 << 20) timeoutMs = 60_000 } - chunkBytes := 32 + chunkBytes := 256 << 10 if value := os.Getenv("RQ_REMOTE_CHUNK_BYTES"); value != "" { parsed, err := strconv.Atoi(value) if err == nil && parsed > 0 { From c4d9e5a7018b77bab92cb4ca405b68434d610b57 Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 17:20:20 +0000 Subject: [PATCH 29/33] Clarify Remote Queries byte shift sizes --- .../impl-agent/remote_query_execute_test.go | 13 +++++++++---- .../remotequeries/impl/remote_query_execute.go | 12 ++++++------ .../live_agent_ipc_par_loop_test.go | 18 +++++++++--------- 3 files changed, 24 insertions(+), 19 deletions(-) diff --git a/comp/api/grpcserver/impl-agent/remote_query_execute_test.go b/comp/api/grpcserver/impl-agent/remote_query_execute_test.go index d7d04e4133a9..d32359adb723 100644 --- a/comp/api/grpcserver/impl-agent/remote_query_execute_test.go +++ b/comp/api/grpcserver/impl-agent/remote_query_execute_test.go @@ -73,11 +73,16 @@ func TestRemoteQueryExecuteRequestFromProtoPreservesCopyStream(t *testing.T) { func TestRemoteQueryIPCStreamCoalescerFlushesDataAtFourMiB(t *testing.T) { stream := &captureRemoteQueryExecuteStreamServer{} coalescer := newRemoteQueryIPCStreamCoalescer(stream) + const ( + threeMiB = 3 << 20 // 3 MiB. + twoMiB = 2 << 20 // 2 MiB. + fiveMiB = 5 << 20 // 5 MiB. + ) require.NoError(t, coalescer.Send(check.RemoteQueryStreamEvent{Type: "metadata", MetadataJSON: `{"operation":"copy_stream","format":"csv"}`})) - require.NoError(t, coalescer.Send(check.RemoteQueryStreamEvent{Type: "data", MetadataJSON: `{"sequence":1,"offset":0,"bytes":3145728}`, Payload: make([]byte, 3<<20)})) + require.NoError(t, coalescer.Send(check.RemoteQueryStreamEvent{Type: "data", MetadataJSON: `{"sequence":1,"offset":0,"bytes":3145728}`, Payload: make([]byte, threeMiB)})) assert.Len(t, stream.chunks, 1, "data below 4MiB should be coalesced before secure IPC send") - require.NoError(t, coalescer.Send(check.RemoteQueryStreamEvent{Type: "data", MetadataJSON: `{"sequence":2,"offset":3145728,"bytes":2097152}`, Payload: make([]byte, 2<<20)})) + require.NoError(t, coalescer.Send(check.RemoteQueryStreamEvent{Type: "data", MetadataJSON: `{"sequence":2,"offset":3145728,"bytes":2097152}`, Payload: make([]byte, twoMiB)})) require.Len(t, stream.chunks, 2, "crossing 4MiB should flush one coalesced data event") firstData := stream.chunks[1].GetEvent().GetData() require.NotNil(t, firstData) @@ -90,8 +95,8 @@ func TestRemoteQueryIPCStreamCoalescerFlushesDataAtFourMiB(t *testing.T) { secondData := stream.chunks[2].GetEvent().GetData() require.NotNil(t, secondData) assert.Equal(t, uint64(remoteQuerySecureIPCDataFlushBytes), secondData.GetOffset()) - assert.Equal(t, uint64((5<<20)-remoteQuerySecureIPCDataFlushBytes), secondData.GetBytes()) - assert.Len(t, secondData.GetPayload(), (5<<20)-remoteQuerySecureIPCDataFlushBytes) + assert.Equal(t, uint64(fiveMiB-remoteQuerySecureIPCDataFlushBytes), secondData.GetBytes()) + assert.Len(t, secondData.GetPayload(), fiveMiB-remoteQuerySecureIPCDataFlushBytes) } func TestRemoteQueryStreamEventFromCheckEventPreservesBinaryPayload(t *testing.T) { diff --git a/comp/remotequeries/impl/remote_query_execute.go b/comp/remotequeries/impl/remote_query_execute.go index 049f75d479cf..006946d72865 100644 --- a/comp/remotequeries/impl/remote_query_execute.go +++ b/comp/remotequeries/impl/remote_query_execute.go @@ -33,12 +33,12 @@ const ( const remoteQueryBinaryPayloadProofQuery = "SELECT decode('00ff80', 'hex') AS payload" var remoteQueryLargePayloadProofQueries = map[string]int{ - "SELECT repeat('x', 1048576) AS payload": 1 << 20, - "SELECT repeat('x', 2097152) AS payload": 2 << 20, - "SELECT repeat('x', 4194304) AS payload": 4 << 20, - "SELECT repeat('x', 8388608) AS payload": 8 << 20, - "SELECT repeat('x', 16777216) AS payload": 16 << 20, - "SELECT repeat('x', 33554432) AS payload": 32 << 20, + "SELECT repeat('x', 1048576) AS payload": 1 << 20, // 1 MiB. + "SELECT repeat('x', 2097152) AS payload": 2 << 20, // 2 MiB. + "SELECT repeat('x', 4194304) AS payload": 4 << 20, // 4 MiB. + "SELECT repeat('x', 8388608) AS payload": 8 << 20, // 8 MiB. + "SELECT repeat('x', 16777216) AS payload": 16 << 20, // 16 MiB. + "SELECT repeat('x', 33554432) AS payload": 32 << 20, // 32 MiB. } type remoteQueryStreamRunner interface { diff --git a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go index 3efa7961c8f1..9d245c0e89d9 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go @@ -39,12 +39,12 @@ const ( ) var remoteQueriesLargePayloadProofQueries = map[string]int{ - "SELECT repeat('x', 1048576) AS payload": 1 << 20, - "SELECT repeat('x', 2097152) AS payload": 2 << 20, - "SELECT repeat('x', 4194304) AS payload": 4 << 20, - "SELECT repeat('x', 8388608) AS payload": 8 << 20, - "SELECT repeat('x', 16777216) AS payload": 16 << 20, - "SELECT repeat('x', 33554432) AS payload": 32 << 20, + "SELECT repeat('x', 1048576) AS payload": 1 << 20, // 1 MiB. + "SELECT repeat('x', 2097152) AS payload": 2 << 20, // 2 MiB. + "SELECT repeat('x', 4194304) AS payload": 4 << 20, // 4 MiB. + "SELECT repeat('x', 8388608) AS payload": 8 << 20, // 8 MiB. + "SELECT repeat('x', 16777216) AS payload": 16 << 20, // 16 MiB. + "SELECT repeat('x', 33554432) AS payload": 32 << 20, // 32 MiB. } func TestRemoteQueriesActionRunsThroughLivePARLoopWithRealAgentIPC(t *testing.T) { @@ -163,13 +163,13 @@ func remoteQueriesProofQueryFromEnv() string { } func remoteQueriesProofCopyLimits(query string) map[string]interface{} { - maxBytes := 4 << 10 + maxBytes := 4 << 10 // 4 KiB. timeoutMs := 5_000 if payloadBytes, ok := remoteQueriesLargePayloadBytes(query); ok { - maxBytes = payloadBytes + (1 << 20) + maxBytes = payloadBytes + (1 << 20) // Add 1 MiB headroom. timeoutMs = 60_000 } - chunkBytes := 256 << 10 + chunkBytes := 256 << 10 // 256 KiB. if value := os.Getenv("RQ_REMOTE_CHUNK_BYTES"); value != "" { parsed, err := strconv.Atoi(value) if err == nil && parsed > 0 { From 0081379afc18460787dc15e7ebbefcef5b1df85a Mon Sep 17 00:00:00 2001 From: nubtron Date: Thu, 21 May 2026 19:07:26 +0000 Subject: [PATCH 30/33] Fix Remote Queries streaming proof assertions --- .../live_agent_ipc_par_loop_test.go | 40 ++----------------- .../fused-local-par-agent-postgres-proof.sh | 29 ++------------ 2 files changed, 6 insertions(+), 63 deletions(-) diff --git a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go index 9d245c0e89d9..7eb01cab8fdd 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go @@ -124,11 +124,11 @@ func TestRemoteQueriesActionRunsThroughLivePARLoopWithRealAgentIPC(t *testing.T) require.True(t, result.Success) require.Equal(t, taskID, result.TaskID) assert.Equal(t, "SUCCEEDED", result.Outputs["status"]) - require.Contains(t, result.Outputs, "rows") + require.Contains(t, result.Outputs, "data") - rows, ok := result.Outputs["rows"].([]interface{}) + data, ok := result.Outputs["data"].(string) require.True(t, ok) - assertRemoteQueriesProofRows(t, proofQuery, rows) + assertRemoteQueriesProofCopyData(t, proofQuery, data) resultEvidence, err := json.Marshal(summarizeRemoteQueriesProofPayload(result.Outputs)) require.NoError(t, err) @@ -230,40 +230,6 @@ func assertRemoteQueriesProofCopyData(t *testing.T, query string, data string) { } } -func assertRemoteQueriesProofRows(t *testing.T, query string, rows []interface{}) { - t.Helper() - - switch query { - case remoteQueriesFixtureTableProofQuery: - require.Len(t, rows, 2) - firstRow, ok := rows[0].(map[string]interface{}) - require.True(t, ok) - assert.Equal(t, "Beautiful city of lights", firstRow["city"]) - assert.Equal(t, "France", firstRow["country"]) - secondRow, ok := rows[1].(map[string]interface{}) - require.True(t, ok) - assert.Equal(t, "New York", secondRow["city"]) - assert.Equal(t, "USA", secondRow["country"]) - case remoteQueriesSeedProofQuery: - require.Len(t, rows, 1) - firstRow, ok := rows[0].(map[string]interface{}) - require.True(t, ok) - assert.Equal(t, float64(1), firstRow["value"]) - default: - expectedPayloadBytes, ok := remoteQueriesLargePayloadBytes(query) - if ok { - require.Len(t, rows, 1) - firstRow, ok := rows[0].(map[string]interface{}) - require.True(t, ok) - payload, ok := firstRow["payload"].(string) - require.True(t, ok, "payload field must be a string") - assert.Len(t, payload, expectedPayloadBytes) - return - } - require.FailNowf(t, "unsupported proof query", "%s=%q must use a bridge-allowlisted proof query", remoteQueriesProofQueryOverrideEnv, query) - } -} - func summarizeRemoteQueriesProofPayload(value interface{}) interface{} { switch typed := value.(type) { case map[string]interface{}: diff --git a/test/remotequeries/fused-local-par-agent-postgres-proof.sh b/test/remotequeries/fused-local-par-agent-postgres-proof.sh index e593ccae1ea7..2ea5a23348cb 100755 --- a/test/remotequeries/fused-local-par-agent-postgres-proof.sh +++ b/test/remotequeries/fused-local-par-agent-postgres-proof.sh @@ -308,7 +308,7 @@ PY local token token=$(cat "$TMP_ROOT/run/auth_token") - log "Preflight real Agent IPC HTTP execute endpoint (dev evidence only)" + log "Preflight real Agent IPC HTTP execute endpoint is disabled; Remote Queries execution requires AgentSecure COPY streaming" local status status=$(curl -sS -k -o "$TMP_ROOT/results/agent-execute-preflight.body" -w '%{http_code}' \ -H "Authorization: Bearer ${token}" \ @@ -319,33 +319,10 @@ PY cat "$TMP_ROOT/results/agent-execute-preflight.body" printf '\n' - if [[ "$status" != "200" ]]; then - echo "FAIL: expected Agent execute preflight HTTP 200, got $status" >&2 + if [[ "$status" != "400" ]]; then + echo "FAIL: expected Agent execute preflight HTTP 400 for disabled inline execution, got $status" >&2 exit 1 fi - RQ_REMOTE_QUERY="$RQ_REMOTE_QUERY" python3 - "$TMP_ROOT/results/agent-execute-preflight.body" <<'PY' -import json -import os -import sys - -with open(sys.argv[1], encoding="utf-8") as f: - body = json.load(f) - -if body.get("status") != "SUCCEEDED": - raise SystemExit(f"Agent execute preflight response status was not SUCCEEDED: {body}") - -rows = body.get("rows") -query = os.environ["RQ_REMOTE_QUERY"] -if query == "SELECT city, country FROM cities ORDER BY city": - expected = [ - {"city": "Beautiful city of lights", "country": "France"}, - {"city": "New York", "country": "USA"}, - ] -else: - expected = [{"value": 1}] -if rows != expected: - raise SystemExit(f"Agent execute preflight rows did not match expected fixture data: rows={rows!r} expected={expected!r}") -PY if grep -Eq 'password|token|secret' "$TMP_ROOT/results/agent-execute-preflight.body"; then echo "FAIL: Agent execute preflight response contained credential-shaped text" >&2 exit 1 From b4607db962ad60fd18a2e23c5ed782790767f4bb Mon Sep 17 00:00:00 2001 From: nubtron Date: Fri, 22 May 2026 13:12:08 +0000 Subject: [PATCH 31/33] Add Bazel files for remote queries packages --- comp/api/grpcserver/impl-agent/server.go | 5 +- comp/remotequeries/fx/BUILD.bazel | 13 + comp/remotequeries/impl/BUILD.bazel | 45 +++ .../impl/remote_query_execute.go | 24 +- comp/remotequeries/impl/remote_query_match.go | 21 +- .../impl/remote_query_par_poc.go | 11 +- comp/remotequeries/impl/remote_query_test.go | 2 +- go.mod | 1 + pkg/collector/python/check.go | 12 +- .../bundles/remotequeries/BUILD.bazel | 256 ++++++++++++++++++ .../bundles/remotequeries/execute.go | 63 +---- .../live_agent_ipc_par_loop_test.go | 8 +- .../standalone_par_process_proof_test.go | 2 +- 13 files changed, 372 insertions(+), 91 deletions(-) create mode 100644 comp/remotequeries/fx/BUILD.bazel create mode 100644 comp/remotequeries/impl/BUILD.bazel create mode 100644 pkg/privateactionrunner/bundles/remotequeries/BUILD.bazel diff --git a/comp/api/grpcserver/impl-agent/server.go b/comp/api/grpcserver/impl-agent/server.go index 60cf34a1a136..f347743a4ada 100644 --- a/comp/api/grpcserver/impl-agent/server.go +++ b/comp/api/grpcserver/impl-agent/server.go @@ -10,7 +10,6 @@ import ( "context" "encoding/json" "errors" - "fmt" "strconv" "strings" "time" @@ -284,7 +283,7 @@ func (s *serverSecure) WorkloadFilterEvaluate(ctx context.Context, req *pb.Workl return s.workloadfilterServer.WorkloadFilterEvaluate(ctx, req) } -func (s *serverSecure) RemoteQueryExecute(_ context.Context, req *pb.RemoteQueryExecuteRequest) (*pb.RemoteQueryExecuteResponse, error) { +func (s *serverSecure) RemoteQueryExecute(_ context.Context, _ *pb.RemoteQueryExecuteRequest) (*pb.RemoteQueryExecuteResponse, error) { return remoteQueryExecuteErrorResponse(remotequeriesimpl.RemoteQueryStatusInvalidRequest, "remote queries require RemoteQueryExecuteStream with operation copy_stream"), nil } @@ -492,7 +491,7 @@ func remoteQueryExecuteRequestFromProto(req *pb.RemoteQueryExecuteRequest) (remo DBName: req.GetTarget().GetDbname(), } if req.GetOperation() != "copy_stream" { - return remotequeriesimpl.RemoteQueryExecuteRequest{}, fmt.Errorf("operation must be copy_stream") + return remotequeriesimpl.RemoteQueryExecuteRequest{}, errors.New("operation must be copy_stream") } return remotequeriesimpl.NewRemoteQueryCopyStreamExecuteRequest(req.GetIntegration(), target, req.GetQuery(), req.GetFormat(), remoteQueryCopyLimitsFromProto(req.GetCopyLimits())) } diff --git a/comp/remotequeries/fx/BUILD.bazel b/comp/remotequeries/fx/BUILD.bazel new file mode 100644 index 000000000000..be2c2428cefa --- /dev/null +++ b/comp/remotequeries/fx/BUILD.bazel @@ -0,0 +1,13 @@ +load("@rules_go//go:def.bzl", "go_library") + +go_library( + name = "fx", + srcs = ["fx.go"], + importpath = "github.com/DataDog/datadog-agent/comp/remotequeries/fx", + visibility = ["//visibility:public"], + deps = [ + "//comp/remotequeries/impl", + "//pkg/util/fxutil", + "@org_uber_go_fx//:fx", + ], +) diff --git a/comp/remotequeries/impl/BUILD.bazel b/comp/remotequeries/impl/BUILD.bazel new file mode 100644 index 000000000000..652370c86b7b --- /dev/null +++ b/comp/remotequeries/impl/BUILD.bazel @@ -0,0 +1,45 @@ +load("@rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "impl", + srcs = [ + "remote_query_execute.go", + "remote_query_match.go", + "remote_query_par_poc.go", + ], + importpath = "github.com/DataDog/datadog-agent/comp/remotequeries/impl", + visibility = ["//visibility:public"], + deps = [ + "//comp/api/api/def", + "//comp/core/config", + "//comp/core/ipc/def", + "//comp/core/ipc/httphelpers", + "@//comp/collector/collector", + "@//pkg/collector/check", + "@in_gopkg_yaml_v3//:yaml_v3", + "@org_uber_go_fx//:fx", + ], +) + +go_test( + name = "impl_test", + srcs = [ + "remote_query_par_poc_test.go", + "remote_query_test.go", + ], + embed = [":impl"], + gotags = ["test"], + deps = [ + "//comp/core/ipc/def", + "//comp/core/ipc/mock", + "@//comp/collector/collector", + "@//comp/core/autodiscovery/integration", + "@//comp/core/diagnose/def", + "@//pkg/aggregator/sender", + "@//pkg/collector/check", + "@//pkg/collector/check/id", + "@//pkg/collector/check/stats", + "@com_github_stretchr_testify//assert", + "@com_github_stretchr_testify//require", + ], +) diff --git a/comp/remotequeries/impl/remote_query_execute.go b/comp/remotequeries/impl/remote_query_execute.go index 006946d72865..303b474cf36f 100644 --- a/comp/remotequeries/impl/remote_query_execute.go +++ b/comp/remotequeries/impl/remote_query_execute.go @@ -146,16 +146,16 @@ func NewRemoteQueryCopyStreamExecuteRequest(integration string, target RemoteQue return RemoteQueryExecuteRequest{}, err } if query == "" { - return RemoteQueryExecuteRequest{}, fmt.Errorf("query is required") + return RemoteQueryExecuteRequest{}, errors.New("query is required") } if !isRemoteQueryAllowedProofQuery(query) { - return RemoteQueryExecuteRequest{}, fmt.Errorf("query is not allowed") + return RemoteQueryExecuteRequest{}, errors.New("query is not allowed") } if format == "" { format = "csv" } if format != "csv" && format != "binary" { - return RemoteQueryExecuteRequest{}, fmt.Errorf("format must be csv or binary") + return RemoteQueryExecuteRequest{}, errors.New("format must be csv or binary") } var parsedLimits *remoteQueryExecuteCopyLimits if limits != nil { @@ -193,8 +193,8 @@ const ( ) // NewRemoteQueryExecuteRequest rejects legacy inline Remote Queries requests. -func NewRemoteQueryExecuteRequest(integration string, target RemoteQueryExecuteTarget, query string, limits *RemoteQueryExecuteLimits) (RemoteQueryExecuteRequest, error) { - return RemoteQueryExecuteRequest{}, fmt.Errorf("operation must be copy_stream") +func NewRemoteQueryExecuteRequest(_ string, _ RemoteQueryExecuteTarget, _ string, _ *RemoteQueryExecuteLimits) (RemoteQueryExecuteRequest, error) { + return RemoteQueryExecuteRequest{}, errors.New("operation must be copy_stream") } type remoteQueryExecuteRequest struct { @@ -314,10 +314,10 @@ func parseExecuteRequest(r *http.Request) (remoteQueryExecuteRequest, string, er } if wireReq.Query == "" { - return remoteQueryExecuteRequest{}, "", fmt.Errorf("query is required") + return remoteQueryExecuteRequest{}, "", errors.New("query is required") } if !isRemoteQueryAllowedProofQuery(wireReq.Query) { - return remoteQueryExecuteRequest{}, "", fmt.Errorf("query is not allowed") + return remoteQueryExecuteRequest{}, "", errors.New("query is not allowed") } limits, err := parseExecuteLimits(wireReq.Limits) @@ -330,19 +330,19 @@ func parseExecuteRequest(r *http.Request) (remoteQueryExecuteRequest, string, er } if wireReq.Operation != "copy_stream" { - return remoteQueryExecuteRequest{}, "", fmt.Errorf("operation must be copy_stream") + return remoteQueryExecuteRequest{}, "", errors.New("operation must be copy_stream") } if wireReq.Format == "" { wireReq.Format = "csv" } if wireReq.Format != "csv" && wireReq.Format != "binary" { - return remoteQueryExecuteRequest{}, "", fmt.Errorf("format must be csv or binary") + return remoteQueryExecuteRequest{}, "", errors.New("format must be csv or binary") } req := remoteQueryExecuteRequest{Integration: integration, Operation: wireReq.Operation, Target: target, Query: wireReq.Query, Format: wireReq.Format, Limits: limits, CopyLimits: copyLimits} requestJSON, err := marshalExecuteRequest(req) if err != nil { - return remoteQueryExecuteRequest{}, "", fmt.Errorf("malformed JSON request") + return remoteQueryExecuteRequest{}, "", errors.New("malformed JSON request") } return req, requestJSON, nil } @@ -440,7 +440,7 @@ func parseRequiredPositiveInt(value *int, name string) (int, error) { return *value, nil } -func (s *RemoteQueryExecuteService) Execute(req RemoteQueryExecuteRequest) RemoteQueryExecuteResult { +func (s *RemoteQueryExecuteService) Execute(_ RemoteQueryExecuteRequest) RemoteQueryExecuteResult { return remoteQueryExecuteErrorResult(http.StatusBadRequest, statusInvalidRequest, "remote queries require operation copy_stream and the streaming executor") } @@ -534,7 +534,7 @@ func remoteQueryExecuteRequestFromInternal(req remoteQueryExecuteRequest) Remote func marshalExecuteRequest(req remoteQueryExecuteRequest) (string, error) { if req.Operation != "copy_stream" { - return "", fmt.Errorf("operation must be copy_stream") + return "", errors.New("operation must be copy_stream") } format := req.Format if format == "" { diff --git a/comp/remotequeries/impl/remote_query_match.go b/comp/remotequeries/impl/remote_query_match.go index 078757ceae04..7b41fd0657c2 100644 --- a/comp/remotequeries/impl/remote_query_match.go +++ b/comp/remotequeries/impl/remote_query_match.go @@ -10,7 +10,6 @@ import ( "bytes" "encoding/json" "errors" - "fmt" "io" "mime" "net/http" @@ -172,7 +171,7 @@ func parseMatchRequest(r *http.Request) (remoteQueryMatchRequest, error) { func parseIntegration(integration string) (string, error) { integration = strings.ToLower(strings.TrimSpace(integration)) if integration == "" { - return "", fmt.Errorf("integration is required") + return "", errors.New("integration is required") } if !integrationNamePattern.MatchString(integration) { return "", invalidRequestError("integration contains invalid characters") @@ -213,12 +212,12 @@ func (t *remoteQueryTargetRequestJSON) UnmarshalJSON(data []byte) error { func parseTarget(target *remoteQueryTargetRequestJSON) (remoteQueryTarget, error) { if target == nil { - return remoteQueryTarget{}, fmt.Errorf("target is required") + return remoteQueryTarget{}, errors.New("target is required") } host := normalizeHost(target.Host) if host == "" { - return remoteQueryTarget{}, fmt.Errorf("target.host is required") + return remoteQueryTarget{}, errors.New("target.host is required") } port, err := parseRequiredPort(target.Port) @@ -227,7 +226,7 @@ func parseTarget(target *remoteQueryTargetRequestJSON) (remoteQueryTarget, error } if target.DBName == "" { - return remoteQueryTarget{}, fmt.Errorf("target.dbname is required") + return remoteQueryTarget{}, errors.New("target.dbname is required") } return remoteQueryTarget{Host: host, Port: port, DBName: target.DBName}, nil @@ -235,10 +234,10 @@ func parseTarget(target *remoteQueryTargetRequestJSON) (remoteQueryTarget, error func parseRequiredPort(port *int) (int, error) { if port == nil { - return 0, fmt.Errorf("target.port is required") + return 0, errors.New("target.port is required") } if *port < 1 || *port > 65535 { - return 0, fmt.Errorf("target.port is out of range") + return 0, errors.New("target.port is out of range") } return *port, nil } @@ -275,15 +274,15 @@ func parseJSONRequestError(err error) error { if errors.As(err, &typeErr) { switch typeErr.Field { case "port", "target.port": - return fmt.Errorf("target.port must be an integer") + return errors.New("target.port must be an integer") case "target": return errTargetMustBeObject case "maxRows", "limits.maxRows": - return fmt.Errorf("limits.maxRows must be an integer") + return errors.New("limits.maxRows must be an integer") case "maxBytes", "limits.maxBytes": - return fmt.Errorf("limits.maxBytes must be an integer") + return errors.New("limits.maxBytes must be an integer") case "timeoutMs", "limits.timeoutMs": - return fmt.Errorf("limits.timeoutMs must be an integer") + return errors.New("limits.timeoutMs must be an integer") case "limits": return errLimitsMustBeObject } diff --git a/comp/remotequeries/impl/remote_query_par_poc.go b/comp/remotequeries/impl/remote_query_par_poc.go index f5993caf186b..d8fc3374a4cf 100644 --- a/comp/remotequeries/impl/remote_query_par_poc.go +++ b/comp/remotequeries/impl/remote_query_par_poc.go @@ -9,6 +9,7 @@ import ( "bytes" "context" "encoding/json" + "errors" "fmt" "io" @@ -62,10 +63,10 @@ func NewRemoteQueryPARHarness(client remoteQueryPARIPCClient, endpointURL string // Execute sends a credential-free target/query request through the injected Agent IPC HTTP client. func (h *RemoteQueryPARHarness) Execute(ctx context.Context, inputs RemoteQueryPARInputs) (RemoteQueryPARResult, error) { if h == nil || h.client == nil { - return RemoteQueryPARResult{}, fmt.Errorf("remote query PAR harness requires an IPC client") + return RemoteQueryPARResult{}, errors.New("remote query PAR harness requires an IPC client") } if h.endpointURL == "" { - return RemoteQueryPARResult{}, fmt.Errorf("remote query PAR harness requires an endpoint URL") + return RemoteQueryPARResult{}, errors.New("remote query PAR harness requires an endpoint URL") } payload, err := json.Marshal(inputs) @@ -82,7 +83,7 @@ func (h *RemoteQueryPARHarness) Execute(ctx context.Context, inputs RemoteQueryP } if postErr != nil { if len(body) > 0 { - return RemoteQueryPARResult{}, fmt.Errorf("remote query IPC request failed with undecodable response") + return RemoteQueryPARResult{}, errors.New("remote query IPC request failed with undecodable response") } return RemoteQueryPARResult{}, fmt.Errorf("remote query IPC request failed: %w", postErr) } @@ -91,7 +92,7 @@ func (h *RemoteQueryPARHarness) Execute(ctx context.Context, inputs RemoteQueryP func decodeRemoteQueryPARResponse(body []byte) (RemoteQueryPARResult, error) { if len(body) == 0 { - return RemoteQueryPARResult{}, fmt.Errorf("empty remote query response") + return RemoteQueryPARResult{}, errors.New("empty remote query response") } decoder := json.NewDecoder(bytes.NewReader(body)) @@ -102,7 +103,7 @@ func decodeRemoteQueryPARResponse(body []byte) (RemoteQueryPARResult, error) { return RemoteQueryPARResult{}, fmt.Errorf("decode remote query response: %w", err) } if result.Status == "" { - return RemoteQueryPARResult{}, fmt.Errorf("remote query response missing status") + return RemoteQueryPARResult{}, errors.New("remote query response missing status") } result.Raw = append(result.Raw[:0], body...) return result, nil diff --git a/comp/remotequeries/impl/remote_query_test.go b/comp/remotequeries/impl/remote_query_test.go index b50496af3051..784f052fd955 100644 --- a/comp/remotequeries/impl/remote_query_test.go +++ b/comp/remotequeries/impl/remote_query_test.go @@ -222,7 +222,7 @@ type fakeCollector struct { func (f fakeCollector) RunCheck(inner check.Check) (checkid.ID, error) { return inner.ID(), nil } func (f fakeCollector) StopCheck(checkid.ID) error { return nil } -func (f fakeCollector) MapOverChecks(cb func([]check.Info)) {} +func (f fakeCollector) MapOverChecks(_ func([]check.Info)) {} func (f fakeCollector) GetChecks() []check.Check { return f.checks } func (f fakeCollector) ReloadAllCheckInstances(string, []check.Check) ([]checkid.ID, error) { return nil, nil diff --git a/go.mod b/go.mod index c374a7c50db6..79da821294b0 100644 --- a/go.mod +++ b/go.mod @@ -965,6 +965,7 @@ require ( github.com/DataDog/datadog-agent/pkg/trace/traceutil v0.77.0-devel.0.20260213154712-e02b9359151a github.com/DataDog/datadog-agent/pkg/util/cgroups v0.64.0-rc.3 github.com/DataDog/datadog-agent/pkg/util/kubernetes/apiserver/common/namespace v0.77.0-devel.0.20260211235139-a5361978c2b6 + github.com/DataDog/datadog-agent/test/fakeintake v0.0.0-00010101000000-000000000000 github.com/DataDog/ddtrivy v0.0.0-20260115083325-07614fb0b8d5 github.com/DataDog/rshell v0.0.13 github.com/aws/aws-sdk-go-v2/service/secretsmanager v1.41.0 diff --git a/pkg/collector/python/check.go b/pkg/collector/python/check.go index 271c7c830730..59354358119f 100644 --- a/pkg/collector/python/check.go +++ b/pkg/collector/python/check.go @@ -44,8 +44,8 @@ char *getStringAddr(char **array, unsigned int idx); extern int remoteQueryStreamEmitBridge(const char *event_type, const char *metadata_json, const uint8_t *payload, size_t payload_len, void *userdata); int run_remote_query_stream(rtloader_t *, rtloader_pyobject_t *check, const char *integration, const char *request_json, int (*emit)(const char *, const char *, const uint8_t *, size_t, void *), void *userdata); -static inline int call_run_remote_query_stream(rtloader_t *rtloader, rtloader_pyobject_t *check, const char *integration, const char *request_json, void *userdata) { - return run_remote_query_stream(rtloader, check, integration, request_json, remoteQueryStreamEmitBridge, userdata); +static inline int call_run_remote_query_stream(rtloader_t *rtloader, rtloader_pyobject_t *check, const char *integration, const char *request_json, uintptr_t userdata) { + return run_remote_query_stream(rtloader, check, integration, request_json, remoteQueryStreamEmitBridge, (void *)userdata); } static inline void call_free(void* ptr) { @@ -169,10 +169,10 @@ func (c *PythonCheck) RunSimple() error { func (c *PythonCheck) RunRemoteQueryStream(integration string, requestJSON string, emit func(checkbase.RemoteQueryStreamEvent) error) error { integration = strings.ToLower(strings.TrimSpace(integration)) if integration == "" { - return fmt.Errorf("integration is required") + return errors.New("integration is required") } if emit == nil { - return fmt.Errorf("emit callback is required") + return errors.New("emit callback is required") } gstate, err := newStickyLock() @@ -192,12 +192,12 @@ func (c *PythonCheck) RunRemoteQueryStream(integration string, requestJSON strin h := cgo.NewHandle(emit) defer h.Delete() - ok := C.call_run_remote_query_stream(rtloader, c.instance, cIntegration, cRequestJSON, unsafe.Pointer(h)) + ok := C.call_run_remote_query_stream(rtloader, c.instance, cIntegration, cRequestJSON, C.uintptr_t(h)) if ok == 0 { if err := getRtLoaderError(); err != nil { return err } - return fmt.Errorf("an error occurred while running remote query stream") + return errors.New("an error occurred while running remote query stream") } return nil } diff --git a/pkg/privateactionrunner/bundles/remotequeries/BUILD.bazel b/pkg/privateactionrunner/bundles/remotequeries/BUILD.bazel new file mode 100644 index 000000000000..c6be2fc8c22b --- /dev/null +++ b/pkg/privateactionrunner/bundles/remotequeries/BUILD.bazel @@ -0,0 +1,256 @@ +load("@rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "remotequeries", + srcs = [ + "entrypoint.go", + "execute.go", + "ipc_client.go", + ], + importpath = "github.com/DataDog/datadog-agent/pkg/privateactionrunner/bundles/remotequeries", + visibility = ["//visibility:public"], + deps = [ + "//pkg/api/security/cert", + "//pkg/config/setup", + "//pkg/privateactionrunner/libs/privateconnection", + "//pkg/privateactionrunner/types", + "//pkg/privateactionrunner/util", + "//pkg/proto/pbgo/core", + "//pkg/util/grpc", + "//pkg/util/system", + "@org_golang_google_grpc//:grpc", + ], +) + +go_test( + name = "remotequeries_test", + srcs = [ + "execute_test.go", + "live_agent_ipc_par_loop_test.go", + "live_par_loop_test.go", + "standalone_par_process_proof_test.go", + ], + embed = [":remotequeries"], + gotags = ["test"], + deps = [ + "//pkg/privateactionrunner/libs/privateconnection", + "//pkg/privateactionrunner/types", + "//pkg/privateactionrunner/util", + "//pkg/proto/pbgo/core", + "@com_github_stretchr_testify//assert", + "@com_github_stretchr_testify//require", + "@org_golang_google_grpc//:grpc", + ] + select({ + "@rules_go//go/platform:aix": [ + "//pkg/config/setup", + "//pkg/privateactionrunner/adapters/config", + "//pkg/privateactionrunner/observability", + "//pkg/privateactionrunner/opms", + "@//pkg/privateactionrunner/adapters/constants", + "@//pkg/privateactionrunner/runners", + "@//pkg/privateactionrunner/task-verifier", + "@com_github_datadog_datadog_agent_test_fakeintake//client", + "@com_github_datadog_datadog_agent_test_fakeintake//server", + "@com_github_datadog_datadog_go_v5//statsd", + "@io_k8s_apimachinery//pkg/util/sets", + "@org_golang_google_protobuf//types/known/structpb", + ], + "@rules_go//go/platform:android": [ + "//pkg/config/setup", + "//pkg/privateactionrunner/adapters/config", + "//pkg/privateactionrunner/observability", + "//pkg/privateactionrunner/opms", + "@//pkg/privateactionrunner/adapters/constants", + "@//pkg/privateactionrunner/runners", + "@//pkg/privateactionrunner/task-verifier", + "@com_github_datadog_datadog_agent_test_fakeintake//client", + "@com_github_datadog_datadog_agent_test_fakeintake//server", + "@com_github_datadog_datadog_go_v5//statsd", + "@io_k8s_apimachinery//pkg/util/sets", + "@org_golang_google_protobuf//types/known/structpb", + ], + "@rules_go//go/platform:darwin": [ + "//pkg/config/setup", + "//pkg/privateactionrunner/adapters/config", + "//pkg/privateactionrunner/observability", + "//pkg/privateactionrunner/opms", + "@//pkg/privateactionrunner/adapters/constants", + "@//pkg/privateactionrunner/runners", + "@//pkg/privateactionrunner/task-verifier", + "@com_github_datadog_datadog_agent_test_fakeintake//client", + "@com_github_datadog_datadog_agent_test_fakeintake//server", + "@com_github_datadog_datadog_go_v5//statsd", + "@io_k8s_apimachinery//pkg/util/sets", + "@org_golang_google_protobuf//types/known/structpb", + ], + "@rules_go//go/platform:dragonfly": [ + "//pkg/config/setup", + "//pkg/privateactionrunner/adapters/config", + "//pkg/privateactionrunner/observability", + "//pkg/privateactionrunner/opms", + "@//pkg/privateactionrunner/adapters/constants", + "@//pkg/privateactionrunner/runners", + "@//pkg/privateactionrunner/task-verifier", + "@com_github_datadog_datadog_agent_test_fakeintake//client", + "@com_github_datadog_datadog_agent_test_fakeintake//server", + "@com_github_datadog_datadog_go_v5//statsd", + "@io_k8s_apimachinery//pkg/util/sets", + "@org_golang_google_protobuf//types/known/structpb", + ], + "@rules_go//go/platform:freebsd": [ + "//pkg/config/setup", + "//pkg/privateactionrunner/adapters/config", + "//pkg/privateactionrunner/observability", + "//pkg/privateactionrunner/opms", + "@//pkg/privateactionrunner/adapters/constants", + "@//pkg/privateactionrunner/runners", + "@//pkg/privateactionrunner/task-verifier", + "@com_github_datadog_datadog_agent_test_fakeintake//client", + "@com_github_datadog_datadog_agent_test_fakeintake//server", + "@com_github_datadog_datadog_go_v5//statsd", + "@io_k8s_apimachinery//pkg/util/sets", + "@org_golang_google_protobuf//types/known/structpb", + ], + "@rules_go//go/platform:illumos": [ + "//pkg/config/setup", + "//pkg/privateactionrunner/adapters/config", + "//pkg/privateactionrunner/observability", + "//pkg/privateactionrunner/opms", + "@//pkg/privateactionrunner/adapters/constants", + "@//pkg/privateactionrunner/runners", + "@//pkg/privateactionrunner/task-verifier", + "@com_github_datadog_datadog_agent_test_fakeintake//client", + "@com_github_datadog_datadog_agent_test_fakeintake//server", + "@com_github_datadog_datadog_go_v5//statsd", + "@io_k8s_apimachinery//pkg/util/sets", + "@org_golang_google_protobuf//types/known/structpb", + ], + "@rules_go//go/platform:ios": [ + "//pkg/config/setup", + "//pkg/privateactionrunner/adapters/config", + "//pkg/privateactionrunner/observability", + "//pkg/privateactionrunner/opms", + "@//pkg/privateactionrunner/adapters/constants", + "@//pkg/privateactionrunner/runners", + "@//pkg/privateactionrunner/task-verifier", + "@com_github_datadog_datadog_agent_test_fakeintake//client", + "@com_github_datadog_datadog_agent_test_fakeintake//server", + "@com_github_datadog_datadog_go_v5//statsd", + "@io_k8s_apimachinery//pkg/util/sets", + "@org_golang_google_protobuf//types/known/structpb", + ], + "@rules_go//go/platform:js": [ + "//pkg/config/setup", + "//pkg/privateactionrunner/adapters/config", + "//pkg/privateactionrunner/observability", + "//pkg/privateactionrunner/opms", + "@//pkg/privateactionrunner/adapters/constants", + "@//pkg/privateactionrunner/runners", + "@//pkg/privateactionrunner/task-verifier", + "@com_github_datadog_datadog_agent_test_fakeintake//client", + "@com_github_datadog_datadog_agent_test_fakeintake//server", + "@com_github_datadog_datadog_go_v5//statsd", + "@io_k8s_apimachinery//pkg/util/sets", + "@org_golang_google_protobuf//types/known/structpb", + ], + "@rules_go//go/platform:linux": [ + "//pkg/config/setup", + "//pkg/privateactionrunner/adapters/config", + "//pkg/privateactionrunner/observability", + "//pkg/privateactionrunner/opms", + "@//pkg/privateactionrunner/adapters/constants", + "@//pkg/privateactionrunner/runners", + "@//pkg/privateactionrunner/task-verifier", + "@com_github_datadog_datadog_agent_test_fakeintake//client", + "@com_github_datadog_datadog_agent_test_fakeintake//server", + "@com_github_datadog_datadog_go_v5//statsd", + "@io_k8s_apimachinery//pkg/util/sets", + "@org_golang_google_protobuf//types/known/structpb", + ], + "@rules_go//go/platform:netbsd": [ + "//pkg/config/setup", + "//pkg/privateactionrunner/adapters/config", + "//pkg/privateactionrunner/observability", + "//pkg/privateactionrunner/opms", + "@//pkg/privateactionrunner/adapters/constants", + "@//pkg/privateactionrunner/runners", + "@//pkg/privateactionrunner/task-verifier", + "@com_github_datadog_datadog_agent_test_fakeintake//client", + "@com_github_datadog_datadog_agent_test_fakeintake//server", + "@com_github_datadog_datadog_go_v5//statsd", + "@io_k8s_apimachinery//pkg/util/sets", + "@org_golang_google_protobuf//types/known/structpb", + ], + "@rules_go//go/platform:openbsd": [ + "//pkg/config/setup", + "//pkg/privateactionrunner/adapters/config", + "//pkg/privateactionrunner/observability", + "//pkg/privateactionrunner/opms", + "@//pkg/privateactionrunner/adapters/constants", + "@//pkg/privateactionrunner/runners", + "@//pkg/privateactionrunner/task-verifier", + "@com_github_datadog_datadog_agent_test_fakeintake//client", + "@com_github_datadog_datadog_agent_test_fakeintake//server", + "@com_github_datadog_datadog_go_v5//statsd", + "@io_k8s_apimachinery//pkg/util/sets", + "@org_golang_google_protobuf//types/known/structpb", + ], + "@rules_go//go/platform:osx": [ + "//pkg/config/setup", + "//pkg/privateactionrunner/adapters/config", + "//pkg/privateactionrunner/observability", + "//pkg/privateactionrunner/opms", + "@//pkg/privateactionrunner/adapters/constants", + "@//pkg/privateactionrunner/runners", + "@//pkg/privateactionrunner/task-verifier", + "@com_github_datadog_datadog_agent_test_fakeintake//client", + "@com_github_datadog_datadog_agent_test_fakeintake//server", + "@com_github_datadog_datadog_go_v5//statsd", + "@io_k8s_apimachinery//pkg/util/sets", + "@org_golang_google_protobuf//types/known/structpb", + ], + "@rules_go//go/platform:plan9": [ + "//pkg/config/setup", + "//pkg/privateactionrunner/adapters/config", + "//pkg/privateactionrunner/observability", + "//pkg/privateactionrunner/opms", + "@//pkg/privateactionrunner/adapters/constants", + "@//pkg/privateactionrunner/runners", + "@//pkg/privateactionrunner/task-verifier", + "@com_github_datadog_datadog_agent_test_fakeintake//client", + "@com_github_datadog_datadog_agent_test_fakeintake//server", + "@com_github_datadog_datadog_go_v5//statsd", + "@io_k8s_apimachinery//pkg/util/sets", + "@org_golang_google_protobuf//types/known/structpb", + ], + "@rules_go//go/platform:qnx": [ + "//pkg/config/setup", + "//pkg/privateactionrunner/adapters/config", + "//pkg/privateactionrunner/observability", + "//pkg/privateactionrunner/opms", + "@//pkg/privateactionrunner/adapters/constants", + "@//pkg/privateactionrunner/runners", + "@//pkg/privateactionrunner/task-verifier", + "@com_github_datadog_datadog_agent_test_fakeintake//client", + "@com_github_datadog_datadog_agent_test_fakeintake//server", + "@com_github_datadog_datadog_go_v5//statsd", + "@io_k8s_apimachinery//pkg/util/sets", + "@org_golang_google_protobuf//types/known/structpb", + ], + "@rules_go//go/platform:solaris": [ + "//pkg/config/setup", + "//pkg/privateactionrunner/adapters/config", + "//pkg/privateactionrunner/observability", + "//pkg/privateactionrunner/opms", + "@//pkg/privateactionrunner/adapters/constants", + "@//pkg/privateactionrunner/runners", + "@//pkg/privateactionrunner/task-verifier", + "@com_github_datadog_datadog_agent_test_fakeintake//client", + "@com_github_datadog_datadog_agent_test_fakeintake//server", + "@com_github_datadog_datadog_go_v5//statsd", + "@io_k8s_apimachinery//pkg/util/sets", + "@org_golang_google_protobuf//types/known/structpb", + ], + "//conditions:default": [], + }), +) diff --git a/pkg/privateactionrunner/bundles/remotequeries/execute.go b/pkg/privateactionrunner/bundles/remotequeries/execute.go index 73e3d8c11529..989ba74aa5be 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/execute.go +++ b/pkg/privateactionrunner/bundles/remotequeries/execute.go @@ -9,7 +9,7 @@ import ( "bytes" "context" "encoding/json" - "fmt" + "errors" "io" "time" "unicode/utf8" @@ -77,20 +77,20 @@ func (a *ExecuteAction) Run( inputs, err := types.ExtractInputs[ExecuteInputs](task) if err != nil { return nil, util.DefaultActionErrorWithDisplayError( - fmt.Errorf("invalid remote query action inputs"), + errors.New("invalid remote query action inputs"), "invalid remote query action inputs", ) } if a == nil || a.newBridgeClient == nil { - return nil, util.DefaultActionError(fmt.Errorf("remote query action requires an Agent IPC client")) + return nil, util.DefaultActionError(errors.New("remote query action requires an Agent IPC client")) } client, err := a.newBridgeClient() if err != nil { return nil, util.DefaultActionErrorWithDisplayError(err, "remote query action could not create an Agent IPC client") } if client == nil { - return nil, util.DefaultActionError(fmt.Errorf("remote query action requires an AgentSecure client")) + return nil, util.DefaultActionError(errors.New("remote query action requires an AgentSecure client")) } rpcStart := time.Now() @@ -144,7 +144,7 @@ func remoteQueryExecuteRequestFromInputs(inputs ExecuteInputs) *pb.RemoteQueryEx func remoteQueryExecuteOutputFromStream(stream grpc.ServerStreamingClient[pb.RemoteQueryExecuteChunk]) (map[string]interface{}, error) { if stream == nil { - return nil, fmt.Errorf("remote query response stream missing") + return nil, errors.New("remote query response stream missing") } typedStreamEvents := make([]map[string]interface{}, 0) @@ -164,16 +164,16 @@ func remoteQueryExecuteOutputFromStream(stream grpc.ServerStreamingClient[pb.Rem return nil, err } if chunk == nil { - return nil, fmt.Errorf("remote query response stream returned nil chunk") + return nil, errors.New("remote query response stream returned nil chunk") } if firstChunkAt.IsZero() { firstChunkAt = time.Now() } if chunk.GetChunkIndex() != expectedChunkIndex { - return nil, fmt.Errorf("remote query response stream chunk index mismatch") + return nil, errors.New("remote query response stream chunk index mismatch") } if seenFinal { - return nil, fmt.Errorf("remote query response stream sent chunk after final") + return nil, errors.New("remote query response stream sent chunk after final") } if event := chunk.GetEvent(); event != nil { streamEvent, err := remoteQueryStreamEventFromProto(event) @@ -190,7 +190,7 @@ func remoteQueryExecuteOutputFromStream(stream grpc.ServerStreamingClient[pb.Rem } typedStreamEvents = append(typedStreamEvents, streamEvent) } else if !chunk.GetFinal() { - return nil, fmt.Errorf("remote query response stream chunk missing typed event") + return nil, errors.New("remote query response stream chunk missing typed event") } seenFinal = chunk.GetFinal() if seenFinal { @@ -199,10 +199,10 @@ func remoteQueryExecuteOutputFromStream(stream grpc.ServerStreamingClient[pb.Rem expectedChunkIndex++ } if !seenFinal { - return nil, fmt.Errorf("remote query response stream missing final chunk") + return nil, errors.New("remote query response stream missing final chunk") } if len(typedStreamEvents) == 0 { - return nil, fmt.Errorf("remote query response stream missing typed events") + return nil, errors.New("remote query response stream missing typed events") } output, err := remoteQueryExecuteOutputFromTypedEvents(typedStreamEvents) if err != nil { @@ -283,7 +283,7 @@ func remoteQueryStreamEventFromProto(event *pb.RemoteQueryExecuteStreamEvent) (m out["attributes"] = e.Error.GetAttributes() } default: - return nil, fmt.Errorf("remote query stream response contained unknown event") + return nil, errors.New("remote query stream response contained unknown event") } return out, nil } @@ -314,11 +314,11 @@ func remoteQueryExecuteOutputFromTypedEvents(events []map[string]interface{}) (m "error": map[string]interface{}{"code": code, "message": message}, }, nil } - return nil, fmt.Errorf("remote query stream response missing final event") + return nil, errors.New("remote query stream response missing final event") } status, _ := finalEvent["status"].(string) if status == "" { - return nil, fmt.Errorf("remote query stream final event missing status") + return nil, errors.New("remote query stream final event missing status") } dataBytes := data.Bytes() output := map[string]interface{}{ @@ -332,41 +332,6 @@ func remoteQueryExecuteOutputFromTypedEvents(events []map[string]interface{}) (m return output, nil } -func remoteQueryExecuteOutputFromProto(resp *pb.RemoteQueryExecuteResponse) (map[string]interface{}, error) { - if resp == nil || resp.GetStatus() == "" { - return nil, fmt.Errorf("remote query response missing status") - } - - output := map[string]interface{}{"status": resp.GetStatus()} - if resp.GetError() != nil { - output["error"] = map[string]interface{}{ - "code": resp.GetError().GetCode(), - "message": resp.GetError().GetMessage(), - } - } - if len(resp.GetColumns()) > 0 { - columns := make([]interface{}, 0, len(resp.GetColumns())) - for _, column := range resp.GetColumns() { - columns = append(columns, column.AsMap()) - } - output["columns"] = columns - } - if len(resp.GetRows()) > 0 { - rows := make([]interface{}, 0, len(resp.GetRows())) - for _, row := range resp.GetRows() { - rows = append(rows, row.AsMap()) - } - output["rows"] = rows - } - if resp.GetTruncated() { - output["truncated"] = true - } - if resp.GetStats() != nil { - output["stats"] = resp.GetStats().AsMap() - } - return output, nil -} - func normalizeRemoteQueryOutput(value interface{}) interface{} { switch v := value.(type) { case map[string]interface{}: diff --git a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go index 7eb01cab8fdd..b2d15f5f8d9a 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go @@ -15,6 +15,7 @@ import ( "fmt" "os" "strconv" + "strings" "testing" "time" @@ -283,11 +284,12 @@ func writeFusedEvidence(t *testing.T, path string, lines []string) { if path == "" { return } - payload := "" + var payload strings.Builder for _, line := range lines { - payload += line + "\n" + payload.WriteString(line) + payload.WriteByte('\n') } - require.NoError(t, os.WriteFile(path, []byte(payload), 0o600)) + require.NoError(t, os.WriteFile(path, []byte(payload.String()), 0o600)) } func getenvRequired(t *testing.T, name string) string { diff --git a/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go b/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go index 0736e7336740..455c4d4f41ea 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go @@ -138,7 +138,7 @@ func TestRemoteQueriesActionRunsThroughStandalonePARProcessWithRealAgentIPC(t *t writeFusedEvidence(t, getenvOptional("RQ_STANDALONE_EVIDENCE_FILE"), []string{ fmt.Sprintf("standalone private-action-runner process pid=%d", parPID), - fmt.Sprintf("separate Agent process pid=%s", agentPID), + "separate Agent process pid=" + agentPID, fmt.Sprintf("fakeintake task enqueued: task_id=%s action_fqn=%s inputs=%s", taskID, fqn, requestEvidence), "standalone PAR process dequeued the fakeintake OPMS task and invoked the registered action", fmt.Sprintf("real AgentSecure IPC called by standalone PAR: 127.0.0.1:%d RemoteQueryExecuteStream", cmdPortInt), From 11f3f1810b5669794000d11026ac1b069e365e36 Mon Sep 17 00:00:00 2001 From: nubtron Date: Fri, 22 May 2026 14:34:26 +0000 Subject: [PATCH 32/33] Keep remote queries Bazel deps scoped --- comp/remotequeries/impl/BUILD.bazel | 16 +- .../impl/remote_query_execute.go | 7 +- comp/remotequeries/impl/remote_query_match.go | 13 +- comp/remotequeries/impl/remote_query_test.go | 10 +- pkg/collector/check/BUILD.bazel | 1 + .../bundles/remotequeries/BUILD.bazel | 221 +----------------- .../live_agent_ipc_par_loop_test.go | 2 +- .../remotequeries/live_par_loop_test.go | 2 +- .../standalone_par_process_proof_test.go | 2 +- rtloader/include/datadog_agent_rtloader.h | 3 +- rtloader/include/rtloader_types.h | 3 +- rtloader/rtloader/api.cpp | 6 +- rtloader/three/three.cpp | 128 +++++----- 13 files changed, 97 insertions(+), 317 deletions(-) diff --git a/comp/remotequeries/impl/BUILD.bazel b/comp/remotequeries/impl/BUILD.bazel index 652370c86b7b..92798db4bee7 100644 --- a/comp/remotequeries/impl/BUILD.bazel +++ b/comp/remotequeries/impl/BUILD.bazel @@ -14,8 +14,7 @@ go_library( "//comp/core/config", "//comp/core/ipc/def", "//comp/core/ipc/httphelpers", - "@//comp/collector/collector", - "@//pkg/collector/check", + "//pkg/collector/check", "@in_gopkg_yaml_v3//:yaml_v3", "@org_uber_go_fx//:fx", ], @@ -30,15 +29,14 @@ go_test( embed = [":impl"], gotags = ["test"], deps = [ + "//comp/core/autodiscovery/integration", + "//comp/core/diagnose/def", "//comp/core/ipc/def", "//comp/core/ipc/mock", - "@//comp/collector/collector", - "@//comp/core/autodiscovery/integration", - "@//comp/core/diagnose/def", - "@//pkg/aggregator/sender", - "@//pkg/collector/check", - "@//pkg/collector/check/id", - "@//pkg/collector/check/stats", + "//pkg/aggregator/sender", + "//pkg/collector/check", + "//pkg/collector/check/id", + "//pkg/collector/check/stats", "@com_github_stretchr_testify//assert", "@com_github_stretchr_testify//require", ], diff --git a/comp/remotequeries/impl/remote_query_execute.go b/comp/remotequeries/impl/remote_query_execute.go index 303b474cf36f..64c23fba86e7 100644 --- a/comp/remotequeries/impl/remote_query_execute.go +++ b/comp/remotequeries/impl/remote_query_execute.go @@ -14,7 +14,6 @@ import ( "net/http" api "github.com/DataDog/datadog-agent/comp/api/api/def" - "github.com/DataDog/datadog-agent/comp/collector/collector" "github.com/DataDog/datadog-agent/pkg/collector/check" ) @@ -87,18 +86,18 @@ func NewRemoteQueryExecuteEndpointProvider(reqs Requires) api.AgentEndpointProvi type remoteQueryExecuteHandler struct { service *RemoteQueryExecuteService - collector collector.Component + collector remoteQueryCollector enabled bool } // RemoteQueryExecuteService executes credential-free Remote Queries requests through loaded checks. type RemoteQueryExecuteService struct { - collector collector.Component + collector remoteQueryCollector enabled bool } // NewRemoteQueryExecuteService creates the shared executor used by the HTTP POC endpoint and AgentSecure RPC. -func NewRemoteQueryExecuteService(collector collector.Component, enabled bool) *RemoteQueryExecuteService { +func NewRemoteQueryExecuteService(collector remoteQueryCollector, enabled bool) *RemoteQueryExecuteService { return &RemoteQueryExecuteService{collector: collector, enabled: enabled} } diff --git a/comp/remotequeries/impl/remote_query_match.go b/comp/remotequeries/impl/remote_query_match.go index 7b41fd0657c2..9c65020a9c64 100644 --- a/comp/remotequeries/impl/remote_query_match.go +++ b/comp/remotequeries/impl/remote_query_match.go @@ -20,7 +20,6 @@ import ( "gopkg.in/yaml.v3" api "github.com/DataDog/datadog-agent/comp/api/api/def" - "github.com/DataDog/datadog-agent/comp/collector/collector" "github.com/DataDog/datadog-agent/comp/core/config" "github.com/DataDog/datadog-agent/pkg/collector/check" ) @@ -43,7 +42,7 @@ type Requires struct { fx.In Cfg config.Component - Collector collector.Component + Collector remoteQueryCollector } // NewRemoteQueryMatchEndpointProvider registers the remote query match endpoint on the internal Agent API. @@ -56,10 +55,16 @@ func NewRemoteQueryMatchEndpointProvider(reqs Requires) api.AgentEndpointProvide } type remoteQueryMatchHandler struct { - collector collector.Component + collector remoteQueryCollector enabled bool } +// remoteQueryCollector is the narrow collector surface Remote Queries needs. +// Keep this local so the POC endpoint does not force Bazel onboarding for the full collector component package. +type remoteQueryCollector interface { + GetChecks() []check.Check +} + type matchResponse struct { Status string `json:"status"` MatchedCount int `json:"matched_count"` @@ -313,7 +318,7 @@ func (h *remoteQueryMatchHandler) findMatches(integration string, target remoteQ return findIntegrationMatches(h.collector, integration, target) } -func findIntegrationMatches(collector collector.Component, integration string, target remoteQueryTarget) []integrationCheckMatch { +func findIntegrationMatches(collector remoteQueryCollector, integration string, target remoteQueryTarget) []integrationCheckMatch { checks := collector.GetChecks() matches := make([]integrationCheckMatch, 0, 1) for _, chk := range checks { diff --git a/comp/remotequeries/impl/remote_query_test.go b/comp/remotequeries/impl/remote_query_test.go index 784f052fd955..ad3dbf8b0b3f 100644 --- a/comp/remotequeries/impl/remote_query_test.go +++ b/comp/remotequeries/impl/remote_query_test.go @@ -15,7 +15,6 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - "github.com/DataDog/datadog-agent/comp/collector/collector" "github.com/DataDog/datadog-agent/comp/core/autodiscovery/integration" diagnose "github.com/DataDog/datadog-agent/comp/core/diagnose/def" "github.com/DataDog/datadog-agent/pkg/aggregator/sender" @@ -220,14 +219,7 @@ type fakeCollector struct { checks []check.Check } -func (f fakeCollector) RunCheck(inner check.Check) (checkid.ID, error) { return inner.ID(), nil } -func (f fakeCollector) StopCheck(checkid.ID) error { return nil } -func (f fakeCollector) MapOverChecks(_ func([]check.Info)) {} -func (f fakeCollector) GetChecks() []check.Check { return f.checks } -func (f fakeCollector) ReloadAllCheckInstances(string, []check.Check) ([]checkid.ID, error) { - return nil, nil -} -func (f fakeCollector) AddEventReceiver(collector.EventReceiver) {} +func (f fakeCollector) GetChecks() []check.Check { return f.checks } type fakeCheck struct { name string diff --git a/pkg/collector/check/BUILD.bazel b/pkg/collector/check/BUILD.bazel index c7b16430cf14..b68529223ced 100644 --- a/pkg/collector/check/BUILD.bazel +++ b/pkg/collector/check/BUILD.bazel @@ -9,6 +9,7 @@ go_library( "jmx.go", "loader.go", "metadata.go", + "remote_query_stream.go", "retry.go", ], importpath = "github.com/DataDog/datadog-agent/pkg/collector/check", diff --git a/pkg/privateactionrunner/bundles/remotequeries/BUILD.bazel b/pkg/privateactionrunner/bundles/remotequeries/BUILD.bazel index c6be2fc8c22b..76e82806edbb 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/BUILD.bazel +++ b/pkg/privateactionrunner/bundles/remotequeries/BUILD.bazel @@ -24,12 +24,7 @@ go_library( go_test( name = "remotequeries_test", - srcs = [ - "execute_test.go", - "live_agent_ipc_par_loop_test.go", - "live_par_loop_test.go", - "standalone_par_process_proof_test.go", - ], + srcs = ["execute_test.go"], embed = [":remotequeries"], gotags = ["test"], deps = [ @@ -40,217 +35,5 @@ go_test( "@com_github_stretchr_testify//assert", "@com_github_stretchr_testify//require", "@org_golang_google_grpc//:grpc", - ] + select({ - "@rules_go//go/platform:aix": [ - "//pkg/config/setup", - "//pkg/privateactionrunner/adapters/config", - "//pkg/privateactionrunner/observability", - "//pkg/privateactionrunner/opms", - "@//pkg/privateactionrunner/adapters/constants", - "@//pkg/privateactionrunner/runners", - "@//pkg/privateactionrunner/task-verifier", - "@com_github_datadog_datadog_agent_test_fakeintake//client", - "@com_github_datadog_datadog_agent_test_fakeintake//server", - "@com_github_datadog_datadog_go_v5//statsd", - "@io_k8s_apimachinery//pkg/util/sets", - "@org_golang_google_protobuf//types/known/structpb", - ], - "@rules_go//go/platform:android": [ - "//pkg/config/setup", - "//pkg/privateactionrunner/adapters/config", - "//pkg/privateactionrunner/observability", - "//pkg/privateactionrunner/opms", - "@//pkg/privateactionrunner/adapters/constants", - "@//pkg/privateactionrunner/runners", - "@//pkg/privateactionrunner/task-verifier", - "@com_github_datadog_datadog_agent_test_fakeintake//client", - "@com_github_datadog_datadog_agent_test_fakeintake//server", - "@com_github_datadog_datadog_go_v5//statsd", - "@io_k8s_apimachinery//pkg/util/sets", - "@org_golang_google_protobuf//types/known/structpb", - ], - "@rules_go//go/platform:darwin": [ - "//pkg/config/setup", - "//pkg/privateactionrunner/adapters/config", - "//pkg/privateactionrunner/observability", - "//pkg/privateactionrunner/opms", - "@//pkg/privateactionrunner/adapters/constants", - "@//pkg/privateactionrunner/runners", - "@//pkg/privateactionrunner/task-verifier", - "@com_github_datadog_datadog_agent_test_fakeintake//client", - "@com_github_datadog_datadog_agent_test_fakeintake//server", - "@com_github_datadog_datadog_go_v5//statsd", - "@io_k8s_apimachinery//pkg/util/sets", - "@org_golang_google_protobuf//types/known/structpb", - ], - "@rules_go//go/platform:dragonfly": [ - "//pkg/config/setup", - "//pkg/privateactionrunner/adapters/config", - "//pkg/privateactionrunner/observability", - "//pkg/privateactionrunner/opms", - "@//pkg/privateactionrunner/adapters/constants", - "@//pkg/privateactionrunner/runners", - "@//pkg/privateactionrunner/task-verifier", - "@com_github_datadog_datadog_agent_test_fakeintake//client", - "@com_github_datadog_datadog_agent_test_fakeintake//server", - "@com_github_datadog_datadog_go_v5//statsd", - "@io_k8s_apimachinery//pkg/util/sets", - "@org_golang_google_protobuf//types/known/structpb", - ], - "@rules_go//go/platform:freebsd": [ - "//pkg/config/setup", - "//pkg/privateactionrunner/adapters/config", - "//pkg/privateactionrunner/observability", - "//pkg/privateactionrunner/opms", - "@//pkg/privateactionrunner/adapters/constants", - "@//pkg/privateactionrunner/runners", - "@//pkg/privateactionrunner/task-verifier", - "@com_github_datadog_datadog_agent_test_fakeintake//client", - "@com_github_datadog_datadog_agent_test_fakeintake//server", - "@com_github_datadog_datadog_go_v5//statsd", - "@io_k8s_apimachinery//pkg/util/sets", - "@org_golang_google_protobuf//types/known/structpb", - ], - "@rules_go//go/platform:illumos": [ - "//pkg/config/setup", - "//pkg/privateactionrunner/adapters/config", - "//pkg/privateactionrunner/observability", - "//pkg/privateactionrunner/opms", - "@//pkg/privateactionrunner/adapters/constants", - "@//pkg/privateactionrunner/runners", - "@//pkg/privateactionrunner/task-verifier", - "@com_github_datadog_datadog_agent_test_fakeintake//client", - "@com_github_datadog_datadog_agent_test_fakeintake//server", - "@com_github_datadog_datadog_go_v5//statsd", - "@io_k8s_apimachinery//pkg/util/sets", - "@org_golang_google_protobuf//types/known/structpb", - ], - "@rules_go//go/platform:ios": [ - "//pkg/config/setup", - "//pkg/privateactionrunner/adapters/config", - "//pkg/privateactionrunner/observability", - "//pkg/privateactionrunner/opms", - "@//pkg/privateactionrunner/adapters/constants", - "@//pkg/privateactionrunner/runners", - "@//pkg/privateactionrunner/task-verifier", - "@com_github_datadog_datadog_agent_test_fakeintake//client", - "@com_github_datadog_datadog_agent_test_fakeintake//server", - "@com_github_datadog_datadog_go_v5//statsd", - "@io_k8s_apimachinery//pkg/util/sets", - "@org_golang_google_protobuf//types/known/structpb", - ], - "@rules_go//go/platform:js": [ - "//pkg/config/setup", - "//pkg/privateactionrunner/adapters/config", - "//pkg/privateactionrunner/observability", - "//pkg/privateactionrunner/opms", - "@//pkg/privateactionrunner/adapters/constants", - "@//pkg/privateactionrunner/runners", - "@//pkg/privateactionrunner/task-verifier", - "@com_github_datadog_datadog_agent_test_fakeintake//client", - "@com_github_datadog_datadog_agent_test_fakeintake//server", - "@com_github_datadog_datadog_go_v5//statsd", - "@io_k8s_apimachinery//pkg/util/sets", - "@org_golang_google_protobuf//types/known/structpb", - ], - "@rules_go//go/platform:linux": [ - "//pkg/config/setup", - "//pkg/privateactionrunner/adapters/config", - "//pkg/privateactionrunner/observability", - "//pkg/privateactionrunner/opms", - "@//pkg/privateactionrunner/adapters/constants", - "@//pkg/privateactionrunner/runners", - "@//pkg/privateactionrunner/task-verifier", - "@com_github_datadog_datadog_agent_test_fakeintake//client", - "@com_github_datadog_datadog_agent_test_fakeintake//server", - "@com_github_datadog_datadog_go_v5//statsd", - "@io_k8s_apimachinery//pkg/util/sets", - "@org_golang_google_protobuf//types/known/structpb", - ], - "@rules_go//go/platform:netbsd": [ - "//pkg/config/setup", - "//pkg/privateactionrunner/adapters/config", - "//pkg/privateactionrunner/observability", - "//pkg/privateactionrunner/opms", - "@//pkg/privateactionrunner/adapters/constants", - "@//pkg/privateactionrunner/runners", - "@//pkg/privateactionrunner/task-verifier", - "@com_github_datadog_datadog_agent_test_fakeintake//client", - "@com_github_datadog_datadog_agent_test_fakeintake//server", - "@com_github_datadog_datadog_go_v5//statsd", - "@io_k8s_apimachinery//pkg/util/sets", - "@org_golang_google_protobuf//types/known/structpb", - ], - "@rules_go//go/platform:openbsd": [ - "//pkg/config/setup", - "//pkg/privateactionrunner/adapters/config", - "//pkg/privateactionrunner/observability", - "//pkg/privateactionrunner/opms", - "@//pkg/privateactionrunner/adapters/constants", - "@//pkg/privateactionrunner/runners", - "@//pkg/privateactionrunner/task-verifier", - "@com_github_datadog_datadog_agent_test_fakeintake//client", - "@com_github_datadog_datadog_agent_test_fakeintake//server", - "@com_github_datadog_datadog_go_v5//statsd", - "@io_k8s_apimachinery//pkg/util/sets", - "@org_golang_google_protobuf//types/known/structpb", - ], - "@rules_go//go/platform:osx": [ - "//pkg/config/setup", - "//pkg/privateactionrunner/adapters/config", - "//pkg/privateactionrunner/observability", - "//pkg/privateactionrunner/opms", - "@//pkg/privateactionrunner/adapters/constants", - "@//pkg/privateactionrunner/runners", - "@//pkg/privateactionrunner/task-verifier", - "@com_github_datadog_datadog_agent_test_fakeintake//client", - "@com_github_datadog_datadog_agent_test_fakeintake//server", - "@com_github_datadog_datadog_go_v5//statsd", - "@io_k8s_apimachinery//pkg/util/sets", - "@org_golang_google_protobuf//types/known/structpb", - ], - "@rules_go//go/platform:plan9": [ - "//pkg/config/setup", - "//pkg/privateactionrunner/adapters/config", - "//pkg/privateactionrunner/observability", - "//pkg/privateactionrunner/opms", - "@//pkg/privateactionrunner/adapters/constants", - "@//pkg/privateactionrunner/runners", - "@//pkg/privateactionrunner/task-verifier", - "@com_github_datadog_datadog_agent_test_fakeintake//client", - "@com_github_datadog_datadog_agent_test_fakeintake//server", - "@com_github_datadog_datadog_go_v5//statsd", - "@io_k8s_apimachinery//pkg/util/sets", - "@org_golang_google_protobuf//types/known/structpb", - ], - "@rules_go//go/platform:qnx": [ - "//pkg/config/setup", - "//pkg/privateactionrunner/adapters/config", - "//pkg/privateactionrunner/observability", - "//pkg/privateactionrunner/opms", - "@//pkg/privateactionrunner/adapters/constants", - "@//pkg/privateactionrunner/runners", - "@//pkg/privateactionrunner/task-verifier", - "@com_github_datadog_datadog_agent_test_fakeintake//client", - "@com_github_datadog_datadog_agent_test_fakeintake//server", - "@com_github_datadog_datadog_go_v5//statsd", - "@io_k8s_apimachinery//pkg/util/sets", - "@org_golang_google_protobuf//types/known/structpb", - ], - "@rules_go//go/platform:solaris": [ - "//pkg/config/setup", - "//pkg/privateactionrunner/adapters/config", - "//pkg/privateactionrunner/observability", - "//pkg/privateactionrunner/opms", - "@//pkg/privateactionrunner/adapters/constants", - "@//pkg/privateactionrunner/runners", - "@//pkg/privateactionrunner/task-verifier", - "@com_github_datadog_datadog_agent_test_fakeintake//client", - "@com_github_datadog_datadog_agent_test_fakeintake//server", - "@com_github_datadog_datadog_go_v5//statsd", - "@io_k8s_apimachinery//pkg/util/sets", - "@org_golang_google_protobuf//types/known/structpb", - ], - "//conditions:default": [], - }), + ], ) diff --git a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go index b2d15f5f8d9a..66730d96b50e 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go @@ -3,7 +3,7 @@ // This product includes software developed at Datadog (https://www.datadoghq.com/). // Copyright 2026-present Datadog, Inc. -//go:build !windows +//go:build remotequeries_live && !windows package com_datadoghq_remotequeries_test diff --git a/pkg/privateactionrunner/bundles/remotequeries/live_par_loop_test.go b/pkg/privateactionrunner/bundles/remotequeries/live_par_loop_test.go index 35c1e6f454cb..cb5057c4406a 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/live_par_loop_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/live_par_loop_test.go @@ -3,7 +3,7 @@ // This product includes software developed at Datadog (https://www.datadoghq.com/). // Copyright 2026-present Datadog, Inc. -//go:build !windows +//go:build remotequeries_live && !windows package com_datadoghq_remotequeries_test diff --git a/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go b/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go index 455c4d4f41ea..daf07b91cb52 100644 --- a/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go +++ b/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go @@ -3,7 +3,7 @@ // This product includes software developed at Datadog (https://www.datadoghq.com/). // Copyright 2026-present Datadog, Inc. -//go:build !windows +//go:build remotequeries_live && !windows package com_datadoghq_remotequeries_test diff --git a/rtloader/include/datadog_agent_rtloader.h b/rtloader/include/datadog_agent_rtloader.h index a469e078c700..41e82d1881a9 100644 --- a/rtloader/include/datadog_agent_rtloader.h +++ b/rtloader/include/datadog_agent_rtloader.h @@ -190,7 +190,8 @@ DATADOG_AGENT_RTLOADER_API int get_check_deprecated(rtloader_t *rtloader, rtload */ DATADOG_AGENT_RTLOADER_API char *run_check(rtloader_t *, rtloader_pyobject_t *check); -/*! \fn int run_remote_query_stream(rtloader_t *, rtloader_pyobject_t *check, const char *integration, const char *request_json, remote_query_stream_emit_cb emit, void *userdata) +/*! \fn int run_remote_query_stream(rtloader_t *, rtloader_pyobject_t *check, const char *integration, const char + *request_json, remote_query_stream_emit_cb emit, void *userdata) \brief Runs the integration remote query streaming helper for a check instance. \param rtloader_t A rtloader_t * pointer to the RtLoader instance. \param check A rtloader_pyobject_t * pointer to the check instance we wish to use. diff --git a/rtloader/include/rtloader_types.h b/rtloader/include/rtloader_types.h index 641782b4a97c..392b8cd482d0 100644 --- a/rtloader/include/rtloader_types.h +++ b/rtloader/include/rtloader_types.h @@ -38,7 +38,8 @@ typedef enum rtloader_gilstate_e { typedef void *(*rtloader_malloc_t)(size_t); typedef void (*rtloader_free_t)(void *); -typedef int (*remote_query_stream_emit_cb)(const char *event_type, const char *metadata_json, const uint8_t *payload, size_t payload_len, void *userdata); +typedef int (*remote_query_stream_emit_cb)(const char *event_type, const char *metadata_json, const uint8_t *payload, + size_t payload_len, void *userdata); typedef enum { DATADOG_AGENT_RTLOADER_GAUGE = 0, diff --git a/rtloader/rtloader/api.cpp b/rtloader/rtloader/api.cpp index 25b94b5ca919..a7882adbacc6 100644 --- a/rtloader/rtloader/api.cpp +++ b/rtloader/rtloader/api.cpp @@ -282,9 +282,9 @@ int run_remote_query_stream(rtloader_t *rtloader, rtloader_pyobject_t *check, co const char *request_json, remote_query_stream_emit_cb emit, void *userdata) { return AS_TYPE(RtLoader, rtloader) - ->runRemoteQueryStream(AS_TYPE(RtLoaderPyObject, check), integration, request_json, emit, userdata) - ? 1 - : 0; + ->runRemoteQueryStream(AS_TYPE(RtLoaderPyObject, check), integration, request_json, emit, userdata) + ? 1 + : 0; } void cancel_check(rtloader_t *rtloader, rtloader_pyobject_t *check) diff --git a/rtloader/three/three.cpp b/rtloader/three/three.cpp index cec202ceac5a..6f09292cc729 100644 --- a/rtloader/three/three.cpp +++ b/rtloader/three/three.cpp @@ -502,76 +502,76 @@ char *Three::runCheck(RtLoaderPyObject *check) } namespace { -std::string normalizeRemoteQueryIntegration(const char *integration) -{ - if (integration == NULL) { - return ""; - } - - std::string normalized(integration); - normalized.erase(normalized.begin(), std::find_if(normalized.begin(), normalized.end(), [](unsigned char ch) { - return !std::isspace(ch); - })); - normalized.erase(std::find_if(normalized.rbegin(), normalized.rend(), [](unsigned char ch) { - return !std::isspace(ch); - }).base(), - normalized.end()); - std::transform(normalized.begin(), normalized.end(), normalized.begin(), [](unsigned char ch) { - return static_cast(std::tolower(ch)); - }); - return normalized; -} + std::string normalizeRemoteQueryIntegration(const char *integration) + { + if (integration == NULL) { + return ""; + } -bool isValidRemoteQueryIntegration(const std::string &integration) -{ - if (integration.empty()) { - return false; + std::string normalized(integration); + normalized.erase(normalized.begin(), std::find_if(normalized.begin(), normalized.end(), [](unsigned char ch) { + return !std::isspace(ch); + })); + normalized.erase( + std::find_if(normalized.rbegin(), normalized.rend(), [](unsigned char ch) { return !std::isspace(ch); }) + .base(), + normalized.end()); + std::transform(normalized.begin(), normalized.end(), normalized.begin(), + [](unsigned char ch) { return static_cast(std::tolower(ch)); }); + return normalized; } - return std::all_of(integration.begin(), integration.end(), [](unsigned char ch) { - return std::islower(ch) || std::isdigit(ch) || ch == '_'; - }); -} - -struct RemoteQueryStreamEmitContext { - remote_query_stream_emit_cb emit; - void *userdata; -}; -PyObject *remoteQueryStreamEmit(PyObject *self, PyObject *args) -{ - RemoteQueryStreamEmitContext *ctx = static_cast(PyCapsule_GetPointer(self, "remote_query_stream_emit")); - if (ctx == NULL || ctx->emit == NULL) { - PyErr_SetString(PyExc_RuntimeError, "remote query stream emit callback is unavailable"); - return NULL; + bool isValidRemoteQueryIntegration(const std::string &integration) + { + if (integration.empty()) { + return false; + } + return std::all_of(integration.begin(), integration.end(), + [](unsigned char ch) { return std::islower(ch) || std::isdigit(ch) || ch == '_'; }); } - const char *event_type = NULL; - const char *metadata_json = NULL; - PyObject *payload = NULL; - if (!PyArg_ParseTuple(args, "ssO:remote_query_stream_emit", &event_type, &metadata_json, &payload)) { - return NULL; - } - if (!PyBytes_Check(payload)) { - PyErr_SetString(PyExc_TypeError, "remote query stream payload must be bytes"); - return NULL; - } + struct RemoteQueryStreamEmitContext { + remote_query_stream_emit_cb emit; + void *userdata; + }; - char *payload_bytes = NULL; - Py_ssize_t payload_len = 0; - if (PyBytes_AsStringAndSize(payload, &payload_bytes, &payload_len) != 0) { - return NULL; - } - int emit_result = ctx->emit(event_type, metadata_json, reinterpret_cast(payload_bytes), static_cast(payload_len), ctx->userdata); - if (emit_result != 0) { - PyErr_SetString(PyExc_RuntimeError, "remote query stream emit callback failed"); - return NULL; - } + PyObject *remoteQueryStreamEmit(PyObject *self, PyObject *args) + { + RemoteQueryStreamEmitContext *ctx + = static_cast(PyCapsule_GetPointer(self, "remote_query_stream_emit")); + if (ctx == NULL || ctx->emit == NULL) { + PyErr_SetString(PyExc_RuntimeError, "remote query stream emit callback is unavailable"); + return NULL; + } - Py_RETURN_NONE; -} + const char *event_type = NULL; + const char *metadata_json = NULL; + PyObject *payload = NULL; + if (!PyArg_ParseTuple(args, "ssO:remote_query_stream_emit", &event_type, &metadata_json, &payload)) { + return NULL; + } + if (!PyBytes_Check(payload)) { + PyErr_SetString(PyExc_TypeError, "remote query stream payload must be bytes"); + return NULL; + } + + char *payload_bytes = NULL; + Py_ssize_t payload_len = 0; + if (PyBytes_AsStringAndSize(payload, &payload_bytes, &payload_len) != 0) { + return NULL; + } + int emit_result = ctx->emit(event_type, metadata_json, reinterpret_cast(payload_bytes), + static_cast(payload_len), ctx->userdata); + if (emit_result != 0) { + PyErr_SetString(PyExc_RuntimeError, "remote query stream emit callback failed"); + return NULL; + } + + Py_RETURN_NONE; + } -PyMethodDef remoteQueryStreamEmitMethod = {"remote_query_stream_emit", remoteQueryStreamEmit, METH_VARARGS, - "Emit a remote query stream event."}; + PyMethodDef remoteQueryStreamEmitMethod + = { "remote_query_stream_emit", remoteQueryStreamEmit, METH_VARARGS, "Emit a remote query stream event." }; } // namespace bool Three::runRemoteQueryStream(RtLoaderPyObject *check, const char *integration, const char *request_json, @@ -595,7 +595,7 @@ bool Three::runRemoteQueryStream(RtLoaderPyObject *check, const char *integratio PyObject *emit_func = NULL; PyObject *result = NULL; std::string module_name = "datadog_checks." + normalized_integration + ".remote_query"; - RemoteQueryStreamEmitContext ctx{emit, userdata}; + RemoteQueryStreamEmitContext ctx{ emit, userdata }; bool ok = false; remote_query_module = PyImport_ImportModule(module_name.c_str()); @@ -634,7 +634,7 @@ bool Three::runRemoteQueryStream(RtLoaderPyObject *check, const char *integratio } ok = true; - done: +done: Py_XDECREF(result); Py_XDECREF(emit_func); Py_XDECREF(capsule); From 9277a1cd6f2c6ad9fb343ab4db94a42b3f6e79b7 Mon Sep 17 00:00:00 2001 From: nubtron Date: Fri, 22 May 2026 15:39:01 +0000 Subject: [PATCH 33/33] Add Remote Queries release note --- ...eries-agent-backed-execution-48ef4b660d8ab523.yaml | 11 +++++++++++ 1 file changed, 11 insertions(+) create mode 100644 releasenotes/notes/remote-queries-agent-backed-execution-48ef4b660d8ab523.yaml diff --git a/releasenotes/notes/remote-queries-agent-backed-execution-48ef4b660d8ab523.yaml b/releasenotes/notes/remote-queries-agent-backed-execution-48ef4b660d8ab523.yaml new file mode 100644 index 000000000000..95d138c4b811 --- /dev/null +++ b/releasenotes/notes/remote-queries-agent-backed-execution-48ef4b660d8ab523.yaml @@ -0,0 +1,11 @@ +# Each section from every release note are combined when the +# CHANGELOG.rst is rendered. So the text needs to be worded so that +# it does not depend on any information only available in another +# section. This may mean repeating some details, but each section +# must be readable independently of the other. +# +# Each section note must be formatted as reStructuredText. +--- +features: + - | + Add Agent-backed Remote Queries execution plumbing for Private Action Runner proof-of-concept workflows.