Fix race condition creating duplicate data sources on fresh databases

documentation/changelog.rst

@@ -69,6 +69,7 @@ Infrastructure / Support
 
 Bugfixes
 -----------
+* Fix a race condition where concurrent jobs on a fresh database could create duplicate data sources (their uniqueness constraint was ineffective for rows with NULL values, such as scheduler sources), after which every subsequent scheduling job failed with ``MultipleResultsFound``; uniqueness is now enforced NULL-safely at the database level (existing duplicates are cleaned up in the migration), the get-or-create logic recovers from losing an insert race, and lookups tolerate pre-existing duplicates by deterministically using the oldest source [see `PR #2359 <https://www.github.com/FlexMeasures/flexmeasures/pull/2359>`_]
 * Scheduling jobs no longer print ``Job ... made schedule.`` before ``scheduler.compute()`` runs (only after a successful schedule) [see `PR #2342 <https://www.github.com/FlexMeasures/flexmeasures/pull/2342>`_]
 * ``flexmeasures add user --roles`` now correctly accepts a comma-separated list of roles (and repeated ``--roles`` options) instead of creating one role whose name contains commas [see `PR #2339 <https://www.github.com/FlexMeasures/flexmeasures/pull/2339>`_]
 * Raise a clear ``ValueError`` when a flex-model references a missing sensor ID instead of ``AttributeError: 'NoneType' object has no attribute 'asset_id'`` [see `PR #2343 <https://www.github.com/FlexMeasures/flexmeasures/pull/2343>`_]

flexmeasures/data/migrations/versions/c9d4f7a21e0b_enforce_null_safe_data_source_uniqueness.py

@@ -0,0 +1,149 @@
+"""Enforce NULL-safe uniqueness of data sources
+
+The data_source table had a unique constraint over
+(name, user_id, account_id, model, version, attributes_hash), but most of
+these columns are nullable and PostgreSQL treats NULLs as distinct values,
+so the constraint never fired for rows with NULLs in any of these columns.
+Notably, scheduler and forecaster sources have no user or account, so
+concurrent get-or-create calls (e.g. several workers computing their first
+schedules against a fresh database) could insert duplicate rows. Every
+subsequent lookup then failed with MultipleResultsFound, wedging all
+scheduling jobs.
+
+This migration:
+
+1. Deduplicates existing data_source rows that are identical in all key
+   columns (treating NULLs as equal). The row with the lowest id is kept,
+   and timed_belief and annotation rows are repointed to it. In the corner
+   case where both a kept and a duplicate source recorded a belief with the
+   same primary key coordinates (sensor, event start, belief horizon,
+   cumulative probability), the kept source's belief wins and the duplicate
+   source's belief is dropped (they are duplicate recordings by the same
+   logical source).
+2. Replaces the NULL-blind unique constraint with a unique expression index
+   that coalesces NULLs to sentinel values which cannot occur in real data
+   (negative ids, empty strings, an empty bytes hash).
+
+Revision ID: c9d4f7a21e0b
+Revises: 4b0f2e9c1a6d
+Create Date: 2026-07-27 12:00:00.000000
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy import text
+
+
+# revision identifiers, used by Alembic.
+revision = "c9d4f7a21e0b"
+down_revision = "4b0f2e9c1a6d"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    bind = op.get_bind()
+    # Legacy tables (pre-timed_belief data model) that may still exist on
+    # long-lived databases and reference data_source; they share the same
+    # primary key layout: (datetime, sensor_id, horizon, data_source_id).
+    legacy_tables = [
+        table
+        for table in ("power", "price", "weather")
+        if sa.inspect(bind).has_table(table)
+    ]
+
+    # 1. Deduplicate: find groups of rows that are identical in all key columns.
+    #    GROUP BY conveniently treats NULLs as equal, matching the semantics of
+    #    the NULL-safe unique index we are about to create.
+    duplicate_groups = bind.execute(
+        text(
+            "SELECT min(id) AS keep_id, array_agg(id ORDER BY id) AS ids "
+            "FROM data_source "
+            "GROUP BY name, user_id, account_id, model, version, attributes_hash "
+            "HAVING count(*) > 1"
+        )
+    ).fetchall()
+    for keep_id, ids in duplicate_groups:
+        for dupe_id in ids:
+            if dupe_id == keep_id:
+                continue
+            # Repoint beliefs to the kept source, except where the kept source
+            # already recorded a belief with the same primary key coordinates.
+            bind.execute(
+                text(
+                    "UPDATE timed_belief tb SET source_id = :keep_id "
+                    "WHERE tb.source_id = :dupe_id "
+                    "AND NOT EXISTS ("
+                    "    SELECT 1 FROM timed_belief tb2 "
+                    "    WHERE tb2.source_id = :keep_id "
+                    "    AND tb2.sensor_id = tb.sensor_id "
+                    "    AND tb2.event_start = tb.event_start "
+                    "    AND tb2.belief_horizon = tb.belief_horizon "
+                    "    AND tb2.cumulative_probability = tb.cumulative_probability"
+                    ")"
+                ),
+                {"keep_id": keep_id, "dupe_id": dupe_id},
+            )
+            # Any beliefs still pointing to the duplicate source collide with
+            # beliefs of the kept source: drop them in favour of the latter.
+            bind.execute(
+                text("DELETE FROM timed_belief WHERE source_id = :dupe_id"),
+                {"dupe_id": dupe_id},
+            )
+            bind.execute(
+                text(
+                    "UPDATE annotation SET source_id = :keep_id WHERE source_id = :dupe_id"
+                ),
+                {"keep_id": keep_id, "dupe_id": dupe_id},
+            )
+            # Do the same for legacy tables that may still reference data_source.
+            for table in legacy_tables:
+                bind.execute(
+                    text(
+                        f"UPDATE {table} t SET data_source_id = :keep_id "  # nosec B608
+                        f"WHERE t.data_source_id = :dupe_id "
+                        f"AND NOT EXISTS ("
+                        f"    SELECT 1 FROM {table} t2 "
+                        f"    WHERE t2.data_source_id = :keep_id "
+                        f"    AND t2.datetime = t.datetime "
+                        f"    AND t2.sensor_id = t.sensor_id "
+                        f"    AND t2.horizon = t.horizon"
+                        f")"
+                    ),
+                    {"keep_id": keep_id, "dupe_id": dupe_id},
+                )
+                bind.execute(
+                    text(
+                        f"DELETE FROM {table} WHERE data_source_id = :dupe_id"
+                    ),  # nosec B608
+                    {"dupe_id": dupe_id},
+                )
+            bind.execute(
+                text("DELETE FROM data_source WHERE id = :dupe_id"),
+                {"dupe_id": dupe_id},
+            )
+
+    # 2. Replace the NULL-blind unique constraint with a NULL-safe unique index.
+    op.drop_constraint("data_source_name_key", "data_source", type_="unique")
+    op.execute(
+        "CREATE UNIQUE INDEX data_source_nullsafe_uniqueness_idx ON data_source "
+        "(name, coalesce(user_id, -1), coalesce(account_id, -1), "
+        "coalesce(model, ''), coalesce(version, ''), "
+        "coalesce(attributes_hash, '\\x'::bytea))"
+    )
+
+
+def downgrade():
+    """Restore the previous (NULL-blind) unique constraint.
+
+    The deduplication of step 1 of the upgrade is intentionally not reversed:
+    the removed rows were duplicates, and the previous get-or-create logic
+    works fine (better, even) without them.
+    """
+    op.drop_index("data_source_nullsafe_uniqueness_idx", table_name="data_source")
+    op.create_unique_constraint(
+        "data_source_name_key",
+        "data_source",
+        ["name", "user_id", "account_id", "model", "version", "attributes_hash"],
+    )

flexmeasures/data/models/data_sources.py

@@ -5,7 +5,7 @@
 import json
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, ClassVar
-from sqlalchemy import select
+from sqlalchemy import select, text
 from sqlalchemy.ext.mutable import MutableDict
 from sqlalchemy.dialects.postgresql import JSONB
 
@@ -273,8 +273,22 @@ class DataSource(db.Model, tb.BeliefSourceDBMixin):
 
     __tablename__ = "data_source"
     __table_args__ = (
-        db.UniqueConstraint(
-            "name", "user_id", "account_id", "model", "version", "attributes_hash"
+        # Enforce uniqueness of (name, user_id, account_id, model, version, attributes_hash).
+        # A plain UniqueConstraint over these columns does not actually prevent duplicates,
+        # because most of them are nullable and PostgreSQL treats NULLs as distinct values.
+        # For example, scheduler sources have no user or account, so concurrent get-or-create
+        # calls against a fresh database used to be able to insert identical rows.
+        # We therefore use a unique expression index that coalesces NULLs to sentinel values
+        # which cannot occur in real data (negative ids, empty strings, an empty bytes hash).
+        db.Index(
+            "data_source_nullsafe_uniqueness_idx",
+            text("name"),
+            text("coalesce(user_id, -1)"),
+            text("coalesce(account_id, -1)"),
+            text("coalesce(model, '')"),
+            text("coalesce(version, '')"),
+            text("coalesce(attributes_hash, '\\x'::bytea)"),
+            unique=True,
         ),
     )
 

flexmeasures/data/services/data_sources.py

@@ -4,6 +4,8 @@
 
 from flask import current_app
 from sqlalchemy import select
+from sqlalchemy.exc import IntegrityError
+from sqlalchemy.sql import Select
 from typing import Type, TypeVar
 
 from flexmeasures import Account, Source, User
@@ -16,6 +18,52 @@
 DG = TypeVar("DG", bound=DataGenerator)
 
 
+def get_first_matching_source(query: Select) -> DataSource | None:
+    """Return the matching data source with the lowest id, or None if there is no match.
+
+    Tolerating multiple matches is a form of defense in depth:
+    while uniqueness of data sources is enforced at the database level,
+    a database may already contain (near-)duplicate rows from before that enforcement
+    (for example, rows created by concurrently running jobs on a fresh database,
+    or rows that differ only in fields the caller did not filter on, such as attributes).
+    Rather than letting such rows fail every lookup with MultipleResultsFound,
+    we deterministically pick the oldest row and log a warning.
+    """
+    sources = db.session.scalars(query.order_by(DataSource.id)).all()
+    if len(sources) > 1:
+        current_app.logger.warning(
+            f"Found {len(sources)} data sources matching a lookup for one (the oldest is {sources[0]}); "
+            f"using the one with the lowest id ({sources[0].id})."
+        )
+    return sources[0] if sources else None
+
+
+def insert_source_race_safely(new_source: DataSource, query: Select) -> DataSource:
+    """Insert a new data source, returning the winning row if we lose an insert race.
+
+    The insert happens within a SAVEPOINT, so that losing a race against a concurrent
+    session creating the same source (e.g. parallel workers scheduling against a fresh
+    database, whose get-or-create logic all found no source yet) doesn't poison the
+    enclosing transaction. Committing the savepoint flushes, which assigns an id so
+    that the new source can be referenced in the current db session.
+
+    :param new_source: the (not yet added) data source to insert
+    :param query:      the query with which to re-fetch the winning row,
+                       should our insert hit a uniqueness conflict
+    """
+    try:
+        with db.session.begin_nested():
+            db.session.add(new_source)
+    except IntegrityError:
+        # We lost the race: another session created the same source concurrently
+        # (the savepoint was rolled back). Fetch the winning row instead.
+        winner = get_first_matching_source(query)
+        if winner is None:
+            raise
+        return winner
+    return new_source
+
+
 def get_or_create_source(
     source: User | str,
     source_type: str | None = None,
@@ -44,7 +92,7 @@ def get_or_create_source(
         query = query.filter(DataSource.name == source)
     else:
         raise TypeError("source should be of type User or str")
-    _source = db.session.execute(query).scalar_one_or_none()
+    _source = get_first_matching_source(query)
     if not _source:
         if is_user(source):
             _source = DataSource(user=source, model=model, version=version)
@@ -60,9 +108,8 @@ def get_or_create_source(
                 account=account,
             )
         current_app.logger.info(f"Setting up {_source} as new data source...")
-        db.session.add(_source)
+        _source = insert_source_race_safely(_source, query)
         if flush:
-            # assigns id so that we can reference the new object in the current db session
             db.session.flush()
     return _source
 

flexmeasures/data/tests/test_data_source.py

@@ -8,7 +8,7 @@
 import numpy as np
 import pandas as pd
 import timely_beliefs as tb
-from sqlalchemy import insert
+from sqlalchemy import func, insert, select
 
 from flexmeasures.data.models.data_sources import keep_latest_version, DataSource
 from flexmeasures.data.models.generic_assets import GenericAsset, GenericAssetType
@@ -519,3 +519,97 @@ def group_key(index_tuple):
             pd.testing.assert_frame_equal(
                 pd.DataFrame(result), pd.DataFrame(expected), check_like=False
             )
+
+
+def test_get_or_create_source_survives_insert_race(db, app, monkeypatch):
+    """Losing an insert race for a new data source must return the winning row.
+
+    On a fresh database, concurrent workers (e.g. running their first-ever
+    scheduling jobs) used to be able to insert duplicate DataSource rows,
+    because each worker's initial lookup found nothing yet.  We simulate the
+    losing worker by patching its initial lookup to find nothing while the row
+    actually exists: its INSERT must then hit the DB-level uniqueness index,
+    roll back to a savepoint and re-fetch the winner, instead of either
+    creating a duplicate or poisoning the session.
+    """
+    from flexmeasures.data.services import data_sources as data_sources_service
+    from flexmeasures.data.services.data_sources import get_or_create_source
+
+    source_info = dict(source_type="scheduler", model="RaceTestScheduler", version="1")
+    winner = get_or_create_source("test-race", **source_info)
+
+    real_fetch = data_sources_service.get_first_matching_source
+    calls = {"n": 0}
+
+    def miss_on_first_lookup(query):
+        calls["n"] += 1
+        if calls["n"] == 1:
+            # Simulate the race: the other worker's row is not seen by our lookup
+            return None
+        return real_fetch(query)
+
+    monkeypatch.setattr(
+        data_sources_service, "get_first_matching_source", miss_on_first_lookup
+    )
+
+    loser = get_or_create_source("test-race", **source_info)
+
+    assert loser.id == winner.id
+    assert calls["n"] == 2, "the IntegrityError path should have re-fetched the winner"
+    num_sources = db.session.scalar(
+        select(func.count())
+        .select_from(DataSource)
+        .filter_by(
+            name="test-race", type="scheduler", model="RaceTestScheduler", version="1"
+        )
+    )
+    assert num_sources == 1, "no duplicate row should have been created"
+
+
+def test_source_lookups_tolerate_duplicates(db, app):
+    """Pre-existing (near-)duplicate sources must not fail lookups with MultipleResultsFound.
+
+    Sources that differ only in their attributes are legitimate separate rows, but a
+    lookup that doesn't filter on attributes matches all of them. Such a lookup should
+    deterministically return the oldest row instead of raising, so that databases which
+    already contain duplicates (created before uniqueness was enforced at the DB level)
+    degrade gracefully instead of failing every scheduling job.
+    """
+    from flexmeasures.data.services.data_sources import get_or_create_source
+    from flexmeasures.data.utils import get_data_source
+
+    source_info = dict(source_type="scheduler", model="DupeScheduler", version="2")
+    source_1 = get_or_create_source("test-dupes", attributes={"a": 1}, **source_info)
+    source_2 = get_or_create_source("test-dupes", attributes={"a": 2}, **source_info)
+    assert source_1.id != source_2.id
+    oldest_id = min(source_1.id, source_2.id)
+
+    # Lookup without an attributes filter matches both rows; this used to raise
+    found = get_or_create_source("test-dupes", **source_info)
+    assert found.id == oldest_id
+
+    # Same for the lower-level get_data_source utility
+    found = get_data_source(
+        "test-dupes",
+        data_source_model="DupeScheduler",
+        data_source_version="2",
+        data_source_type="scheduler",
+    )
+    assert found.id == oldest_id
+
+
+def test_exact_duplicate_sources_rejected_by_db(db, app):
+    """The DB must reject exact duplicates even when the unique key columns hold NULLs.
+
+    The previous UniqueConstraint was NULL-blind (PostgreSQL treats NULLs as
+    distinct), so rows like scheduler sources - which have no user or account -
+    could be duplicated freely. The NULL-safe unique index must reject them.
+    """
+    from sqlalchemy.exc import IntegrityError
+
+    kwargs = dict(name="test-unique", type="scheduler", model="X", version="3")
+    db.session.add(DataSource(**kwargs))
+    db.session.flush()
+    with pytest.raises(IntegrityError):
+        with db.session.begin_nested():  # keep the outer transaction usable
+            db.session.add(DataSource(**kwargs))