Skip to content

Add Postgres remote query COPY executor with optional intake upload - #23499

Draft
nubtron wants to merge 27 commits into
masterfrom
nubtron/remote-queries-poc
Draft

Add Postgres remote query COPY executor with optional intake upload#23499
nubtron wants to merge 27 commits into
masterfrom
nubtron/remote-queries-poc

Conversation

@nubtron

@nubtron nubtron commented Apr 28, 2026

Copy link
Copy Markdown
Collaborator

What does this PR do?

Adds a proof-of-concept Postgres remote query executor for the remote queries research path, with two result modes:

  • without resultDelivery, COPY bytes stream inline as COPY events through the existing Agent result path;
  • with resultDelivery, the integration uploads bounded COPY bytes directly to its-agent-intake over HTTP using a multipart contract, then returns a receipt.

The executor includes:

  • an explicit check registry abstraction with exact {host, port, dbname} matching and fail-closed target selection;
  • credential-shaped field rejection in request payloads and a strict proof-query policy;
  • binary-safe COPY streaming with configurable chunk/row/byte limits and timeout;
  • direct multipart intake upload with a server-owned part size, COPY-chunk aggregation into parts, finalize/abort, idempotent part retries, and receipt mapping;
  • a 10 GiB total upload cap and a 5-minute HTTP timeout;
  • unit tests and a targeted Docker-backed integration test.

Verified locally in this worktree with ddev --no-interactive test -s postgres (lint passed), ddev validate config -s postgres and ddev validate models -s postgres (valid / in sync), and ddev --no-interactive test postgres -- tests/test_remote_query.py (149 unit tests passed). The live integration test could not run here because a shared long-running proof container already holds host port 5432; the consolidated tree is byte-identical to the hosted #24926 tip, which already ran the integration test in its own CI.

Motivation

This consolidates the former draft stack into a single review surface:

Both child drafts are closed in favor of this PR; their branches remain available. The stack is linear, so this PR advances nubtron/remote-queries-poc to the hosted #24926 tree by fast-forward, preserving the existing signed commits without rewriting history. No intermediate relay or chunk-upload variants are reintroduced; the final tree keeps inline COPY streaming when resultDelivery is absent and direct multipart intake upload when it is present.

This is intentionally a proof executor, not a production remote-query implementation. It proves the Agent-local Python execution seam once an explicit registry supplies the matching check, plus the optional direct intake upload path.

Review checklist (to be filled by reviewers)

  • Feature or bugfix MUST have appropriate tests (unit, integration, e2e)
  • Add qa/required if this PR needs QA validation, or qa/skip-qa if it does not. Exactly one of the two is required.
  • If you need to backport this PR to another branch, you can add the backport/<branch-name> label to the PR and it will automatically open a backport PR once this one is merged

@nubtron nubtron added the qa/skip-qa Automatically skip this PR for the next QA label Apr 28, 2026
@nubtron nubtron self-assigned this Apr 28, 2026
@nubtron
nubtron force-pushed the nubtron/remote-queries-poc branch from ce610ea to 0f5123e Compare April 28, 2026 15:55
@datadog-prod-us1-6

datadog-prod-us1-6 Bot commented Apr 28, 2026

Copy link
Copy Markdown

Tests  Code Coverage

All CI checks and tests passed. Datadog automation helped this PR pass.

🎉 All green!

🧪 All tests passed
❄️ No new flaky tests detected

🔄 Datadog retried 1 test - 1 passed on retry View in Datadog

🎯 Code Coverage (details)
Patch Coverage: 95.72%
Overall Coverage: 94.38% (+5.78%)

This comment will be updated automatically if new data arrives.
🔗 Commit SHA: 3c202f6 | Docs | View more details | Give us feedback!

@nubtron
nubtron force-pushed the nubtron/remote-queries-poc branch from 0f5123e to 1f512f2 Compare April 28, 2026 16:02
@codecov

codecov Bot commented Apr 28, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 90.33333% with 29 lines in your changes missing coverage. Please review.
✅ Project coverage is 89.35%. Comparing base (18c80aa) to head (92ac25c).
⚠️ Report is 3 commits behind head on master.

Additional details and impacted files
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@nubtron
nubtron force-pushed the nubtron/remote-queries-poc branch from 30da8a7 to 92ac25c Compare April 30, 2026 13:31
@cit-pr-commenter-54b7da

cit-pr-commenter-54b7da Bot commented Aug 19, 2026

Copy link
Copy Markdown

evalya-impact-summary

evalya impact analysis
Impact analysis: 0 selected, 0 skipped (of 0 test tasks)
Publish tasks:   1 (always emitted)
Diff (6 files):
  postgres/changelog.d/23499.added
  postgres/changelog.d/23499.added.1
  postgres/changelog.d/23499.fixed
  postgres/datadog_checks/postgres/remote_query.py
  postgres/tests/test_remote_query.py
  postgres/tests/test_remote_query_integration.py

Debug a specific task: evalya plan impact --path <path> --task <task>

Learn more about CI impact filtering

Extend the Postgres remote query COPY stream with an optional
resultDelivery instruction (mode POC_PUBLIC_CHUNKED_UPLOAD). When
present, the integration feeds bounded COPY bytes through the
existing native emit callback to the Agent-owned upload transport:
chunks are sized by resultDelivery.chunkBytes, capped by
resultDelivery.maxBytes, and carry a SHA-256 digest; the STARTED
event carries the sanitized upload-session handle (uploadId only) and
the final event carries a compact provisional uploadReceipt aligned
to the Agent-owned receipt shape. Python never receives the upload
baseUrl, scoped token, or API/app keys and performs no HTTP upload;
the Agent Go side retains those secrets. Omitting resultDelivery
keeps the inline streaming behavior unchanged.
…tion

Replace the emit-callback upload bridge with direct HTTP upload to
its-agent-intake from the Postgres Python integration. When resultDelivery is
present, the integration PUTs bounded COPY chunks straight to the intake over
HTTP using the base URL and scoped upload token forwarded through the native
bridge plus the org API key and POC application key read from Agent config
via datadog_agent.get_config.

Bulk chunk bytes no longer traverse the native emit bridge, AgentSecure, PAR,
or AP action output; only metadata/final/error events cross back. The
integration owns chunking, per-chunk sha256, the HTTP PUT with retry on
transient 408/429/5xx, finalize, and best-effort abort on failure.

RemoteQueryResultDelivery now accepts and requires baseUrl and token.
execute_agent_rpc_stream_copy routes upload-mode requests to a new
_execute_upload_stream that drives the existing chunk iterator but sinks data
events to an injectable _UploadClient (production uses requests; tests inject
a fake). Omitting resultDelivery keeps the inline streaming path unchanged.
…ad proof

The Postgres COPY bridge now uploads chunks directly to its-agent-intake over
HTTP instead of feeding bytes through the native emit callback. Replace the
obsolete emit-bridge proof (which asserted data events arrive at the native
callback with sha256 and one-chunk-in-flight) with a direct-HTTP upload
proof.

The new proof drives the real _execute_upload_stream with an injectable
FakeUploadClient and asserts:

- Bulk bytes go directly to the intake over HTTP (fake put_chunk calls) in 1 MiB
  bounds, not through the emit callback; only metadata/final/error cross back.
- One-chunk backpressure: COPY reads are fenced during each HTTP upload
  (lockstep reads_at_put == 1..N), so the full payload is never buffered ahead.
- Per-chunk SHA-256 matches the payload and the aggregate matches the
  incremental reference at 8 MiB and 32 MiB.
- Finalize is called once and the camelCase receipt surfaces under the
  snake_case upload_receipt key.
- HTTP failure aborts the session and emits an error event; maxBytes is
  enforced at MiB scale.

The 243021 task reference is preserved as the deterministic upload session id.
… contract

Replace the Postgres Remote Queries chunk-object/manifest uploader with a
sequential multipart HTTP client using the POC_PUBLIC_MULTIPART_UPLOAD
contract. The integration aggregates COPY output into at most one configured
partBytes buffer and uploads contiguous 1-based /parts/{part_number} requests
carrying X-DD-Part-SHA256/Bytes/Rows headers, retries the same part
idempotently, finalizes once, and returns the small objectPath/partCount
receipt. A final short part is allowed and a zero-row result finalizes with
zero parts. Overall stream SHA-256, maxBytes enforcement, abort behavior,
credential redaction, and no bulk Agent output are preserved.

Delete the old POC_PUBLIC_CHUNKED_UPLOAD mode and the chunkBytes/manifestPath/
chunkCount aliases and compatibility. Add tests for the exact HTTP contract,
idempotent retry, fail-closed rejection, multipart sizing with a final short
part, and zero-row finalization.
…gregate chunks into parts

partBytes (the multipart part size) and limits.chunkBytes (the COPY streaming chunk
size) are distinct concepts. The worker injects an 8 MiB partBytes while the COPY
chunk size defaults to 1 MiB, so the previous cross-field clamp
(resultDelivery.partBytes <= limits.chunkBytes) rejected every normal multipart
dispatch. Remove that clamp; keep resultDelivery.maxBytes <= limits.maxBytes.

The COPY stream now always emits limits.chunkBytes-sized events. The upload client
aggregates those events into partBytes-sized parts via a bounded _MultipartBuffer (at
most one part buffered), computes each part's SHA-256 over the aggregated body, and
uploads contiguous 1-based parts. The STARTED metadata reports the COPY chunkBytes and
the resultDelivery partBytes separately; the provisional receipt reports the part
count (ceil(totalBytes/partBytes)), not the chunk count. The no-compat multipart
contract is preserved.

Add regressions proving partBytes may exceed limits.chunkBytes and that the upload
client aggregates COPY chunks into one partBytes part (with a final short part),
plus that each part carries the SHA-256 of its aggregated body.
…d 5-minute HTTP timeout

Raise the POC multipart upload capacity and timeouts to match the plan's larger
result envelope. resultDelivery.maxBytes is now accepted up to a 10 GiB server-owned
maximum (and still must not exceed limits.maxBytes), and partBytes up to 128 MiB,
independent of the 1 MiB COPY read chunk. The COPY stream still emits
limits.chunkBytes-sized events that the upload client aggregates into partBytes parts,
so memory stays bounded to one multipart part.

Change the per-upload HTTP timeout from a 60-second scalar to an explicit
(connect, read) tuple with a 10-second connect timeout and a 5-minute read timeout so
a slow large-part upload is not cut short while a stuck connect fails fast. The timeout
is now plumbed through _RequestsUploadClient into _upload_with_retry (the constructor
parameter was previously stored but unused); retry count and backoff stay bounded.

Add regressions proving the 10 GiB maxBytes boundary is accepted when the extraction
cap matches and rejected one byte past it, the 128 MiB partBytes boundary, that many
1 MiB COPY reads aggregate into 64 MiB parts without materializing the full result,
and that the HTTP timeout is the 5-minute (connect, read) tuple. No test allocates
10 GiB. Same-part retry, final short part, maxBytes fail closed, abort, checksums, and
no bulk Agent output are preserved.
The intake finalize route serializes the canonical final path as the
snake-case object_path field, but _intake_receipt_to_camel read the
obsolete object_key, producing an empty AP receipt objectPath on the
multipart finalization path. Map object_path to objectPath instead.

Remote Queries is greenfield, so the object_key alias is not preserved.
Update the finalize fixtures to object_path and add focused unit tests
covering the mapping and the dropped alias.
@nubtron nubtron changed the title Add Postgres remote query proof executor Add Postgres remote query COPY executor with optional intake upload Aug 27, 2026
…receipt

Full-object readback is debug-only/default-off, so the intake finalize response may omit or return an empty aggregate sha256. Omit the field from the emitted upload receipt when absent or empty, forward it verbatim when present and valid (64-char hex), and fail closed with an invalid_receipt error when a present value is malformed. Per-part X-DD-Part-SHA256 behavior is unchanged.
@dd-octo-sts

dd-octo-sts Bot commented Aug 28, 2026

Copy link
Copy Markdown
Contributor

Validation Report

All 21 validations passed.

Show details
Validation Description Status
agent-reqs Verify check versions match the Agent requirements file
ci Validate CI configuration and code coverage settings
codeowners Validate every integration has a CODEOWNERS entry
config Validate default configuration files against spec.yaml
dep Verify dependency pins are consistent and Agent-compatible
http Validate integrations use the HTTP wrapper correctly
imports Validate check imports do not use deprecated modules
integration-style Validate check code style conventions
jmx-metrics Validate JMX metrics definition files and config
labeler Validate PR labeler config matches integration directories
legacy-signature Validate no integration uses the legacy Agent check signature
license-headers Validate Python files have proper license headers
licenses Validate third-party license attribution list
metadata Validate metadata.csv metric definitions
models Validate configuration data models match spec.yaml
openmetrics Validate OpenMetrics integrations disable the metric limit
package Validate Python package metadata and naming
qa-label Validate the pull request declares whether it needs QA for the next Agent release
readmes Validate README files have required sections
saved-views Validate saved view JSON file structure and fields
version Validate version consistency between package and changelog

View full run

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

integration/postgres qa/skip-qa Automatically skip this PR for the next QA

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant