Upgrade from psycopg2 to psycopg3

dev/environment

@@ -9,7 +9,7 @@ AWS_ACCESS_KEY_ID=foo
 AWS_SECRET_ACCESS_KEY=foo
 BROKER_URL=sqs://localstack:4566/?region=us-east-1&queue_name_prefix=warehouse-dev
 
-DATABASE_URL=postgresql://postgres@db/warehouse
+DATABASE_URL=postgresql+psycopg://postgres@db/warehouse
 
 ELASTICSEARCH_URL=http://elasticsearch:9200/development
 

docs/dev/development/getting-started.rst

@@ -427,7 +427,7 @@ compilation errors due to your system not including libraries
 or binaries required by some of Warehouse's dependencies.
 
 An example of such dependency is
-`psycopg2 <https://www.psycopg.org/docs/install.html#prerequisites>`_
+`psycopg <https://www.psycopg.org/psycopg3/docs/basic/install.html#local-installation>`_
 which requires PostgreSQL binaries and will fail if not present.
 
 If there's a specific use case you think requires development outside

requirements/lint.in

@@ -16,7 +16,6 @@ types-first
 types-html5lib
 types-itsdangerous
 types-passlib
-types-psycopg2
 types-python-slugify
 types-pytz
 types-redis

requirements/lint.txt

@@ -310,10 +310,6 @@ types-passlib==1.7.7.12 \
     --hash=sha256:6abbf2400a8f1cba48639753e3a034af507a765489bb070974d7f68d9ceef883 \
     --hash=sha256:7a4df64b53c2746f804aa29fb361974e5894e0df30ff18cf60b9518696ffc9d3
     # via -r requirements/lint.in
-types-psycopg2==2.9.21.11 \
-    --hash=sha256:7a323d7744bc8a882fb5a6f63448e903fc70d3dc0d6da9ec1f9c6c4dc10a7102 \
-    --hash=sha256:d5077eacf90e61db8c0b8eea2fdc9d4a97d7aaa16865fb4bd7034a7571520b4d
-    # via -r requirements/lint.in
 types-pyopenssl==23.2.0.2 \
     --hash=sha256:19536aa3debfbe25a918cf0d898e9f5fbbe6f3594a429da7914bf331deb1b342 \
     --hash=sha256:6a010dac9ecd42b582d7dd2cc3e9e40486b79b3b64bb2fffba1474ff96af906d

requirements/main.in

@@ -39,7 +39,7 @@ paginate_sqlalchemy
 passlib>=1.6.4
 pip-api
 premailer
-psycopg2
+psycopg[c]
 pycurl
 pydantic
 pyqrcode

requirements/main.txt

@@ -1209,21 +1209,13 @@ protobuf==4.23.4 \
     #   googleapis-common-protos
     #   grpcio-status
     #   proto-plus
-psycopg2==2.9.6 \
-    --hash=sha256:11aca705ec888e4f4cea97289a0bf0f22a067a32614f6ef64fcf7b8bfbc53744 \
-    --hash=sha256:1861a53a6a0fd248e42ea37c957d36950da00266378746588eab4f4b5649e95f \
-    --hash=sha256:2362ee4d07ac85ff0ad93e22c693d0f37ff63e28f0615a16b6635a645f4b9214 \
-    --hash=sha256:36c941a767341d11549c0fbdbb2bf5be2eda4caf87f65dfcd7d146828bd27f39 \
-    --hash=sha256:53f4ad0a3988f983e9b49a5d9765d663bbe84f508ed655affdb810af9d0972ad \
-    --hash=sha256:869776630c04f335d4124f120b7fb377fe44b0a7645ab3c34b4ba42516951889 \
-    --hash=sha256:a8ad4a47f42aa6aec8d061fdae21eaed8d864d4bb0f0cade5ad32ca16fcd6258 \
-    --hash=sha256:b81fcb9ecfc584f661b71c889edeae70bae30d3ef74fa0ca388ecda50b1222b7 \
-    --hash=sha256:d24ead3716a7d093b90b27b3d73459fe8cd90fd7065cf43b3c40966221d8c394 \
-    --hash=sha256:ded2faa2e6dfb430af7713d87ab4abbfc764d8d7fb73eafe96a24155f906ebf5 \
-    --hash=sha256:f15158418fd826831b28585e2ab48ed8df2d0d98f502a2b4fe619e7d5ca29011 \
-    --hash=sha256:f75001a1cbbe523e00b0ef896a5a1ada2da93ccd752b7636db5a99bc57c44494 \
-    --hash=sha256:f7a7a5ee78ba7dc74265ba69e010ae89dae635eea0e97b055fb641a01a31d2b1
+psycopg[c]==3.1.9 \
+    --hash=sha256:ab400f207a8c120bafdd8077916d8f6c0106e809401378708485b016508c30c9 \
+    --hash=sha256:fbbac339274d8733ee70ba9822297af3e8871790a26e967b5ea53e30a4b74dcc
     # via -r requirements/main.in
+psycopg-c==3.1.9 \
+    --hash=sha256:d160b45b0ee1eb05d78a81538c2bc6868bacb5f421b7190ed65d4681e4552455
+    # via psycopg
 pyasn1==0.5.0 \
     --hash=sha256:87a2121042a1ac9358cabcaf1d07680ff97ee6404333bacca15f76aa8ad01a57 \
     --hash=sha256:97b7290ca68e62a832558ec3976f15cbf911bf5d7c7039d8b861c2a0ece69fde
@@ -1530,6 +1522,7 @@ typing-extensions==4.7.1 \
     # via
     #   alembic
     #   limits
+    #   psycopg
     #   pydantic
     #   sqlalchemy
 tzdata==2023.3 \

requirements/tests.in

@@ -3,7 +3,7 @@ factory_boy
 freezegun
 pretend
 pytest>=3.0.0
-pytest-postgresql>=3.1.3,<4.0.0
+pytest-postgresql>=3.1.3,<6.0.0
 pytest-socket
 pytz
 responses>=0.5.1

requirements/tests.txt

@@ -207,16 +207,20 @@ psutil==5.9.5 \
     --hash=sha256:c607bb3b57dc779d55e1554846352b4e358c10fff3abf3514a7a6601beebdb30 \
     --hash=sha256:ea8518d152174e1249c4f2a1c89e3e6065941df2fa13a1ab45327716a23c2b48
     # via mirakuru
+psycopg==3.1.9 \
+    --hash=sha256:ab400f207a8c120bafdd8077916d8f6c0106e809401378708485b016508c30c9 \
+    --hash=sha256:fbbac339274d8733ee70ba9822297af3e8871790a26e967b5ea53e30a4b74dcc
+    # via pytest-postgresql
 pytest==7.4.0 \
     --hash=sha256:78bf16451a2eb8c7a2ea98e32dc119fd2aa758f1d5d66dbf0a59d69a3969df32 \
     --hash=sha256:b4bf8c45bd59934ed84001ad51e11b4ee40d40a1229d2c79f9c592b0a3f6bd8a
     # via
     #   -r requirements/tests.in
     #   pytest-postgresql
     #   pytest-socket
-pytest-postgresql==3.1.3 \
-    --hash=sha256:05b87a192741511f5171e0300689a531a2a48b4483c69ae2b5f565d3e429b1d5 \
-    --hash=sha256:3649bcac5a0cd0d2cc1470a1087739990d402e2e910d53265ac486321a833898
+pytest-postgresql==5.0.0 \
+    --hash=sha256:22edcbafab8995ee85b8d948ddfaad4f70c2c7462303d7477ecd2f77fc9d15bd \
+    --hash=sha256:6e8f0773b57c9b8975b6392c241b7b81b7018f32079a533f368f2fbda732ecd3
     # via -r requirements/tests.in
 pytest-socket==0.6.0 \
     --hash=sha256:363c1d67228315d4fc7912f1aabfd570de29d0e3db6217d61db5728adacd7138 \
@@ -294,6 +298,10 @@ types-pyyaml==6.0.12.11 \
     --hash=sha256:7d340b19ca28cddfdba438ee638cd4084bde213e501a3978738543e27094775b \
     --hash=sha256:a461508f3096d1d5810ec5ab95d7eeecb651f3a15b71959999988942063bf01d
     # via responses
+typing-extensions==4.7.1 \
+    --hash=sha256:440d5dd3af93b060174bf433bccd69b0babc3b15b1a8dca43789fd7f61514b36 \
+    --hash=sha256:b75ddc264f0ba5615db7ba217daeb99701ad295353c45f9e95963337ceeeffb2
+    # via psycopg
 urllib3==1.26.16 \
     --hash=sha256:8d36afa7616d8ab714608411b4a3b13e58f463aee519024578e062e141dce20f \
     --hash=sha256:8f135f6502756bde6b2a9b28989df5fbe87c9970cecaa69041edcce7f0589b14
@@ -312,3 +320,9 @@ webtest==3.0.0 \
     --hash=sha256:2a001a9efa40d2a7e5d9cd8d1527c75f41814eb6afce2c3d207402547b1e5ead \
     --hash=sha256:54bd969725838d9861a9fa27f8d971f79d275d94ae255f5c501f53bb6d9929eb
     # via -r requirements/tests.in
+
+# The following packages are considered to be unsafe in a requirements file:
+setuptools==68.0.0 \
+    --hash=sha256:11e52c67415a381d10d6b462ced9cfb97066179f0e871399e006c4ab101fc85f \
+    --hash=sha256:baf1fdb41c6da4cd2eae722e135500da913332ab3f2f5c7d33af9b492acb5235
+    # via pytest-postgresql

tests/conftest.py

@@ -28,7 +28,7 @@
 import webtest as _webtest
 
 from jinja2 import Environment, FileSystemLoader
-from psycopg2.errors import InvalidCatalogName
+from psycopg.errors import InvalidCatalogName
 from pyramid.i18n import TranslationString
 from pyramid.static import ManifestCacheBuster
 from pyramid_jinja2 import IJinja2Environment
@@ -242,7 +242,7 @@ def database(request):
     def drop_database():
         janitor.drop()
 
-    return f"postgresql://{pg_user}@{pg_host}:{pg_port}/{pg_db}"
+    return f"postgresql+psycopg://{pg_user}@{pg_host}:{pg_port}/{pg_db}"
 
 
 class MockManifestCacheBuster(ManifestCacheBuster):

tests/unit/cli/test_db.py

@@ -33,21 +33,6 @@
 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)
@@ -65,8 +50,6 @@ 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 _compare_alembic_locks(connection.execute.calls)
     assert alembic_branches.calls == [pretend.call(alembic_config)]
 
 
@@ -87,8 +70,6 @@ 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 _compare_alembic_locks(connection.execute.calls)
     assert alembic_current.calls == [pretend.call(alembic_config)]
 
 
@@ -109,8 +90,6 @@ 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 _compare_alembic_locks(connection.execute.calls)
     assert alembic_downgrade.calls == [pretend.call(alembic_config, "-1")]
 
 
@@ -139,8 +118,6 @@ 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 _compare_alembic_locks(connection.execute.calls)
     assert alembic_heads.calls == [pretend.call(alembic_config, **ekwargs)]
 
 
@@ -161,8 +138,6 @@ 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 _compare_alembic_locks(connection.execute.calls)
     assert alembic_history.calls == [pretend.call(alembic_config, "foo:bar")]
 
 
@@ -202,8 +177,6 @@ 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 _compare_alembic_locks(connection.execute.calls)
     assert alembic_merge.calls == [pretend.call(alembic_config, *eargs, **ekwargs)]
 
 
@@ -260,8 +233,6 @@ 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 _compare_alembic_locks(connection.execute.calls)
     assert alembic_revision.calls == [pretend.call(alembic_config, **ekwargs)]
 
 
@@ -282,8 +253,6 @@ 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 _compare_alembic_locks(connection.execute.calls)
     assert alembic_show.calls == [pretend.call(alembic_config, "foo")]
 
 
@@ -304,8 +273,6 @@ 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 _compare_alembic_locks(connection.execute.calls)
     assert alembic_stamp.calls == [pretend.call(alembic_config, "foo")]
 
 
@@ -326,8 +293,6 @@ 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 _compare_alembic_locks(connection.execute.calls)
     assert alembic_upgrade.calls == [pretend.call(alembic_config, "foo")]
 
 

tests/unit/test_db.py

@@ -14,7 +14,7 @@
 
 import alembic.config
 import pretend
-import psycopg2.extensions
+import psycopg
 import pytest
 import sqlalchemy
 import venusian
@@ -114,7 +114,7 @@ def config_cls():
 
 def test_raises_db_available_error(pyramid_services, metrics):
     def raiser():
-        raise OperationalError("foo", {}, psycopg2.OperationalError())
+        raise OperationalError("foo", {}, psycopg.OperationalError())
 
     engine = pretend.stub(connect=raiser)
     request = pretend.stub(
@@ -199,7 +199,6 @@ def test_create_session_read_only_mode(
 
     connection = pretend.stub(
         connection=pretend.stub(
-            get_transaction_status=lambda: pretend.stub(),
             set_session=lambda **kw: None,
             rollback=lambda: None,
         ),

tests/unit/utils/test_wsgi.py

@@ -13,6 +13,8 @@
 import pretend
 import pytest
 
+from sqlalchemy import type_coerce
+from sqlalchemy.dialects.postgresql import INET
 from sqlalchemy.exc import NoResultFound
 
 from warehouse.ip_addresses.models import IpAddress
@@ -196,16 +198,22 @@ def test_ip_address_exists(db_request):
 
 def test_ip_address_created(db_request):
     with pytest.raises(NoResultFound):
-        db_request.db.query(IpAddress).filter_by(ip_address="192.0.2.69").one()
+        db_request.db.query(IpAddress).filter_by(
+            ip_address=type_coerce("192.0.2.69", INET)
+        ).one()
 
     db_request.environ["GEOIP_CITY"] = "Anytown, ST"
     db_request.remote_addr = "192.0.2.69"
     db_request.remote_addr_hashed = "deadbeef"
 
     wsgi._ip_address(db_request)
 
-    ip_address = db_request.db.query(IpAddress).filter_by(ip_address="192.0.2.69").one()
-    assert ip_address.ip_address == "192.0.2.69"
+    ip_address = (
+        db_request.db.query(IpAddress)
+        .filter_by(ip_address=type_coerce("192.0.2.69", INET))
+        .one()
+    )
+    assert str(ip_address.ip_address) == "192.0.2.69"
     assert ip_address.hashed_ip_address == "deadbeef"
     assert ip_address.geoip_info == {"city": "Anytown, ST"}
 

warehouse/accounts/views.py

@@ -834,7 +834,7 @@ def _error(message):
     try:
         email = (
             request.db.query(Email)
-            .filter(Email.id == data["email.id"], Email.user == request.user)
+            .filter(Email.id == int(data["email.id"]), Email.user == request.user)
             .one()
         )
     except NoResultFound:

warehouse/admin/bans.py

@@ -10,6 +10,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from sqlalchemy import type_coerce
+from sqlalchemy.dialects.postgresql import INET
+
 from warehouse.accounts.interfaces import IUserService
 from warehouse.events.models import IpAddress
 
@@ -18,10 +21,10 @@ class Bans:
     def __init__(self, request):
         self.request = request
 
-    def by_ip(self, ip_address):
+    def by_ip(self, ip_address: str) -> bool:
         banned = (
             self.request.db.query(IpAddress)
-            .filter_by(ip_address=ip_address, is_banned=True)
+            .filter_by(ip_address=type_coerce(ip_address, INET), is_banned=True)
             .one_or_none()
         )
         if banned is not None:

warehouse/banners/views.py

@@ -29,7 +29,7 @@ def list_banner_messages(request):
     if banner_id:
         query = request.db.query(Banner).filter(Banner.id == banner_id)
     else:
-        today = str(datetime.date.today())
+        today = datetime.date.today()
         query = request.db.query(Banner).filter(
             (Banner.active == True) & (Banner.end >= today)  # noqa
         )