From 01043188cd74f19654685d96e189191dc7f0e537 Mon Sep 17 00:00:00 2001
From: Jahnvi Thakkar <jathakkar@microsoft.com>
Date: Tue, 2 Jun 2026 11:32:22 +0530
Subject: [PATCH] Fix exception pickle/unpickle round-trip (#587)

Add __reduce__ to the base Exception class and ConnectionStringParseError
so that pickle.loads(pickle.dumps(exc)) reconstructs the exception with
the correct constructor arguments.

Previously, super().__init__(self.message) stored only one combined string
in self.args, but __init__ requires two positional arguments (driver_error,
ddbc_error). This caused TypeError on unpickle, breaking multiprocessing,
concurrent.futures, celery, and copy.deepcopy.

Fixes #587
---
 mssql_python/exceptions.py   |  6 +++
 tests/test_006_exceptions.py | 73 ++++++++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+)

diff --git a/mssql_python/exceptions.py b/mssql_python/exceptions.py
index f2285bce5..b4c8e3ce4 100644
--- a/mssql_python/exceptions.py
+++ b/mssql_python/exceptions.py
@@ -30,6 +30,9 @@ def __init__(self, errors: list) -> None:
         message = "Connection string parsing failed:\n  " + "\n  ".join(errors)
         super().__init__(message)
 
+    def __reduce__(self):
+        return (self.__class__, (self.errors,))
+
 
 class Exception(builtins.Exception):
     """
@@ -47,6 +50,9 @@ def __init__(self, driver_error: str, ddbc_error: str) -> None:
             self.message = f"Driver Error: {self.driver_error}"
         super().__init__(self.message)
 
+    def __reduce__(self):
+        return (self.__class__, (self.driver_error, self.ddbc_error))
+
 
 class Warning(Exception):
     """
diff --git a/tests/test_006_exceptions.py b/tests/test_006_exceptions.py
index c763ed556..b0d117761 100644
--- a/tests/test_006_exceptions.py
+++ b/tests/test_006_exceptions.py
@@ -451,3 +451,76 @@ def test_truncate_error_message_return_paths():
         # If the exception handling worked, it would have been caught
         # and the function would return the original message (line 531)
         pass
+
+
+# ---------------------------------------------------------------------------
+# Pickle / unpickle round-trip tests
+# ---------------------------------------------------------------------------
+
+
+def test_exception_pickle_roundtrip():
+    """All DB-API exception subclasses must survive a pickle round-trip."""
+    import pickle
+    import copy
+
+    exception_classes = [
+        Warning,
+        Error,
+        InterfaceError,
+        DatabaseError,
+        DataError,
+        OperationalError,
+        IntegrityError,
+        InternalError,
+        ProgrammingError,
+        NotSupportedError,
+    ]
+
+    for cls in exception_classes:
+        original = cls("driver msg", "ddbc msg")
+
+        # pickle round-trip
+        restored = pickle.loads(pickle.dumps(original))
+
+        assert type(restored) is cls, f"{cls.__name__}: type mismatch after unpickle"
+        assert restored.driver_error == "driver msg", f"{cls.__name__}: driver_error mismatch"
+        assert restored.ddbc_error == "ddbc msg", f"{cls.__name__}: ddbc_error mismatch"
+        assert str(restored) == str(original), f"{cls.__name__}: str() mismatch"
+
+        # copy.deepcopy also uses __reduce__
+        deep = copy.deepcopy(original)
+        assert type(deep) is cls
+        assert deep.driver_error == "driver msg"
+
+
+def test_exception_pickle_empty_ddbc_error():
+    """Exceptions with empty ddbc_error should also round-trip cleanly."""
+    import pickle
+
+    original = ProgrammingError("cursor is closed", "")
+    restored = pickle.loads(pickle.dumps(original))
+
+    assert type(restored) is ProgrammingError
+    assert restored.driver_error == "cursor is closed"
+    assert restored.ddbc_error == ""
+    assert str(restored) == str(original)
+
+
+def test_connection_string_parse_error_pickle_roundtrip():
+    """ConnectionStringParseError should survive a pickle round-trip."""
+    import pickle
+    import copy
+
+    errors = ["Unknown keyword: foo", "Missing value for: bar"]
+    original = ConnectionStringParseError(errors)
+
+    restored = pickle.loads(pickle.dumps(original))
+
+    assert type(restored) is ConnectionStringParseError
+    assert restored.errors == errors
+    assert str(restored) == str(original)
+
+    # copy.deepcopy
+    deep = copy.deepcopy(original)
+    assert type(deep) is ConnectionStringParseError
+    assert deep.errors == errors