fix(dagster-postgres): use INSERT ON CONFLICT in initialize_concurrency_limit_to_default

Fixes #33552

The base-class `initialize_concurrency_limit_to_default` uses an INSERT-
then-catch-IntegrityError-then-UPDATE pattern. On PostgreSQL, a failed
INSERT aborts the entire transaction, so the fallback UPDATE fails with
`InFailedSqlTransaction`. This crashes every run that uses `pool=` when
the concurrency key already exists (race condition or retry).

Override both `initialize_concurrency_limit_to_default` and
`_reconcile_concurrency_limits_from_slots` in `PostgresEventLogStorage`
to use `INSERT ... ON CONFLICT` — the same atomic upsert pattern already
used by `store_asset_event` and `add_dynamic_partitions` in the same
class. The legacy-schema path also moves `_allocate_concurrency_slots`
inside the transaction context (fixing a pre-existing use-after-close).

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: iaintbyte

python_modules/dagster/dagster_tests/storage_tests/utils/event_log_storage.py

@@ -5863,6 +5863,41 @@ def test_changing_default_concurrency_key(
 
         assert storage.get_concurrency_info("foo").slot_count == 0
 
+    def test_default_concurrency_idempotent(
+        self, storage: EventLogStorage, instance: DagsterInstance
+    ):
+        """Calling initialize_concurrency_limit_to_default twice with the same key must not
+        crash (regression test for PostgreSQL UniqueViolation / InFailedSqlTransaction)."""
+        assert storage
+        if not storage.supports_global_concurrency_limits:
+            pytest.skip("storage does not support global op concurrency")
+
+        if not self.can_set_concurrency_defaults():
+            pytest.skip("storage does not support setting global op concurrency defaults")
+
+        self.set_default_op_concurrency(instance, storage, 5)
+
+        assert storage.initialize_concurrency_limit_to_default("bar")
+        assert storage.get_concurrency_info("bar").slot_count == 5
+
+        # Second call with same default — must be idempotent, not crash
+        assert storage.initialize_concurrency_limit_to_default("bar")
+        assert storage.get_concurrency_info("bar").slot_count == 5
+
+        # Change the default and re-initialize — limit must propagate
+        self.set_default_op_concurrency(instance, storage, 3)
+        assert storage.initialize_concurrency_limit_to_default("bar")
+        assert storage.get_concurrency_info("bar").slot_count == 3
+
+        # Explicitly set via set_concurrency_slots (user-managed), then
+        # re-initialize from default — default must reclaim the row
+        storage.set_concurrency_slots("bar", 10)
+        assert storage.get_concurrency_info("bar").slot_count == 10
+
+        self.set_default_op_concurrency(instance, storage, 4)
+        assert storage.initialize_concurrency_limit_to_default("bar")
+        assert storage.get_concurrency_info("bar").slot_count == 4
+
     def test_asset_checks(
         self,
         storage: EventLogStorage,

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

@@ -19,6 +19,7 @@
     SqlEventLogStorageMetadata,
     SqlEventLogStorageTable,
 )
+from dagster._core.storage.event_log.schema import ConcurrencyLimitsTable, ConcurrencySlotsTable
 from dagster._core.storage.event_log.base import EventLogCursor
 from dagster._core.storage.event_log.migration import ASSET_KEY_INDEX_COLS
 from dagster._core.storage.event_log.polling_event_watcher import SqlPollingEventWatcher
@@ -314,6 +315,98 @@ def store_asset_event(self, event: EventLogEntry, event_id: int) -> None:
                 query = query.on_conflict_do_nothing()
             conn.execute(query)
 
+    def _reconcile_concurrency_limits_from_slots(self) -> None:
+        if not self.has_concurrency_limits_table:
+            return
+
+        if not self._has_rows(ConcurrencySlotsTable) or self._has_rows(ConcurrencyLimitsTable):
+            return
+
+        with self.index_transaction() as conn:
+            rows = conn.execute(
+                db_select(
+                    [
+                        ConcurrencySlotsTable.c.concurrency_key,
+                        db.func.count().label("count"),
+                    ]
+                )
+                .where(
+                    ConcurrencySlotsTable.c.deleted == False,  # noqa: E712
+                )
+                .group_by(
+                    ConcurrencySlotsTable.c.concurrency_key,
+                )
+            ).fetchall()
+            if rows:
+                conn.execute(
+                    db_dialects.postgresql.insert(ConcurrencyLimitsTable)
+                    .values(
+                        [
+                            {
+                                "concurrency_key": row[0],
+                                "limit": row[1],
+                            }
+                            for row in rows
+                        ]
+                    )
+                    .on_conflict_do_nothing(),
+                )
+
+    def initialize_concurrency_limit_to_default(self, concurrency_key: str) -> bool:
+        if not self.has_concurrency_limits_table:
+            return False
+
+        self._reconcile_concurrency_limits_from_slots()
+
+        if not self.has_instance:
+            return False
+
+        default_limit = self._instance.global_op_concurrency_default_limit
+        has_default_pool_limit_col = self.has_default_pool_limit_col
+
+        if has_default_pool_limit_col:
+            with self.index_transaction() as conn:
+                if default_limit is None:
+                    conn.execute(
+                        ConcurrencyLimitsTable.delete().where(
+                            ConcurrencyLimitsTable.c.concurrency_key == concurrency_key,
+                        )
+                    )
+                    self._allocate_concurrency_slots(conn, concurrency_key, 0)
+                else:
+                    stmt = db_dialects.postgresql.insert(ConcurrencyLimitsTable).values(
+                        concurrency_key=concurrency_key,
+                        limit=default_limit,
+                        using_default_limit=True,
+                    )
+                    stmt = stmt.on_conflict_do_update(
+                        index_elements=[ConcurrencyLimitsTable.c.concurrency_key],
+                        set_={
+                            "limit": stmt.excluded.limit,
+                            "using_default_limit": stmt.excluded.using_default_limit,
+                        },
+                        where=db.or_(
+                            ConcurrencyLimitsTable.c.limit != stmt.excluded.limit,
+                            ConcurrencyLimitsTable.c.using_default_limit
+                            != stmt.excluded.using_default_limit,
+                        ),
+                    )
+                    conn.execute(stmt)
+                    self._allocate_concurrency_slots(conn, concurrency_key, default_limit)
+            return True
+
+        if default_limit is None:
+            return False
+
+        with self.index_transaction() as conn:
+            conn.execute(
+                db_dialects.postgresql.insert(ConcurrencyLimitsTable)
+                .values(concurrency_key=concurrency_key, limit=default_limit)
+                .on_conflict_do_nothing(),
+            )
+            self._allocate_concurrency_slots(conn, concurrency_key, default_limit)
+        return True
+
     def add_dynamic_partitions(
         self, partitions_def_name: str, partition_keys: Sequence[str]
     ) -> None:

python_modules/libraries/dagster-postgres/dagster_postgres_tests/test_event_log.py

@@ -45,6 +45,18 @@ def event_log_storage(self, instance):
         assert isinstance(event_log_storage, PostgresEventLogStorage)
         yield event_log_storage
 
+    def can_set_concurrency_defaults(self):
+        return True
+
+    def set_default_op_concurrency(self, instance, storage, limit):
+        settings = dict(instance._settings) if instance._settings else {}  # noqa: SLF001
+        concurrency = dict(settings.get("concurrency", {}))
+        pools = dict(concurrency.get("pools", {}))
+        pools["default_limit"] = limit
+        concurrency["pools"] = pools
+        settings["concurrency"] = concurrency
+        instance._settings = settings  # noqa: SLF001
+
     def can_wipe_asset_partitions(self) -> bool:
         return False