scheme&table queries

Author: vgvoleg

ydb_dbapi/connection.py

@@ -1,10 +1,17 @@
 from typing import (
     Any,
+    List,
     NamedTuple,
     Optional,
 )
+import posixpath
 
 import ydb
+from ydb.retries import retry_operation_async
+
+from .utils import (
+    handle_ydb_errors,
+)
 
 from .errors import (
     # InterfaceError,
@@ -38,6 +45,8 @@ def __init__(
             "ydb_table_path_prefix", ""
         )
 
+        self.driver: ydb.aio.Driver = self.conn_kwargs.pop("ydb_driver", None)
+
         self.session_pool: ydb.aio.QuerySessionPool = self.conn_kwargs.pop(
             "ydb_session_pool", None
         )
@@ -125,6 +134,48 @@ def get_isolation_level(self) -> str:
         else:
             raise NotSupportedError(f"{self.tx_mode.name} is not supported")
 
+    @handle_ydb_errors
+    async def describe(self, table_path: str) -> ydb.TableSchemeEntry:
+        abs_table_path = posixpath.join(
+            self.database, self.table_path_prefix, table_path
+        )
+        return self.driver.table_client.describe_table(abs_table_path)
+
+    @handle_ydb_errors
+    async def check_exists(self, table_path: str) -> ydb.SchemeEntry:
+        abs_table_path = posixpath.join(
+            self.database, self.table_path_prefix, table_path
+        )
+        return await self._check_path_exists(abs_table_path)
+
+    @handle_ydb_errors
+    async def get_table_names(self) -> List[str]:
+        abs_dir_path = posixpath.join(self.database, self.table_path_prefix)
+        names = await self._get_table_names(abs_dir_path)
+        return [posixpath.relpath(path, abs_dir_path) for path in names]
+
+    async def _check_path_exists(self, table_path: str) -> ydb.SchemeEntry:
+        try:
+            await retry_operation_async(
+                self.driver.scheme_client.describe_path, table_path
+            )
+            return True
+        except ydb.SchemeError:
+            return False
+
+    async def _get_table_names(self, abs_dir_path: str) -> List[str]:
+        directory = await retry_operation_async(
+            self.driver.scheme_client.list_directory, abs_dir_path
+        )
+        result = []
+        for child in directory.children:
+            child_abs_path = posixpath.join(abs_dir_path, child.name)
+            if child.is_table():
+                result.append(child_abs_path)
+            elif child.is_directory() and not child.name.startswith("."):
+                result.extend(self.get_table_names(child_abs_path))
+        return result
+
 
 async def connect() -> Connection:
     return Connection()

ydb_dbapi/utils.py

@@ -0,0 +1,49 @@
+import functools
+import ydb
+
+from .errors import (
+    DatabaseError,
+    DataError,
+    IntegrityError,
+    InternalError,
+    NotSupportedError,
+    OperationalError,
+    ProgrammingError,
+)
+
+
+def handle_ydb_errors(func):
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except (ydb.issues.AlreadyExists, ydb.issues.PreconditionFailed) as e:
+            raise IntegrityError(e.message, original_error=e) from e
+        except (ydb.issues.Unsupported, ydb.issues.Unimplemented) as e:
+            raise NotSupportedError(e.message, original_error=e) from e
+        except (ydb.issues.BadRequest, ydb.issues.SchemeError) as e:
+            raise ProgrammingError(e.message, original_error=e) 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, original_error=e) from e
+        except ydb.issues.GenericError as e:
+            raise DataError(e.message, original_error=e) from e
+        except ydb.issues.InternalError as e:
+            raise InternalError(e.message, original_error=e) from e
+        except ydb.Error as e:
+            raise DatabaseError(e.message, original_error=e) from e
+        except Exception as e:
+            raise DatabaseError("Failed to execute query") from e
+
+    return wrapper