diff --git a/tests/unit/cli/test_db.py b/tests/unit/cli/test_db.py index c8ce709390a1..5e6aa1be10eb 100644 --- a/tests/unit/cli/test_db.py +++ b/tests/unit/cli/test_db.py @@ -33,6 +33,21 @@ from warehouse.cli.db.upgrade import upgrade +def _compare_alembic_locks(calls: list[pretend.call]) -> bool: + sql = [] + for t in calls: + assert len(t.args) == 1 + assert len(t.kwargs) == 0 + + tc = t.args[0] + assert isinstance(tc, sqlalchemy.sql.expression.TextClause) + sql.append(tc.text) + return sql == [ + "SELECT pg_advisory_lock(hashtext('alembic'))", + "SELECT pg_advisory_unlock(hashtext('alembic'))", + ] + + def test_branches_command(monkeypatch, cli, pyramid_config): alembic_branches = pretend.call_recorder(lambda config: None) monkeypatch.setattr(alembic.command, "branches", alembic_branches) @@ -51,10 +66,7 @@ def test_branches_command(monkeypatch, cli, pyramid_config): result = cli.invoke(branches, obj=pyramid_config) assert result.exit_code == 0 assert alembic_config.attributes == {"connection": connection} - assert connection.execute.calls == [ - pretend.call("SELECT pg_advisory_lock(hashtext('alembic'))"), - pretend.call("SELECT pg_advisory_unlock(hashtext('alembic'))"), - ] + assert _compare_alembic_locks(connection.execute.calls) assert alembic_branches.calls == [pretend.call(alembic_config)] @@ -76,10 +88,7 @@ def test_current_command(monkeypatch, cli, pyramid_config): result = cli.invoke(current, obj=pyramid_config) assert result.exit_code == 0 assert alembic_config.attributes == {"connection": connection} - assert connection.execute.calls == [ - pretend.call("SELECT pg_advisory_lock(hashtext('alembic'))"), - pretend.call("SELECT pg_advisory_unlock(hashtext('alembic'))"), - ] + assert _compare_alembic_locks(connection.execute.calls) assert alembic_current.calls == [pretend.call(alembic_config)] @@ -101,10 +110,7 @@ def test_downgrade_command(monkeypatch, cli, pyramid_config): result = cli.invoke(downgrade, ["--", "-1"], obj=pyramid_config) assert result.exit_code == 0 assert alembic_config.attributes == {"connection": connection} - assert connection.execute.calls == [ - pretend.call("SELECT pg_advisory_lock(hashtext('alembic'))"), - pretend.call("SELECT pg_advisory_unlock(hashtext('alembic'))"), - ] + assert _compare_alembic_locks(connection.execute.calls) assert alembic_downgrade.calls == [pretend.call(alembic_config, "-1")] @@ -134,10 +140,7 @@ def test_heads_command(monkeypatch, cli, pyramid_config, args, ekwargs): result = cli.invoke(heads, args, obj=pyramid_config) assert result.exit_code == 0 assert alembic_config.attributes == {"connection": connection} - assert connection.execute.calls == [ - pretend.call("SELECT pg_advisory_lock(hashtext('alembic'))"), - pretend.call("SELECT pg_advisory_unlock(hashtext('alembic'))"), - ] + assert _compare_alembic_locks(connection.execute.calls) assert alembic_heads.calls == [pretend.call(alembic_config, **ekwargs)] @@ -159,10 +162,7 @@ def test_history_command(monkeypatch, cli, pyramid_config): result = cli.invoke(history, ["foo:bar"], obj=pyramid_config) assert result.exit_code == 0 assert alembic_config.attributes == {"connection": connection} - assert connection.execute.calls == [ - pretend.call("SELECT pg_advisory_lock(hashtext('alembic'))"), - pretend.call("SELECT pg_advisory_unlock(hashtext('alembic'))"), - ] + assert _compare_alembic_locks(connection.execute.calls) assert alembic_history.calls == [pretend.call(alembic_config, "foo:bar")] @@ -203,10 +203,7 @@ def test_merge_command(monkeypatch, cli, pyramid_config, args, eargs, ekwargs): result = cli.invoke(merge, args, obj=pyramid_config) assert result.exit_code == 0 assert alembic_config.attributes == {"connection": connection} - assert connection.execute.calls == [ - pretend.call("SELECT pg_advisory_lock(hashtext('alembic'))"), - pretend.call("SELECT pg_advisory_unlock(hashtext('alembic'))"), - ] + assert _compare_alembic_locks(connection.execute.calls) assert alembic_merge.calls == [pretend.call(alembic_config, *eargs, **ekwargs)] @@ -264,10 +261,7 @@ def test_revision_command(monkeypatch, cli, pyramid_config, args, ekwargs): result = cli.invoke(revision, args, obj=pyramid_config) assert result.exit_code == 0 assert alembic_config.attributes == {"connection": connection} - assert connection.execute.calls == [ - pretend.call("SELECT pg_advisory_lock(hashtext('alembic'))"), - pretend.call("SELECT pg_advisory_unlock(hashtext('alembic'))"), - ] + assert _compare_alembic_locks(connection.execute.calls) assert alembic_revision.calls == [pretend.call(alembic_config, **ekwargs)] @@ -289,10 +283,7 @@ def test_show_command(monkeypatch, cli, pyramid_config): result = cli.invoke(show, ["foo"], obj=pyramid_config) assert result.exit_code == 0 assert alembic_config.attributes == {"connection": connection} - assert connection.execute.calls == [ - pretend.call("SELECT pg_advisory_lock(hashtext('alembic'))"), - pretend.call("SELECT pg_advisory_unlock(hashtext('alembic'))"), - ] + assert _compare_alembic_locks(connection.execute.calls) assert alembic_show.calls == [pretend.call(alembic_config, "foo")] @@ -314,10 +305,7 @@ def test_stamp_command(monkeypatch, cli, pyramid_config): result = cli.invoke(stamp, ["foo"], obj=pyramid_config) assert result.exit_code == 0 assert alembic_config.attributes == {"connection": connection} - assert connection.execute.calls == [ - pretend.call("SELECT pg_advisory_lock(hashtext('alembic'))"), - pretend.call("SELECT pg_advisory_unlock(hashtext('alembic'))"), - ] + assert _compare_alembic_locks(connection.execute.calls) assert alembic_stamp.calls == [pretend.call(alembic_config, "foo")] @@ -339,10 +327,7 @@ def test_upgrade_command(monkeypatch, cli, pyramid_config): result = cli.invoke(upgrade, ["foo"], obj=pyramid_config) assert result.exit_code == 0 assert alembic_config.attributes == {"connection": connection} - assert connection.execute.calls == [ - pretend.call("SELECT pg_advisory_lock(hashtext('alembic'))"), - pretend.call("SELECT pg_advisory_unlock(hashtext('alembic'))"), - ] + assert _compare_alembic_locks(connection.execute.calls) assert alembic_upgrade.calls == [pretend.call(alembic_config, "foo")] diff --git a/warehouse/cli/db/__init__.py b/warehouse/cli/db/__init__.py index 8c0fdf5ee7aa..08368cddc36b 100644 --- a/warehouse/cli/db/__init__.py +++ b/warehouse/cli/db/__init__.py @@ -12,6 +12,8 @@ import contextlib +from sqlalchemy import text + from warehouse.cli import warehouse @@ -20,7 +22,7 @@ def alembic_lock(engine, alembic_config): with engine.begin() as connection: # Attempt to acquire the alembic lock, this will wait until the lock # has been acquired allowing multiple commands to wait for each other. - connection.execute("SELECT pg_advisory_lock(hashtext('alembic'))") + connection.execute(text("SELECT pg_advisory_lock(hashtext('alembic'))")) try: # Tell Alembic use our current connection instead of creating it's @@ -31,7 +33,7 @@ def alembic_lock(engine, alembic_config): yield alembic_config finally: # Finally we need to release the lock we've acquired. - connection.execute("SELECT pg_advisory_unlock(hashtext('alembic'))") + connection.execute(text("SELECT pg_advisory_unlock(hashtext('alembic'))")) @warehouse.group() # pragma: no branch diff --git a/warehouse/migrations/env.py b/warehouse/migrations/env.py index 23a05d64ba40..4cc737aab941 100644 --- a/warehouse/migrations/env.py +++ b/warehouse/migrations/env.py @@ -49,16 +49,28 @@ def run_migrations_online(): url = options.pop("url") connectable = create_engine(url, poolclass=pool.NullPool) - with connectable.connect() as connection: - connection.execute(text("SET statement_timeout = 5000")) - connection.execute(text("SET lock_timeout = 4000")) + with connectable.connect() as connection: + connection.execute(text("SET statement_timeout = 5000")) + connection.execute(text("SET lock_timeout = 4000")) + context.configure( + connection=connection, + target_metadata=db.metadata, + compare_server_default=True, + transaction_per_migration=True, + ) + with context.begin_transaction(): + context.run_migrations() + else: context.configure( - connection=connection, + connection=connectable, target_metadata=db.metadata, compare_server_default=True, transaction_per_migration=True, ) + context.execute(text("SET statement_timeout = 5000")) + context.execute(text("SET lock_timeout = 4000")) + with context.begin_transaction(): context.run_migrations() diff --git a/warehouse/migrations/versions/4490777c984f_migrate_existing_data_for_release_is_.py b/warehouse/migrations/versions/4490777c984f_migrate_existing_data_for_release_is_.py index 7a94448429e2..115b0b8df68d 100644 --- a/warehouse/migrations/versions/4490777c984f_migrate_existing_data_for_release_is_.py +++ b/warehouse/migrations/versions/4490777c984f_migrate_existing_data_for_release_is_.py @@ -55,7 +55,7 @@ def upgrade(): """ ) ) - conn.execute("COMMIT") + conn.execute(sa.text("COMMIT")) op.alter_column( "releases",