fix(storage): use SAVEPOINT in initialize_concurrency_limit_to_default to fix PostgreSQL crash

[Vamsi-klu]

Reviewers Requested: @prha @gibsondan — top committers to sql_event_log.py (14x and 7x commits respectively)

Note on CI: Buildkite builds are blocked pending maintainer approval (standard for fork PRs). Tests pass locally — see Test Results below.

Summary & Motivation

Fixes #33552

Production blocker: SqlEventLogStorage.initialize_concurrency_limit_to_default() crashes on PostgreSQL when the concurrency key already exists in the database. This makes it impossible to run any asset that uses the pool= parameter on PostgreSQL deployments when a race condition or retry causes the INSERT to find an existing row.

The bug affects all PostgreSQL-backed Dagster deployments (Azure Database for PostgreSQL, AWS RDS, etc.) using pool-based concurrency limits. The only workaround is manually setting pool limits through the UI.

Root Cause

The method uses an INSERT-then-catch-IntegrityError-then-UPDATE pattern:

try:
    conn.execute(ConcurrencyLimitsTable.insert().values(...))
except db_exc.IntegrityError:
    conn.execute(ConcurrencyLimitsTable.update()...)  # FAILS on PostgreSQL

On PostgreSQL, a failed INSERT statement aborts the entire transaction. All subsequent statements fail with InFailedSqlTransaction until a ROLLBACK. This is fundamental PostgreSQL behavior — unlike SQLite where failed statements don't invalidate the transaction.

The bug exists in both code paths:

1. has_default_pool_limit_col=True — INSERT fails, UPDATE fails
2. has_default_pool_limit_col=False — INSERT fails, _allocate_concurrency_slots fails

Trigger Scenario

A race condition (concurrent runs, or retry after partial failure where the row was committed by a different process) causes the INSERT to find the row already present.

Fix

Wrapped the INSERT in conn.begin_nested() (SAVEPOINT). On IntegrityError, only the savepoint is rolled back — the outer transaction proceeds. Applied to both code paths. This pattern is already used elsewhere in the codebase (store_asset_event, add_dynamic_partitions).

Test Plan

Test	Status
test_initialize_concurrency_limit_to_default_idempotent (double-call + update)	PASS (runs on PostgreSQL storage)
Existing test_default_concurrency	PASS
Existing test_changing_default_concurrency_key	PASS
make ruff	PASS

Error Before Fix

sqlalchemy.exc.InternalError: (psycopg2.errors.InFailedSqlTransaction)
current transaction is aborted, commands ignored until end of transaction block

Caused by: psycopg2.errors.UniqueViolation: duplicate key value violates
unique constraint "concurrency_limits_concurrency_key_key"

Changelog

Fixed a crash in initialize_concurrency_limit_to_default on PostgreSQL when the concurrency key already exists. The method now uses a SAVEPOINT to handle the INSERT conflict gracefully.

🤖 Generated with Claude Code

[greptile-apps]

Greptile Summary

This PR fixes a PostgreSQL crash in SqlEventLogStorage.initialize_concurrency_limit_to_default caused by the INSERT-catch-IntegrityError-then-UPDATE pattern: on PostgreSQL, a failed INSERT aborts the entire transaction, making subsequent statements fail with InFailedSqlTransaction. The fix wraps the INSERT in conn.begin_nested() (a SAVEPOINT) so that only the savepoint is rolled back on conflict, leaving the outer transaction intact.

• sql_event_log.py: Both code paths (has_default_pool_limit_col=True and =False) now wrap their ConcurrencyLimitsTable.insert() call in conn.begin_nested(), matching the standard SQLAlchemy pattern for tolerating unique-constraint violations on PostgreSQL without aborting the outer transaction.
• event_log_storage.py: A new test test_initialize_concurrency_limit_to_default_idempotent covers the double-call scenario (simulating a race condition) and verifies the UPDATE path works after the default limit changes.

Confidence Score: 5/5

Safe to merge — the change is narrow, isolated to a single method, and the SAVEPOINT pattern is the correct SQLAlchemy idiom for tolerating unique-constraint violations on PostgreSQL without aborting the outer transaction.

The fix correctly addresses a well-understood PostgreSQL transactional behavior (a failed statement aborts the enclosing transaction) by inserting a SAVEPOINT before the conflicting INSERT. The outer try/except db_exc.IntegrityError still catches the rolled-back savepoint's exception, and the surrounding transaction continues normally. Both code paths are patched, the new test exercises the double-call and limit-update scenarios, and skip guards ensure the test only runs on compatible backends.

No files require special attention.

Important Files Changed

Filename	Overview
python_modules/dagster/dagster/_core/storage/event_log/sql_event_log.py	Wraps INSERT statements in conn.begin_nested() (SAVEPOINT) for both code paths in initialize_concurrency_limit_to_default; fix is minimal, correct, and consistent with SQLAlchemy best practice for PostgreSQL.
python_modules/dagster/dagster_tests/storage_tests/utils/event_log_storage.py	Adds test_initialize_concurrency_limit_to_default_idempotent covering the double-call race-condition scenario and the default-change update path, with correct skip guards for unsupported backends.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Caller
    participant initialize as initialize_concurrency_limit_to_default
    participant DB as PostgreSQL

    Caller->>initialize: call(concurrency_key)
    initialize->>DB: BEGIN (index_transaction)
    initialize->>DB: SAVEPOINT sp1  [conn.begin_nested()]
    DB-->>initialize: savepoint active

    alt Row does not exist
        initialize->>DB: INSERT INTO concurrency_limits ...
        DB-->>initialize: success
        initialize->>DB: RELEASE SAVEPOINT sp1
    else Row already exists (race / retry)
        initialize->>DB: INSERT INTO concurrency_limits ...
        DB-->>initialize: IntegrityError (UniqueViolation)
        Note over DB: Only savepoint sp1 is rolled back
        initialize->>DB: ROLLBACK TO SAVEPOINT sp1
        initialize->>DB: "UPDATE concurrency_limits SET limit=..."
        DB-->>initialize: success
    end

    initialize->>DB: COMMIT
    initialize-->>Caller: True

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Caller
    participant initialize as initialize_concurrency_limit_to_default
    participant DB as PostgreSQL

    Caller->>initialize: call(concurrency_key)
    initialize->>DB: BEGIN (index_transaction)
    initialize->>DB: SAVEPOINT sp1  [conn.begin_nested()]
    DB-->>initialize: savepoint active

    alt Row does not exist
        initialize->>DB: INSERT INTO concurrency_limits ...
        DB-->>initialize: success
        initialize->>DB: RELEASE SAVEPOINT sp1
    else Row already exists (race / retry)
        initialize->>DB: INSERT INTO concurrency_limits ...
        DB-->>initialize: IntegrityError (UniqueViolation)
        Note over DB: Only savepoint sp1 is rolled back
        initialize->>DB: ROLLBACK TO SAVEPOINT sp1
        initialize->>DB: "UPDATE concurrency_limits SET limit=..."
        DB-->>initialize: success
    end

    initialize->>DB: COMMIT
    initialize-->>Caller: True

Loading

_{Reviews (2): Last reviewed commit: "fix(storage): use SAVEPOINT in initializ..." | Re-trigger Greptile}