mikebridge commented on code in PR #44009:
URL: https://github.com/apache/superset/pull/44009#discussion_r4020984789


##########
superset/versioning/changes/listener.py:
##########
@@ -395,7 +396,101 @@ def _persist_buffered_records(
         incr_capture_error("bulk_insert")
 
 
-def register_change_record_listener() -> None:  # noqa: C901
+def finalize_change_records(session: Session) -> None:
+    """Build and persist the transaction's change records at commit time.
+
+    Module-level (rather than a closure inside the registration function)
+    so the capture write path can be exercised directly by unit tests
+    against an isolated session; it depends only on the session and the
+    module helpers, never on the registered entity classes.
+    """
+    if session.in_nested_transaction() or session.info.get(_FINALIZING_KEY):
+        return
+
+    session.info[_FINALIZING_KEY] = True
+    # Measures the FINALIZE stage only: the timer starts after the flush,
+    # which excludes the transaction's own write cost but also excludes
+    # capture_initial_states' per-entity pre-state SELECTs (those are timed
+    # as their own ``capture_initial_states`` stage in before_flush) — and
+    # runs through every capture step and early return. Every commit on the
+    # session emits a sample, including commits touching no versioned
+    # entity, because the whole-listener overhead is exactly what the
+    # kill-switch removes; a flush that raises emits nothing.
+    start: float | None = None
+    try:
+        session.flush()
+        start = perf_counter()
+        initial_states: dict[tuple[str, int], tuple[Any, dict[str, Any]]] = (
+            session.info.get(_INITIAL_STATES_KEY, {})
+        )
+        buffer = _build_scalar_buffer(initial_states)
+
+        try:
+            tx_id = _current_transaction_id(session)
+        except Exception:  # pylint: disable=broad-except
+            logger.exception("version_changes: transaction lookup failed")
+            incr_capture_error("transaction_lookup")
+            return
+        if tx_id is None:
+            return
+
+        _stamp_action_kind_on_transaction(session, tx_id)
+        _append_child_records_to_buffer(session, tx_id, buffer)
+        _inject_action_meta_record(session, buffer)
+
+        if buffer:
+            _persist_buffered_records(session, tx_id, buffer)
+    finally:
+        session.info.pop(_FINALIZING_KEY, None)
+        if start is not None:
+            emit_capture_timing("finalize", (perf_counter() - start) * 1000.0)
+
+
+def _capture_initial_states(
+    session: Session, versioned_classes: tuple[type, ...]
+) -> None:
+    """The ``before_flush`` capture stage: retain each dirty versioned entity's
+    pre-flush database state for the final diff.
+
+    Timed as its own metric stage. The per-entity pre-state SELECTs issued
+    here are the capture cost that scales with the number of dirty versioned
+    entities — on a bulk edit plausibly the dominant cost the kill-switch
+    removes — and they run before the flush, outside ``finalize``'s timer. A
+    sample is emitted only when at least one versioned entity was captured,
+    so the many unrelated autoflushes do not flood the series with empty
+    samples; together with ``finalize`` the two stages cover the whole
+    listener. Module-level (not the registered closure) so it is
+    unit-testable without ``db.session``.
+    """
+    initial_states: dict[tuple[str, int], tuple[Any, dict[str, Any]]] = (
+        session.info.setdefault(_INITIAL_STATES_KEY, {})
+    )
+    start = perf_counter()
+    captured = 0
+    try:
+        for obj in list(session.dirty):
+            if isinstance(obj, versioned_classes):
+                # Count captures, not candidates: an entity already retained
+                # from an earlier flush returns early without a SELECT, and a
+                # sample for it would be near-zero noise diluting the upper
+                # percentiles the series is alerted on.
+                before = len(initial_states)
+                _capture_dirty_entity_initial_state(session, obj, 
initial_states)
+                captured += len(initial_states) - before
+    except Exception:  # pylint: disable=broad-except
+        # Twin of the transaction-lookup guard in finalize: a versioning bug
+        # must never break a user's save, so a raise in the per-entity capture
+        # is logged and counted rather than propagated out of before_flush.
+        logger.exception("version_changes: initial-state capture failed")
+        incr_capture_error("capture_initial_states")

Review Comment:
   The head now catches each capture-call failure separately, and the 
three-entity regression test retains entities one and three when two fails. 
Could we regard that case as addressed and handle the earlier identity lookup 
in the newer thread?



##########
superset/versioning/metrics.py:
##########
@@ -52,3 +52,37 @@ def incr_capture_error(stage: str) -> None:
         
stats_logger_manager.instance.incr(f"{_CAPTURE_METRIC_PREFIX}.{stage}.error")
     except Exception:  # pylint: disable=broad-except
         logger.exception("versioning: failed to emit capture-error metric")
+
+
+def emit_capture_timing(stage: str, duration_ms: float) -> None:
+    """Emit the write-path latency for one capture *stage*, in milliseconds.
+
+    The documented recovery lever for capture trouble is the
+    ``ENABLE_VERSIONING_CAPTURE`` kill-switch, flipped on save-path
+    slowdown — the ``superset.versioning.capture.<stage>.latency`` series
+    are the signal an operator alerts on before flipping it. There are two
+    stages, and TOGETHER they cover the cost the kill-switch removes:
+    ``capture_initial_states`` (the before-flush per-entity pre-state reads,
+    which scale with the number of dirty versioned entities — on a bulk
+    edit the dominant cost — sampled only when at least one was captured)

Review Comment:
   Good catch—the metric already samples attempted reads, including batches 
that retain no states. Could we take your wording so the metrics documentation 
matches the listener and its tests?



##########
superset/versioning/changes/listener.py:
##########
@@ -173,22 +174,34 @@ def build_action_headline(
 _REGISTERED_SENTINEL = "_versioning_change_listener_registered"
 
 
+def _uncaptured_identity(
+    obj: Any,
+    initial_states: dict[tuple[str, int], tuple[Any, dict[str, Any]]],
+) -> tuple[str, int] | None:
+    """Return the entity identity only when its pre-state needs a read."""
+    entity_kind: str | None = ENTITY_KIND_BY_CLASS_NAME.get(type(obj).__name__)
+    if entity_kind is None:
+        return None
+    entity_id: int | None = getattr(obj, "id", None)
+    if entity_id is None:
+        return None
+    key: tuple[str, int] = (entity_kind, entity_id)
+    if key in initial_states:
+        return None
+    return key
+
+
 def _capture_dirty_entity_initial_state(

Review Comment:
   The wrapper is indeed test-only, although eligibility and deduplication 
already share `_uncaptured_identity`. Could we remove the unused wrapper and 
move the first-state-retention test onto `_capture_initial_states` instead of 
dropping that coverage?



##########
superset/versioning/changes/listener.py:
##########
@@ -395,7 +408,115 @@ def _persist_buffered_records(
         incr_capture_error("bulk_insert")
 
 
-def register_change_record_listener() -> None:  # noqa: C901
+def finalize_change_records(session: Session) -> None:
+    """Build and persist the transaction's change records at commit time.
+
+    Module-level (rather than a closure inside the registration function)
+    so the capture write path can be exercised directly by unit tests
+    against an isolated session; it depends only on the session and the
+    module helpers, never on the registered entity classes.
+    """
+    if session.in_nested_transaction() or session.info.get(_FINALIZING_KEY):
+        return
+
+    session.info[_FINALIZING_KEY] = True
+    # Measures the FINALIZE stage only: the timer starts after the flush,
+    # which excludes the transaction's own write cost but also excludes
+    # capture_initial_states' per-entity pre-state SELECTs (those are timed
+    # as their own ``capture_initial_states`` stage in before_flush) — and
+    # runs through every capture step and early return. Every commit on the
+    # session emits a sample, including commits touching no versioned
+    # entity, because the whole-listener overhead is exactly what the
+    # kill-switch removes; a flush that raises emits nothing.
+    start: float | None = None
+    try:
+        session.flush()
+        start = perf_counter()
+        initial_states: dict[tuple[str, int], tuple[Any, dict[str, Any]]] = (
+            session.info.get(_INITIAL_STATES_KEY, {})
+        )
+        buffer = _build_scalar_buffer(initial_states)
+
+        try:
+            tx_id = _current_transaction_id(session)
+        except Exception:  # pylint: disable=broad-except
+            logger.exception("version_changes: transaction lookup failed")
+            incr_capture_error("transaction_lookup")
+            return
+        if tx_id is None:
+            return
+
+        _stamp_action_kind_on_transaction(session, tx_id)
+        _append_child_records_to_buffer(session, tx_id, buffer)
+        _inject_action_meta_record(session, buffer)
+
+        if buffer:
+            _persist_buffered_records(session, tx_id, buffer)
+    finally:
+        session.info.pop(_FINALIZING_KEY, None)
+        if start is not None:
+            emit_capture_timing("finalize", (perf_counter() - start) * 1000.0)
+
+
+def _capture_initial_states(
+    session: Session, versioned_classes: tuple[type, ...]
+) -> None:
+    """Retain dirty versioned entities' pre-flush states for the final diff.
+
+    Timed as its own metric stage. The per-entity pre-state SELECTs issued
+    here are the capture cost that scales with the number of dirty versioned
+    entities — on a bulk edit plausibly the dominant cost the kill-switch
+    removes — and they run before the flush, outside ``finalize``'s timer. A
+    sample is emitted only when at least one pre-state read was attempted,
+    including failed reads that return no state,
+    so the many unrelated autoflushes do not flood the series with empty
+    samples; together with ``finalize`` the two stages cover the whole
+    listener. Module-level (not the registered closure) so it is
+    unit-testable without ``db.session``.
+    """
+    initial_states: dict[tuple[str, int], tuple[Any, dict[str, Any]]] = (
+        session.info.setdefault(_INITIAL_STATES_KEY, {})
+    )
+    start: float = perf_counter()
+    attempted: bool = False
+    obj: Any
+    try:
+        for obj in list(session.dirty):
+            if isinstance(obj, versioned_classes):
+                key: tuple[str, int] | None = _uncaptured_identity(obj, 
initial_states)
+                if key is None:
+                    continue

Review Comment:
   Good catch—identity lookup still sits outside the per-entity handler, so its 
failure reaches the outer guard and skips later entities. Could we take your 
suggested guard and add the three-entity identity-access regression case?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to