diff --git a/cmd/agent/subcommands/run/command.go b/cmd/agent/subcommands/run/command.go index c6d2af666f63..ae6c178088a7 100644 --- a/cmd/agent/subcommands/run/command.go +++ b/cmd/agent/subcommands/run/command.go @@ -45,6 +45,7 @@ import ( networkconfigmanagementfx "github.com/DataDog/datadog-agent/comp/networkconfigmanagement/fx" traceroute "github.com/DataDog/datadog-agent/comp/networkpath/traceroute/def" remotetraceroute "github.com/DataDog/datadog-agent/comp/networkpath/traceroute/fx-remote" + 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" ssistatusfx "github.com/DataDog/datadog-agent/comp/updater/ssistatus/fx" @@ -579,6 +580,7 @@ func getSharedFxOption() fx.Option { remoteagentregistryfx.Module(), haagentfx.Module(), doqueryactionsfx.Module(), + remotequeriesfx.Module(), metricscompressorfx.Module(), diagnosefx.Module(), ipcfx.ModuleReadWrite(), diff --git a/comp/api/grpcserver/impl-agent/grpc.go b/comp/api/grpcserver/impl-agent/grpc.go index aaa31e321a7f..1bc8ab297d60 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/def" 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 { @@ -137,6 +139,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 @@ -149,6 +152,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, @@ -167,6 +171,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..d32359adb723 --- /dev/null +++ b/comp/api/grpcserver/impl-agent/remote_query_execute_test.go @@ -0,0 +1,125 @@ +// 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" + + "google.golang.org/grpc" + + "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" +) + +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 TestRemoteQueryExecuteRejectsUnaryInlineMode(t *testing.T) { + resp, err := (&serverSecure{}).RemoteQueryExecute(context.Background(), &pb.RemoteQueryExecuteRequest{}) + + require.NoError(t, err) + assert.Equal(t, "invalid_request", resp.GetStatus()) + require.NotNil(t, resp.GetError()) + assert.Equal(t, "invalid_request", resp.GetError().GetCode()) + assert.Contains(t, resp.GetError().GetMessage(), "RemoteQueryExecuteStream") +} + +func TestRemoteQueryExecuteStreamReturnsSanitizedUnavailableWhenServiceMissing(t *testing.T) { + stream := &captureRemoteQueryExecuteStreamServer{} + err := (&serverSecure{}).RemoteQueryExecuteStream(&pb.RemoteQueryExecuteRequest{}, stream) + + require.NoError(t, err) + 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) { + 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 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, 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, 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) + 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(fiveMiB-remoteQuerySecureIPCDataFlushBytes), secondData.GetBytes()) + assert.Len(t, secondData.GetPayload(), fiveMiB-remoteQuerySecureIPCDataFlushBytes) +} + +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()) +} + +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 c689fdbefc4c..09e46ebaf24a 100644 --- a/comp/api/grpcserver/impl-agent/server.go +++ b/comp/api/grpcserver/impl-agent/server.go @@ -6,14 +6,19 @@ package agentimpl import ( + "bytes" "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 +38,8 @@ import ( "github.com/DataDog/datadog-agent/comp/metadata/host/impl/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" + "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" @@ -60,6 +67,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 +282,443 @@ 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, _ *pb.RemoteQueryExecuteRequest) (*pb.RemoteQueryExecuteResponse, error) { + 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 remoteQueryExecuteStreamError(remotequeriesimpl.RemoteQueryStatusExecutorUnavailable, "remote query executor is unavailable", stream) + } + + execReq, err := remoteQueryExecuteRequestFromProto(req) + if err != nil { + return remoteQueryExecuteStreamError(remotequeriesimpl.RemoteQueryStatusInvalidRequest, err.Error(), stream) + } + + coalescer := newRemoteQueryIPCStreamCoalescer(stream) + result := s.remoteQueries.ExecuteStream(execReq, coalescer.Send) + if result.Error != nil { + if err := coalescer.Flush(); err != nil { + return err + } + return remoteQueryExecuteStreamErrorAt(result.Error.Code, result.Error.Message, stream, coalescer.NextChunkIndex()) + } + if err := coalescer.Flush(); err != nil { + return err + } + 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 + + 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, start: time.Now()} +} + +func (c *remoteQueryIPCStreamCoalescer) NextChunkIndex() int32 { + return c.chunkIndex +} + +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 + } + data := protoEvent.GetData() + if data == nil { + if err := c.Flush(); err != nil { + return err + } + 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 { + 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)), + }}, + } + 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) + c.dataOffset += uint64(len(payload)) + c.dataChunks++ + if c.data.Len() == 0 { + c.dataStarted = false + } + return nil +} + +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 0, err + } + duration := time.Since(start) + c.sendCalls++ + c.sendDuration += duration + if duration > c.maxSendDuration { + c.maxSendDuration = duration + } + c.chunkIndex++ + 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) { + target := remotequeriesimpl.RemoteQueryExecuteTarget{ + Host: req.GetTarget().GetHost(), + Port: int(req.GetTarget().GetPort()), + DBName: req.GetTarget().GetDbname(), + } + if req.GetOperation() != "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())) +} + +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 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: chunkIndex, + Event: &pb.RemoteQueryExecuteStreamEvent{Event: &pb.RemoteQueryExecuteStreamEvent_Error{Error: &pb.RemoteQueryStreamError{ + Code: code, + Message: message, + }}}, + }); err != nil { + return err + } + return stream.Send(&pb.RemoteQueryExecuteChunk{ChunkIndex: chunkIndex + 1, Final: true}) +} + +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, + 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/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/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/fx/fx.go b/comp/remotequeries/fx/fx.go new file mode 100644 index 000000000000..51919ed2df44 --- /dev/null +++ b/comp/remotequeries/fx/fx.go @@ -0,0 +1,21 @@ +// 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/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.NewRemoteQueryMatchEndpointProvider), + fx.Provide(remotequeriesimpl.NewRemoteQueryExecuteEndpointProvider), + ) +} diff --git a/comp/remotequeries/impl/BUILD.bazel b/comp/remotequeries/impl/BUILD.bazel new file mode 100644 index 000000000000..92798db4bee7 --- /dev/null +++ b/comp/remotequeries/impl/BUILD.bazel @@ -0,0 +1,43 @@ +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", + "//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/autodiscovery/integration", + "//comp/core/diagnose/def", + "//comp/core/ipc/def", + "//comp/core/ipc/mock", + "//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 new file mode 100644 index 000000000000..64c23fba86e7 --- /dev/null +++ b/comp/remotequeries/impl/remote_query_execute.go @@ -0,0 +1,582 @@ +// 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" + "errors" + "fmt" + "io" + "net/http" + + api "github.com/DataDog/datadog-agent/comp/api/api/def" + "github.com/DataDog/datadog-agent/pkg/collector/check" +) + +const ( + // 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" + + remoteQueryProofSeedQuery = "SELECT 1 AS value" + remoteQueryFixtureTableProofQuery = "SELECT city, country FROM cities ORDER BY city" + + statusExecutorUnavailable = "executor_unavailable" +) + +const remoteQueryBinaryPayloadProofQuery = "SELECT decode('00ff80', 'hex') AS payload" + +var remoteQueryLargePayloadProofQueries = map[string]int{ + "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 { + RunRemoteQueryStream(integration string, requestJSON string, emit func(check.RemoteQueryStreamEvent) error) error +} + +func isRemoteQueryAllowedProofQuery(query string) bool { + switch query { + case remoteQueryProofSeedQuery, remoteQueryFixtureTableProofQuery, remoteQueryBinaryPayloadProofQuery: + return true + default: + _, ok := remoteQueryLargePayloadProofQueries[query] + return ok + } +} + +type remoteQueryCheckUnwrapper interface { + Unwrap() check.Check +} + +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{ + service: NewRemoteQueryExecuteService(reqs.Collector, reqs.Cfg.GetBool(RemoteQueriesExecuteEnabledConfig)), + } + return api.NewAgentEndpointProvider(h.handle, RemoteQueryExecuteEndpointPath, http.MethodPost) +} + +type remoteQueryExecuteHandler struct { + service *RemoteQueryExecuteService + collector remoteQueryCollector + enabled bool +} + +// RemoteQueryExecuteService executes credential-free Remote Queries requests through loaded checks. +type RemoteQueryExecuteService struct { + collector remoteQueryCollector + enabled bool +} + +// NewRemoteQueryExecuteService creates the shared executor used by the HTTP POC endpoint and AgentSecure RPC. +func NewRemoteQueryExecuteService(collector remoteQueryCollector, 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 +} + +// 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{}, errors.New("query is required") + } + if !isRemoteQueryAllowedProofQuery(query) { + return RemoteQueryExecuteRequest{}, errors.New("query is not allowed") + } + if format == "" { + format = "csv" + } + if format != "csv" && format != "binary" { + return RemoteQueryExecuteRequest{}, errors.New("format must be csv or binary") + } + 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. +type RemoteQueryExecuteError struct { + Code string + Message string +} + +// RemoteQueryExecuteResult is the service result. +type RemoteQueryExecuteResult struct { + HTTPStatus int + Status string + Error *RemoteQueryExecuteError +} + +const ( + // RemoteQueryStatusInvalidRequest reports a malformed or disallowed request. + RemoteQueryStatusInvalidRequest = statusInvalidRequest + // RemoteQueryStatusExecutorUnavailable reports an unavailable matched executor or bridge dependency. + RemoteQueryStatusExecutorUnavailable = statusExecutorUnavailable +) + +// NewRemoteQueryExecuteRequest rejects legacy inline Remote Queries requests. +func NewRemoteQueryExecuteRequest(_ string, _ RemoteQueryExecuteTarget, _ string, _ *RemoteQueryExecuteLimits) (RemoteQueryExecuteRequest, error) { + return RemoteQueryExecuteRequest{}, errors.New("operation must be copy_stream") +} + +type remoteQueryExecuteRequest struct { + Integration string + Operation string + Target remoteQueryTarget + Query string + Format string + Limits *remoteQueryExecuteLimits + CopyLimits *remoteQueryExecuteCopyLimits +} + +type remoteQueryExecuteRequestJSON struct { + 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 { + MaxRows *int `json:"maxRows"` + MaxBytes *int `json:"maxBytes"` + TimeoutMs *int `json:"timeoutMs"` +} + +type remoteQueryExecuteLimits struct { + MaxRows int + MaxBytes int + 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 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"` + DBName string `json:"dbname"` +} + +func (h *remoteQueryExecuteHandler) handle(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + + 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, _, err := parseExecuteRequest(r) + if err != nil { + writeExecuteParseError(w, err) + return + } + + 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, `{"status":"SUCCEEDED"}`) +} + +func parseExecuteRequest(r *http.Request) (remoteQueryExecuteRequest, string, error) { + if !isJSONContentType(r.Header.Get("Content-Type")) { + return remoteQueryExecuteRequest{}, "", invalidRequestError("content-type must be application/json") + } + + defer r.Body.Close() + var wireReq remoteQueryExecuteRequestJSON + if err := decodeStrictJSON(r.Body, &wireReq); err != nil { + return remoteQueryExecuteRequest{}, "", parseJSONRequestError(err) + } + + integration, err := parseIntegration(wireReq.Integration) + if err != nil { + return remoteQueryExecuteRequest{}, "", err + } + + target, err := parseTarget(wireReq.Target) + if err != nil { + return remoteQueryExecuteRequest{}, "", err + } + + if wireReq.Query == "" { + return remoteQueryExecuteRequest{}, "", errors.New("query is required") + } + if !isRemoteQueryAllowedProofQuery(wireReq.Query) { + return remoteQueryExecuteRequest{}, "", errors.New("query is not allowed") + } + + limits, err := parseExecuteLimits(wireReq.Limits) + if err != nil { + return remoteQueryExecuteRequest{}, "", err + } + copyLimits, err := parseExecuteCopyLimits(wireReq.CopyLimits) + if err != nil { + return remoteQueryExecuteRequest{}, "", err + } + + if wireReq.Operation != "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{}, "", 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{}, "", errors.New("malformed JSON request") + } + return req, requestJSON, nil +} + +var ( + errLimitsUnknownField = errors.New("limits contains unknown field") + 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 + } + + type limitsAlias remoteQueryExecuteLimitsRequestJSON + var limits limitsAlias + if err := decodeStrictJSON(bytes.NewReader(data), &limits); err != nil { + if isUnknownJSONFieldError(err) { + return errLimitsUnknownField + } + return err + } + *l = remoteQueryExecuteLimitsRequestJSON(limits) + 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 + } + + maxRows, err := parseRequiredPositiveInt(limits.MaxRows, "limits.maxRows") + if err != nil { + return nil, err + } + maxBytes, err := parseRequiredPositiveInt(limits.MaxBytes, "limits.maxBytes") + if err != nil { + return nil, err + } + timeoutMs, err := parseRequiredPositiveInt(limits.TimeoutMs, "limits.timeoutMs") + if err != nil { + return nil, err + } + + return &remoteQueryExecuteLimits{MaxRows: maxRows, MaxBytes: maxBytes, TimeoutMs: timeoutMs}, nil +} + +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", name) + } + return *value, nil +} + +func (s *RemoteQueryExecuteService) Execute(_ RemoteQueryExecuteRequest) RemoteQueryExecuteResult { + 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 remoteQueryExecuteErrorResult(http.StatusBadRequest, statusInvalidRequest, "operation must be copy_stream") + } + 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, + Status: status, + Error: &RemoteQueryExecuteError{Code: status, Message: message}, + } +} + +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" { + return "", errors.New("operation must be 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 +} + +func writeExecuteParseError(w http.ResponseWriter, err error) { + parseErr, ok := err.(requestParseError) + if !ok { + writeExecuteError(w, http.StatusBadRequest, statusInvalidRequest, err.Error()) + return + } + + writeExecuteError(w, http.StatusBadRequest, parseErr.status, parseErr.message) +} + +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/remotequeries/impl/remote_query_match.go b/comp/remotequeries/impl/remote_query_match.go new file mode 100644 index 000000000000..9c65020a9c64 --- /dev/null +++ b/comp/remotequeries/impl/remote_query_match.go @@ -0,0 +1,417 @@ +// 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 ( + "bytes" + "encoding/json" + "errors" + "io" + "mime" + "net/http" + "regexp" + "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/core/config" + "github.com/DataDog/datadog-agent/pkg/collector/check" +) + +const ( + // 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" + + 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. +type Requires struct { + fx.In + + Cfg config.Component + Collector remoteQueryCollector +} + +// 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(RemoteQueriesMatchEnabledConfig), + } + return api.NewAgentEndpointProvider(h.handle, RemoteQueryMatchEndpointPath, http.MethodPost) +} + +type remoteQueryMatchHandler struct { + 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"` + 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 remoteQueryMatchRequest struct { + Integration string + 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 + 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} +} + +var integrationNamePattern = regexp.MustCompile(`^[a-z0-9_]+$`) + +type integrationInstanceTarget struct { + host string + port int + dbname string +} + +func (h *remoteQueryMatchHandler) 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 + } + + req, err := parseMatchRequest(r) + if err != nil { + writeMatchParseError(w, err) + return + } + + matches := h.findMatches(req.Integration, req.Target) + switch len(matches) { + case 0: + 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 integration checks found") + } +} + +func parseMatchRequest(r *http.Request) (remoteQueryMatchRequest, error) { + if !isJSONContentType(r.Header.Get("Content-Type")) { + return remoteQueryMatchRequest{}, invalidRequestError("content-type must be application/json") + } + + defer r.Body.Close() + var wireReq remoteQueryMatchRequestJSON + if err := decodeStrictJSON(r.Body, &wireReq); err != nil { + return remoteQueryMatchRequest{}, parseJSONRequestError(err) + } + + integration, err := parseIntegration(wireReq.Integration) + if err != nil { + return remoteQueryMatchRequest{}, err + } + target, err := parseTarget(wireReq.Target) + if err != nil { + return remoteQueryMatchRequest{}, err + } + return remoteQueryMatchRequest{Integration: integration, Target: target}, nil +} + +func parseIntegration(integration string) (string, error) { + integration = strings.ToLower(strings.TrimSpace(integration)) + if integration == "" { + return "", errors.New("integration is required") + } + if !integrationNamePattern.MatchString(integration) { + return "", invalidRequestError("integration contains invalid characters") + } + return integration, nil +} + +func isJSONContentType(contentType string) bool { + mediaType, _, err := mime.ParseMediaType(contentType) + if err != nil { + return false + } + return mediaType == "application/json" +} + +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 + } + + type targetAlias remoteQueryTargetRequestJSON + var target targetAlias + if err := decodeStrictJSON(bytes.NewReader(data), &target); err != nil { + if isUnknownJSONFieldError(err) { + return errTargetUnknownField + } + return err + } + *t = remoteQueryTargetRequestJSON(target) + return nil +} + +func parseTarget(target *remoteQueryTargetRequestJSON) (remoteQueryTarget, error) { + if target == nil { + return remoteQueryTarget{}, errors.New("target is required") + } + + host := normalizeHost(target.Host) + if host == "" { + return remoteQueryTarget{}, errors.New("target.host is required") + } + + port, err := parseRequiredPort(target.Port) + if err != nil { + return remoteQueryTarget{}, err + } + + if target.DBName == "" { + return remoteQueryTarget{}, errors.New("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, errors.New("target.port is required") + } + if *port < 1 || *port > 65535 { + return 0, errors.New("target.port is out of range") + } + 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 errors.New("target.port must be an integer") + case "target": + return errTargetMustBeObject + case "maxRows", "limits.maxRows": + return errors.New("limits.maxRows must be an integer") + case "maxBytes", "limits.maxBytes": + return errors.New("limits.maxBytes must be an integer") + case "timeoutMs", "limits.timeoutMs": + return errors.New("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 { + host = strings.ToLower(strings.TrimSpace(host)) + return strings.TrimSuffix(host, ".") +} + +type integrationCheckMatch struct { + check check.Check + sanitized sanitizedMatch +} + +func (h *remoteQueryMatchHandler) findMatches(integration string, target remoteQueryTarget) []integrationCheckMatch { + return findIntegrationMatches(h.collector, integration, target) +} + +func findIntegrationMatches(collector remoteQueryCollector, integration string, target remoteQueryTarget) []integrationCheckMatch { + checks := collector.GetChecks() + matches := make([]integrationCheckMatch, 0, 1) + for _, chk := range checks { + if normalizeIntegrationName(chk.String()) != integration { + continue + } + + instanceTarget, ok := parseIntegrationInstanceTarget(chk.InstanceConfig()) + if !ok { + continue + } + + if instanceTarget.host == target.Host && instanceTarget.port == target.Port && instanceTarget.dbname == target.DBName { + matches = append(matches, integrationCheckMatch{ + check: chk, + sanitized: sanitizedMatch{ + Integration: integration, + Loader: chk.Loader(), + ConfigProvider: chk.ConfigProvider(), + }, + }) + } + } + return matches +} + +func normalizeIntegrationName(name string) string { + return strings.ToLower(strings.TrimSpace(name)) +} + +func parseIntegrationInstanceTarget(instanceConfig string) (integrationInstanceTarget, bool) { + var fields map[string]any + if err := yaml.Unmarshal([]byte(instanceConfig), &fields); err != nil || fields == nil { + return integrationInstanceTarget{}, false + } + + host, ok := fields["host"].(string) + if !ok { + return integrationInstanceTarget{}, false + } + host = normalizeHost(host) + if host == "" { + return integrationInstanceTarget{}, false + } + + port, ok := yamlInt(fields["port"]) + if !ok || port < 1 || port > 65535 { + return integrationInstanceTarget{}, false + } + + dbname, ok := fields["dbname"].(string) + if !ok || dbname == "" { + return integrationInstanceTarget{}, false + } + + return integrationInstanceTarget{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 writeMatchParseError(w http.ResponseWriter, err error) { + parseErr, ok := err.(requestParseError) + if !ok { + writeMatchResponse(w, http.StatusBadRequest, statusInvalidRequest, 0, nil, err.Error()) + return + } + + 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) { + 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/remotequeries/impl/remote_query_par_poc.go b/comp/remotequeries/impl/remote_query_par_poc.go new file mode 100644 index 000000000000..d8fc3374a4cf --- /dev/null +++ b/comp/remotequeries/impl/remote_query_par_poc.go @@ -0,0 +1,110 @@ +// 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" + "errors" + "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"` + 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. +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{}, errors.New("remote query PAR harness requires an IPC client") + } + if h.endpointURL == "" { + return RemoteQueryPARResult{}, errors.New("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{}, errors.New("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{}, errors.New("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{}, 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_par_poc_test.go b/comp/remotequeries/impl/remote_query_par_poc_test.go new file mode 100644 index 000000000000..5ec718fd2268 --- /dev/null +++ b/comp/remotequeries/impl/remote_query_par_poc_test.go @@ -0,0 +1,144 @@ +// 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 TestRemoteQueryPARHarnessUsesCredentialFreeIPCPostShape(t *testing.T) { + client := &capturePostClient{response: []byte(`{"status":"SUCCEEDED","rows":[{"value":1}]}`)} + harness := NewRemoteQueryPARHarness(client, "https://localhost:5001"+AgentRemoteQueryExecuteEndpointPath) + + result, err := harness.Execute(context.Background(), RemoteQueryPARInputs{ + Integration: "postgres", + Operation: "copy_stream", + Format: "csv", + Target: remoteQueryTargetJSON{Host: "localhost", Port: 5432, DBName: "postgres"}, + Query: remoteQueryProofSeedQuery, + 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","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 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) + server := ipc.NewMockServer(ipc.HTTPMiddleware(mux)) + harness := NewRemoteQueryPARHarness(ipc.GetClient(), server.URL+AgentRemoteQueryExecuteEndpointPath) + + 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, statusInvalidRequest, result.Status) + require.NotNil(t, result.Error) + assert.Contains(t, result.Error.Message, "streaming executor") + assert.NotContains(t, string(result.Raw), "datastore-secret") +} + +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(AgentRemoteQueryExecuteEndpointPath, handler.handle) + server := ipc.NewMockServer(ipc.HTTPMiddleware(mux)) + harness := NewRemoteQueryPARHarness(ipc.GetClient(), server.URL+AgentRemoteQueryExecuteEndpointPath) + + tests := []struct { + name string + inputs RemoteQueryPARInputs + wantStatus string + wantCode string + }{ + { + name: "target not found", + inputs: RemoteQueryPARInputs{ + Integration: "postgres", + Operation: "copy_stream", + Format: "csv", + Target: remoteQueryTargetJSON{Host: "localhost", Port: 5432, DBName: "other"}, + Query: remoteQueryProofSeedQuery, + }, + 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", + }, + 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 _ remoteQueryPARIPCClient = (*capturePostClient)(nil) diff --git a/comp/remotequeries/impl/remote_query_test.go b/comp/remotequeries/impl/remote_query_test.go new file mode 100644 index 000000000000..ad3dbf8b0b3f --- /dev/null +++ b/comp/remotequeries/impl/remote_query_test.go @@ -0,0 +1,531 @@ +// 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/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: `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"extra":true}`, + wantError: "request contains unknown field", + }, + { + name: "unknown target field", + body: `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres","extra":true}}`, + wantError: "target contains unknown 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 unknown field", + }, + { + name: "credential-like target field is unknown", + body: `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres","username":"alice"}}`, + wantError: "target contains unknown field", + }, + { + name: "non-integer port", + body: `{"integration":"postgres","target":{"host":"localhost","port":5432.1,"dbname":"postgres"}}`, + wantError: "target.port must be an integer", + }, + { + name: "string port", + body: `{"integration":"postgres","target":{"host":"localhost","port":"5432","dbname":"postgres"}}`, + wantError: "target.port must be an integer", + }, + { + name: "missing dbname", + body: `{"integration":"postgres","target":{"host":"localhost","port":5432}}`, + wantError: "target.dbname is required", + }, + { + name: "malformed JSON", + 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, RemoteQueryMatchEndpointPath, 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, RemoteQueryMatchEndpointPath, strings.NewReader( + `{"integration":"postgres","target":{"host":" LocalHost. ","port":5432,"dbname":"Postgres"}}`, + )) + req.Header.Set("Content-Type", "application/json; charset=utf-8") + + parsed, err := parseMatchRequest(req) + require.NoError(t, err) + assert.Equal(t, "postgres", parsed.Integration) + assert.Equal(t, remoteQueryTarget{Host: "localhost", Port: 5432, DBName: "Postgres"}, parsed.Target) +} + +func TestParseMatchRequestRejectsInvalidIntegration(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, RemoteQueryMatchEndpointPath, strings.NewReader( + `{"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, "integration contains invalid characters", err.Error()) +} + +func TestRemoteQueryMatchHandlerDisabled(t *testing.T) { + handler := &remoteQueryMatchHandler{enabled: false, collector: fakeCollector{}} + + 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 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, `{"integration":"postgres","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 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, `{"integration":"postgres","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 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, `{"integration":"postgres","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 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"}}`) + + assert.Equal(t, http.StatusBadRequest, recorder.Code) + body := recorder.Body.String() + assert.Contains(t, body, `"status":"invalid_request"`) + assert.Contains(t, body, "target contains unknown field") + assert.NotContains(t, body, "postgres://secret-value@example/db") + assert.NotContains(t, body, "secret-value") +} + +func TestRemoteQueryMatchHandlerRejectsInvalidIntegration(t *testing.T) { + handler := &remoteQueryMatchHandler{enabled: true, collector: fakeCollector{}} + + recorder := callMatchHandler(handler, `{"integration":"my-sql","target":{"host":"localhost","port":3306,"dbname":"mysql"}}`) + + 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") +} + +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() + + 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 *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) + return recorder +} + +type fakeCollector struct { + checks []check.Check +} + +func (f fakeCollector) GetChecks() []check.Check { return f.checks } + +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 } + +func TestParseExecuteRequestValidatesStrictShape(t *testing.T) { + tests := []struct { + name string + body string + wantError string + }{ + { + name: "unknown top level field", + body: `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value","extra":true}`, + wantError: "request contains unknown 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 unknown field", + }, + { + name: "unknown target field", + body: `{"integration":"postgres","target":{"host":"localhost","port":5432,"dbname":"postgres","extra":true},"query":"SELECT 1 AS value"}`, + wantError: "target contains unknown 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: "target contains unknown field", + }, + { + name: "non-exact query", + 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: `{"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-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: "limits contains unknown field", + }, + { + name: "string maxRows", + 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: `{"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, RemoteQueryExecuteEndpointPath, 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 TestParseExecuteRequestNormalizesAndMarshalsCopyStreamExecutorJSON(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, RemoteQueryExecuteEndpointPath, strings.NewReader( + `{"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, `{"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") +} + +func TestParseExecuteRequestRejectsInvalidIntegration(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, RemoteQueryExecuteEndpointPath, strings.NewReader( + `{"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, "integration contains invalid characters", err.Error()) +} + +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") + + _, _, 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","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, `{"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 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) { + handler := &remoteQueryExecuteHandler{enabled: false, collector: fakeCollector{}} + + 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 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","operation":"copy_stream","format":"csv","target":{"host":"LOCALHOST.","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) + + assert.Equal(t, http.StatusBadRequest, recorder.Code) + assert.Contains(t, recorder.Body.String(), "streaming executor") + 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: []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 []check.RemoteQueryStreamEvent + result := service.ExecuteStream(req, func(event check.RemoteQueryStreamEvent) 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.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{}} + + recorder := callExecuteHandler(handler, `{"integration":"my-sql","target":{"host":"localhost","port":3306,"dbname":"mysql"},"query":"SELECT 1 AS value"}`) + + 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") +} + +func TestRemoteQueryExecuteHandlerNoMatchAndAmbiguous(t *testing.T) { + t.Run("no match", func(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: secret-value\n"}}, + }}} + + 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.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") + }) + + t.Run("ambiguous", func(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: 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, `{"integration":"postgres","operation":"copy_stream","format":"csv","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) + + 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") + }) +} + +func TestRemoteQueryExecuteHandlerUnsupportedAndRunnerErrorAreSanitized(t *testing.T) { + t.Run("unsupported", func(t *testing.T) { + 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, `{"integration":"postgres","operation":"copy_stream","format":"csv","target":{"host":"localhost","port":5432,"dbname":"postgres"},"query":"SELECT 1 AS value"}`) + + assert.Equal(t, http.StatusBadRequest, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"status":"invalid_request"`) + assert.NotContains(t, recorder.Body.String(), "secret-value") + }) + + t.Run("runner error", func(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: secret-value\n"}, err: assert.AnError}, + }}} + + 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.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()) + }) +} + +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) + return recorder +} + +type fakeWrappedCheck struct { + check.Check +} + +func (f fakeWrappedCheck) Unwrap() check.Check { + return f.Check +} + +type fakeRunnerCheck struct { + fakeCheck + err error +} + +type fakeStreamRunnerCheck struct { + fakeRunnerCheck + events []check.RemoteQueryStreamEvent + streamSeen string + streamCalls int +} + +func (f *fakeStreamRunnerCheck) RunRemoteQueryStream(integration string, requestJSON string, emit func(check.RemoteQueryStreamEvent) 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/go.mod b/go.mod index 83cb00ada87d..21b94affa03e 100644 --- a/go.mod +++ b/go.mod @@ -954,6 +954,7 @@ require ( github.com/DataDog/datadog-agent/pkg/util/cgroups v0.64.0-rc.3 github.com/DataDog/datadog-agent/pkg/util/hostport v0.0.0-00010101000000-000000000000 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.18 github.com/aptly-dev/aptly v1.6.3-0.20260504093056-0d31298f3709 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/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 66c608e07473..59354358119f 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" @@ -33,12 +34,19 @@ import ( ) /* +#include #include #include "datadog_agent_rtloader.h" #include "rtloader_mem.h" 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, uintptr_t userdata) { + return run_remote_query_stream(rtloader, check, integration, request_json, remoteQueryStreamEmitBridge, (void *)userdata); +} static inline void call_free(void* ptr) { _free(ptr); @@ -157,6 +165,43 @@ func (c *PythonCheck) RunSimple() error { return c.runCheck(false) } +// 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)) + if integration == "" { + return errors.New("integration is required") + } + if emit == nil { + return errors.New("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, C.uintptr_t(h)) + if ok == 0 { + if err := getRtLoaderError(); err != nil { + return err + } + return errors.New("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 9f41fc0fa4bb..5198a2fd425b 100644 --- a/pkg/collector/python/check_test.go +++ b/pkg/collector/python/check_test.go @@ -70,3 +70,15 @@ func TestCheckDiagnosesDeserialization(t *testing.T) { func TestRunAfterCancel(t *testing.T) { testRunAfterCancel(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..1a4bd0e3a600 --- /dev/null +++ b/pkg/collector/python/remote_query_stream.go @@ -0,0 +1,40 @@ +// 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 +#include +*/ +import "C" + +import ( + "runtime/cgo" + "unsafe" + + checkbase "github.com/DataDog/datadog-agent/pkg/collector/check" +) + +//export remoteQueryStreamEmitBridge +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(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 6ff27e7a4398..1fa6f167fc6e 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" ) @@ -104,6 +105,30 @@ char *get_check_diagnoses(rtloader_t *s, rtloader_pyobject_t *check) { return get_check_diagnoses_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 *, 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 = "{\"status\":\"SUCCEEDED\"}"; + uint8_t payload[] = {0x00, 0xff, 0x80}; + if (run_remote_query_stream_return && emit != NULL) { + 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; + } + } + return run_remote_query_stream_return; +} + // // get_check MOCK // @@ -203,6 +228,12 @@ void reset_check_mock() { get_check_diagnoses_return = NULL; get_check_diagnoses_calls = 0; + 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" @@ -668,6 +699,65 @@ func testGetDiagnoses(t *testing.T) { assert.Zero(t, len(diagnoses[1].Remediation)) } +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 []checkbase.RemoteQueryStreamEvent + err = check.RunRemoteQueryStream(" Postgres ", `{"operation":"copy_stream"}`, func(event checkbase.RemoteQueryStreamEvent) error { + events = append(events, event) + return nil + }) + + require.NoError(t, err) + 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) + 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(checkbase.RemoteQueryStreamEvent) 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(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) +} + func testRunAfterCancel(t *testing.T) { mockRtloader(t) diff --git a/pkg/privateactionrunner/bundles/registry.go b/pkg/privateactionrunner/bundles/registry.go index 01ba8947f307..dee2dc4934bc 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(configuration), "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 7d6abe309225..11b1394878fc 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/BUILD.bazel b/pkg/privateactionrunner/bundles/remotequeries/BUILD.bazel new file mode 100644 index 000000000000..76e82806edbb --- /dev/null +++ b/pkg/privateactionrunner/bundles/remotequeries/BUILD.bazel @@ -0,0 +1,39 @@ +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"], + 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", + ], +) diff --git a/pkg/privateactionrunner/bundles/remotequeries/entrypoint.go b/pkg/privateactionrunner/bundles/remotequeries/entrypoint.go new file mode 100644 index 000000000000..f7b61b70a180 --- /dev/null +++ b/pkg/privateactionrunner/bundles/remotequeries/entrypoint.go @@ -0,0 +1,45 @@ +// 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 +} + +var defaultBridgeClientFactory BridgeClientFactory = NewDefaultBridgeClient + +func NewRemoteQueriesBundle() *RemoteQueriesBundle { + return &RemoteQueriesBundle{ + actions: map[string]types.Action{ + 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/execute.go b/pkg/privateactionrunner/bundles/remotequeries/execute.go new file mode 100644 index 000000000000..989ba74aa5be --- /dev/null +++ b/pkg/privateactionrunner/bundles/remotequeries/execute.go @@ -0,0 +1,357 @@ +// 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" + "errors" + "io" + "time" + "unicode/utf8" + + "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" +) + +// 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. +type BridgeClientFactory func() (BridgeClient, error) + +type ExecuteAction struct { + newBridgeClient BridgeClientFactory +} + +func NewExecuteAction(newBridgeClient BridgeClientFactory) *ExecuteAction { + return &ExecuteAction{newBridgeClient: newBridgeClient} +} + +type ExecuteInputs struct { + 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 { + 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"` +} + +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, + _ *privateconnection.PrivateCredentials, +) (interface{}, error) { + actionStart := time.Now() + inputs, err := types.ExtractInputs[ExecuteInputs](task) + if err != nil { + return nil, util.DefaultActionErrorWithDisplayError( + errors.New("invalid remote query action inputs"), + "invalid remote query action inputs", + ) + } + + if a == nil || a.newBridgeClient == nil { + 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(errors.New("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 +} + +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), + 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), + } + } + 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 +} + +func remoteQueryExecuteOutputFromStream(stream grpc.ServerStreamingClient[pb.RemoteQueryExecuteChunk]) (map[string]interface{}, error) { + if stream == nil { + return nil, errors.New("remote query response stream missing") + } + + 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 { + chunk, err := stream.Recv() + if err == io.EOF { + break + } + if err != nil { + return nil, err + } + if chunk == nil { + return nil, errors.New("remote query response stream returned nil chunk") + } + if firstChunkAt.IsZero() { + firstChunkAt = time.Now() + } + if chunk.GetChunkIndex() != expectedChunkIndex { + return nil, errors.New("remote query response stream chunk index mismatch") + } + if seenFinal { + return nil, errors.New("remote query response stream sent chunk after final") + } + if event := chunk.GetEvent(); event != nil { + streamEvent, err := remoteQueryStreamEventFromProto(event) + 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, errors.New("remote query response stream chunk missing typed event") + } + seenFinal = chunk.GetFinal() + if seenFinal { + finalChunkAt = time.Now() + } + expectedChunkIndex++ + } + if !seenFinal { + return nil, errors.New("remote query response stream missing final chunk") + } + if len(typedStreamEvents) == 0 { + return nil, errors.New("remote query response stream missing typed events") + } + 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) { + 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, errors.New("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 errorEvent 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 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, errors.New("remote query stream response missing final event") + } + status, _ := finalEvent["status"].(string) + if status == "" { + return nil, errors.New("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 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 new file mode 100644 index 000000000000..4e8d0ef77856 --- /dev/null +++ b/pkg/privateactionrunner/bundles/remotequeries/execute_test.go @@ -0,0 +1,211 @@ +// 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" + "io" + "testing" + + "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) { + 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", + "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(1024), client.request.GetCopyLimits().GetChunkBytes()) + requestEvidence, err := json.Marshal(client.request) + require.NoError(t, err) + assert.NotContains(t, string(requestEvidence), "secret-value") + 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) { + client := &captureBridgeClient{chunks: []*pb.RemoteQueryExecuteChunk{ + {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 }) + + 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, []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{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 + }) + + 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, 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 { + 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 nil, 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 + } + return &captureRemoteQueryExecuteStream{chunks: c.chunks}, 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/ipc_client.go b/pkg/privateactionrunner/bundles/remotequeries/ipc_client.go new file mode 100644 index 000000000000..14cdab6181db --- /dev/null +++ b/pkg/privateactionrunner/bundles/remotequeries/ipc_client.go @@ -0,0 +1,54 @@ +// 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" + "fmt" + "strconv" + "time" + + "github.com/DataDog/datadog-agent/pkg/api/security/cert" + 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, error) { + cfg := pkgconfigsetup.Datadog() + + clientTLSConfig, _, _, err := cert.FetchIPCCert(cfg) + if err != nil { + return nil, fmt.Errorf("fetch Agent IPC certificate: %w", err) + } + + 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, fmt.Errorf("create AgentSecure client: %w", err) + } + return client, nil +} + +func agentIPCHost() (string, error) { + cfg := pkgconfigsetup.Datadog() + 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) + } + 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 new file mode 100644 index 000000000000..66730d96b50e --- /dev/null +++ b/pkg/privateactionrunner/bundles/remotequeries/live_agent_ipc_par_loop_test.go @@ -0,0 +1,300 @@ +// 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 remotequeries_live && !windows + +package com_datadoghq_remotequeries_test + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "fmt" + "os" + "strconv" + "strings" + "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" + 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{ + "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) { + 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()) + proofQuery := remoteQueriesProofQueryFromEnv() + inputs := map[string]interface{}{ + "integration": "postgres", + "operation": "copy_stream", + "format": "csv", + "target": remoteQueriesPostgresTargetFromEnv(t), + "query": proofQuery, + "copyLimits": remoteQueriesProofCopyLimits(proofQuery), + } + 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 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)) + require.NoError(t, err) + if !result.Success { + 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) + assert.Equal(t, "SUCCEEDED", result.Outputs["status"]) + require.Contains(t, result.Outputs, "data") + + data, ok := result.Outputs["data"].(string) + require.True(t, ok) + assertRemoteQueriesProofCopyData(t, proofQuery, data) + + resultEvidence, err := json.Marshal(summarizeRemoteQueriesProofPayload(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 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", + }) +} + +func getenvOptional(name string) string { + return os.Getenv(name) +} + +func remoteQueriesProofQueryFromEnv() string { + if query := os.Getenv(remoteQueriesProofQueryOverrideEnv); query != "" { + return query + } + return remoteQueriesFixtureTableProofQuery +} + +func remoteQueriesProofCopyLimits(query string) map[string]interface{} { + maxBytes := 4 << 10 // 4 KiB. + timeoutMs := 5_000 + if payloadBytes, ok := remoteQueriesLargePayloadBytes(query); ok { + maxBytes = payloadBytes + (1 << 20) // Add 1 MiB headroom. + timeoutMs = 60_000 + } + chunkBytes := 256 << 10 // 256 KiB. + 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": chunkBytes, + "maxBytes": maxBytes, + "maxRowBytes": 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 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() + + 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: + 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) + } +} + +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() + + 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 == "" { + return + } + var payload strings.Builder + for _, line := range lines { + payload.WriteString(line) + payload.WriteByte('\n') + } + require.NoError(t, os.WriteFile(path, []byte(payload.String()), 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/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..cb5057c4406a --- /dev/null +++ b/pkg/privateactionrunner/bundles/remotequeries/live_par_loop_test.go @@ -0,0 +1,185 @@ +// 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 remotequeries_live && !windows + +package com_datadoghq_remotequeries_test + +import ( + "context" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "encoding/json" + "io" + "testing" + "time" + + "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" + "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" + 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" +) + +func TestRemoteQueriesActionRunsThroughLivePARLoopAndFakeintake(t *testing.T) { + t.Setenv(app.InternalSkipTaskVerificationEnvVar, "true") + + 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() + + 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", + }, + "operation": "copy_stream", + "format": "csv", + "query": "SELECT 1 AS value", + "copyLimits": map[string]interface{}{ + "chunkBytes": 1024, + "maxBytes": 1024, + "maxRowBytes": 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, "events") + + select { + 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, "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()) + assert.Equal(t, "postgres", req.GetTarget().GetDbname()) + case <-time.After(2 * time.Second): + require.FailNow(t, "remote query action did not call the AgentSecure client") + } + + 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 captureAgentSecureClient struct { + requests chan<- *pb.RemoteQueryExecuteRequest + response *pb.RemoteQueryExecuteResponse +} + +func (c *captureAgentSecureClient) RemoteQueryExecute(_ context.Context, req *pb.RemoteQueryExecuteRequest, _ ...grpc.CallOption) (*pb.RemoteQueryExecuteResponse, error) { + 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 + 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 { + 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 new file mode 100644 index 000000000000..daf07b91cb52 --- /dev/null +++ b/pkg/privateactionrunner/bundles/remotequeries/standalone_par_process_proof_test.go @@ -0,0 +1,221 @@ +// 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 remotequeries_live && !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()) + 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, + "copyLimits": remoteQueriesProofCopyLimits(proofQuery), + } + 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 RemoteQueryExecuteStream", cmdPortInt) + require.NoError(t, fakeintakeClient.EnqueuePARTask(taskID, fqn, inputs)) + + result, err := fakeintakeClient.GetPARTaskResult(taskID, remoteQueriesProofResultTimeout(proofQuery)) + require.NoError(t, err) + if !result.Success { + 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) + require.Equal(t, taskID, result.TaskID) + assert.Equal(t, "SUCCEEDED", result.Outputs["status"]) + 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) + } + + resultEvidence, err := json.Marshal(summarizeRemoteQueriesProofPayload(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), + "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), + 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: 120 +`, 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/pkg/proto/datadog/api/v1/BUILD.bazel b/pkg/proto/datadog/api/v1/BUILD.bazel index da552f241c1e..4615dd8e8e18 100644 --- a/pkg/proto/datadog/api/v1/BUILD.bazel +++ b/pkg/proto/datadog/api/v1/BUILD.bazel @@ -15,6 +15,7 @@ proto_library( "//pkg/proto/datadog/workloadfilter:workloadfilter_proto", "//pkg/proto/datadog/workloadmeta:workloadmeta_proto", "@protobuf//:empty_proto", + "@protobuf//:struct_proto", ], ) diff --git a/pkg/proto/datadog/api/v1/api.proto b/pkg/proto/datadog/api/v1/api.proto index 91882e0fc173..34d2cdad9718 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,93 @@ 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 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 { + 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; +} + +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 { + reserved 1; + int32 chunk_index = 2; + bool final = 3; + RemoteQueryExecuteStreamEvent event = 4; +} + service AgentSecure { // subscribes to added, removed, or changed entities in the Tagger // and streams them to clients as events. @@ -71,6 +159,12 @@ 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); + + // 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. 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..f8e3c93a3187 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,956 @@ 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 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"` + 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[3] + 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[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 RemoteQueryExecuteRequest.ProtoReflect.Descriptor instead. +func (*RemoteQueryExecuteRequest) Descriptor() ([]byte, []int) { + return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{3} +} + +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 +} + +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"` + 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[4] + 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[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 RemoteQueryExecuteError.ProtoReflect.Descriptor instead. +func (*RemoteQueryExecuteError) Descriptor() ([]byte, []int) { + return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{4} +} + +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[5] + 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[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 RemoteQueryExecuteResponse.ProtoReflect.Descriptor instead. +func (*RemoteQueryExecuteResponse) Descriptor() ([]byte, []int) { + return file_datadog_api_v1_api_proto_rawDescGZIP(), []int{5} +} + +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 +} + +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"` + 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[11] + 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[11] + 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{11} +} + +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 +} + +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 = "" + "\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\"\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\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" + + "\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\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\"\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\x05eventJ\x04\b\x01\x10\x022Z\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\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" + @@ -46,91 +991,141 @@ 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\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 ( + 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, 15) 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 + (*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 + (*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.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 + 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 + 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() } @@ -145,18 +1140,25 @@ 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: 0, + NumMessages: 15, 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 7fd74127169d..1786843eeac5 100644 --- a/pkg/proto/pbgo/core/api_grpc.pb.go +++ b/pkg/proto/pbgo/core/api_grpc.pb.go @@ -146,6 +146,8 @@ 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_RemoteQueryExecuteStream_FullMethodName = "/datadog.api.v1.AgentSecure/RemoteQueryExecuteStream" AgentSecure_StreamKubeMetadata_FullMethodName = "/datadog.api.v1.AgentSecure/StreamKubeMetadata" ) @@ -185,6 +187,10 @@ 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) + // 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) } @@ -416,9 +422,38 @@ 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) 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 } @@ -471,6 +506,10 @@ 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) + // 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 mustEmbedUnimplementedAgentSecureServer() @@ -537,6 +576,12 @@ func (UnimplementedAgentSecureServer) StreamConfigEvents(*ConfigStreamRequest, g func (UnimplementedAgentSecureServer) WorkloadFilterEvaluate(context.Context, *WorkloadFilterEvaluateRequest) (*WorkloadFilterEvaluateResponse, error) { return nil, status.Error(codes.Unimplemented, "method WorkloadFilterEvaluate not implemented") } +func (UnimplementedAgentSecureServer) RemoteQueryExecute(context.Context, *RemoteQueryExecuteRequest) (*RemoteQueryExecuteResponse, error) { + return nil, status.Error(codes.Unimplemented, "method RemoteQueryExecute not implemented") +} +func (UnimplementedAgentSecureServer) RemoteQueryExecuteStream(*RemoteQueryExecuteRequest, grpc.ServerStreamingServer[RemoteQueryExecuteChunk]) error { + return status.Error(codes.Unimplemented, "method RemoteQueryExecuteStream not implemented") +} func (UnimplementedAgentSecureServer) StreamKubeMetadata(*KubeMetadataStreamRequest, grpc.ServerStreamingServer[KubeMetadataStreamResponse]) error { return status.Error(codes.Unimplemented, "method StreamKubeMetadata not implemented") } @@ -846,6 +891,35 @@ 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_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 { @@ -916,6 +990,10 @@ var AgentSecure_ServiceDesc = grpc.ServiceDesc{ MethodName: "WorkloadFilterEvaluate", Handler: _AgentSecure_WorkloadFilterEvaluate_Handler, }, + { + MethodName: "RemoteQueryExecute", + Handler: _AgentSecure_RemoteQueryExecute_Handler, + }, }, Streams: []grpc.StreamDesc{ { @@ -944,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 1c09277b913a..3715563dc831 100644 --- a/pkg/proto/pbgo/mocks/core/api_mockgen.pb.go +++ b/pkg/proto/pbgo/mocks/core/api_mockgen.pb.go @@ -385,6 +385,46 @@ 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...) +} + +// 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() @@ -731,6 +771,35 @@ 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) +} + +// 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/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. diff --git a/rtloader/include/datadog_agent_rtloader.h b/rtloader/include/datadog_agent_rtloader.h index 8b3b8f8036fe..41e82d1881a9 100644 --- a/rtloader/include/datadog_agent_rtloader.h +++ b/rtloader/include/datadog_agent_rtloader.h @@ -190,6 +190,22 @@ 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) + \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 98c371e42a8c..c7c7af7a4d8a 100644 --- a/rtloader/include/rtloader.h +++ b/rtloader/include/rtloader.h @@ -121,6 +121,19 @@ class RtLoader */ virtual char *runCheck(RtLoaderPyObject *check) = 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..392b8cd482d0 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,6 +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 enum { DATADOG_AGENT_RTLOADER_GAUGE = 0, diff --git a/rtloader/rtloader/api.cpp b/rtloader/rtloader/api.cpp index 0c62375399d6..a7882adbacc6 100644 --- a/rtloader/rtloader/api.cpp +++ b/rtloader/rtloader/api.cpp @@ -278,6 +278,15 @@ char *run_check(rtloader_t *rtloader, rtloader_pyobject_t *check) return AS_TYPE(RtLoader, rtloader)->runCheck(AS_TYPE(RtLoaderPyObject, check)); } +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/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 36906d82d5a7..6f09292cc729 100644 --- a/rtloader/three/three.cpp +++ b/rtloader/three/three.cpp @@ -19,6 +19,7 @@ #include "util.h" #include +#include #include #include @@ -140,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; } } @@ -497,6 +501,149 @@ char *Three::runCheck(RtLoaderPyObject *check) return ret; } +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 == '_'; }); + } + + 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; + } + + 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." }; +} // namespace + +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 ffce489bc4b7..6fc3cbdd1560 100644 --- a/rtloader/three/three.h +++ b/rtloader/three/three.h @@ -65,6 +65,8 @@ class Three : public RtLoader const char *provider_str, RtLoaderPyObject *&check); char *runCheck(RtLoaderPyObject *check); + 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/fused-local-par-agent-postgres-proof.sh b/test/remotequeries/fused-local-par-agent-postgres-proof.sh new file mode 100755 index 000000000000..2ea5a23348cb --- /dev/null +++ b/test/remotequeries/fused-local-par-agent-postgres-proof.sh @@ -0,0 +1,382 @@ +#!/usr/bin/env bash +# 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 -> 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_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. + +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_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_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:-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:-} + +AGENT_PID="" +POSTGRES_COMPOSE_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="" +} + +docker_compose() { + docker compose -f "$POSTGRES_COMPOSE_FILE" -p "$POSTGRES_COMPOSE_PROJECT" "$@" +} + +cleanup_all() { + cleanup_agent + if [[ "$POSTGRES_COMPOSE_STARTED" == "1" ]]; then + docker_compose down --remove-orphans >/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 +} + +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" == "bob" && \ + "$RQ_POSTGRES_PASSWORD" == "bob" ]] +} + +start_postgres_fixture() { + if postgres_is_ready; then + log "Using existing compatible Postgres fixture at $RQ_POSTGRES_HOST:$RQ_POSTGRES_PORT/$RQ_POSTGRES_DBNAME" + return + 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 + + [[ -f "$POSTGRES_COMPOSE_FILE" ]] || { + echo "Postgres compose fixture not found: $POSTGRES_COMPOSE_FILE" >&2 + exit 1 + } + docker compose version >/dev/null 2>&1 || { + echo "docker compose is required to start the integrations-core Postgres fixture" >&2 + exit 1 + } + + 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_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" < "$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 + 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({ + "integration": "postgres", + "target": { + "host": os.environ["RQ_POSTGRES_HOST"], + "port": int(os.environ["RQ_POSTGRES_PORT"]), + "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}, +})) +PY +) + local token + token=$(cat "$TMP_ROOT/run/auth_token") + + 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}" \ + -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" != "400" ]]; then + echo "FAIL: expected Agent execute preflight HTTP 400 for disabled inline execution, got $status" >&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 AgentSecure gRPC 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" \ + 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" +} + +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 AgentSecure IPC evidence are in $TMP_ROOT/results/fused-proof-evidence.txt" +} + +main "$@" 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..cdbb6e0befa3 --- /dev/null +++ b/test/remotequeries/standalone-par-agentsecure-postgres-proof.sh @@ -0,0 +1,517 @@ +#!/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 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. +# +# 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_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 +# 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_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_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_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_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} +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:-} + +AGENT_PID="" +POSTGRES_COMPOSE_STARTED=0 +PROOF_CASE_NAME="" +CASE_RESULTS_DIR="" + +PROOF_CASE_NAMES=( + "seed" + "fixture-city" + "copy-fixture-city" + "copy-binary-payload" + "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 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" + "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)" "$*" +} + +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="" +} + +docker_compose() { + docker compose -f "$POSTGRES_COMPOSE_FILE" -p "$POSTGRES_COMPOSE_PROJECT" "$@" +} + +cleanup_all() { + cleanup_agent + if [[ "$POSTGRES_COMPOSE_STARTED" == "1" ]]; then + docker_compose down --remove-orphans >/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 +} + +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" == "bob" && \ + "$RQ_POSTGRES_PASSWORD" == "bob" ]] +} + +start_postgres_fixture() { + if postgres_is_ready; then + log "Using existing compatible Postgres fixture at $RQ_POSTGRES_HOST:$RQ_POSTGRES_PORT/$RQ_POSTGRES_DBNAME" + return + 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 + + [[ -f "$POSTGRES_COMPOSE_FILE" ]] || { + echo "Postgres compose fixture not found: $POSTGRES_COMPOSE_FILE" >&2 + exit 1 + } + docker compose version >/dev/null 2>&1 || { + echo "docker compose is required to start the integrations-core Postgres fixture" >&2 + exit 1 + } + + 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_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" < 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" \ + "$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 + 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({ + "integration": "postgres", + "target": { + "host": os.environ["RQ_POSTGRES_HOST"], + "port": int(os.environ["RQ_POSTGRES_PORT"]), + "dbname": os.environ["RQ_POSTGRES_DBNAME"], + }, + "query": os.environ["RQ_REMOTE_QUERY"], + "limits": json.loads(os.environ["RQ_LIMITS_JSON"]), +})) +PY +) + local token + token=$(cat "$TMP_ROOT/run/auth_token") + + 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 status + 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 "$status_file" + if [[ -s "$body_file" ]]; then + summarize_json_file "$body_file" "$summary_file" + fi + + if [[ "$status" != "400" ]]; then + echo "FAIL: expected Agent execute preflight HTTP 400 for disabled inline execution, got $status" >&2 + exit 1 + fi + rm -f "$body_file" +} + +run_standalone_go_proof() { + log "[$PROOF_CASE_NAME] Running standalone PAR process -> real AgentSecure gRPC streaming 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="$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" \ + 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" +} + +run_proof_case() { + PROOF_CASE_NAME=$1 + RQ_REMOTE_QUERY=$2 + if [[ "$RQ_REMOTE_OPERATION_WAS_SET" != "1" ]]; then + RQ_REMOTE_OPERATION=copy_stream + fi + if [[ "$RQ_REMOTE_FORMAT_WAS_SET" != "1" ]]; then + RQ_REMOTE_FORMAT=csv + 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" + + 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() { + 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 + run_proof_cases + + log "Sanitized standalone proof evidence" + find "$TMP_ROOT/results" -name 'standalone-proof-evidence.txt' -print -exec cat {} \; + + 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 "$@"