Merge branch 'main' of github.com:FlexMeasures/flexmeasures

Author: nhoening

documentation/changelog.rst

@@ -114,6 +114,7 @@ Infrastructure / Support
 * ``flexmeasures db upgrade`` now runs ``VACUUM ANALYZE`` after upgrading by default, so Postgres has fresh planner statistics right after a migration; opt out with ``--no-vacuum`` [see `PR #2333 <https://www.github.com/FlexMeasures/flexmeasures/pull/2333>`_]
 * Automate Docker Hub image publishing and a PyPI install smoke test on release, add a manually-triggered QA workflow that runs the toy tutorials and HEMS walkthrough against a local Docker Compose stack, and add a helper script to list merged PRs since the last tag [see `PR #2260 <https://www.github.com/FlexMeasures/flexmeasures/pull/2260>`_]
 * Warn hosts when the database schema is not at the latest migration, and skip startup template provisioning until migrations are applied [see `PR #2309 <https://www.github.com/FlexMeasures/flexmeasures/pull/2309>`_]
+* Avoid creating default data while running database maintenance commands such as ``flexmeasures db upgrade`` and ``flexmeasures db-ops restore``, so hosts can upgrade or restore databases without startup provisioning touching an outdated or partially restored schema [see `PR #2428 <https://www.github.com/FlexMeasures/flexmeasures/pull/2428>`_]
 * Add ``FLEXMEASURES_DEFAULT_JOB_TIMEOUT`` and ``FLEXMEASURES_JOB_TIMEOUT`` settings for configuring RQ job timeouts globally and per queue, and log actionable guidance when a forecasting job times out [see `PR #2318 <https://github.com/FlexMeasures/flexmeasures/pull/2318>`_]
 * Stop manual runs of the Docker publishing workflow from overwriting the ``latest`` image tag, and let them opt in to it explicitly [see `PR #2316 <https://www.github.com/FlexMeasures/flexmeasures/pull/2316>`_]
 * Add a pre-commit hook that blocks image files (png, jpg, gif, bmp, tiff, webp, ico, psd) from being committed outside of ``flexmeasures/ui/static/`` and ``documentation/``, to protect the git history from binary bloat; screenshots belong in the ``FlexMeasures/screenshots`` repo instead [see `PR #2315 <https://www.github.com/FlexMeasures/flexmeasures/pull/2315>`_]

flexmeasures/data/__init__.py

@@ -19,6 +19,26 @@
 
 ma: Marshmallow = Marshmallow()
 
+DB_COMMANDS = frozenset(db_cli_group.commands.keys())
+DB_OPS_COMMANDS = frozenset(("dump", "reset", "restore"))
+
+
+def is_running_database_command() -> bool:
+    """Return whether this process is running a database maintenance command."""
+    args = sys.argv[1:]
+    return _has_command_pair(args, "db", DB_COMMANDS) or _has_command_pair(
+        args, "db-ops", DB_OPS_COMMANDS
+    )
+
+
+def _has_command_pair(
+    args: list[str], command_group: str, commands: frozenset[str]
+) -> bool:
+    return any(
+        arg == command_group and i + 1 < len(args) and args[i + 1] in commands
+        for i, arg in enumerate(args)
+    )
+
 
 def _is_running_db_upgrade_command() -> bool:
     """Return whether this process is already running the Alembic upgrade command."""

flexmeasures/data/scripts/data_gen.py

@@ -20,6 +20,18 @@
 from flexmeasures.data.utils import TEMPLATE_COPY_GUIDANCE_PREFIX
 
 
+def _skip_default_data_creation_for_database_command(default_data_name: str) -> bool:
+    """Return whether implicit default data creation should be skipped."""
+    from flexmeasures.data import is_running_database_command
+
+    if is_running_database_command():
+        click.echo(
+            f"Skipping default {default_data_name} creation during database maintenance command."
+        )
+        return True
+    return False
+
+
 def add_default_data_sources(db: SQLAlchemy):
     for source_name, source_type in (
         ("Seita", "demo script"),
@@ -234,6 +246,9 @@ def provision_default_template_assets(db: SQLAlchemy):
     This currently provisions the single-asset starter templates which are
     intended to show up in the asset copy UI.
     """
+    if _skip_default_data_creation_for_database_command("template asset"):
+        return
+
     asset_types = add_default_asset_types(db)
 
     # Battery
@@ -372,6 +387,9 @@ def populate_initial_structure(db: SQLAlchemy):
     """
     Add initially useful structural data.
     """
+    if _skip_default_data_creation_for_database_command("initial structure"):
+        return
+
     click.echo("Populating the database %s with structural data ..." % db.engine)
     add_default_data_sources(db)
     add_default_user_roles(db)

flexmeasures/data/tests/test_template_assets.py

@@ -2,7 +2,16 @@
 
 from flexmeasures.data.models.generic_assets import GenericAsset
 from flexmeasures.data.models.time_series import Sensor
-from flexmeasures.data.scripts.data_gen import provision_default_template_assets
+from flexmeasures.data.scripts.data_gen import (
+    populate_initial_structure,
+    provision_default_template_assets,
+)
+
+
+class _FailingDb:
+    @property
+    def session(self):
+        raise AssertionError("Default data creation should not touch the database.")
 
 
 def test_provision_default_template_assets_creates_single_asset_templates(
@@ -71,3 +80,28 @@ def test_provision_default_template_assets_is_idempotent(fresh_db):
         fresh_db.session.scalar(select(func.count()).select_from(Sensor))
         == sensor_count
     )
+
+
+def test_initial_structure_creation_skips_database_commands(monkeypatch):
+    monkeypatch.setattr("flexmeasures.data.is_running_database_command", lambda: True)
+
+    populate_initial_structure(_FailingDb())
+
+
+def test_template_asset_provisioning_skips_database_commands(fresh_db, monkeypatch):
+    monkeypatch.setattr("flexmeasures.data.is_running_database_command", lambda: True)
+    asset_count = fresh_db.session.scalar(
+        select(func.count()).select_from(GenericAsset)
+    )
+    sensor_count = fresh_db.session.scalar(select(func.count()).select_from(Sensor))
+
+    provision_default_template_assets(fresh_db)
+
+    assert (
+        fresh_db.session.scalar(select(func.count()).select_from(GenericAsset))
+        == asset_count
+    )
+    assert (
+        fresh_db.session.scalar(select(func.count()).select_from(Sensor))
+        == sensor_count
+    )

flexmeasures/data/tests/test_utils.py

@@ -7,6 +7,7 @@
 from flexmeasures.data import db, register_at
 from flexmeasures.data.utils import (
     DatabaseSchemaRevisionStatus,
+    database_schema_has_revision,
     format_database_schema_revision_status,
     get_database_schema_revision_status,
 )
@@ -40,6 +41,25 @@ def get_heads(self) -> tuple[str, ...]:
         return self._heads
 
 
+class _DummyRevision:
+    def __init__(self, revision: str):
+        self.revision = revision
+
+
+class _DummyRevisionMap:
+    def __init__(self, revisions: tuple[str, ...]):
+        self._revisions = revisions
+
+    def iterate_revisions(self, *args, **kwargs):
+        return (_DummyRevision(revision) for revision in self._revisions)
+
+
+class _DummyScriptDirectoryWithRevisionMap(_DummyScriptDirectory):
+    def __init__(self, heads: tuple[str, ...], revisions: tuple[str, ...]):
+        super().__init__(heads)
+        self.revision_map = _DummyRevisionMap(revisions)
+
+
 def test_schema_mismatch_log_record_is_deduplicated(
     app, clean_redis, monkeypatch, caplog
 ):
@@ -181,3 +201,39 @@ def raise_operational_error():
     assert status.expected_heads == ("head-a",)
     assert status.inspection_error is not None
     assert status.is_migrated_to_head is False
+
+
+def test_database_schema_has_revision_when_revision_is_in_current_history(
+    app, monkeypatch
+):
+    monkeypatch.setattr(db.engine, "connect", lambda: _DummyConnection())
+    monkeypatch.setattr(
+        "flexmeasures.data.utils.MigrationContext.configure",
+        lambda connection: _DummyMigrationContext(("head-a",)),
+    )
+    monkeypatch.setattr(
+        "flexmeasures.data.utils.ScriptDirectory.from_config",
+        lambda config: _DummyScriptDirectoryWithRevisionMap(
+            heads=("head-a",), revisions=("head-a", "required-a")
+        ),
+    )
+
+    assert database_schema_has_revision(app, "required-a") is True
+
+
+def test_database_schema_has_revision_false_when_revision_is_not_in_current_history(
+    app, monkeypatch
+):
+    monkeypatch.setattr(db.engine, "connect", lambda: _DummyConnection())
+    monkeypatch.setattr(
+        "flexmeasures.data.utils.MigrationContext.configure",
+        lambda connection: _DummyMigrationContext(("old-a",)),
+    )
+    monkeypatch.setattr(
+        "flexmeasures.data.utils.ScriptDirectory.from_config",
+        lambda config: _DummyScriptDirectoryWithRevisionMap(
+            heads=("head-a",), revisions=("old-a",)
+        ),
+    )
+
+    assert database_schema_has_revision(app, "required-a") is False

flexmeasures/data/utils.py

@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 from alembic.config import Config as AlembicConfig
+from alembic.script.revision import RevisionError
 from alembic.runtime.migration import MigrationContext
 from alembic.script import ScriptDirectory
 from dataclasses import dataclass
@@ -86,6 +87,38 @@ def get_database_schema_revision_status(app) -> DatabaseSchemaRevisionStatus:
     )
 
 
+def database_schema_has_revision(app, required_revision: str) -> bool:
+    """Return whether the connected database includes a specific Alembic revision."""
+    revision_status = get_database_schema_revision_status(app)
+    if (
+        revision_status.inspection_error is not None
+        or not revision_status.current_heads
+    ):
+        return False
+
+    migrate_extension = app.extensions.get("migrate")
+    if migrate_extension is None:
+        return False
+
+    alembic_config = AlembicConfig()
+    alembic_config.set_main_option("script_location", migrate_extension.directory)
+    script = ScriptDirectory.from_config(alembic_config)
+
+    for current_head in revision_status.current_heads:
+        try:
+            revisions = script.revision_map.iterate_revisions(
+                current_head,
+                required_revision,
+                inclusive=True,
+                assert_relative_length=False,
+            )
+        except RevisionError:
+            continue
+        if any(revision.revision == required_revision for revision in revisions):
+            return True
+    return False
+
+
 def format_database_schema_revision_status(
     status: DatabaseSchemaRevisionStatus,
 ) -> str:

flexmeasures/utils/app_utils.py

@@ -21,19 +21,29 @@
     _sentry_filter_notfound,
 )
 
+TEMPLATE_ASSETS_REQUIRED_MIGRATION = "4b0f2e9c1a6d"
+
 
 def provision_default_template_assets_on_startup(app: Flask) -> None:
     """Provision starter template assets when startup settings and schema allow it."""
+    from flexmeasures.data import is_running_database_command
+
+    if is_running_database_command():
+        return
+
     if (
         not app.config.get("FLEXMEASURES_CREATE_TEMPLATE_ASSETS_ON_STARTUP", False)
         or app.testing
         or app.config.get("FLEXMEASURES_ENV") == "documentation"
     ):
         return
 
-    if not getattr(app, "database_schema_is_migrated_to_head", True):
+    from flexmeasures.data.utils import database_schema_has_revision
+
+    if not database_schema_has_revision(app, TEMPLATE_ASSETS_REQUIRED_MIGRATION):
         app.logger.info(
-            "Skipping startup template provisioning because the database schema is not at the Alembic head revision yet."
+            "Skipping startup template provisioning because the database schema is missing the required migration "
+            f"{TEMPLATE_ASSETS_REQUIRED_MIGRATION}."
         )
         return
 

flexmeasures/utils/tests/test_app_utils.py

@@ -15,6 +15,7 @@
 from flexmeasures.data import (
     _is_running_db_upgrade_command,
     _schema_mismatch_deduplication_key,
+    is_running_database_command,
 )
 from flexmeasures.data.utils import DatabaseSchemaRevisionStatus
 from flexmeasures.utils.app_utils import (
@@ -431,7 +432,10 @@ def fail_if_called(db):
     monkeypatch.setitem(
         app.config, "FLEXMEASURES_CREATE_TEMPLATE_ASSETS_ON_STARTUP", True
     )
-    monkeypatch.setattr(app, "database_schema_is_migrated_to_head", False)
+    monkeypatch.setattr("flexmeasures.data.is_running_database_command", lambda: False)
+    monkeypatch.setattr(
+        "flexmeasures.data.utils.database_schema_has_revision", lambda app, rev: False
+    )
     monkeypatch.setattr(
         "flexmeasures.data.scripts.data_gen.provision_default_template_assets",
         fail_if_called,
@@ -443,6 +447,57 @@ def fail_if_called(db):
     assert "Skipping startup template provisioning" in caplog.text
 
 
+def test_provision_default_template_assets_on_startup_skips_database_commands(
+    app, monkeypatch
+):
+    def fail_if_called(db):
+        raise AssertionError("Template provisioning should not run.")
+
+    monkeypatch.setattr(app, "testing", False)
+    monkeypatch.setitem(app.config, "FLEXMEASURES_ENV", "production")
+    monkeypatch.setitem(
+        app.config, "FLEXMEASURES_CREATE_TEMPLATE_ASSETS_ON_STARTUP", True
+    )
+    monkeypatch.setattr("flexmeasures.data.is_running_database_command", lambda: True)
+    monkeypatch.setattr(
+        "flexmeasures.data.scripts.data_gen.provision_default_template_assets",
+        fail_if_called,
+    )
+
+    provision_default_template_assets_on_startup(app)
+
+
+def test_is_running_database_command(monkeypatch):
+    monkeypatch.setattr(
+        "sys.argv",
+        [
+            "/path/to/flexmeasures",
+            "--custom-option-with-value",
+            "some-value",
+            "db-ops",
+            "restore",
+        ],
+    )
+
+    assert is_running_database_command() is True
+
+
+def test_is_running_database_command_false_for_other_cli_groups(monkeypatch):
+    monkeypatch.setattr(
+        "sys.argv", ["/path/to/flexmeasures", "add", "initial-structure"]
+    )
+
+    assert is_running_database_command() is False
+
+
+def test_is_running_database_command_false_for_option_value(monkeypatch):
+    monkeypatch.setattr(
+        "sys.argv", ["/path/to/flexmeasures", "add", "toy-account", "--name", "db"]
+    )
+
+    assert is_running_database_command() is False
+
+
 def test_is_running_db_upgrade_command(monkeypatch):
     monkeypatch.setattr("sys.argv", ["/path/to/flexmeasures", "db", "upgrade", "--sql"])