From 46807dbbe2f249cf9efcde933436940f2580b32d Mon Sep 17 00:00:00 2001 From: "F.N. Claessen" Date: Mon, 27 Jul 2026 13:00:36 +0200 Subject: [PATCH 1/2] fix: prevent duplicate data sources from concurrent get-or-create On a fresh database, concurrent jobs (e.g. several workers computing their first-ever schedules) could each run the get-or-create logic for the same data source, find nothing, and insert duplicate rows. Every subsequent lookup then failed with MultipleResultsFound, permanently wedging all scheduling jobs. The existing unique constraint over (name, user_id, account_id, model, version, attributes_hash) never fired for such rows, because PostgreSQL treats NULLs as distinct and scheduler/forecaster sources have NULL user_id and account_id (and sometimes NULL attributes_hash). This commit: - replaces the NULL-blind unique constraint with a NULL-safe unique expression index (coalescing NULLs to sentinel values), via a migration that first deduplicates existing rows (keeping the lowest id and repointing timed_belief, annotation and legacy power/price/weather rows) - makes both get-or-create code paths (get_or_create_source and get_data_source) insert within a savepoint and re-fetch the winning row on IntegrityError - makes source lookups tolerate remaining (near-)duplicates by deterministically picking the lowest id with a logged warning, instead of raising MultipleResultsFound Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WtuVTVfL4fQ9QSqbLXmAGD Signed-off-by: F.N. Claessen --- ...nforce_null_safe_data_source_uniqueness.py | 149 ++++++++++++++++++ flexmeasures/data/models/data_sources.py | 20 ++- flexmeasures/data/services/data_sources.py | 53 ++++++- flexmeasures/data/tests/test_data_source.py | 96 ++++++++++- flexmeasures/data/utils.py | 32 ++-- 5 files changed, 329 insertions(+), 21 deletions(-) create mode 100644 flexmeasures/data/migrations/versions/c9d4f7a21e0b_enforce_null_safe_data_source_uniqueness.py diff --git a/flexmeasures/data/migrations/versions/c9d4f7a21e0b_enforce_null_safe_data_source_uniqueness.py b/flexmeasures/data/migrations/versions/c9d4f7a21e0b_enforce_null_safe_data_source_uniqueness.py new file mode 100644 index 0000000000..2929d4a11f --- /dev/null +++ b/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"], + ) diff --git a/flexmeasures/data/models/data_sources.py b/flexmeasures/data/models/data_sources.py index 642b51dc66..5b41c84b0c 100644 --- a/flexmeasures/data/models/data_sources.py +++ b/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, ), ) diff --git a/flexmeasures/data/services/data_sources.py b/flexmeasures/data/services/data_sources.py index d8260b022e..b2c6ca601e 100644 --- a/flexmeasures/data/services/data_sources.py +++ b/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 diff --git a/flexmeasures/data/tests/test_data_source.py b/flexmeasures/data/tests/test_data_source.py index 9f9d917835..665457200d 100644 --- a/flexmeasures/data/tests/test_data_source.py +++ b/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)) diff --git a/flexmeasures/data/utils.py b/flexmeasures/data/utils.py index 91f3e1023d..0efbcf2257 100644 --- a/flexmeasures/data/utils.py +++ b/flexmeasures/data/utils.py @@ -121,27 +121,31 @@ def get_data_source( """Make sure we have a data source. Create one if it doesn't exist, and add to session. Meant for scripts that may run for the first time. """ + # Imported here to avoid a circular import at module load time. + from flexmeasures.data.services.data_sources import ( + get_first_matching_source, + insert_source_race_safely, + ) - data_source = db.session.execute( - select(DataSource).filter_by( - name=data_source_name, - model=data_source_model, - version=data_source_version, - type=data_source_type, - ) - ).scalar_one_or_none() + query = select(DataSource).filter_by( + name=data_source_name, + model=data_source_model, + version=data_source_version, + type=data_source_type, + ) + data_source = get_first_matching_source(query) if data_source is None: - data_source = DataSource( + new_source = DataSource( name=data_source_name, model=data_source_model, version=data_source_version, type=data_source_type, ) - db.session.add(data_source) - db.session.flush() # populate the primary key attributes (like id) without committing the transaction - current_app.logger.info( - f'Session updated with new {data_source_type} data source "{data_source.__repr__()}".' - ) + data_source = insert_source_race_safely(new_source, query) + if data_source is new_source: + current_app.logger.info( + f'Session updated with new {data_source_type} data source "{data_source.__repr__()}".' + ) return data_source From eeeb0077a0cda4d91ecf8ca9af731913de986c65 Mon Sep 17 00:00:00 2001 From: "F.N. Claessen" Date: Mon, 27 Jul 2026 13:02:12 +0200 Subject: [PATCH 2/2] docs: changelog entry for data source race fix Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WtuVTVfL4fQ9QSqbLXmAGD Signed-off-by: F.N. Claessen --- documentation/changelog.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/documentation/changelog.rst b/documentation/changelog.rst index b3af30cec6..6861221295 100644 --- a/documentation/changelog.rst +++ b/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 `_] * Scheduling jobs no longer print ``Job ... made schedule.`` before ``scheduler.compute()`` runs (only after a successful schedule) [see `PR #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 `_] * 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 `_]