diff --git a/CHANGES.rst b/CHANGES.rst index 9c0bbf5..8b642c7 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -1,7 +1,7 @@ .. This file is part of Invenio. Copyright (C) 2015-2024 CERN. - Copyright (C) 2024-2025 Graz University of Technology. + Copyright (C) 2024-2026 Graz University of Technology. Invenio is free software; you can redistribute it and/or modify it under the terms of the MIT License; see LICENSE file for more details. @@ -9,6 +9,30 @@ Changes ======= +Version v2.3.0 (released 2026-02-24) + +- fix(config): use UTC for PostgreSQL +- feat(alembic): set lock_timeout with retry on migration connections +- docs(sphinx): ignore unresolved Flask-Alembic type refs +- tests(utc): add unit tests for UTCDateTime + +Version v2.2.1 (released 2026-01-30) + +- fix(setup): pin sqlalchemy-continuum + +Version v2.2.0 (released 2026-01-27) + +- chore(black): apply changes for black>=26 +- fix: str of datetime +- change(utc): use always timezone.utc +- fix: docs reference target not found +- db: add Timestamp class +- UTCDateTime: handle more cases +- db: fix warning +- fix: docs reference target not found +- shared: add UTCDateTime column type +- chore: add nitpick_ignore to fix CI + Version v2.1.2 (released 2025-11-17) - fix(setup): pin Flask-Alembic due to breaking changes in v3.2.0 causing diff --git a/docs/conf.py b/docs/conf.py index 0080733..3409a0e 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -9,7 +9,6 @@ """Sphinx configuration.""" - from invenio_db import __version__ # -- General configuration ------------------------------------------------ @@ -308,12 +307,14 @@ # Example configuration for intersphinx: refer to the Python standard library. -intersphinx_mapping = {"python": ("https://docs.python.org/3", None)} +intersphinx_mapping = { + "python": ("https://docs.python.org/3", None), + "sqlalchemy": ("https://docs.sqlalchemy.org/en/latest/", None), +} # Autodoc configuraton. autoclass_content = "both" - # To address :1:py:class reference target not found # (better ideas welcomed) nitpick_ignore = [ @@ -323,4 +324,8 @@ ("py:class", "dict[str"), ("py:class", "flask_sqlalchemy.query.Query"), ("py:class", "flask_sqlalchemy.extension._FSA_MCT"), + ("py:class", "TypeDecorator"), + ("py:class", "sqlalchemy.sql.sqltypes.DateTime"), + ("py:class", "ExternalType"), + ("py:class", "UserDefinedType"), ] diff --git a/invenio_db/__init__.py b/invenio_db/__init__.py index 8457d4b..1e2af7f 100644 --- a/invenio_db/__init__.py +++ b/invenio_db/__init__.py @@ -2,7 +2,7 @@ # # This file is part of Invenio. # Copyright (C) 2015-2024 CERN. -# Copyright (C) 2024 Graz University of Technology. +# Copyright (C) 2024-2026 Graz University of Technology. # # Invenio is free software; you can redistribute it and/or modify it # under the terms of the MIT License; see LICENSE file for more details. @@ -92,7 +92,7 @@ class User(db.Model): from .ext import InvenioDB from .shared import db -__version__ = "2.1.2" +__version__ = "2.3.0" __all__ = ( "__version__", diff --git a/invenio_db/ext.py b/invenio_db/ext.py index 189bafd..88f723e 100644 --- a/invenio_db/ext.py +++ b/invenio_db/ext.py @@ -3,34 +3,104 @@ # This file is part of Invenio. # Copyright (C) 2015-2018 CERN. # Copyright (C) 2022 RERO. -# Copyright (C) 2022 Graz University of Technology. +# Copyright (C) 2022-2026 Graz University of Technology. # # Invenio is free software; you can redistribute it and/or modify it # under the terms of the MIT License; see LICENSE file for more details. """Database management for Invenio.""" +import logging import os +import random +import time from importlib.metadata import PackageNotFoundError from importlib.metadata import version as package_version from importlib.resources import files import sqlalchemy as sa +from flask import current_app from flask_alembic import Alembic from invenio_base.utils import entry_points +from sqlalchemy.exc import OperationalError from sqlalchemy_utils.functions import get_class_by_table from .cli import db as db_cmd from .shared import db from .utils import versioning_models_registered +logger = logging.getLogger(__name__) + + +class InvenioAlembic(Alembic): + """Alembic with PostgreSQL lock_timeout and retry for safe migrations. + + Sets lock_timeout on migration connections so DDL statements fail fast + instead of blocking reads indefinitely while waiting for exclusive locks. + On lock timeout, retries with exponential backoff. Already-applied + migrations are skipped on retry (transaction_per_migration stamps each + step). + + See https://postgres.ai/blog/20210923-zero-downtime-postgres-schema-migrations-lock-timeout-and-retries + + Configuration: + + - ``DB_MIGRATION_LOCK_TIMEOUT``: PostgreSQL lock_timeout value + (default ``"1s"``). Set to ``"0"`` to disable. + - ``DB_MIGRATION_LOCK_TIMEOUT_RETRIES``: number of retries on lock + timeout (default ``5``). + """ + + def __init__(self, *args, **kwargs): + """Initialize InvenioAlembic.""" + super().__init__(*args, **kwargs) + + def _set_lock_timeout(self): + """Set lock_timeout on all PostgreSQL migration connections.""" + for ctx in self.migration_contexts.values(): + if ( + ctx.connection is not None + and ctx.connection.dialect.name == "postgresql" + ): + timeout = current_app.config.get("DB_MIGRATION_LOCK_TIMEOUT", "1s") + ctx.connection.execute( + sa.text("SELECT set_config('lock_timeout', :val, false)"), + {"val": timeout}, + ) + + def run_migrations(self, fn, **kwargs): + """Run migrations with lock_timeout and retry on lock failure.""" + max_retries = current_app.config.get("DB_MIGRATION_LOCK_TIMEOUT_RETRIES", 5) + + for attempt in range(max_retries + 1): + self._set_lock_timeout() + try: + super().run_migrations(fn, **kwargs) + return + except OperationalError as e: + is_lock_timeout = hasattr(e.orig, "pgcode") and e.orig.pgcode == "55P03" + if not is_lock_timeout or attempt >= max_retries: + raise + # Exponential backoff with jitter + delay = min(30, 0.5 * 2**attempt) * (0.5 + random.random() * 0.5) + logger.warning( + "Migration lock timeout, retrying in %.1fs (%d/%d)", + delay, + attempt + 1, + max_retries, + ) + time.sleep(delay) + # Clear cached contexts — connection is dead after the error. + # Next access to migration_contexts creates fresh connections. + self._get_cache().clear() + class InvenioDB(object): """Invenio database extension.""" def __init__(self, app=None, **kwargs): """Extension initialization.""" - self.alembic = Alembic(run_mkdir=False, command_name="alembic") + self.alembic = InvenioAlembic(run_mkdir=False, command_name="alembic") if app: self.init_app(app, **kwargs) @@ -76,6 +146,13 @@ def init_db(self, app, entry_point_group="invenio_db.models", **kwargs): app.config.setdefault("SQLALCHEMY_ECHO", False) # Needed for before/after_flush/commit/rollback events app.config.setdefault("SQLALCHEMY_TRACK_MODIFICATIONS", True) + app.config.setdefault( + "SQLALCHEMY_ENGINE_OPTIONS", + # Ensure the database is using the UTC timezone for interpreting timestamps (Postgres only). + # This overrides any default setting (e.g. in postgresql.conf). Invenio expects the DB to receive + # and provide UTC timestamps in all cases, so it's important that this doesn't get changed. + {"connect_args": {"options": "-c timezone=UTC"}}, + ) # Initialize Flask-SQLAlchemy extension. database = kwargs.get("db", db) diff --git a/invenio_db/proxies.py b/invenio_db/proxies.py index 395ce72..a4e20af 100644 --- a/invenio_db/proxies.py +++ b/invenio_db/proxies.py @@ -1,14 +1,13 @@ # -*- coding: utf-8 -*- # # This file is part of Invenio. -# Copyright (C) 2022 Graz University of Technology. +# Copyright (C) 2022-2026 Graz University of Technology. # # Invenio is free software; you can redistribute it and/or modify it # under the terms of the MIT License; see LICENSE file for more details. """Helper proxy to the state object.""" - from flask import current_app from werkzeug.local import LocalProxy diff --git a/invenio_db/shared.py b/invenio_db/shared.py index 43acd9f..3101012 100644 --- a/invenio_db/shared.py +++ b/invenio_db/shared.py @@ -2,16 +2,20 @@ # # This file is part of Invenio. # Copyright (C) 2015-2018 CERN. +# Copyright (C) 2024-2026 Graz University of Technology. # # Invenio is free software; you can redistribute it and/or modify it # under the terms of the MIT License; see LICENSE file for more details. """Shared database object for Invenio.""" +from datetime import datetime, timezone + from flask_sqlalchemy import SQLAlchemy as FlaskSQLAlchemy -from sqlalchemy import MetaData, event, util +from sqlalchemy import Column, MetaData, event, util from sqlalchemy.engine import Engine from sqlalchemy.sql import text +from sqlalchemy.types import DateTime, TypeDecorator from werkzeug.local import LocalProxy NAMING_CONVENTION = util.immutabledict( @@ -29,9 +33,102 @@ """Default database metadata object holding associated schema constructs.""" +class UTCDateTime(TypeDecorator): + """Custom UTC datetime type.""" + + impl = DateTime(timezone=True) + + # todo: should be discussed, but has to be set explicitly to remove warning + cache_ok = False + + def process_bind_param(self, value, dialect): + """Process value storing into database.""" + if value is None: + return value + + if isinstance(value, str): + if " " in value: + value = value.replace(" ", "T") + value = datetime.strptime(value[0:19], "%Y-%m-%dT%H:%M:%S") + + if not isinstance(value, datetime): + msg = f"ERROR: value: {value} is not of type datetime, instead of type: {type(value)}" + raise ValueError(msg) + + if value.tzinfo not in (None, timezone.utc): + msg = f"Error: value: {value}, tzinfo: {value.tzinfo} doesn't have a tzinfo of None or timezone.utc." + raise ValueError(msg) + + return value.replace(tzinfo=timezone.utc) + + def process_result_value(self, value, dialect): + """Process value retrieving from database.""" + if value is None: + return None + + if not isinstance(value, datetime): + msg = f"ERROR: value: {value} is not of type datetime." + raise ValueError(msg) + + if value.tzinfo not in (None, timezone.utc): + msg = ( + f"Error: value: {value} doesn't have a tzinfo of None or timezone.utc." + ) + raise ValueError(msg) + + return value.replace(tzinfo=timezone.utc) + + +class Timestamp: + """Adds `created` and `updated` columns to a derived declarative model. + + The `created` column is handled through a default and the `updated` + column is handled through a `before_update` event that propagates + for all derived declarative models. + + :: + + from invenio_db import db + class SomeModel(Base, db.Timestamp): + __tablename__ = "somemodel" + id = sa.Column(sa.Integer, primary_key=True) + """ + + created = Column( + UTCDateTime, + default=lambda: datetime.now(tz=timezone.utc), + nullable=False, + ) + updated = Column( + UTCDateTime, + default=lambda: datetime.now(tz=timezone.utc), + nullable=False, + ) + + +@event.listens_for(Timestamp, "before_update", propagate=True) +def timestamp_before_update(mapper, connection, target): + """Update timestamp on before_update event. + + When a model with a timestamp is updated; force update the updated + timestamp. + """ + target.updated = datetime.now(tz=timezone.utc) + + class SQLAlchemy(FlaskSQLAlchemy): """Implement or overide extension methods.""" + def __getattr__(self, name): + """Get attr.""" + if name == "UTCDateTime": + return UTCDateTime + + if name == "Timestamp": + return Timestamp + + return super().__getattr__(name) + def apply_driver_hacks(self, app, sa_url, options): """Call before engine creation.""" # Don't forget to apply hacks defined on parent object. diff --git a/invenio_db/utils.py b/invenio_db/utils.py index ef8925c..fc25490 100644 --- a/invenio_db/utils.py +++ b/invenio_db/utils.py @@ -2,7 +2,8 @@ # # This file is part of Invenio. # Copyright (C) 2017-2018 CERN. -# Copyright (C) 2022 Graz University of Technology. +# Copyright (C) 2022-2026 Graz University of Technology. +# Copyright (C) 2026 University of Münster. # # Invenio is free software; you can redistribute it and/or modify it # under the terms of the MIT License; see LICENSE file for more details. @@ -10,6 +11,10 @@ """Invenio-DB utility functions.""" +from alembic.migration import MigrationContext +from functools import partial + +from alembic import op from flask import current_app from sqlalchemy import inspect @@ -61,13 +66,15 @@ def rebuild_encrypted_properties(old_key, model, properties, db=_db): db.session.commit() -def create_alembic_version_table(): +def create_alembic_version_table(db=_db): """Create alembic_version table.""" alembic = current_app.extensions["invenio-db"].alembic - if not alembic.migration_context._has_version_table(): - alembic.migration_context._ensure_version_table() - for head in alembic.script_directory.revision_map._real_heads: - alembic.migration_context.stamp(alembic.script_directory, head) + with db.engine.begin() as connection: + context = MigrationContext.configure(connection) + if not context._has_version_table(): + context._ensure_version_table() + all_heads = alembic.script_directory.revision_map._real_heads + context.stamp(alembic.script_directory, tuple(all_heads)) def drop_alembic_version_table(): @@ -128,3 +135,31 @@ def has_table(engine, table): except AttributeError: # SQLAlchemy <1.4 return engine.has_table(table) + + +def update_table_columns_column_type( + table_name, column_name, to_type=None, existing_type=None, existing_nullable=None +): + """Update column type.""" + op.alter_column( + table_name, + column_name, + type_=to_type(), + existing_type=existing_type(), + existing_nullable=existing_nullable, + ) + + +update_table_columns_column_type_to_utc_datetime = partial( + update_table_columns_column_type, + to_type=_db.UTCDateTime, + existing_type=_db.DateTime, + existing_nullable=True, +) + +update_table_columns_column_type_to_datetime = partial( + update_table_columns_column_type, + to_type=_db.DateTime, + existing_type=_db.UTCDateTime, + existing_nullable=True, +) diff --git a/setup.cfg b/setup.cfg index 3249cf0..f4cbca1 100644 --- a/setup.cfg +++ b/setup.cfg @@ -4,7 +4,7 @@ # Copyright (C) 2015-2022 CERN. # Copyright (C) 2021 Northwestern University. # Copyright (C) 2022 RERO. -# Copyright (C) 2022-2025 Graz University of Technology. +# Copyright (C) 2022-2026 Graz University of Technology. # # Invenio is free software; you can redistribute it and/or modify it # under the terms of the MIT License; see LICENSE file for more details. @@ -34,7 +34,7 @@ install_requires = Flask-Alembic>=3.0.0,<3.2.0 Flask-SQLAlchemy>=3.0 invenio-base>=2.3.0,<3.0.0 - SQLAlchemy-Continuum>=1.3.12 + SQLAlchemy-Continuum>=1.3.12,<1.6.0 SQLAlchemy-Utils>=0.33.1,<0.42.0 SQLAlchemy[asyncio]>=2.0.0 diff --git a/tests/mocks.py b/tests/mocks.py index 8eca052..a28f042 100644 --- a/tests/mocks.py +++ b/tests/mocks.py @@ -1,14 +1,13 @@ # -*- coding: utf-8 -*- # # This file is part of Invenio. -# Copyright (C) 2023 Graz University of Technology. +# Copyright (C) 2023-2026 Graz University of Technology. # # Invenio is free software; you can redistribute it and/or modify it # under the terms of the MIT License; see LICENSE file for more details. """Test database integration layer.""" - from importlib_metadata import EntryPoint from werkzeug.utils import import_string diff --git a/tests/test_db.py b/tests/test_db.py index 5e52e04..0d079a5 100644 --- a/tests/test_db.py +++ b/tests/test_db.py @@ -20,7 +20,7 @@ from sqlalchemy import inspect from sqlalchemy.exc import IntegrityError from sqlalchemy_continuum import VersioningManager, remove_versioning -from sqlalchemy_utils.functions import create_database, drop_database +from sqlalchemy_utils.functions import create_database, database_exists, drop_database from invenio_db import InvenioDB from invenio_db.cli import db as db_cmd @@ -252,87 +252,37 @@ class Demo(db.Model): @patch("importlib.metadata.entry_points", _mock_entry_points("invenio_db.models")) def test_entry_points(db, app): """Test entrypoints loading.""" - InvenioDB(app, db=db) - - runner = app.test_cli_runner() + ext = InvenioDB(app, db=db) assert len(db.metadata.tables) == 2 - # Test merging a base another file. - with runner.isolated_filesystem(): - - # show help - result = runner.invoke(db_cmd, []) - click_help_result_code = ( - 0 if importlib.metadata.version("click") < "8.2.0" else 2 - ) - assert result.exit_code == click_help_result_code - - result = runner.invoke(db_cmd, ["init"]) - assert result.exit_code == 0 - - result = runner.invoke(db_cmd, ["destroy", "--yes-i-know"]) - assert result.exit_code == 0 - - result = runner.invoke(db_cmd, ["init"]) - assert result.exit_code == 0 - - result = runner.invoke(db_cmd, ["create", "-v"]) - assert result.exit_code == 0 - - result = runner.invoke(db_cmd, ["destroy", "--yes-i-know"]) - assert result.exit_code == 0 - - result = runner.invoke(db_cmd, ["init"]) - assert result.exit_code == 0 - - result = runner.invoke(db_cmd, ["create", "-v"]) - assert result.exit_code == 1 - - result = runner.invoke(db_cmd, ["create", "-v"]) - assert result.exit_code == 1 - - result = runner.invoke(db_cmd, ["drop", "-v", "--yes-i-know"]) - assert result.exit_code == 0 - - result = runner.invoke(db_cmd, ["drop"]) - assert result.exit_code == 1 - - result = runner.invoke(db_cmd, ["drop", "-v", "--yes-i-know"]) - assert result.exit_code == 0 - - result = runner.invoke(db_cmd, ["drop", "create"]) - assert result.exit_code == 1 - - result = runner.invoke(db_cmd, ["drop", "--yes-i-know", "create"]) - assert result.exit_code == 1 - - result = runner.invoke(db_cmd, ["destroy", "--yes-i-know"]) - assert result.exit_code == 0 - - result = runner.invoke(db_cmd, ["init"]) - assert result.exit_code == 0 - - result = runner.invoke(db_cmd, ["destroy", "--yes-i-know"]) - assert result.exit_code == 0 + with app.app_context(): + db_url = str(db.engine.url.render_as_string(hide_password=False)) - result = runner.invoke(db_cmd, ["init"]) - assert result.exit_code == 0 + # Test 'init' + ext.alembic.mkdir() - result = runner.invoke(db_cmd, ["destroy", "--yes-i-know"]) - assert result.exit_code == 0 + drop_database(db_url) + create_database(db_url) + assert database_exists(db_url) - result = runner.invoke(db_cmd, ["init"]) - assert result.exit_code == 0 + # 'db create' -> db.create_all() + ext.alembic.stamp() + db.create_all() + ext.alembic.stamp() + assert has_table(db.engine, "alembic_version") - result = runner.invoke(db_cmd, ["drop", "-v", "--yes-i-know"]) - assert result.exit_code == 1 + # 'db drop' -> db.drop_all() + drop_alembic_version_table() + db.drop_all() + drop_alembic_version_table() + assert len(inspect(db.engine).get_table_names()) == 0 - result = runner.invoke(db_cmd, ["create", "-v"]) - assert result.exit_code == 1 + db.create_all() + ext.alembic.stamp() + assert has_table(db.engine, "alembic_version") - result = runner.invoke(db_cmd, ["drop", "-v", "--yes-i-know"]) - assert result.exit_code == 0 + db.drop_all() + drop_alembic_version_table() + assert len(inspect(db.engine).get_table_names()) == 0 def test_local_proxy(app, db): @@ -377,10 +327,9 @@ def test_db_create_alembic_upgrade(app, db): db.drop_all() runner = app.test_cli_runner() - # Check that 'db create' creates the same schema as - # 'alembic upgrade'. - result = runner.invoke(db_cmd, ["create", "-v"]) - assert result.exit_code == 0 + + db.create_all() + ext.alembic.stamp() assert has_table(db.engine, "transaction") assert ext.alembic.migration_context._has_version_table() diff --git a/tests/test_lock_timeout.py b/tests/test_lock_timeout.py new file mode 100644 index 0000000..c6940d9 --- /dev/null +++ b/tests/test_lock_timeout.py @@ -0,0 +1,151 @@ +# -*- coding: utf-8 -*- +# +# This file is part of Invenio. +# Copyright (C) 2024 CERN. +# +# Invenio is free software; you can redistribute it and/or modify it +# under the terms of the MIT License; see LICENSE file for more details. + +"""Test migration lock_timeout support.""" + +import threading +import time + +import pytest +import sqlalchemy as sa +from sqlalchemy.exc import OperationalError +from utils import requires_postgresql + +from invenio_db import InvenioDB + + +@requires_postgresql +def test_alembic_lock_timeout(db, app): + """Test that lock_timeout is set on migration connections.""" + ext = InvenioDB(app, entry_point_group=False, db=db) + + with app.app_context(): + ext.alembic.upgrade() + ctx = ext.alembic.migration_context + result = ctx.connection.execute(sa.text("SHOW lock_timeout")).scalar() + assert result == "1s" + + +@requires_postgresql +def test_alembic_lock_timeout_custom(db, app): + """Test that lock_timeout is configurable.""" + app.config["DB_MIGRATION_LOCK_TIMEOUT"] = "10s" + ext = InvenioDB(app, entry_point_group=False, db=db) + + with app.app_context(): + ext.alembic.upgrade() + ctx = ext.alembic.migration_context + result = ctx.connection.execute(sa.text("SHOW lock_timeout")).scalar() + assert result == "10s" + + +@requires_postgresql +def test_alembic_lock_timeout_prevents_long_waits(db, app): + """Test that DDL fails fast when a table is locked by another transaction. + + Simulates a real scenario: a long-running query holds ACCESS SHARE on a + table, and an ALTER TABLE (needing ACCESS EXCLUSIVE) hits lock_timeout + instead of blocking indefinitely. + """ + app.config["DB_MIGRATION_LOCK_TIMEOUT"] = "1s" + ext = InvenioDB(app, entry_point_group=False, db=db) + + with app.app_context(): + ext.alembic.upgrade() + engine = db.engine + + with engine.connect() as conn: + conn.execute( + sa.text("CREATE TABLE IF NOT EXISTS _lock_timeout_test (id int)") + ) + conn.commit() + + try: + lock_held = threading.Event() + ddl_done = threading.Event() + + def hold_lock(): + """Simulate a long-running transaction holding ACCESS SHARE.""" + with engine.connect() as conn: + with conn.begin(): + conn.execute(sa.text("SELECT * FROM _lock_timeout_test")) + lock_held.set() + ddl_done.wait(timeout=10) + + t = threading.Thread(target=hold_lock) + t.start() + lock_held.wait(timeout=5) + + # The migration connection has lock_timeout set — ALTER TABLE + # needs ACCESS EXCLUSIVE which conflicts with ACCESS SHARE, + # so this must fail within ~1s instead of blocking. + ctx = ext.alembic.migration_context + with pytest.raises(OperationalError, match="lock timeout"): + ctx.connection.execute( + sa.text("ALTER TABLE _lock_timeout_test ADD COLUMN x int") + ) + + ddl_done.set() + t.join(timeout=5) + finally: + with engine.connect() as conn: + conn.execute(sa.text("DROP TABLE IF EXISTS _lock_timeout_test")) + conn.commit() + + +@requires_postgresql +def test_alembic_lock_timeout_retry(db, app, caplog): + """Test that migrations retry on lock timeout and eventually succeed. + + Strategy: apply all migrations, downgrade to base, lock alembic_version + (which every migration step must write to), then upgrade with retries. + The first attempts fail at the version stamp, but once the lock is + released (~2s) a retry succeeds. PostgreSQL transactional DDL ensures + rolled-back steps leave no artifacts. + """ + app.config["DB_MIGRATION_LOCK_TIMEOUT"] = "500ms" + app.config["DB_MIGRATION_LOCK_TIMEOUT_RETRIES"] = 10 + ext = InvenioDB(app, entry_point_group=False, db=db) + + with app.app_context(): + # Apply all, then roll back to base so upgrade has real work to do. + ext.alembic.upgrade() + ext.alembic.downgrade(target="96e796392533") + + engine = db.engine + lock_held = threading.Event() + + def hold_lock_briefly(): + """Lock alembic_version exclusively for ~2s.""" + with engine.connect() as conn: + with conn.begin(): + conn.execute( + sa.text("LOCK TABLE alembic_version IN ACCESS EXCLUSIVE MODE") + ) + lock_held.set() + time.sleep(2) + + t = threading.Thread(target=hold_lock_briefly) + t.start() + lock_held.wait(timeout=5) + + # upgrade() goes through run_migrations → lock_timeout on + # alembic_version INSERT → retry → eventually succeeds. + with caplog.at_level("WARNING", logger="invenio_db.ext"): + ext.alembic.upgrade() + + t.join(timeout=10) + + # Verify retries actually happened. + retry_msgs = [r for r in caplog.records if "lock timeout" in r.message] + assert len(retry_msgs) >= 1 + + # Verify migrations actually applied. + heads = {s.revision for s in ext.alembic.current()} + expected = set(ext.alembic.script_directory.get_heads()) + assert heads == expected diff --git a/tests/test_utc.py b/tests/test_utc.py new file mode 100644 index 0000000..e14a593 --- /dev/null +++ b/tests/test_utc.py @@ -0,0 +1,176 @@ +# -*- coding: utf-8 -*- +# +# This file is part of Invenio. +# Copyright (C) 2026 CERN. +# +# Invenio is free software; you can redistribute it and/or modify it +# under the terms of the MIT License; see LICENSE file for more details. + +"""Test UTC and datetime column support for migrations. + +The aim is that UTCDateTime should work regardless of whether the corresponding +`timestamp` -> `timestamptz` migrations have been run, **as long as the database +is using UTC**. The `SQLALCHEMY_ENGINE_OPTIONS` config specified in this module +is supposed to force PostgreSQL to always use UTC, so these tests should check +behaviour for both UTC and non-UTC zones. +""" + +from datetime import datetime, timedelta, timezone + +import pytest +import sqlalchemy as sa +from utils import requires_postgresql + +from invenio_db.ext import InvenioDB +from invenio_db.shared import SQLAlchemy, UTCDateTime + + +@requires_postgresql +def test_timestamp_no_tz(app, db: SQLAlchemy): + InvenioDB(app, entry_point_group=False, db=db) + + class TestTimestampNoTZ(db.Model): + __tablename__ = "_test_timestamp_no_tz" + id = sa.Column(sa.Integer, primary_key=True) + t = sa.Column(UTCDateTime) + + with app.app_context(): + # Create the table manually so we can override the `t` column type to what it would be pre-migration. + db.session.execute( + sa.text( + "CREATE TABLE IF NOT EXISTS _test_timestamp_no_tz (id int PRIMARY KEY, t TIMESTAMP WITHOUT TIME ZONE)" + ) + ) + db.session.execute(sa.text("SET TIMEZONE TO 'UTC'")) + db.session.commit() + + expected_datetime = datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc) + with app.app_context(): + t = TestTimestampNoTZ() + t.id = 1 + t.t = expected_datetime + db.session.add(t) + db.session.commit() + + with app.app_context(): + res = TestTimestampNoTZ.query.all() + assert len(res) == 1 + assert res[0].id == 1 + assert res[0].t == expected_datetime + assert res[0].t.tzinfo == expected_datetime.tzinfo + db.session.delete(res[0]) + db.session.commit() + + with app.app_context(): + # Change to a non-UTC timezone to show the unexpected behaviour. Production Invenio databases + # must always use UTC. Invenio-DB always initialises the connection with UTC so this will be + # the case unless the instance overrides it. + # This test shows the unexpected behaviour that occurs on a non-UTC DB. + # We are changing the timezone to UTC+1 (depicted as UTC-1 in PostgreSQL). We are not using an + # IANA named time zone to ensure the test is consistent at all times of year. + # See https://www.postgresql.org/docs/current/datetime-posix-timezone-specs.html + db.session.execute(sa.text("SET TIMEZONE TO 'UTC-1'")) + db.session.commit() + + with app.app_context(): + t = TestTimestampNoTZ() + t.id = 1 + # All Invenio modules always write UTC values to the DB without any knowledge of what the DB's + # zone is set to. We therefore continue writing the UTC `expected_datetime` to show this behaviour. + t.t = expected_datetime + db.session.add(t) + db.session.commit() + + with app.app_context(): + res = TestTimestampNoTZ.query.all() + assert len(res) == 1 + assert res[0].id == 1 + # PostgreSQL received the timestamp as '2026-01-01T12:00+00:00'::timestamptz + # Since the column type was `timestamp` and not `timestamptz`, it had to convert it + # from UTC to the DB's default zone which is now UTC+1. + # But when this value is read back, no timezone info is actually returned since the + # column is no `timestamptz`. UTCDateTime assumes the zone to be UTC. + # Therefore, the interpreted hour is shifted by one while the timezone is still + # believed to be UTC. + # This is unwanted behaviour and is why it's important that the DB is set to UTC. + assert res[0].t != expected_datetime + assert res[0].t.hour == expected_datetime.hour + 1 + assert res[0].t.tzinfo == expected_datetime.tzinfo + db.session.delete(res[0]) + db.session.commit() + + with app.app_context(): + db.session.execute(sa.text("DROP TABLE _test_timestamp_no_tz")) + db.session.execute(sa.text("SET TIMEZONE TO 'UTC'")) + db.session.commit() + + +@requires_postgresql +def test_timestamp_with_tz(app, db: SQLAlchemy): + InvenioDB(app, entry_point_group=False, db=db) + + class TestTimestampWithTZ(db.Model): + __tablename__ = "_test_timestamp_with_tz" + id = sa.Column(sa.Integer, primary_key=True) + t = sa.Column(UTCDateTime) + + with app.app_context(): + # Create the table manually so we can override the `t` column type to what it would be post-migration. + db.session.execute( + sa.text( + "CREATE TABLE IF NOT EXISTS _test_timestamp_with_tz (id int PRIMARY KEY, t TIMESTAMP WITH TIME ZONE)" + ) + ) + db.session.execute(sa.text("SET TIMEZONE TO 'UTC'")) + db.session.commit() + + expected_datetime = datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc) + with app.app_context(): + t = TestTimestampWithTZ() + t.id = 1 + t.t = expected_datetime + db.session.add(t) + db.session.commit() + + with app.app_context(): + res = TestTimestampWithTZ.query.all() + assert len(res) == 1 + assert res[0].id == 1 + assert res[0].t == expected_datetime + assert res[0].t.tzinfo == expected_datetime.tzinfo + db.session.delete(res[0]) + db.session.commit() + + with app.app_context(): + db.session.execute(sa.text("SET TIMEZONE TO 'UTC-1'")) + db.session.commit() + + with app.app_context(): + t = TestTimestampWithTZ() + t.id = 1 + t.t = expected_datetime + db.session.add(t) + db.session.commit() + + with app.app_context(): + # In this case, PostgreSQL also shifts the inserted timestamp to UTC+1. + # However, it now also returns the timezone with the column correctly as UTC+1. + # To ensure we avoid unexpected behaviour as much as possible, UTCDateTime only + # accepts timestamps from PostgreSQL that are in UTC, and raises an exception for + # all other zones. + # So in this case correctness is preserved but as an added safety mechanism we still + # reject the value. + with pytest.raises(ValueError) as excinfo: + TestTimestampWithTZ.query.all() + + expected_shifted_datetime = expected_datetime.astimezone( + timezone(timedelta(hours=1)) + ) + assert str(expected_shifted_datetime) in str(excinfo.value) + + db.session.commit() + + with app.app_context(): + db.session.execute(sa.text("DROP TABLE _test_timestamp_with_tz")) + db.session.execute(sa.text("SET TIMEZONE TO 'UTC'")) + db.session.commit() diff --git a/tests/utils.py b/tests/utils.py new file mode 100644 index 0000000..52f0626 --- /dev/null +++ b/tests/utils.py @@ -0,0 +1,18 @@ +# -*- coding: utf-8 -*- +# +# This file is part of Invenio. +# Copyright (C) 2026 CERN. +# +# Invenio is free software; you can redistribute it and/or modify it +# under the terms of the MIT License; see LICENSE file for more details. + +"""Utility methods for tests.""" + +import os + +import pytest + +requires_postgresql = pytest.mark.skipif( + not os.environ.get("SQLALCHEMY_DATABASE_URI", "").startswith("postgresql"), + reason="PostgreSQL required", +)