postgres: add deterministic Remote Queries COPY relay proofs

postgres/datadog_checks/postgres/remote_query.py

@@ -45,8 +45,15 @@
     )
 )
 
+# Server-owned maximums for the POC multipart upload path. The backend selects/clamps these,
+# not the caller. The multipart part size is independent of the COPY read chunk size
+# (limits.chunkBytes); it may be up to 128 MiB. maxBytes must not exceed the caller/backend COPY
+# safety cap (limits.maxBytes).
+REMOTE_QUERY_UPLOAD_MAX_PART_BYTES = 128 * 1024 * 1024
+REMOTE_QUERY_UPLOAD_MAX_TOTAL_BYTES = 10 * 1024 * 1024 * 1024
+
 CopyStreamFormat = Literal['csv', 'binary']
-ResultDeliveryMode = Literal['POC_PUBLIC_CHUNKED_UPLOAD']
+ResultDeliveryMode = Literal['POC_PUBLIC_MULTIPART_UPLOAD']
 ResultDeliveryFormat = Literal['csv']
 ResultDeliveryCompression = Literal['none']
 CopyStreamEmit = Callable[[str, str, bytes], None]
@@ -138,7 +145,7 @@ class RemoteQueryResultDelivery(BaseModel):
     HTTP. The Agent forwards the intake base URL and scoped upload token here; the integration
     reads the org API key and POC application key from Agent config via ``datadog_agent.get_config``
     and attaches them to its own HTTP upload requests. The integration performs the HTTP upload
-    itself; bulk chunk bytes never traverse the native emit bridge, AgentSecure, PAR, or AP action
+    itself; bulk part bytes never traverse the native emit bridge, AgentSecure, PAR, or AP action
     output. Omitting ``resultDelivery`` keeps the existing inline streaming behavior unchanged.
     """
 
@@ -148,15 +155,15 @@ class RemoteQueryResultDelivery(BaseModel):
     upload_id: StrictStr = Field(alias='uploadId', min_length=1)
     base_url: StrictStr = Field(alias='baseUrl', min_length=1)
     token: StrictStr = Field(alias='token', min_length=1)
-    chunk_bytes: StrictInt = Field(alias='chunkBytes', ge=1)
-    max_bytes: StrictInt = Field(alias='maxBytes', ge=1)
+    part_bytes: StrictInt = Field(alias='partBytes', ge=1, le=REMOTE_QUERY_UPLOAD_MAX_PART_BYTES)
+    max_bytes: StrictInt = Field(alias='maxBytes', ge=1, le=REMOTE_QUERY_UPLOAD_MAX_TOTAL_BYTES)
     format: ResultDeliveryFormat = 'csv'
     compression: ResultDeliveryCompression = 'none'
 
     @model_validator(mode='after')
-    def validate_chunk_within_max(self) -> 'RemoteQueryResultDelivery':
-        if self.chunk_bytes > self.max_bytes:
-            raise ValueError('chunkBytes must not exceed maxBytes')
+    def validate_part_within_max(self) -> 'RemoteQueryResultDelivery':
+        if self.part_bytes > self.max_bytes:
+            raise ValueError('partBytes must not exceed maxBytes')
         return self
 
 
@@ -175,20 +182,21 @@ class RemoteQueryCopyRequest(BaseModel):
     @model_validator(mode='after')
     def validate_format_consistency(self) -> 'RemoteQueryCopyRequest':
         # When resultDelivery is present, the COPY stream format must match the upload format so the
-        # emitted bytes and the manifest agree. The current contract allows only ``csv`` for upload,
-        # so upload mode is CSV-only; a ``binary`` COPY stream with a ``csv`` upload is rejected.
+        # emitted bytes and the finalized object agree. The current contract allows only ``csv`` for
+        # upload, so upload mode is CSV-only; a ``binary`` COPY stream with a ``csv`` upload is rejected.
         if self.result_delivery is not None and self.format != self.result_delivery.format:
             raise ValueError('format must match resultDelivery.format when resultDelivery is present')
         return self
 
     @model_validator(mode='after')
-    def validate_chunk_within_limits(self) -> 'RemoteQueryCopyRequest':
-        # The upload caps must not widen the caller/backend COPY safety caps. Fail closed so a
+    def validate_max_within_limits(self) -> 'RemoteQueryCopyRequest':
+        # The upload byte cap must not widen the caller/backend COPY safety cap. Fail closed so a
         # backend-injected resultDelivery cannot raise the integration's configured byte ceiling;
-        # equal or smaller upload caps are accepted.
+        # an equal or smaller upload maxBytes is accepted. ``partBytes`` (the multipart part size)
+        # and ``limits.chunkBytes`` (the COPY streaming chunk size) are distinct concepts, so
+        # partBytes may exceed chunkBytes: the COPY stream emits chunkBytes-sized events that the
+        # upload client aggregates into partBytes-sized parts.
         if self.result_delivery is not None:
-            if self.result_delivery.chunk_bytes > self.limits.chunk_bytes:
-                raise ValueError('resultDelivery.chunkBytes must not exceed limits.chunkBytes')
             if self.result_delivery.max_bytes > self.limits.max_bytes:
                 raise ValueError('resultDelivery.maxBytes must not exceed limits.maxBytes')
         return self
@@ -343,21 +351,20 @@ def _dbname_from_check(check: 'PostgreSql') -> str | None:
 
 def _started_metadata(request: RemoteQueryCopyRequest) -> dict[str, Any]:
     result_delivery = request.result_delivery
-    chunk_bytes = result_delivery.chunk_bytes if result_delivery is not None else request.limits.chunk_bytes
     max_bytes = result_delivery.max_bytes if result_delivery is not None else request.limits.max_bytes
     metadata: dict[str, Any] = {
         'status': 'STARTED',
         'format': request.format,
         'operation': request.operation,
-        'chunkBytes': chunk_bytes,
+        'chunkBytes': request.limits.chunk_bytes,
         'maxBytes': max_bytes,
         'maxRowBytes': request.limits.max_row_bytes,
     }
     if result_delivery is not None:
         metadata['resultDelivery'] = {
             'mode': result_delivery.mode,
             'uploadId': result_delivery.upload_id,
-            'chunkBytes': result_delivery.chunk_bytes,
+            'partBytes': result_delivery.part_bytes,
             'maxBytes': result_delivery.max_bytes,
             'format': result_delivery.format,
             'compression': result_delivery.compression,
@@ -370,15 +377,15 @@ def _succeeded_metadata(state: _CopyStreamState, started_at: float, request: Rem
     result_delivery = request.result_delivery
     if result_delivery is not None:
         # Provisional receipt aligned to the Agent-owned uploadReceipt shape
-        # {mode, uploadId, bucketName, manifestPath, totalBytes, totalRows, chunkCount, sha256}.
+        # {mode, uploadId, bucketName, objectPath, totalBytes, totalRows, partCount, sha256}.
         # Python owns only the fields it can compute from the byte stream; the Agent Go side
-        # enriches it with bucketName, manifestPath, totalRows, and the aggregate sha256 after
+        # enriches it with bucketName, objectPath, totalRows, and the aggregate sha256 after
         # it finalizes the upload session.
         metadata['uploadReceipt'] = {
             'mode': result_delivery.mode,
             'uploadId': result_delivery.upload_id,
             'totalBytes': state.bytes_emitted,
-            'chunkCount': state.chunks_emitted,
+            'partCount': _multipart_part_count(state.bytes_emitted, result_delivery.part_bytes),
         }
     return metadata
 
@@ -442,7 +449,7 @@ def _copy_stream_data_events(
 ) -> Iterator[tuple[CopyStreamEvent, _CopyStreamState]]:
     limits = request.limits
     result_delivery = request.result_delivery
-    chunk_bytes = result_delivery.chunk_bytes if result_delivery is not None else limits.chunk_bytes
+    chunk_bytes = limits.chunk_bytes
     max_bytes = result_delivery.max_bytes if result_delivery is not None else limits.max_bytes
     max_row_bytes = limits.max_row_bytes
     timeout_ms = limits.timeout_ms
@@ -608,21 +615,29 @@ def _validation_location(location: tuple[Any, ...]) -> str:
 
 
 # ---------------------------------------------------------------------------
-# Direct chunked upload to its-agent-intake (POC_PUBLIC_CHUNKED_UPLOAD)
+# Direct multipart upload to its-agent-intake (POC_PUBLIC_MULTIPART_UPLOAD)
 #
-# When resultDelivery is present, the integration uploads bounded COPY chunks
+# When resultDelivery is present, the integration uploads bounded COPY parts
 # directly to its-agent-intake over HTTP. The Agent forwards the intake base
 # URL and scoped upload token in resultDelivery, and the integration reads the
 # org API key and POC application key from Agent config via datadog_agent.get_config.
-# Bulk chunk bytes never traverse the native emit bridge, AgentSecure, PAR, or
+# Bulk part bytes never traverse the native emit bridge, AgentSecure, PAR, or
 # AP action output; only the compact final receipt is emitted back.
 
-REMOTE_QUERY_UPLOAD_HTTP_TIMEOUT_SECONDS = 60
 REMOTE_QUERY_UPLOAD_TEST_DRIVE_HEADER = 'test-drive-its-agent-intake-poc'
 REMOTE_QUERY_UPLOAD_TEST_DRIVE_CONFIG_KEY = 'remote_queries.execute.intake_test_drive_selector'
 REMOTE_QUERY_UPLOAD_MAX_RETRIES = 4
 REMOTE_QUERY_UPLOAD_INITIAL_BACKOFF_SECONDS = 0.1
 REMOTE_QUERY_UPLOAD_MAX_BACKOFF_SECONDS = 5.0
+# Per-upload HTTP timeout as an explicit (connect, read) tuple: a short connect timeout and a
+# 5-minute read timeout so a slow part upload (e.g. a large part over a constrained link) is
+# not cut short, while a stuck connect fails fast. Retry count and backoff stay bounded above.
+REMOTE_QUERY_UPLOAD_HTTP_CONNECT_TIMEOUT_SECONDS = 10
+REMOTE_QUERY_UPLOAD_HTTP_READ_TIMEOUT_SECONDS = 300
+REMOTE_QUERY_UPLOAD_HTTP_TIMEOUT = (
+    REMOTE_QUERY_UPLOAD_HTTP_CONNECT_TIMEOUT_SECONDS,
+    REMOTE_QUERY_UPLOAD_HTTP_READ_TIMEOUT_SECONDS,
+)
 
 
 @dataclass(frozen=True)
@@ -636,7 +651,9 @@ class _UploadCredentials:
 
 
 class _UploadClient(Protocol):
-    def put_chunk(self, creds: _UploadCredentials, index: int, payload: bytes, sha256_hex: str, rows: int) -> None: ...
+    def put_part(
+        self, creds: _UploadCredentials, part_number: int, payload: bytes, sha256_hex: str, rows: int
+    ) -> None: ...
 
     def finalize(self, creds: _UploadCredentials) -> Mapping[str, Any]: ...
 
@@ -646,8 +663,8 @@ def abort(self, creds: _UploadCredentials) -> None: ...
 class _RequestsUploadClient:
     """HTTP upload client for its-agent-intake. Imports requests lazily."""
 
-    def __init__(self, timeout_seconds: int = REMOTE_QUERY_UPLOAD_HTTP_TIMEOUT_SECONDS) -> None:
-        self._timeout = timeout_seconds
+    def __init__(self, timeout: tuple[int, int] = REMOTE_QUERY_UPLOAD_HTTP_TIMEOUT) -> None:
+        self._timeout = timeout
 
     def _headers(self, creds: _UploadCredentials, content_type: str | None = None) -> dict[str, str]:
         headers = {
@@ -661,25 +678,25 @@ def _headers(self, creds: _UploadCredentials, content_type: str | None = None) -
             headers[REMOTE_QUERY_UPLOAD_TEST_DRIVE_HEADER] = creds.test_drive_selector
         return headers
 
-    def put_chunk(self, creds: _UploadCredentials, index: int, payload: bytes, sha256_hex: str, rows: int) -> None:
+    def put_part(self, creds: _UploadCredentials, part_number: int, payload: bytes, sha256_hex: str, rows: int) -> None:
         headers = self._headers(creds, 'application/octet-stream')
-        headers['X-DD-Chunk-SHA256'] = sha256_hex
-        headers['X-DD-Chunk-Bytes'] = str(len(payload))
-        headers['X-DD-Chunk-Rows'] = str(rows)
-        url = '{}/uploads/{}/chunks/{}'.format(creds.base_url.rstrip('/'), creds.upload_id, index)
-        _upload_with_retry('PUT', url, headers, payload)
+        headers['X-DD-Part-SHA256'] = sha256_hex
+        headers['X-DD-Part-Bytes'] = str(len(payload))
+        headers['X-DD-Part-Rows'] = str(rows)
+        url = '{}/uploads/{}/parts/{}'.format(creds.base_url.rstrip('/'), creds.upload_id, part_number)
+        _upload_with_retry('PUT', url, headers, payload, self._timeout)
 
     def finalize(self, creds: _UploadCredentials) -> Mapping[str, Any]:
         headers = self._headers(creds, 'application/json')
         url = '{}/uploads/{}/finalize'.format(creds.base_url.rstrip('/'), creds.upload_id)
-        _status, body = _upload_with_retry('POST', url, headers, b'{}')
+        _status, body = _upload_with_retry('POST', url, headers, b'{}', self._timeout)
         return json.loads(body.decode('utf-8'))
 
     def abort(self, creds: _UploadCredentials) -> None:
         headers = self._headers(creds, 'application/json')
         url = '{}/uploads/{}/abort'.format(creds.base_url.rstrip('/'), creds.upload_id)
         try:
-            _upload_with_retry('POST', url, headers, b'{}')
+            _upload_with_retry('POST', url, headers, b'{}', self._timeout)
         except _CopyStreamFailure:
             LOGGER.debug('Remote query upload abort failed (best-effort)', exc_info=True)
 
@@ -688,16 +705,20 @@ def _is_transient_upload_status(status: int) -> bool:
     return status == 408 or status == 429 or status >= 500
 
 
-def _upload_with_retry(method: str, url: str, headers: Mapping[str, str], body: bytes) -> tuple[int, bytes]:
+def _upload_with_retry(
+    method: str,
+    url: str,
+    headers: Mapping[str, str],
+    body: bytes,
+    timeout: tuple[int, int] = REMOTE_QUERY_UPLOAD_HTTP_TIMEOUT,
+) -> tuple[int, bytes]:
     import requests  # lazy: only the POC upload path needs it
 
     backoff = REMOTE_QUERY_UPLOAD_INITIAL_BACKOFF_SECONDS
     last_err: Any = None
     for attempt in range(REMOTE_QUERY_UPLOAD_MAX_RETRIES + 1):
         try:
-            resp = requests.request(
-                method, url, headers=dict(headers), data=body, timeout=REMOTE_QUERY_UPLOAD_HTTP_TIMEOUT_SECONDS
-            )
+            resp = requests.request(method, url, headers=dict(headers), data=body, timeout=timeout)
         except requests.exceptions.RequestException as e:
             last_err = e
         else:
@@ -721,7 +742,7 @@ def _upload_with_retry(method: str, url: str, headers: Mapping[str, str], body:
 
 def _is_upload_request(request: Mapping[str, Any]) -> bool:
     delivery = request.get('resultDelivery')
-    return isinstance(delivery, Mapping) and delivery.get('mode') == 'POC_PUBLIC_CHUNKED_UPLOAD'
+    return isinstance(delivery, Mapping) and delivery.get('mode') == 'POC_PUBLIC_MULTIPART_UPLOAD'
 
 
 def _get_agent_config(key: str) -> str:
@@ -756,15 +777,23 @@ def _count_newlines(payload: bytes) -> int:
     return payload.count(b'\n')
 
 
+def _multipart_part_count(total_bytes: int, part_bytes: int) -> int:
+    # Number of multipart parts for ``total_bytes`` aggregated at ``part_bytes``: full parts of
+    # exactly partBytes plus a final short part, or zero parts for an empty result.
+    if total_bytes <= 0:
+        return 0
+    return (total_bytes + part_bytes - 1) // part_bytes
+
+
 def _intake_receipt_to_camel(resp: Mapping[str, Any]) -> dict[str, Any]:
     return {
-        'mode': resp.get('mode', 'POC_PUBLIC_CHUNKED_UPLOAD'),
+        'mode': resp.get('mode', 'POC_PUBLIC_MULTIPART_UPLOAD'),
         'uploadId': resp.get('upload_id', ''),
         'bucketName': resp.get('bucket_name', ''),
-        'manifestPath': resp.get('manifest_key', ''),
+        'objectPath': resp.get('object_path', ''),
         'totalBytes': resp.get('total_bytes', 0),
         'totalRows': resp.get('total_rows', 0),
-        'chunkCount': resp.get('chunk_count', 0),
+        'partCount': resp.get('part_count', 0),
         'sha256': resp.get('sha256', ''),
     }
 
@@ -778,9 +807,52 @@ def _safe_abort(client: _UploadClient, creds: _UploadCredentials) -> None:
         LOGGER.debug('Remote query upload abort failed (best-effort)', exc_info=True)
 
 
-def _upload_one_chunk(client: _UploadClient, creds: _UploadCredentials, event: CopyStreamEvent) -> None:
-    metadata = event.metadata
-    client.put_chunk(creds, metadata['sequence'], event.payload, metadata['sha256'], _count_newlines(event.payload))
+@dataclass(frozen=True)
+class _MultipartPart:
+    part_number: int
+    payload: bytes
+
+
+class _MultipartBuffer:
+    """Aggregate COPY chunk bytes into server-clamped multipart parts.
+
+    The COPY stream emits ``limits.chunkBytes``-sized chunks; this buffer accumulates them into
+    ``partBytes``-sized parts (the multipart part size), which may exceed the COPY chunk size.
+    At most one part is buffered at a time. Part numbers are contiguous and 1-based to match
+    provider multipart conventions.
+    """
+
+    def __init__(self, part_bytes: int) -> None:
+        self._part_bytes = part_bytes
+        self._pending = bytearray()
+        self._next_part_number = 1
+
+    def extend(self, data: bytes) -> None:
+        self._pending.extend(data)
+
+    def _take(self) -> _MultipartPart:
+        payload = bytes(self._pending[: self._part_bytes])
+        del self._pending[: self._part_bytes]
+        part = _MultipartPart(self._next_part_number, payload)
+        self._next_part_number += 1
+        return part
+
+    def full_parts(self) -> Iterator[_MultipartPart]:
+        while len(self._pending) >= self._part_bytes:
+            yield self._take()
+
+    def flush_final(self) -> _MultipartPart | None:
+        if not self._pending:
+            return None
+        return self._take()
+
+
+def _upload_one_part(client: _UploadClient, creds: _UploadCredentials, part: _MultipartPart) -> None:
+    # The part SHA-256 is over the aggregated part body (all chunks combined), not the
+    # individual COPY chunks, so the intake can verify the part it receives.
+    client.put_part(
+        creds, part.part_number, part.payload, hashlib.sha256(part.payload).hexdigest(), _count_newlines(part.payload)
+    )
 
 
 def _finalize_upload(
@@ -801,7 +873,7 @@ def _execute_upload_stream(
     emit: CopyStreamEmit,
     http_client: _UploadClient | None = None,
 ) -> None:
-    """Upload COPY chunks directly to its-agent-intake; emit only metadata/final/error events."""
+    """Upload COPY parts directly to its-agent-intake; emit only metadata/final/error events."""
     creds = _resolve_upload_credentials(request)
     if not creds.api_key or not creds.app_key:
         _emit_copy_event(
@@ -814,11 +886,25 @@ def _execute_upload_stream(
         return
     client = http_client if http_client is not None else _default_upload_client()
     events = iter_agent_rpc_stream_copy_events(request, StaticPostgresCheckRegistry([check]))
+    buffer: _MultipartBuffer | None = None
     try:
         for event in events:
-            if event.event_type == 'data':
-                _upload_one_chunk(client, creds, event)
+            if event.event_type == 'metadata':
+                # The STARTED metadata carries the backend-validated partBytes; the COPY stream
+                # emits chunkBytes-sized events that this client aggregates into partBytes parts.
+                delivery_meta = event.metadata.get('resultDelivery') or {}
+                buffer = _MultipartBuffer(int(delivery_meta.get('partBytes') or 0))
+                _emit_copy_event(emit, event)
+            elif event.event_type == 'data':
+                if buffer is not None:
+                    buffer.extend(event.payload)
+                    for part in buffer.full_parts():
+                        _upload_one_part(client, creds, part)
             elif event.event_type == 'final':
+                if buffer is not None:
+                    final_part = buffer.flush_final()
+                    if final_part is not None:
+                        _upload_one_part(client, creds, final_part)
                 _finalize_upload(client, creds, event, emit)
             elif event.event_type == 'error':
                 _safe_abort(client, creds)