Add original error to dbapi exceptions

Author: LuckySting

ydb_sqlalchemy/__init__.py

@@ -1 +1,2 @@
 from .dbapi import IsolationLevel  # noqa: F401
+from .sqlalchemy import Upsert, types, upsert  # noqa: F401

ydb_sqlalchemy/dbapi/connection.py

@@ -162,7 +162,7 @@ def _create_driver(self):
         try:
             self._maybe_await(driver.wait, timeout=5, fail_fast=True)
         except ydb.Error as e:
-            raise InterfaceError(e.message, e.issues, e.status) from e
+            raise InterfaceError(e.message, original_error=e) from e
         except Exception as e:
             self._maybe_await(driver.stop)
             raise InterfaceError(f"Failed to connect to YDB, details {driver.discovery_debug_details()}") from e

ydb_sqlalchemy/dbapi/cursor.py

@@ -41,11 +41,11 @@ def wrapper(*args, **kwargs):
         try:
             return func(*args, **kwargs)
         except (ydb.issues.AlreadyExists, ydb.issues.PreconditionFailed) as e:
-            raise IntegrityError(e.message, e.issues, e.status) from e
+            raise IntegrityError(e.message, original_error=e) from e
         except (ydb.issues.Unsupported, ydb.issues.Unimplemented) as e:
-            raise NotSupportedError(e.message, e.issues, e.status) from e
+            raise NotSupportedError(e.message, original_error=e) from e
         except (ydb.issues.BadRequest, ydb.issues.SchemeError) as e:
-            raise ProgrammingError(e.message, e.issues, e.status) from e
+            raise ProgrammingError(e.message, original_error=e) from e
         except (
             ydb.issues.TruncatedResponseError,
             ydb.issues.ConnectionError,
@@ -59,13 +59,13 @@ def wrapper(*args, **kwargs):
             ydb.issues.SessionExpired,
             ydb.issues.SessionPoolEmpty,
         ) as e:
-            raise OperationalError(e.message, e.issues, e.status) from e
+            raise OperationalError(e.message, original_error=e) from e
         except ydb.issues.GenericError as e:
-            raise DataError(e.message, e.issues, e.status) from e
+            raise DataError(e.message, original_error=e) from e
         except ydb.issues.InternalError as e:
-            raise InternalError(e.message, e.issues, e.status) from e
+            raise InternalError(e.message, original_error=e) from e
         except ydb.Error as e:
-            raise DatabaseError(e.message, e.issues, e.status) from e
+            raise DatabaseError(e.message, original_error=e) from e
         except Exception as e:
             raise DatabaseError("Failed to execute query") from e
 
@@ -214,7 +214,7 @@ def _rows_iterable(self, chunks_iterable: ydb.convert.ResultSets):
                     #  of this PEP to return a sequence: https://www.python.org/dev/peps/pep-0249/#fetchmany
                     yield row[::]
         except ydb.Error as e:
-            raise DatabaseError(e.message, e.issues, e.status) from e
+            raise DatabaseError(e.message, original_error=e) from e
 
     def _ensure_prefetched(self):
         if self.rows is not None and self._rows_prefetched is None:

ydb_sqlalchemy/dbapi/errors.py

@@ -1,15 +1,27 @@
+from typing import Optional, List
+
+import ydb
+from google.protobuf.message import Message
+
+
 class Warning(Exception):
     pass
 
 
 class Error(Exception):
-    def __init__(self, message, issues=None, status=None):
+    def __init__(
+        self,
+        message: str,
+        original_error: Optional[ydb.Error] = None,
+    ):
         super(Error, self).__init__(message)
 
-        pretty_issues = _pretty_issues(issues)
-        self.issues = issues
-        self.message = pretty_issues or message
-        self.status = status
+        self.original_error = original_error
+        if original_error:
+            pretty_issues = _pretty_issues(original_error.issues)
+            self.issues = original_error.issues
+            self.message = pretty_issues or message
+            self.status = original_error.status
 
 
 class InterfaceError(Error):
@@ -44,7 +56,7 @@ class NotSupportedError(DatabaseError):
     pass
 
 
-def _pretty_issues(issues):
+def _pretty_issues(issues: List[Message]) -> str:
     if issues is None:
         return None
 
@@ -56,7 +68,7 @@ def _pretty_issues(issues):
     return "\n" + "\n".join(children_messages)
 
 
-def _get_messages(issue, max_depth=100, indent=2, depth=0, root=False):
+def _get_messages(issue: Message, max_depth: int = 100, indent: int = 2, depth: int = 0, root: bool = False) -> str:
     if depth >= max_depth:
         return None