Make extract_postgres_error work with Psycopg 3.1

Author: Photonios

psqlextra/_version.py

@@ -1 +1 @@
-__version__ = "2.0.4"
+__version__ = "2.0.9rc3+swen.4"

psqlextra/error.py

@@ -1,11 +1,21 @@
-from typing import Optional
-
-import psycopg2
+from typing import Optional, Union
 
 from django import db
 
+try:
+    from psycopg2 import Error as Psycopg2Error
+except ImportError:
+    Psycopg2Error = None
+
+try:
+    from psycopg import Error as Psycopg3Error
+except ImportError:
+    Psycopg3Error = None
 
-def extract_postgres_error(error: db.Error) -> Optional[psycopg2.Error]:
+
+def extract_postgres_error(
+    error: db.Error,
+) -> Optional[Union["Psycopg2Error", "Psycopg3Error"]]:
     """Extracts the underlying :see:psycopg2.Error from the specified Django
     database error.
 
@@ -14,7 +24,30 @@ def extract_postgres_error(error: db.Error) -> Optional[psycopg2.Error]:
     the cause of the error.
     """
 
-    if not isinstance(error.__cause__, psycopg2.Error):
+    if (Psycopg2Error and not isinstance(error.__cause__, Psycopg2Error)) and (
+        Psycopg3Error and not isinstance(error.__cause__, Psycopg3Error)
+    ):
         return None
 
     return error.__cause__
+
+
+def extract_postgres_error_code(error: db.Error) -> Optional[str]:
+    """Extracts the underlying Postgres error code.
+
+    As per PEP-249, Django wraps all database errors in its own
+    exception. We can extract the underlying database error by examaning
+    the cause of the error.
+    """
+
+    cause = error.__cause__
+    if not cause:
+        return None
+
+    if Psycopg2Error and isinstance(cause, Psycopg2Error):
+        return cause.pgcode
+
+    if Psycopg3Error and isinstance(cause, Psycopg3Error):
+        return cause.sqlstate
+
+    return None

tests/test_schema.py

@@ -5,7 +5,7 @@
 from django.db import InternalError, ProgrammingError, connection
 from psycopg2 import errorcodes
 
-from psqlextra.error import extract_postgres_error
+from psqlextra.error import extract_postgres_error_code
 from psqlextra.schema import PostgresSchema, postgres_temporary_schema
 
 
@@ -92,8 +92,8 @@ def test_postgres_schema_delete_and_create():
     with pytest.raises(InternalError) as exc_info:
         schema = PostgresSchema.delete_and_create(schema.name)
 
-    pg_error = extract_postgres_error(exc_info.value)
-    assert pg_error.pgcode == errorcodes.DEPENDENT_OBJECTS_STILL_EXIST
+    pg_error = extract_postgres_error_code(exc_info.value)
+    assert pg_error == errorcodes.DEPENDENT_OBJECTS_STILL_EXIST
 
     # Verify that the schema and table still exist
     assert _does_schema_exist(schema.name)
@@ -112,8 +112,8 @@ def test_postgres_schema_delete_and_create():
             cursor.execute("SELECT * FROM test.bla")
             assert cursor.fetchone() == ("hello",)
 
-    pg_error = extract_postgres_error(exc_info.value)
-    assert pg_error.pgcode == errorcodes.UNDEFINED_TABLE
+    pg_error = extract_postgres_error_code(exc_info.value)
+    assert pg_error == errorcodes.UNDEFINED_TABLE
 
 
 def test_postgres_schema_delete():
@@ -134,8 +134,8 @@ def test_postgres_schema_delete_not_empty():
     with pytest.raises(InternalError) as exc_info:
         schema.delete()
 
-    pg_error = extract_postgres_error(exc_info.value)
-    assert pg_error.pgcode == errorcodes.DEPENDENT_OBJECTS_STILL_EXIST
+    pg_error = extract_postgres_error_code(exc_info.value)
+    assert pg_error == errorcodes.DEPENDENT_OBJECTS_STILL_EXIST
 
 
 def test_postgres_schema_delete_cascade_not_empty():
@@ -176,8 +176,8 @@ def test_postgres_temporary_schema_not_empty():
                     f"CREATE TABLE {schema.name}.mytable AS SELECT 'hello world'"
                 )
 
-    pg_error = extract_postgres_error(exc_info.value)
-    assert pg_error.pgcode == errorcodes.DEPENDENT_OBJECTS_STILL_EXIST
+    pg_error = extract_postgres_error_code(exc_info.value)
+    assert pg_error == errorcodes.DEPENDENT_OBJECTS_STILL_EXIST
 
 
 def test_postgres_temporary_schema_not_empty_cascade():