fix(events): preserve native JSON batch parameters

Author: cofin

sqlspec/adapters/psqlpy/events/backend.py

@@ -311,8 +311,8 @@ async def _publish_durable(
                     {
                         "event_id": event_id,
                         "channel": channel,
-                        "payload_json": to_json(payload),
-                        "metadata_json": to_json(metadata) if metadata is not None else None,
+                        "payload_json": payload,
+                        "metadata_json": metadata,
                         "status": "pending",
                         "available_at": now,
                         "lease_expires_at": None,

sqlspec/extensions/events/_queue.py

@@ -14,7 +14,7 @@
 from sqlspec.extensions.events._names import normalize_queue_table_name
 from sqlspec.extensions.events._payload import parse_event_timestamp
 from sqlspec.utils.logging import get_logger
-from sqlspec.utils.serializers import from_json, to_json
+from sqlspec.utils.serializers import from_json
 from sqlspec.utils.uuids import uuid4
 
 if TYPE_CHECKING:
@@ -190,8 +190,8 @@ def _batch_insert_parameters(
             records.append({
                 "event_id": event_id,
                 "channel": channel,
-                "payload_json": to_json(payload),
-                "metadata_json": to_json(metadata) if metadata is not None else None,
+                "payload_json": payload,
+                "metadata_json": metadata,
                 "status": _PENDING_STATUS,
                 "available_at": now,
                 "lease_expires_at": None,

tests/integration/adapters/asyncpg/extensions/events/test_listen_notify.py

@@ -66,12 +66,24 @@ async def _handler(message: Any) -> None:
         event_id = await channel.publish("alerts", {"action": "hybrid_async"})
         await _wait_for_message(received)
 
+    batch_ids = await channel.publish_many([
+        ("alerts", {"index": 1}, None),
+        ("alerts", {"index": 2}, {"source": "batch"}),
+    ])
+    for _ in range(_MAX_POLL_ATTEMPTS):
+        if len(received) >= 3:
+            break
+        await asyncio.sleep(_POLL_INTERVAL)
+
     await channel.stop_listener(listener.id)
 
-    assert received, "listener did not receive message"
+    assert len(received) >= 3, "listener did not receive the single event and complete batch"
     message = received[0]
     assert message.event_id == event_id
     assert message.payload["action"] == "hybrid_async"
+    assert [message.event_id for message in received[1:3]] == batch_ids
+    assert [message.payload["index"] for message in received[1:3]] == [1, 2]
+    assert received[2].metadata == {"source": "batch"}
 
     await channel.shutdown()
     if config.connection_instance:

tests/integration/adapters/psqlpy/extensions/events/test_listen_notify.py

@@ -100,12 +100,24 @@ async def _handler(message: Any) -> None:
             if received:
                 break
             await asyncio.sleep(0.05)
+
+        batch_ids = await channel.publish_many([
+            ("alerts", {"index": 1}, None),
+            ("alerts", {"index": 2}, {"source": "batch"}),
+        ])
+        for _ in range(200):
+            if len(received) >= 3:
+                break
+            await asyncio.sleep(0.05)
         await channel.stop_listener(listener.id)
 
-        assert received, "listener did not receive message"
+        assert len(received) >= 3, "listener did not receive the single event and complete batch"
         message = received[0]
         assert message.event_id == event_id
         assert message.payload["action"] == "hybrid"
+        assert [message.event_id for message in received[1:3]] == batch_ids
+        assert [message.payload["index"] for message in received[1:3]] == [1, 2]
+        assert received[2].metadata == {"source": "batch"}
     finally:
         backend = getattr(channel, "_backend", None)
         if backend and hasattr(backend, "shutdown"):

tests/integration/adapters/psycopg/extensions/events/test_listen_notify.py

@@ -172,10 +172,22 @@ async def _handler(message: Any) -> None:
         if received:
             break
         await asyncio.sleep(0.05)
+
+    batch_ids = await channel.publish_many([
+        ("alerts", {"index": 1}, None),
+        ("alerts", {"index": 2}, {"source": "batch"}),
+    ])
+    for _ in range(200):
+        if len(received) >= 3:
+            break
+        await asyncio.sleep(0.05)
     await channel.stop_listener(listener.id)
 
-    assert received, "listener did not receive message"
+    assert len(received) >= 3, "listener did not receive the single event and complete batch"
     message = received[0]
 
     assert message.event_id == event_id
     assert message.payload["action"] == "hybrid-async"
+    assert [message.event_id for message in received[1:3]] == batch_ids
+    assert [message.payload["index"] for message in received[1:3]] == [1, 2]
+    assert received[2].metadata == {"source": "batch"}

tests/unit/extensions/test_events/test_batch_publish.py

@@ -573,8 +573,9 @@ async def test_async_postgres_hybrid_single_and_batch_preserve_empty_metadata(
     batch_backend = backend_factory(batch_config)
     await batch_backend.publish_many([("events", {"value": 1}, {})])
 
-    assert single_config.driver.execute_calls[0].parameters["metadata_json"] == to_json({})
-    assert batch_config.driver.execute_many_calls[0][1][0]["metadata_json"] == to_json({})
+    expected_single_metadata = {} if config_type is _PsqlpyConfig else to_json({})
+    assert single_config.driver.execute_calls[0].parameters["metadata_json"] == expected_single_metadata
+    assert batch_config.driver.execute_many_calls[0][1][0]["metadata_json"] == {}
 
 
 def test_psycopg_sync_hybrid_publish_many_bulk_inserts_and_marks_each_channel_once() -> None:
@@ -630,7 +631,7 @@ def test_psycopg_sync_hybrid_single_and_batch_preserve_empty_metadata() -> None:
     batch_backend.publish_many([("events", {"value": 1}, {})])
 
     assert single_config.driver.execute_calls[0].parameters["metadata_json"] == to_json({})
-    assert batch_config.driver.execute_many_calls[0][1][0]["metadata_json"] == to_json({})
+    assert batch_config.driver.execute_many_calls[0][1][0]["metadata_json"] == {}
 
 
 def _notify_event_ids(parameters: "list[Any]") -> list[str]:

tests/unit/extensions/test_events/test_queue.py

@@ -103,13 +103,15 @@ def test_table_queue_empty_poll_backoff_state_is_bounded(tmp_path: Any) -> None:
 
 
 def test_table_queue_batch_records_have_deterministic_created_order() -> None:
-    _, records = SyncTableEventQueue._batch_insert_parameters([
-        ("events", {"index": index}, None) for index in range(3)
-    ])
+    payloads = [{"index": index} for index in range(3)]
+    metadata = {"source": "test"}
+    _, records = SyncTableEventQueue._batch_insert_parameters([("events", payload, metadata) for payload in payloads])
 
     created_at = [record["created_at"] for record in records]
     assert created_at == sorted(created_at)
     assert len(set(created_at)) == 3
+    assert [record["payload_json"] for record in records] == payloads
+    assert all(record["metadata_json"] is metadata for record in records)
 
 
 def test_table_event_queue_default_table_name(tmp_path) -> None: