Make Postgres remote query COPY stream binary-safe

Author: nubtron

postgres/datadog_checks/postgres/remote_query.py

@@ -27,6 +27,7 @@
 RemoteQueryCopySql = Literal[
     'SELECT 1 AS value',
     '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",
@@ -36,9 +37,8 @@
     "SELECT i, repeat('x', 1000) AS payload FROM generate_series(1, 3000) AS i",
 ]
 
-CopyStreamFormat = Literal['csv']
-CopyStreamEvent = Mapping[str, Any]
-CopyStreamEmit = Callable[[CopyStreamEvent], None]
+CopyStreamFormat = Literal['csv', 'binary']
+CopyStreamEmit = Callable[[str, str, bytes], None]
 
 if TYPE_CHECKING:
     from datadog_checks.postgres import PostgreSql
@@ -131,6 +131,13 @@ class _CopyStreamState:
     bytes_emitted: int = 0
 
 
+@dataclass(frozen=True)
+class CopyStreamEvent:
+    event_type: str
+    metadata: Mapping[str, Any]
+    payload: bytes = b''
+
+
 class _CopyStreamFailure(Exception):
     def __init__(self, code: str, message: str, retryable: bool = False):
         self.code = code
@@ -164,29 +171,33 @@ def execute_agent_rpc_stream_copy(
     try:
         request = json.loads(request_json)
     except (TypeError, ValueError):
-        emit(
+        _emit_copy_event(
+            emit,
             _stream_failed_event(
                 'invalid_request', 'Invalid remote query request: request_json must be a valid JSON object.'
-            )
+            ),
         )
         return
 
     if not isinstance(request, Mapping):
-        emit(
-            _stream_failed_event('invalid_request', 'Invalid remote query request: request_json must be a JSON object.')
+        _emit_copy_event(
+            emit,
+            _stream_failed_event(
+                'invalid_request', 'Invalid remote query request: request_json must be a JSON object.'
+            ),
         )
         return
 
     events = iter_agent_rpc_stream_copy_events(request, StaticPostgresCheckRegistry([check]))
     try:
         for event in events:
-            emit(event)
+            _emit_copy_event(emit, event)
     except BaseException:
         events.close()
         raise
 
 
-def iter_agent_rpc_stream_copy_events(request: Any, registry: PostgresCheckRegistry) -> Iterator[dict[str, Any]]:
+def iter_agent_rpc_stream_copy_events(request: Any, registry: PostgresCheckRegistry) -> Iterator[CopyStreamEvent]:
     """Yield COPY byte-stream events for unit tests and callback adaptation."""
     started_at = time.monotonic()
     try:
@@ -298,7 +309,7 @@ def _execute_safe_query(
 
 def _iter_copy_stream_events(
     check: 'PostgreSql', request: RemoteQueryCopyRequest, started_at: float
-) -> Iterator[dict[str, Any]]:
+) -> Iterator[CopyStreamEvent]:
     db_pool = getattr(check, 'db_pool', None)
     if db_pool is None:
         yield _stream_failed_event(
@@ -316,15 +327,17 @@ def _iter_copy_stream_events(
         )
         return
 
-    yield {
-        'type': 'metadata',
-        'status': 'STARTED',
-        'format': request.format,
-        'operation': request.operation,
-        'chunkBytes': request.limits.chunk_bytes,
-        'maxBytes': request.limits.max_bytes,
-        'maxRowBytes': request.limits.max_row_bytes,
-    }
+    yield CopyStreamEvent(
+        'metadata',
+        {
+            'status': 'STARTED',
+            'format': request.format,
+            'operation': request.operation,
+            'chunkBytes': request.limits.chunk_bytes,
+            'maxBytes': request.limits.max_bytes,
+            'maxRowBytes': request.limits.max_row_bytes,
+        },
+    )
 
     state = _CopyStreamState()
     error: _CopyStreamFailure | None = None
@@ -351,16 +364,15 @@ def _iter_copy_stream_events(
         )
         return
 
-    yield {
-        'type': 'final',
-        'status': 'SUCCEEDED',
-        'stats': _copy_stream_stats(state, started_at, request.format),
-    }
+    yield CopyStreamEvent(
+        'final',
+        {'status': 'SUCCEEDED', 'stats': _copy_stream_stats(state, started_at, request.format)},
+    )
 
 
 def _copy_stream_data_events(
     check: 'PostgreSql', request: RemoteQueryCopyRequest, state: _CopyStreamState, started_at: float
-) -> Iterator[tuple[dict[str, Any], _CopyStreamState]]:
+) -> Iterator[tuple[CopyStreamEvent, _CopyStreamState]]:
     limits = request.limits
     deadline = started_at + (limits.timeout_ms / 1000)
     copy_sql = _copy_stdout_sql(request.query, request.format)
@@ -413,10 +425,11 @@ def _copy_stream_data_events(
 
 
 def _copy_stdout_sql(query: str, stream_format: CopyStreamFormat) -> str:
-    # CSV is the only initial COPY byte-stream format. It is compact, standard, and preserves raw COPY bytes.
-    if stream_format != 'csv':
-        raise _CopyStreamFailure('invalid_request', 'Unsupported COPY stream format.')
-    return f'COPY ({query}) TO STDOUT WITH (FORMAT CSV)'
+    if stream_format == 'csv':
+        return f'COPY ({query}) TO STDOUT WITH (FORMAT CSV)'
+    if stream_format == 'binary':
+        return f'COPY ({query}) TO STDOUT WITH (FORMAT BINARY)'
+    raise _CopyStreamFailure('invalid_request', 'Unsupported COPY stream format.')
 
 
 def _set_statement_timeout(cursor: Any, timeout_ms: int) -> str | None:
@@ -445,9 +458,17 @@ def _raise_if_timed_out(deadline: float) -> None:
         raise _CopyStreamFailure('timeout', 'COPY stream exceeded timeoutMs.', retryable=True)
 
 
-def _copy_data_event(data: bytearray, state: _CopyStreamState) -> tuple[dict[str, Any], _CopyStreamState]:
+def _copy_data_event(data: bytearray, state: _CopyStreamState) -> tuple[CopyStreamEvent, _CopyStreamState]:
     payload = bytes(data)
-    event = {'type': 'data', 'sequence': state.sequence, 'data': payload, 'bytes': len(payload)}
+    event = CopyStreamEvent(
+        'data',
+        {
+            'sequence': state.sequence,
+            'offset': state.bytes_emitted,
+            'bytes': len(payload),
+        },
+        payload,
+    )
     next_state = _CopyStreamState(
         sequence=state.sequence + 1,
         chunks_emitted=state.chunks_emitted + 1,
@@ -475,17 +496,20 @@ def _stream_failed_event(
     retryable: bool = False,
     stats: Mapping[str, Any] | None = None,
     elapsed_ms: int | None = None,
-) -> dict[str, Any]:
-    event = {
-        'type': 'final',
+) -> CopyStreamEvent:
+    metadata = {
         'status': 'FAILED',
         'error': {'code': code, 'message': message, 'retryable': retryable},
     }
     if stats is not None:
-        event['stats'] = dict(stats)
+        metadata['stats'] = dict(stats)
     elif elapsed_ms is not None:
-        event['stats'] = {'elapsedMs': elapsed_ms}
-    return event
+        metadata['stats'] = {'elapsedMs': elapsed_ms}
+    return CopyStreamEvent('error', metadata)
+
+
+def _emit_copy_event(emit: CopyStreamEmit, event: CopyStreamEvent) -> None:
+    emit(event.event_type, json.dumps(event.metadata, default=str), event.payload)
 
 
 def _response_columns(description: Sequence[Any], rows: Sequence[Sequence[Any]]) -> list[dict[str, str]]:

postgres/tests/test_remote_query.py

@@ -433,16 +433,24 @@ def collect_copy_events(request, check):
     return list(iter_agent_rpc_stream_copy_events(request, StaticPostgresCheckRegistry([check])))
 
 
+def event_metadata(event):
+    return event.metadata
+
+
+def event_payload(event):
+    return event.payload
+
+
 def test_copy_stream_requires_explicit_operation_before_pool_access():
     pool = FakePool(copy_blocks=[b'1\n'])
     request = valid_copy_request()
     request.pop('operation')
 
     events = collect_copy_events(request, make_check(pool=pool))
 
-    assert events[-1]['status'] == 'FAILED'
-    assert events[-1]['error']['code'] == 'invalid_request'
-    assert 'operation' in events[-1]['error']['message']
+    assert event_metadata(events[-1])['status'] == 'FAILED'
+    assert event_metadata(events[-1])['error']['code'] == 'invalid_request'
+    assert 'operation' in event_metadata(events[-1])['error']['message']
     assert pool.requested_dbnames == []
 
 
@@ -452,9 +460,9 @@ def test_copy_stream_rejects_unknown_fields_without_echoing_secrets(caplog):
 
     events = collect_copy_events(request, make_check(pool=pool))
 
-    assert events[-1]['status'] == 'FAILED'
-    assert events[-1]['error']['code'] == 'invalid_request'
-    assert 'password' in events[-1]['error']['message']
+    assert event_metadata(events[-1])['status'] == 'FAILED'
+    assert event_metadata(events[-1])['error']['code'] == 'invalid_request'
+    assert 'password' in event_metadata(events[-1])['error']['message']
     assert 'SECRET_DO_NOT_LOG' not in str(events)
     assert 'SECRET_DO_NOT_LOG' not in caplog.text
     assert pool.requested_dbnames == []
@@ -466,9 +474,9 @@ def test_copy_stream_rejects_non_copy_allowlisted_queries_before_pool_access():
 
     events = collect_copy_events(request, make_check(pool=pool))
 
-    assert events[-1]['status'] == 'FAILED'
-    assert events[-1]['error']['code'] == 'invalid_request'
-    assert 'query' in events[-1]['error']['message']
+    assert event_metadata(events[-1])['status'] == 'FAILED'
+    assert event_metadata(events[-1])['error']['code'] == 'invalid_request'
+    assert 'query' in event_metadata(events[-1])['error']['message']
     assert pool.requested_dbnames == []
 
 
@@ -478,33 +486,46 @@ def test_copy_stream_uses_connection_pool_and_emits_chunked_copy_bytes():
 
     events = collect_copy_events(valid_copy_request(), check)
 
-    assert events[0]['type'] == 'metadata'
-    assert events[0]['operation'] == 'copy_stream'
-    assert events[0]['format'] == 'csv'
-    data_events = [event for event in events if event['type'] == 'data']
-    assert [event['sequence'] for event in data_events] == [0, 1, 2]
-    assert [event['data'] for event in data_events] == [b'abcdefgh', b'ijklmnop', b'qr']
-    assert [event['bytes'] for event in data_events] == [8, 8, 2]
-    assert events[-1]['type'] == 'final'
-    assert events[-1]['status'] == 'SUCCEEDED'
-    assert events[-1]['stats']['bytesEmitted'] == 18
-    assert events[-1]['stats']['chunksEmitted'] == 3
+    assert events[0].event_type == 'metadata'
+    assert event_metadata(events[0])['operation'] == 'copy_stream'
+    assert event_metadata(events[0])['format'] == 'csv'
+    data_events = [event for event in events if event.event_type == 'data']
+    assert [event_metadata(event)['sequence'] for event in data_events] == [0, 1, 2]
+    assert [event_metadata(event)['offset'] for event in data_events] == [0, 8, 16]
+    assert [event_payload(event) for event in data_events] == [b'abcdefgh', b'ijklmnop', b'qr']
+    assert [event_metadata(event)['bytes'] for event in data_events] == [8, 8, 2]
+    assert events[-1].event_type == 'final'
+    assert event_metadata(events[-1])['status'] == 'SUCCEEDED'
+    assert event_metadata(events[-1])['stats']['bytesEmitted'] == 18
+    assert event_metadata(events[-1])['stats']['chunksEmitted'] == 3
     assert pool.requested_dbnames == ['datadog_test']
     assert pool.closed_copies == 1
 
 
+def test_copy_stream_data_event_payload_preserves_arbitrary_bytes():
+    arbitrary_bytes = b'\x00\xff\x80abc\n'
+    pool = FakePool(copy_blocks=[arbitrary_bytes])
+
+    events = collect_copy_events(valid_copy_request(), make_check(pool=pool))
+
+    data_events = [event for event in events if event.event_type == 'data']
+    assert len(data_events) == 1
+    assert event_payload(data_events[0]) == arbitrary_bytes
+    assert isinstance(event_payload(data_events[0]), bytes)
+
+
 def test_copy_stream_enforces_max_bytes_without_exceeding_limit():
     pool = FakePool(copy_blocks=[b'abcdefgh', b'ijklmnop'])
     request = valid_copy_request(limits={'chunkBytes': 8, 'maxBytes': 10, 'maxRowBytes': 32, 'timeoutMs': 5000})
 
     events = collect_copy_events(request, make_check(pool=pool))
 
-    data_events = [event for event in events if event['type'] == 'data']
-    assert [event['data'] for event in data_events] == [b'abcdefgh', b'ij']
-    assert sum(event['bytes'] for event in data_events) == 10
-    assert events[-1]['status'] == 'FAILED'
-    assert events[-1]['error']['code'] == 'max_bytes_exceeded'
-    assert events[-1]['stats']['bytesEmitted'] == 10
+    data_events = [event for event in events if event.event_type == 'data']
+    assert [event_payload(event) for event in data_events] == [b'abcdefgh', b'ij']
+    assert sum(event_metadata(event)['bytes'] for event in data_events) == 10
+    assert event_metadata(events[-1])['status'] == 'FAILED'
+    assert event_metadata(events[-1])['error']['code'] == 'max_bytes_exceeded'
+    assert event_metadata(events[-1])['stats']['bytesEmitted'] == 10
     assert pool.closed_copies == 1
 
 
@@ -513,32 +534,41 @@ def test_copy_stream_enforces_max_row_bytes_after_copy_block_arrives():
 
     events = collect_copy_events(valid_copy_request(), make_check(pool=pool))
 
-    assert [event['data'] for event in events if event['type'] == 'data'] == []
-    assert events[-1]['status'] == 'FAILED'
-    assert events[-1]['error']['code'] == 'max_row_bytes_exceeded'
-    assert 'row granularity' in events[-1]['error']['message']
+    assert [event_payload(event) for event in events if event.event_type == 'data'] == []
+    assert event_metadata(events[-1])['status'] == 'FAILED'
+    assert event_metadata(events[-1])['error']['code'] == 'max_row_bytes_exceeded'
+    assert 'row granularity' in event_metadata(events[-1])['error']['message']
     assert pool.closed_copies == 1
 
 
-def test_agent_rpc_stream_copy_adapts_iterator_to_callback():
-    pool = FakePool(copy_blocks=[b'1\n'])
+def test_agent_rpc_stream_copy_adapts_iterator_to_binary_safe_callback():
+    arbitrary_bytes = b'\x00\xff\x80abc\n'
+    pool = FakePool(copy_blocks=[arbitrary_bytes])
     events = []
 
-    execute_agent_rpc_stream_copy(json.dumps(valid_copy_request()), make_check(pool=pool), events.append)
+    execute_agent_rpc_stream_copy(
+        json.dumps(valid_copy_request()), make_check(pool=pool), lambda *event: events.append(event)
+    )
 
-    assert [event['type'] for event in events] == ['metadata', 'data', 'final']
-    assert events[1]['data'] == b'1\n'
-    assert events[-1]['status'] == 'SUCCEEDED'
+    assert [event[0] for event in events] == ['metadata', 'data', 'final']
+    assert json.loads(events[1][1])['bytes'] == len(arbitrary_bytes)
+    assert events[1][2] == arbitrary_bytes
+    assert isinstance(events[1][2], bytes)
+    assert json.loads(events[-1][1])['status'] == 'SUCCEEDED'
 
 
 def test_agent_rpc_stream_copy_rejects_malformed_json_without_echoing_input(caplog):
     pool = FakePool(copy_blocks=[b'1\n'])
     events = []
 
-    execute_agent_rpc_stream_copy('{"password": "SECRET_DO_NOT_LOG"', make_check(pool=pool), events.append)
+    execute_agent_rpc_stream_copy(
+        '{"password": "SECRET_DO_NOT_LOG"', make_check(pool=pool), lambda *event: events.append(event)
+    )
 
-    assert events[-1]['status'] == 'FAILED'
-    assert events[-1]['error']['code'] == 'invalid_request'
+    metadata = json.loads(events[-1][1])
+    assert events[-1][0] == 'error'
+    assert metadata['status'] == 'FAILED'
+    assert metadata['error']['code'] == 'invalid_request'
     assert 'SECRET_DO_NOT_LOG' not in str(events)
     assert 'SECRET_DO_NOT_LOG' not in caplog.text
     assert pool.requested_dbnames == []
@@ -548,13 +578,13 @@ def test_agent_rpc_stream_copy_closes_copy_when_callback_raises():
     pool = FakePool(copy_blocks=[b'12345678', b'abcdef'])
     events = []
 
-    def emit(event):
-        events.append(event)
-        if event['type'] == 'data':
+    def emit(event_type, metadata_json, payload):
+        events.append((event_type, metadata_json, payload))
+        if event_type == 'data':
             raise RuntimeError('stop streaming')
 
     with pytest.raises(RuntimeError, match='stop streaming'):
         execute_agent_rpc_stream_copy(json.dumps(valid_copy_request()), make_check(pool=pool), emit)
 
-    assert [event['type'] for event in events] == ['metadata', 'data']
+    assert [event[0] for event in events] == ['metadata', 'data']
     assert pool.closed_copies == 1