diff --git a/ydb_sqlalchemy/__init__.py b/ydb_sqlalchemy/__init__.py index 2e5fbab..7e39278 100644 --- a/ydb_sqlalchemy/__init__.py +++ b/ydb_sqlalchemy/__init__.py @@ -1 +1,2 @@ from .dbapi import IsolationLevel # noqa: F401 +from .sqlalchemy import Upsert, types, upsert # noqa: F401 diff --git a/ydb_sqlalchemy/dbapi/connection.py b/ydb_sqlalchemy/dbapi/connection.py index a90239f..df73c4c 100644 --- a/ydb_sqlalchemy/dbapi/connection.py +++ b/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 diff --git a/ydb_sqlalchemy/dbapi/cursor.py b/ydb_sqlalchemy/dbapi/cursor.py index 9d74424..27e6593 100644 --- a/ydb_sqlalchemy/dbapi/cursor.py +++ b/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: diff --git a/ydb_sqlalchemy/dbapi/errors.py b/ydb_sqlalchemy/dbapi/errors.py index a67c0ce..70b55eb 100644 --- a/ydb_sqlalchemy/dbapi/errors.py +++ b/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