From ea2555cd5969844528b012f16f1be81216fa494f Mon Sep 17 00:00:00 2001
From: Andreas Motl <andreas.motl@crate.io>
Date: Thu, 28 Sep 2023 15:39:29 +0200
Subject: [PATCH] Improve exception handling: Properly raise `IntegrityError`
 exceptions

... when receiving `DuplicateKeyException` errors from CrateDB, instead
of the more general `ProgrammingError`.
---
 CHANGES.txt                   |  2 ++
 src/crate/client/http.py      | 13 +++++++++++++
 src/crate/client/test_http.py | 26 ++++++++++++++++++++++++--
 3 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/CHANGES.txt b/CHANGES.txt
index 4fbb10b9..808430c7 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -15,6 +15,8 @@ Unreleased
   ``subjectAltName`` attribute will be used.
 - SQLAlchemy: Improve DDL compiler to ignore foreign key and uniqueness
   constraints
+- DBAPI: Properly raise ``IntegrityError`` exceptions instead of
+  ``ProgrammingError``, when CrateDB raises a ``DuplicateKeyException``.
 
 .. _urllib3 v2.0 migration guide: https://urllib3.readthedocs.io/en/latest/v2-migration-guide.html
 .. _urllib3 v2.0 roadmap: https://urllib3.readthedocs.io/en/stable/v2-roadmap.html
diff --git a/src/crate/client/http.py b/src/crate/client/http.py
index 0cce7bda..1318cca2 100644
--- a/src/crate/client/http.py
+++ b/src/crate/client/http.py
@@ -56,6 +56,7 @@
     BlobLocationNotFoundException,
     DigestNotFoundException,
     ProgrammingError,
+    IntegrityError,
 )
 
 
@@ -191,6 +192,18 @@ def _ex_to_message(ex):
 
 
 def _raise_for_status(response):
+    """
+    Properly raise `IntegrityError` exceptions for CrateDB's `DuplicateKeyException` errors.
+    """
+    try:
+        return _raise_for_status_real(response)
+    except ProgrammingError as ex:
+        if "DuplicateKeyException" in ex.message:
+            raise IntegrityError(ex.message, error_trace=ex.error_trace) from ex
+        raise
+
+
+def _raise_for_status_real(response):
     """ make sure that only crate.exceptions are raised that are defined in
     the DB-API specification """
     message = ''
diff --git a/src/crate/client/test_http.py b/src/crate/client/test_http.py
index de287834..fe20d386 100644
--- a/src/crate/client/test_http.py
+++ b/src/crate/client/test_http.py
@@ -44,8 +44,7 @@
 from setuptools.ssl_support import find_ca_bundle
 
 from .http import Client, CrateJsonEncoder, _get_socket_opts, _remove_certs_for_non_https
-from .exceptions import ConnectionError, ProgrammingError
-
+from .exceptions import ConnectionError, ProgrammingError, IntegrityError
 
 REQUEST = 'crate.client.http.Server.request'
 CA_CERT_PATH = find_ca_bundle()
@@ -91,6 +90,17 @@ def bad_bulk_response():
     return r
 
 
+def duplicate_key_exception():
+    r = fake_response(409, 'Conflict')
+    r.data = json.dumps({
+        "error": {
+            "code": 4091,
+            "message": "DuplicateKeyException[A document with the same primary key exists already]"
+        }
+    }).encode()
+    return r
+
+
 def fail_sometimes(*args, **kwargs):
     if random.randint(1, 100) % 10 == 0:
         raise urllib3.exceptions.MaxRetryError(None, '/_sql', '')
@@ -302,6 +312,18 @@ def test_uuid_serialization(self, request):
         self.assertEqual(data['args'], [str(uid)])
         client.close()
 
+    @patch(REQUEST, fake_request(duplicate_key_exception()))
+    def test_duplicate_key_error(self):
+        """
+        Verify that an `IntegrityError` is raised on duplicate key errors,
+        instead of the more general `ProgrammingError`.
+        """
+        client = Client(servers="localhost:4200")
+        with self.assertRaises(IntegrityError) as cm:
+            client.sql('INSERT INTO testdrive (foo) VALUES (42)')
+        self.assertEqual(cm.exception.message,
+                         "DuplicateKeyException[A document with the same primary key exists already]")
+
 
 @patch(REQUEST, fail_sometimes)
 class ThreadSafeHttpClientTest(TestCase):