Merge pull request #2551 from arc53/collisons-in-messages

fix: classify FK violations in message journal to skip pointless seq-retry

application/streaming/message_journal.py

@@ -59,6 +59,27 @@ def _strip_null_bytes(value: Any) -> Any:
     return value
 
 
+# Postgres SQLSTATE for a foreign-key violation. ``message_events`` has
+# an FK to ``conversation_messages(id)``, so a 23503 means that parent
+# row was never committed — the event can't be journaled and a
+# ``sequence_no`` retry is pointless churn. A 23505 (PK collision), by
+# contrast, is recoverable by rewriting at a fresh seq.
+_FK_VIOLATION_SQLSTATE = "23503"
+
+
+def _is_foreign_key_violation(exc: IntegrityError) -> bool:
+    """Return ``True`` when ``exc`` wraps a Postgres FK violation (23503).
+
+    ``IntegrityError.orig`` is the underlying driver error, which carries
+    the SQLSTATE on ``.sqlstate`` (psycopg3) or ``.pgcode`` (psycopg2).
+    When neither is present (an error we can't classify) this returns
+    ``False`` so the caller keeps its existing seq-retry behavior.
+    """
+    orig = getattr(exc, "orig", None)
+    sqlstate = getattr(orig, "sqlstate", None) or getattr(orig, "pgcode", None)
+    return sqlstate == _FK_VIOLATION_SQLSTATE
+
+
 def record_event(
     message_id: str,
     sequence_no: int,
@@ -108,59 +129,75 @@ def record_event(
                 message_id, sequence_no, event_type, materialised_payload
             )
         journal_committed = True
-    except IntegrityError:
+    except IntegrityError as exc:
-        # Composite-PK collision on (message_id, sequence_no). Most
+        if _is_foreign_key_violation(exc):
-        # likely cause is a stale ``latest_sequence_no`` seed on a
+            # No ``conversation_messages`` parent row — the event
-        # continuation retry — the route read MAX(seq) from a separate
+            # references a message that was never committed. A
-        # connection before another writer committed past it. Look up
+            # ``sequence_no`` retry can't fix that, so log the real
-        # the live latest and retry once with latest+1 so the event is
+            # cause and drop instead of churning the readonly probe +
-        # not silently lost. Bounded to a single retry — if two
+            # retry (which would just FK-fail again).
-        # writers keep racing in lockstep the route-level retry will
+            logger.warning(
-        # converge them across attempts.
+                "record_event: no conversation_messages row for "
-        try:
+                "message_id=%s (foreign-key violation); dropping "
-            with db_readonly() as conn:
+                "(seq=%s type=%s). The parent message row must be "
-                latest = MessageEventsRepository(conn).latest_sequence_no(
+                "committed before its events are journaled.",
-                    message_id
+                message_id,
-                )
+                sequence_no,
-            materialised_seq = (latest if latest is not None else -1) + 1
+                event_type,
-            with db_session() as conn:
+            )
-                MessageEventsRepository(conn).record(
+        else:
+            # Composite-PK collision on (message_id, sequence_no). Most
+            # likely cause is a stale ``latest_sequence_no`` seed on a
+            # continuation retry — the route read MAX(seq) from a
+            # separate connection before another writer committed past
+            # it. Look up the live latest and retry once with latest+1
+            # so the event is not silently lost. Bounded to a single
+            # retry — if two writers keep racing in lockstep the
+            # route-level retry will converge them across attempts.
+            try:
+                with db_readonly() as conn:
+                    latest = MessageEventsRepository(conn).latest_sequence_no(
+                        message_id
+                    )
+                materialised_seq = (latest if latest is not None else -1) + 1
+                with db_session() as conn:
+                    MessageEventsRepository(conn).record(
+                        message_id,
+                        materialised_seq,
+                        event_type,
+                        materialised_payload,
+                    )
+                journal_committed = True
+                logger.info(
+                    "record_event: collision at seq=%s recovered → wrote at "
+                    "seq=%s message_id=%s type=%s",
+                    sequence_no,
+                    materialised_seq,
                     message_id,
+                    event_type,
+                )
+            except IntegrityError:
+                # Second collision under the same retry — give up and
+                # log. The route's nonlocal counter will continue at
+                # ``sequence_no+1`` on the next emit; the next call may
+                # land cleanly past the contended window.
+                logger.warning(
+                    "record_event: IntegrityError persists after seq+1 "
+                    "retry; dropping. message_id=%s original_seq=%s "
+                    "retry_seq=%s type=%s",
+                    message_id,
+                    sequence_no,
                     materialised_seq,
                     event_type,
-                    materialised_payload,
                 )
-            journal_committed = True
+            except Exception:
-            logger.info(
+                logger.exception(
-                "record_event: collision at seq=%s recovered → wrote at "
+                    "record_event: retry path failed unexpectedly "
-                "seq=%s message_id=%s type=%s",
+                    "(message_id=%s seq=%s type=%s)",
-                sequence_no,
+                    message_id,
-                materialised_seq,
+                    sequence_no,
-                message_id,
+                    event_type,
-                event_type,
+                )
-            )
-        except IntegrityError:
-            # Second collision under the same retry — give up and log.
-            # The route's nonlocal counter will continue at
-            # ``sequence_no+1`` on the next emit; the next call may
-            # land cleanly past the contended window.
-            logger.warning(
-                "record_event: IntegrityError persists after seq+1 retry; "
-                "dropping. message_id=%s original_seq=%s retry_seq=%s "
-                "type=%s",
-                message_id,
-                sequence_no,
-                materialised_seq,
-                event_type,
-            )
-        except Exception:
-            logger.exception(
-                "record_event: retry path failed unexpectedly "
-                "(message_id=%s seq=%s type=%s)",
-                message_id,
-                sequence_no,
-                event_type,
-            )
     except Exception:
         logger.exception(
             "message_events INSERT failed: message_id=%s seq=%s type=%s",
@@ -297,7 +334,21 @@ class BatchedJournalWriter:
                 MessageEventsRepository(conn).bulk_record(
                     self._message_id, pending
                 )
-        except IntegrityError:
+        except IntegrityError as exc:
+            if _is_foreign_key_violation(exc):
+                # The whole batch references a conversation_messages
+                # parent that doesn't exist — every per-row retry would
+                # FK-fail too, so drop the batch once instead of N
+                # pointless re-attempts. The buffer is already cleared.
+                logger.warning(
+                    "BatchedJournalWriter: no conversation_messages row "
+                    "for message_id=%s (foreign-key violation); dropping "
+                    "%d event(s). The parent message row must be committed "
+                    "before its events are journaled.",
+                    self._message_id,
+                    len(pending),
+                )
+                return
             logger.info(
                 "BatchedJournalWriter: bulk INSERT collided for "
                 "message_id=%s n=%d; falling back to per-row writes",

tests/test_message_journal.py

@@ -18,6 +18,21 @@ from application.streaming.message_journal import (
 )
 
 
+def _integrity_error(sqlstate: str):
+    """Build an ``IntegrityError`` whose ``.orig`` reports ``sqlstate``.
+
+    Mirrors how psycopg3 surfaces a Postgres error code on the DBAPI
+    exception that SQLAlchemy wraps: ``23503`` = foreign-key violation
+    (the ``conversation_messages`` parent row is missing), ``23505`` =
+    unique/PK collision (a duplicate ``sequence_no``).
+    """
+    from sqlalchemy.exc import IntegrityError
+
+    orig = Exception("pg error")
+    orig.sqlstate = sqlstate
+    return IntegrityError("stmt", {}, orig)
+
+
 @pytest.mark.unit
 class TestRecordEvent:
     def test_invalid_args_return_false(self):
@@ -245,6 +260,85 @@ class TestRecordEvent:
             # serve it on a future reconnect anyway.
             mock_topic.publish.assert_called_once()
 
+    def test_foreign_key_violation_drops_without_seq_retry(self):
+        """A 23503 (foreign-key violation) means the
+        ``conversation_messages`` parent row was never committed — a
+        ``sequence_no`` retry can never land the row, so the write is
+        dropped immediately without probing ``latest_sequence_no`` or
+        re-attempting. The live publish still fires (best-effort tail).
+        """
+        repo_first = MagicMock(name="repo_first")
+        repo_first.record.side_effect = _integrity_error("23503")
+        repo_readonly = MagicMock(name="repo_readonly")
+        repo_retry = MagicMock(name="repo_retry")
+
+        repo_instances = iter([repo_first, repo_readonly, repo_retry])
+
+        with patch(
+            "application.streaming.message_journal.db_session"
+        ) as mock_session, patch(
+            "application.streaming.message_journal.db_readonly"
+        ) as mock_readonly, patch(
+            "application.streaming.message_journal.MessageEventsRepository",
+            side_effect=lambda conn: next(repo_instances),
+        ), patch(
+            "application.streaming.message_journal.Topic"
+        ) as mock_topic_cls:
+            mock_session.return_value.__enter__.return_value = MagicMock()
+            mock_readonly.return_value.__enter__.return_value = MagicMock()
+            mock_topic = MagicMock()
+            mock_topic_cls.return_value = mock_topic
+
+            result = record_event("msg-1", 3, "answer", {"text": "hi"})
+
+            assert result is False
+            repo_first.record.assert_called_once_with(
+                "msg-1", 3, "answer", {"text": "hi"}
+            )
+            # No seq+1 retry: the readonly probe and retry insert never run.
+            mock_readonly.assert_not_called()
+            repo_readonly.latest_sequence_no.assert_not_called()
+            repo_retry.record.assert_not_called()
+            # Live tail still gets the frame at the original seq.
+            mock_topic.publish.assert_called_once()
+            wire = json.loads(mock_topic.publish.call_args[0][0])
+            assert wire["sequence_no"] == 3
+
+    def test_pk_collision_still_retries_with_seq_plus_one(self):
+        """An explicit 23505 (unique/PK collision) keeps the existing
+        recover-by-retry behavior: probe ``latest_sequence_no`` and
+        rewrite at latest+1. Regression guard pinning the FK-vs-PK split.
+        """
+        repo_first = MagicMock(name="repo_first")
+        repo_first.record.side_effect = _integrity_error("23505")
+        repo_readonly = MagicMock(name="repo_readonly")
+        repo_readonly.latest_sequence_no.return_value = 7
+        repo_retry = MagicMock(name="repo_retry")
+
+        repo_instances = iter([repo_first, repo_readonly, repo_retry])
+
+        with patch(
+            "application.streaming.message_journal.db_session"
+        ) as mock_session, patch(
+            "application.streaming.message_journal.db_readonly"
+        ) as mock_readonly, patch(
+            "application.streaming.message_journal.MessageEventsRepository",
+            side_effect=lambda conn: next(repo_instances),
+        ), patch(
+            "application.streaming.message_journal.Topic"
+        ) as mock_topic_cls:
+            mock_session.return_value.__enter__.return_value = MagicMock()
+            mock_readonly.return_value.__enter__.return_value = MagicMock()
+            mock_topic_cls.return_value = MagicMock()
+
+            result = record_event("msg-1", 3, "answer", {"text": "hi"})
+
+            assert result is True
+            repo_readonly.latest_sequence_no.assert_called_once_with("msg-1")
+            repo_retry.record.assert_called_once_with(
+                "msg-1", 8, "answer", {"text": "hi"}
+            )
+
     def test_payload_with_null_byte_stripped_before_insert(self):
         """``\\x00`` in a string value would crash JSONB INSERT — strip
         at the journal boundary so the row lands and reconnecting
@@ -522,6 +616,62 @@ class TestBatchedJournalWriter:
             assert per_row_repo.record.call_args_list[0].args[1] == 0
             assert per_row_repo.record.call_args_list[1].args[1] == 1
 
+    def test_bulk_foreign_key_violation_drops_batch_without_per_row(self):
+        """A 23503 (foreign-key violation) on the bulk INSERT means the
+        whole batch references a ``conversation_messages`` parent that
+        doesn't exist — every per-row retry would FK-fail too. The
+        batch is dropped once without the per-row fallback (no latest
+        probe, no re-attempts, no publish), and the buffer is cleared.
+        """
+        session_cm, readonly_cm, repo_cls_p, topic_cls_p = self._patch_io()
+        with session_cm as mock_session, readonly_cm as mock_readonly, repo_cls_p as mock_repo_cls, topic_cls_p as mock_topic_cls:
+            mock_session.return_value.__enter__.return_value = MagicMock()
+            mock_readonly.return_value.__enter__.return_value = MagicMock()
+
+            bulk_repo = MagicMock(name="bulk_repo")
+            bulk_repo.bulk_record.side_effect = _integrity_error("23503")
+            per_row_repo = MagicMock(name="per_row_repo")
+            mock_repo_cls.side_effect = [bulk_repo, per_row_repo]
+
+            mock_topic = MagicMock()
+            mock_topic_cls.return_value = mock_topic
+
+            writer = BatchedJournalWriter("msg-1", batch_size=2)
+            writer.record(0, "answer", {"text": "a"})
+            writer.record(1, "answer", {"text": "b"})
+
+            bulk_repo.bulk_record.assert_called_once()
+            # No per-row fallback: the parent is missing for the whole batch.
+            mock_readonly.assert_not_called()
+            per_row_repo.record.assert_not_called()
+            mock_topic.publish.assert_not_called()
+            # Buffer cleared so close() is a no-op and memory stays bounded.
+            assert writer._buffer == []
+
+    def test_bulk_collision_23505_still_falls_back_to_per_row(self):
+        """An explicit 23505 (unique/PK collision) keeps the per-row
+        fallback — the parent exists, only a ``sequence_no`` collided.
+        Regression guard for the FK-vs-PK split at the bulk boundary.
+        """
+        session_cm, readonly_cm, repo_cls_p, topic_cls_p = self._patch_io()
+        with session_cm as mock_session, readonly_cm as mock_readonly, repo_cls_p as mock_repo_cls, topic_cls_p as mock_topic_cls:
+            mock_session.return_value.__enter__.return_value = MagicMock()
+            mock_readonly.return_value.__enter__.return_value = MagicMock()
+
+            bulk_repo = MagicMock(name="bulk_repo")
+            bulk_repo.bulk_record.side_effect = _integrity_error("23505")
+            per_row_repo = MagicMock(name="per_row_repo")
+            mock_repo_cls.side_effect = [bulk_repo, per_row_repo, per_row_repo]
+
+            mock_topic_cls.return_value = MagicMock()
+
+            writer = BatchedJournalWriter("msg-1", batch_size=2)
+            writer.record(0, "answer", {"text": "a"})
+            writer.record(1, "answer", {"text": "b"})
+
+            bulk_repo.bulk_record.assert_called_once()
+            assert per_row_repo.record.call_count == 2
+
     def test_flush_clears_buffer_even_on_total_failure(self):
         """A flush that fails with a non-IntegrityError exception must
         still clear the buffer — leaving rows would grow memory