Unwrap SQLAlchemy IntegrityErrors (#18891)

* Unwrap SQLAlchemy errors

* Add a test

* Update tests accordingly

Author: di

tests/unit/forklift/test_legacy.py

@@ -15,12 +15,12 @@
 from unittest import mock
 
 import pretend
+import psycopg
 import pytest
 
 from pypi_attestations import Attestation, Envelope, VerificationMaterial
 from pyramid.httpexceptions import HTTPBadRequest, HTTPForbidden, HTTPTooManyRequests
 from sqlalchemy import and_, exists
-from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import joinedload
 from trove_classifiers import classifiers
 from webob.multidict import MultiDict
@@ -4789,7 +4789,7 @@ def test_all_valid_classifiers_can_be_created(self, db_request):
     @pytest.mark.parametrize("parent_classifier", ["private", "Private", "PrIvAtE"])
     def test_private_classifiers_cannot_be_created(self, db_request, parent_classifier):
         db_request.db.add(Classifier(classifier=f"{parent_classifier} :: Foo"))
-        with pytest.raises(IntegrityError):
+        with pytest.raises(psycopg.errors.CheckViolation):
             db_request.db.commit()
 
     def test_equivalent_version_one_release(self, pyramid_config, db_request):

tests/unit/ip_addresses/test_models.py

@@ -1,9 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 
+import psycopg
 import pytest
 
 from sqlalchemy import sql
-from sqlalchemy.exc import IntegrityError
 
 from warehouse.ip_addresses.models import BanReason
 
@@ -46,5 +46,5 @@ def test_invalid_transformed(self, db_request):
         ],
     )
     def test_ban_data_constraint(self, db_request, kwargs):
-        with pytest.raises(IntegrityError):
+        with pytest.raises(psycopg.errors.CheckViolation):
             DBIpAddressFactory(**kwargs)

tests/unit/oidc/models/test_github.py

@@ -1,8 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import pretend
+import psycopg
 import pytest
-import sqlalchemy
 
 from tests.common.db.oidc import GitHubPublisherFactory, PendingGitHubPublisherFactory
 from warehouse.oidc import errors
@@ -676,7 +676,7 @@ def test_github_publisher_duplicates_cant_be_created(self, db_request):
         )
         db_request.db.add(publisher2)
 
-        with pytest.raises(sqlalchemy.exc.IntegrityError):
+        with pytest.raises(psycopg.errors.UniqueViolation):
             db_request.db.commit()
 
     @pytest.mark.parametrize(

tests/unit/oidc/models/test_gitlab.py

@@ -1,8 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import pretend
+import psycopg
 import pytest
-import sqlalchemy
 
 from tests.common.db.oidc import GitLabPublisherFactory, PendingGitLabPublisherFactory
 from warehouse.oidc import errors
@@ -710,7 +710,7 @@ def test_gitlab_publisher_duplicates_cant_be_created(self, db_request):
         )
         db_request.db.add(publisher2)
 
-        with pytest.raises(sqlalchemy.exc.IntegrityError):
+        with pytest.raises(psycopg.errors.UniqueViolation):
             db_request.db.commit()
 
     @pytest.mark.parametrize(

tests/unit/organizations/test_models.py

@@ -3,13 +3,13 @@
 import datetime
 
 import pretend
+import psycopg
 import pytest
 
 from freezegun import freeze_time
 from pyramid.authorization import Allow
 from pyramid.httpexceptions import HTTPPermanentRedirect
 from pyramid.location import lineage
-from sqlalchemy.exc import IntegrityError
 
 from warehouse.authnz import Permissions
 from warehouse.organizations.models import (
@@ -724,7 +724,7 @@ def test_unique_constraint(self, db_session):
         )
 
         # Attempt to create duplicate - should raise IntegrityError
-        with pytest.raises(IntegrityError):
+        with pytest.raises(psycopg.errors.UniqueViolation):
             DBOrganizationOIDCIssuerFactory.create(
                 organization=organization,
                 issuer_type=OIDCIssuerType.GitLab,

tests/unit/test_db.py

@@ -11,7 +11,7 @@
 import zope.sqlalchemy
 
 from sqlalchemy import event
-from sqlalchemy.exc import OperationalError
+from sqlalchemy.exc import DBAPIError, OperationalError
 
 from warehouse import db
 from warehouse.admin.flags import AdminFlag, AdminFlagValue
@@ -22,6 +22,7 @@
     _configure_alembic,
     _create_session,
     includeme,
+    unwrap_dbapi_exceptions,
 )
 
 
@@ -241,3 +242,33 @@ class FakeRegistry(dict):
         )
     ]
     assert config.registry["sqlalchemy.engine"] is engine
+
+
+def test_unwrap_dbapi_exceptions():
+    original_exception = psycopg.OperationalError()
+    sqlalchemy_exception = DBAPIError("foo", {}, original_exception)
+    context = pretend.stub(
+        sqlalchemy_exception=sqlalchemy_exception,
+        original_exception=original_exception,
+    )
+
+    with pytest.raises(psycopg.OperationalError) as e:
+        unwrap_dbapi_exceptions(context)
+
+    assert e.value is original_exception
+
+
+def test_unwrap_dbapi_exceptions_no_op():
+    # Not a DBAPIError
+    context = pretend.stub(
+        sqlalchemy_exception=OperationalError("foo", {}, None),
+        original_exception=None,
+    )
+    unwrap_dbapi_exceptions(context)
+
+    # No original exception
+    context = pretend.stub(
+        sqlalchemy_exception=DBAPIError("foo", {}, None),
+        original_exception=None,
+    )
+    unwrap_dbapi_exceptions(context)

warehouse/db.py

@@ -16,7 +16,7 @@
 from pyramid.renderers import JSON
 from sqlalchemy import event, func, inspect
 from sqlalchemy.dialects.postgresql import UUID as PG_UUID
-from sqlalchemy.exc import IntegrityError, OperationalError
+from sqlalchemy.exc import DBAPIError, IntegrityError, OperationalError
 from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, sessionmaker
 
 from warehouse.metrics import IMetricsService
@@ -166,6 +166,19 @@ def cleanup(request):
     return session
 
 
+@event.listens_for(sqlalchemy.engine.Engine, "handle_error")
+def unwrap_dbapi_exceptions(context):
+    """
+    Listens for SQLAlchemy errors and raises the original
+    DBAPI (e.g., psycopg) exception instead.
+    """
+    if (
+        isinstance(context.sqlalchemy_exception, DBAPIError)
+        and context.original_exception
+    ):
+        raise context.original_exception from context.sqlalchemy_exception
+
+
 def includeme(config):
     # Add a directive to get an alembic configuration.
     config.add_directive("alembic_config", _configure_alembic)

warehouse/organizations/services.py

@@ -2,8 +2,9 @@
 
 import datetime
 
+from psycopg.errors import UniqueViolation
 from sqlalchemy import delete, func, orm, select
-from sqlalchemy.exc import IntegrityError, NoResultFound
+from sqlalchemy.exc import NoResultFound
 from zope.interface import implementer
 
 from warehouse.accounts.models import TermsOfServiceEngagement, User
@@ -518,7 +519,7 @@ def rename_organization(self, organization_id, name):
         try:
             self.db.flush()  # flush db now so organization.normalized_name available
             self.add_catalog_entry(organization_id)
-        except IntegrityError:
+        except UniqueViolation:
             raise ValueError(f'Organization name "{name}" has been used')
 
         return organization