Skip to content

fix: update migrations code for sqlalchemy 2.0 #13948

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jun 15, 2023
Merged

fix: update migrations code for sqlalchemy 2.0 #13948

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1b9c171
fix: update migrations code for sqlalchemy 2.0
miketheman Jun 15, 2023
ce8dd7e
test: update test suite to account for TextClause
miketheman Jun 15, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
65 changes: 25 additions & 40 deletions tests/unit/cli/test_db.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)]


Expand All @@ -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)]


Expand All @@ -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")]


Expand Down Expand Up @@ -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)]


Expand All @@ -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")]


Expand Down Expand Up @@ -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)]


Expand Down Expand Up @@ -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)]


Expand All @@ -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")]


Expand All @@ -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")]


Expand All @@ -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")]


Expand Down
6 changes: 4 additions & 2 deletions warehouse/cli/db/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@

import contextlib

from sqlalchemy import text

from warehouse.cli import warehouse


Expand All @@ -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
Expand All @@ -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
Expand Down
20 changes: 16 additions & 4 deletions warehouse/migrations/env.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down
2 changes: 1 addition & 1 deletion warehouse/migrations/versions/4490777c984f_migrate_existing_data_for_release_is_.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@ def upgrade():
"""
)
)
conn.execute("COMMIT")
conn.execute(sa.text("COMMIT"))
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This surfaced during a make initdb vs running migrations. I chose to preserve the text vs conn.commit() since the function has other semantics applicable to the running connection.


op.alter_column(
"releases",
Expand Down