From 5a9c6d510c4587a402c0a678f7e708b280c9bf17 Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Fri, 11 Feb 2022 20:27:40 +0100
Subject: [PATCH] Set Error.sqlcode when unknown codes are received

Close #225.
---
 docs/news.rst             |  1 +
 psycopg/psycopg/errors.py |  4 ++++
 tests/test_errors.py      | 21 +++++++++++++++++++++
 3 files changed, 26 insertions(+)

diff --git a/docs/news.rst b/docs/news.rst
index 1e6cd9437..c09133611 100644
--- a/docs/news.rst
+++ b/docs/news.rst
@@ -24,6 +24,7 @@ Psycopg 3.1 (unreleased)
 Psycopg 3.0.9 (unreleased)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^
 
+- Set `Error.sqlstate` when an unknown code is received (:ticket:`#225`).
 - Add the `!tzdata` package as a dependency on Windows in order to handle time
   zones (:ticket:`#223`).
 
diff --git a/psycopg/psycopg/errors.py b/psycopg/psycopg/errors.py
index 0e0db5204..66ded14ee 100644
--- a/psycopg/psycopg/errors.py
+++ b/psycopg/psycopg/errors.py
@@ -59,6 +59,10 @@ class Error(Exception):
         self._info = info
         self._encoding = encoding
 
+        # Handle sqlstate codes for which we don't have a class.
+        if not self.sqlstate and info:
+            self.sqlstate = self.diag.sqlstate
+
     @property
     def diag(self) -> "Diagnostic":
         """
diff --git a/tests/test_errors.py b/tests/test_errors.py
index aa61546ae..5ae3b51b9 100644
--- a/tests/test_errors.py
+++ b/tests/test_errors.py
@@ -239,3 +239,24 @@ def test_query_context(conn):
     assert exc.value.diag.message_primary in s
     assert "ERROR" not in s
     assert not s.endswith("\n")
+
+
+def test_unknown_sqlstate(conn):
+    code = "PXX99"
+    with pytest.raises(KeyError):
+        e.lookup(code)
+
+    with pytest.raises(e.ProgrammingError) as excinfo:
+        conn.execute(
+            f"""
+            do $$begin
+            raise exception 'made up code' using errcode = '{code}';
+            end$$ language plpgsql
+            """
+        )
+    exc = excinfo.value
+    assert exc.diag.sqlstate == code
+    assert exc.sqlstate == code
+    # Survives pickling too
+    pexc = pickle.loads(pickle.dumps(exc))
+    assert pexc.sqlstate == code
-- 
2.47.2