Group event log index writes in transactions (#25660)

## Summary & Motivation
As we prepare to write more derived tables on event log writes, we
should group the writes into transactions.

## Test Plan
BK

Internal-RevId: 1d689b77cfa6a046a573c4f0939006bf1886e5e4

Author: prha

python_modules/dagster/dagster/_core/storage/event_log/sql_event_log.py

@@ -229,15 +229,26 @@ def _event_to_row(self, event: EventLogEntry) -> dict[str, Any]:
 
     def has_asset_key_col(self, column_name: str) -> bool:
         with self.index_connection() as conn:
-            column_names = [x.get("name") for x in db.inspect(conn).get_columns(AssetKeyTable.name)]
-            return column_name in column_names
+            return self._has_asset_key_col_on_connection(conn, column_name)
+
+    def _has_asset_key_col_on_connection(self, conn: Connection, column_name: str) -> bool:
+        column_names = [x.get("name") for x in db.inspect(conn).get_columns(AssetKeyTable.name)]
+        return column_name in column_names
 
     def has_asset_key_index_cols(self) -> bool:
         return self.has_asset_key_col("last_materialization_timestamp")
 
+    def _has_asset_key_index_cols_on_connection(self, conn: Connection) -> bool:
+        return self._has_asset_key_col_on_connection(conn, "last_materialization_timestamp")
+
     def store_asset_event(self, event: EventLogEntry, event_id: int):
         check.inst_param(event, "event", EventLogEntry)
+        check.int_param(event_id, "event_id")
+
+        with self.index_transaction() as conn:
+            self._store_asset_event(conn, event, event_id)
 
+    def _store_asset_event(self, conn: Connection, event: EventLogEntry, event_id: int) -> None:
         if not (event.dagster_event and event.dagster_event.asset_key):
             return
 
@@ -253,7 +264,9 @@ def store_asset_event(self, event: EventLogEntry, event_id: int):
         #
         # https://github.com/dagster-io/dagster/issues/3945
 
-        values = self._get_asset_entry_values(event, event_id, self.has_asset_key_index_cols())
+        values = self._get_asset_entry_values(
+            event, event_id, self._has_asset_key_index_cols_on_connection(conn)
+        )
         if not values:
             return
         insert_statement = AssetKeyTable.insert().values(
@@ -267,11 +280,10 @@ def store_asset_event(self, event: EventLogEntry, event_id: int):
             )
         )
 
-        with self.index_connection() as conn:
-            try:
-                conn.execute(insert_statement)
-            except db_exc.IntegrityError:
-                conn.execute(update_statement)
+        try:
+            conn.execute(insert_statement)
+        except db_exc.IntegrityError:
+            conn.execute(update_statement)
 
     def _get_asset_entry_values(
         self, event: EventLogEntry, event_id: int, has_asset_key_index_cols: bool
@@ -337,6 +349,12 @@ def store_asset_event_tags(
         check.sequence_param(events, "events", EventLogEntry)
         check.sequence_param(event_ids, "event_ids", int)
 
+        with self.index_transaction() as conn:
+            self._store_asset_event_tags(conn, events, event_ids)
+
+    def _store_asset_event_tags(
+        self, conn: Connection, events: Sequence[EventLogEntry], event_ids: Sequence[int]
+    ) -> None:
         all_values = [
             dict(
                 event_id=event_id,
@@ -353,8 +371,7 @@ def store_asset_event_tags(
         # migration to create the table. On read, we will throw an error if the table does not
         # exist.
         if len(all_values) > 0 and self.has_table(AssetEventTagsTable.name):
-            with self.index_connection() as conn:
-                conn.execute(AssetEventTagsTable.insert(), all_values)
+            conn.execute(AssetEventTagsTable.insert(), all_values)
 
     def _tags_for_asset_event(self, event: EventLogEntry) -> Mapping[str, str]:
         tags = {}
@@ -388,14 +405,14 @@ def store_event(self, event: EventLogEntry) -> None:
             and event.dagster_event_type in ASSET_EVENTS
             and event.dagster_event.asset_key  # type: ignore
         ):
-            self.store_asset_event(event, event_id)
-
             if event_id is None:
                 raise DagsterInvariantViolationError(
                     "Cannot store asset event tags for null event id."
                 )
 
-            self.store_asset_event_tags([event], [event_id])
+            with self.index_transaction() as conn:
+                self._store_asset_event_tags(conn, [event], [event_id])
+                self._store_asset_event(conn, event, event_id)
 
         if event.is_dagster_event and event.dagster_event_type in ASSET_CHECK_EVENTS:
             self.store_asset_check_event(event, event_id)

python_modules/dagster/dagster/_core/storage/event_log/sqlite/sqlite_event_log.py

@@ -275,20 +275,17 @@ def store_event(self, event: EventLogEntry) -> None:
             event_id = None
 
             # mirror the event in the cross-run index database
-            with (
-                self.index_connection() as conn,
-                db_result(conn, insert_event_statement) as result,
-            ):
-                event_id = result.inserted_primary_key[0]
-
-            self.store_asset_event(event, event_id)
+            with self.index_transaction() as conn:
+                with db_result(conn, insert_event_statement) as result:
+                    event_id = result.inserted_primary_key[0]
 
-            if event_id is None:
-                raise DagsterInvariantViolationError(
-                    "Cannot store asset event tags for null event id."
-                )
+                if event_id is None:
+                    raise DagsterInvariantViolationError(
+                        "Cannot store asset event tags for null event id."
+                    )
 
-            self.store_asset_event_tags([event], [event_id])
+                self._store_asset_event_tags(conn, [event], [event_id])
+                self._store_asset_event(conn, event, event_id)
 
         if event.is_dagster_event and event.dagster_event_type in ASSET_CHECK_EVENTS:
             # mirror the event in the cross-run index database

python_modules/libraries/dagster-mysql/dagster_mysql/event_log/event_log.py

@@ -149,33 +149,42 @@ def get_server_version(self) -> str | None:
         return cast("str", row[0])
 
     def store_asset_event(self, event: EventLogEntry, event_id: int) -> None:
+        check.inst_param(event, "event", EventLogEntry)
+        check.int_param(event_id, "event_id")
+
+        with self.index_transaction() as conn:
+            self._store_asset_event(conn, event, event_id)
+
+    def _store_asset_event(self, conn: Connection, event: EventLogEntry, event_id: int) -> None:
         # last_materialization_timestamp is updated upon observation, materialization, materialization_planned
         # See SqlEventLogStorage.store_asset_event method for more details
 
+        if not (event.dagster_event and event.dagster_event.asset_key):
+            return
+
         values = self._get_asset_entry_values(
             event, event_id, self.has_secondary_index(ASSET_KEY_INDEX_COLS)
         )
-        with self.index_connection() as conn:
-            if values:
+        if values:
+            conn.execute(
+                db_dialects.mysql.insert(AssetKeyTable)
+                .values(
+                    asset_key=event.dagster_event.asset_key.to_string(),
+                    **values,
+                )
+                .on_duplicate_key_update(
+                    **values,
+                )
+            )
+        else:
+            try:
                 conn.execute(
-                    db_dialects.mysql.insert(AssetKeyTable)
-                    .values(
-                        asset_key=event.dagster_event.asset_key.to_string(),  # type: ignore  # (possible none)
-                        **values,
-                    )
-                    .on_duplicate_key_update(
-                        **values,
+                    db_dialects.mysql.insert(AssetKeyTable).values(
+                        asset_key=event.dagster_event.asset_key.to_string(),
                     )
                 )
-            else:
-                try:
-                    conn.execute(
-                        db_dialects.mysql.insert(AssetKeyTable).values(
-                            asset_key=event.dagster_event.asset_key.to_string(),  # type: ignore  # (possible none)
-                        )
-                    )
-                except db_exc.IntegrityError:
-                    pass
+            except db_exc.IntegrityError:
+                pass
 
     def _connect(self) -> ContextManager[Connection]:
         return create_mysql_connection(self._engine, __file__, "event log")

python_modules/libraries/dagster-postgres/dagster_postgres/event_log/event_log.py

@@ -209,14 +209,14 @@ def store_event(self, event: EventLogEntry) -> None:
             and event.dagster_event_type in ASSET_EVENTS
             and event.dagster_event.asset_key  # type: ignore
         ):
-            self.store_asset_event(event, event_id)
-
             if event_id is None:
                 raise DagsterInvariantViolationError(
                     "Cannot store asset event tags for null event id."
                 )
 
-            self.store_asset_event_tags([event], [event_id])
+            with self.index_transaction() as conn:
+                self._store_asset_event_tags(conn, [event], [event_id])
+                self._store_asset_event(conn, event, event_id)
 
         if event.is_dagster_event and event.dagster_event_type in ASSET_CHECK_EVENTS:
             self.store_asset_check_event(event, event_id)
@@ -250,20 +250,26 @@ def store_event_batch(self, events: Sequence[EventLogEntry]) -> None:
                 )
                 event_ids = [cast("int", row[0]) for row in result.fetchall()]
 
-            # We only update the asset table with the last event
-            self.store_asset_event(events[-1], event_ids[-1])
-
             if any(event_id is None for event_id in event_ids):
                 raise DagsterInvariantViolationError(
                     "Cannot store asset event tags for null event id."
                 )
 
-            self.store_asset_event_tags(events, event_ids)
+            with self.index_transaction() as conn:
+                self._store_asset_event_tags(conn, events, event_ids)
+                # We only update the asset table with the last event.
+                self._store_asset_event(conn, events[-1], event_ids[-1])
         else:
             return super().store_event_batch(events)
 
     def store_asset_event(self, event: EventLogEntry, event_id: int) -> None:
         check.inst_param(event, "event", EventLogEntry)
+        check.int_param(event_id, "event_id")
+
+        with self.index_transaction() as conn:
+            self._store_asset_event(conn, event, event_id)
+
+    def _store_asset_event(self, conn: Connection, event: EventLogEntry, event_id: int) -> None:
         if not (event.dagster_event and event.dagster_event.asset_key):
             return
 
@@ -300,19 +306,18 @@ def store_asset_event(self, event: EventLogEntry, event_id: int) -> None:
         values = self._get_asset_entry_values(
             event, event_id, self.has_secondary_index(ASSET_KEY_INDEX_COLS)
         )
-        with self.index_connection() as conn:
-            query = db_dialects.postgresql.insert(AssetKeyTable).values(
-                asset_key=event.dagster_event.asset_key.to_string(),
-                **values,
+        query = db_dialects.postgresql.insert(AssetKeyTable).values(
+            asset_key=event.dagster_event.asset_key.to_string(),
+            **values,
+        )
+        if values:
+            query = query.on_conflict_do_update(
+                index_elements=[AssetKeyTable.c.asset_key],
+                set_=dict(**values),
             )
-            if values:
-                query = query.on_conflict_do_update(
-                    index_elements=[AssetKeyTable.c.asset_key],
-                    set_=dict(**values),
-                )
-            else:
-                query = query.on_conflict_do_nothing()
-            conn.execute(query)
+        else:
+            query = query.on_conflict_do_nothing()
+        conn.execute(query)
 
     def add_dynamic_partitions(
         self, partitions_def_name: str, partition_keys: Sequence[str]