Merge pull request #19 from ydb-platform/dbapi-excp-mapping

dbapi: ydb driver exception mapping

Author: Valeria1235

test_dbapi/test_dbapi.py

@@ -15,7 +15,7 @@ def test_connection(connection):
         cur.execute("DROP TABLE foo", context={"isddl": True})
 
     assert not connection.check_exists("/local/foo")
-    with pytest.raises(dbapi.DatabaseError):
+    with pytest.raises(dbapi.ProgrammingError):
         connection.describe("/local/foo")
 
     cur.execute("CREATE TABLE foo(id Int64 NOT NULL, PRIMARY KEY (id))", context={"isddl": True})
@@ -107,3 +107,34 @@ def test_cursor(connection):
 
     cur.close()
     cur2.close()
+
+
+def test_errors(connection):
+    with pytest.raises(dbapi.InterfaceError):
+        dbapi.connect("localhost:2136", database="/local666")
+
+    cur = connection.cursor()
+
+    with suppress(dbapi.DatabaseError):
+        cur.execute("DROP TABLE test", context={"isddl": True})
+
+    with pytest.raises(dbapi.DataError):
+        cur.execute("SELECT 18446744073709551616")
+
+    with pytest.raises(dbapi.DataError):
+        cur.execute("SELECT * FROM 拉屎")
+
+    with pytest.raises(dbapi.DataError):
+        cur.execute("SELECT floor(5 / 2)")
+
+    with pytest.raises(dbapi.ProgrammingError):
+        cur.execute("SELECT * FROM test")
+
+    cur.execute("CREATE TABLE test(id Int64, PRIMARY KEY (id))", context={"isddl": True})
+
+    cur.execute("INSERT INTO test(id) VALUES(1)")
+    with pytest.raises(dbapi.IntegrityError):
+        cur.execute("INSERT INTO test(id) VALUES(1)")
+
+    cur.execute("DROP TABLE test", context={"isddl": True})
+    cur.close()

ydb_sqlalchemy/dbapi/connection.py

@@ -2,7 +2,7 @@
 
 import ydb
 from .cursor import Cursor
-from .errors import DatabaseError
+from .errors import InterfaceError, ProgrammingError, DatabaseError
 
 
 class Connection:
@@ -19,10 +19,12 @@ def describe(self, table_path):
         full_path = posixpath.join(self.database, table_path)
         try:
             return self.pool.retry_operation_sync(lambda cli: cli.describe_table(full_path))
+        except ydb.issues.SchemeError as e:
+            raise ProgrammingError(e.message, e.issues, e.status) from e
         except ydb.Error as e:
-            raise DatabaseError(e.message, e.issues, e.status)
-        except Exception:
-            raise DatabaseError(f"Failed to describe table {table_path}")
+            raise DatabaseError(e.message, e.issues, e.status) from e
+        except Exception as e:
+            raise DatabaseError(f"Failed to describe table {table_path}") from e
 
     def check_exists(self, table_path):
         try:
@@ -61,8 +63,8 @@ def _create_driver(endpoint, database, **conn_kwargs):
         try:
             driver.wait(timeout=5, fail_fast=True)
         except ydb.Error as e:
-            raise DatabaseError(e.message, e.issues, e.status)
-        except Exception:
+            raise InterfaceError(e.message, e.issues, e.status) from e
+        except Exception as e:
             driver.stop()
-            raise DatabaseError(f"Failed to connect to YDB, details {driver.discovery_debug_details()}")
+            raise InterfaceError(f"Failed to connect to YDB, details {driver.discovery_debug_details()}") from e
         return driver

ydb_sqlalchemy/dbapi/cursor.py

@@ -7,7 +7,15 @@
 from typing import Optional, Dict, Any
 
 import ydb
-from .errors import IntegrityError, DatabaseError, ProgrammingError
+from .errors import (
+    InternalError,
+    IntegrityError,
+    DataError,
+    DatabaseError,
+    ProgrammingError,
+    OperationalError,
+    NotSupportedError,
+)
 
 
 logger = logging.getLogger(__name__)
@@ -107,6 +115,28 @@ def _execute_in_pool(cli):
                     return cli.transaction().execute(prepared_query, sql_params, commit_tx=True)
             except (ydb.issues.AlreadyExists, ydb.issues.PreconditionFailed) as e:
                 raise IntegrityError(e.message, e.issues, e.status) from e
+            except (ydb.issues.Unsupported, ydb.issues.Unimplemented) as e:
+                raise NotSupportedError(e.message, e.issues, e.status) from e
+            except (ydb.issues.BadRequest, ydb.issues.SchemeError) as e:
+                raise ProgrammingError(e.message, e.issues, e.status) from e
+            except (
+                ydb.issues.TruncatedResponseError,
+                ydb.issues.ConnectionError,
+                ydb.issues.Aborted,
+                ydb.issues.Unavailable,
+                ydb.issues.Overloaded,
+                ydb.issues.Undetermined,
+                ydb.issues.Timeout,
+                ydb.issues.Cancelled,
+                ydb.issues.SessionBusy,
+                ydb.issues.SessionExpired,
+                ydb.issues.SessionPoolEmpty,
+            ) as e:
+                raise OperationalError(e.message, e.issues, e.status) from e
+            except ydb.issues.GenericError as e:
+                raise DataError(e.message, e.issues, e.status) from e
+            except ydb.issues.InternalError as e:
+                raise InternalError(e.message, e.issues, e.status) from e
             except ydb.Error as e:
                 raise DatabaseError(e.message, e.issues, e.status) from e
 
@@ -144,7 +174,7 @@ def _rows_iterable(self, chunks_iterable):
                     #  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)
+            raise DatabaseError(e.message, e.issues, e.status) from e
 
     def _ensure_prefetched(self):
         if self.rows is not None and self._rows_prefetched is None: