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 `_] 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