ydb-sqlalchemy 0.0.1b4__py2.py3-none-any.whl

ydb_sqlalchemy/dbapi/cursor.py

@@ -0,0 +1,304 @@
+import collections.abc
+import dataclasses
+import functools
+import itertools
+import logging
+from typing import Any, Dict, List, Mapping, Optional, Sequence, Union
+
+import ydb
+import ydb.aio
+from sqlalchemy import util
+
+from .errors import (
+    DatabaseError,
+    DataError,
+    IntegrityError,
+    InternalError,
+    NotSupportedError,
+    OperationalError,
+    ProgrammingError,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def get_column_type(type_obj: Any) -> str:
+    return str(ydb.convert.type_to_native(type_obj))
+
+
+@dataclasses.dataclass
+class YdbQuery:
+    yql_text: str
+    parameters_types: Dict[str, Union[ydb.PrimitiveType, ydb.AbstractTypeBuilder]] = dataclasses.field(
+        default_factory=dict
+    )
+    is_ddl: bool = False
+
+
+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
+
+
+class Cursor:
+    def __init__(
+        self,
+        session_pool: Union[ydb.SessionPool, ydb.aio.SessionPool],
+        tx_mode: ydb.AbstractTransactionModeBuilder,
+        tx_context: Optional[ydb.BaseTxContext] = None,
+    ):
+        self.session_pool = session_pool
+        self.tx_mode = tx_mode
+        self.tx_context = tx_context
+        self.description = None
+        self.arraysize = 1
+        self.rows = None
+        self._rows_prefetched = None
+
+    @_handle_ydb_errors
+    def describe_table(self, abs_table_path: str) -> ydb.TableDescription:
+        return self._retry_operation_in_pool(self._describe_table, abs_table_path)
+
+    def check_exists(self, table_path: str) -> bool:
+        try:
+            self._retry_operation_in_pool(self._describe_path, table_path)
+            return True
+        except ydb.SchemeError:
+            return False
+
+    @_handle_ydb_errors
+    def get_table_names(self) -> List[str]:
+        directory: ydb.Directory = self._retry_operation_in_pool(self._list_directory)
+        return [child.name for child in directory.children if child.is_table()]
+
+    def execute(self, operation: YdbQuery, parameters: Optional[Mapping[str, Any]] = None):
+        if operation.is_ddl or not operation.parameters_types:
+            query = operation.yql_text
+            is_ddl = operation.is_ddl
+        else:
+            query = ydb.DataQuery(operation.yql_text, operation.parameters_types)
+            is_ddl = operation.is_ddl
+
+        logger.info("execute sql: %s, params: %s", query, parameters)
+        if is_ddl:
+            chunks = self._execute_ddl(query)
+        else:
+            chunks = self._execute_dml(query, parameters)
+
+        rows = self._rows_iterable(chunks)
+        # Prefetch the description:
+        try:
+            first_row = next(rows)
+        except StopIteration:
+            pass
+        else:
+            rows = itertools.chain((first_row,), rows)
+        if self.rows is not None:
+            rows = itertools.chain(self.rows, rows)
+
+        self.rows = rows
+
+    @_handle_ydb_errors
+    def _execute_dml(
+        self, query: Union[ydb.DataQuery, str], parameters: Optional[Mapping[str, Any]] = None
+    ) -> ydb.convert.ResultSets:
+        prepared_query = query
+        if isinstance(query, str) and parameters:
+            if self.tx_context:
+                prepared_query = self._run_operation_in_session(self._prepare, query)
+            else:
+                prepared_query = self._retry_operation_in_pool(self._prepare, query)
+
+        if self.tx_context:
+            return self._run_operation_in_tx(self._execute_in_tx, prepared_query, parameters)
+
+        return self._retry_operation_in_pool(self._execute_in_session, self.tx_mode, prepared_query, parameters)
+
+    @_handle_ydb_errors
+    def _execute_ddl(self, query: str) -> ydb.convert.ResultSets:
+        return self._retry_operation_in_pool(self._execute_scheme, query)
+
+    @staticmethod
+    def _execute_scheme(session: ydb.Session, query: str) -> ydb.convert.ResultSets:
+        return session.execute_scheme(query)
+
+    @staticmethod
+    def _describe_table(session: ydb.Session, abs_table_path: str) -> ydb.TableDescription:
+        return session.describe_table(abs_table_path)
+
+    @staticmethod
+    def _describe_path(session: ydb.Session, table_path: str) -> ydb.SchemeEntry:
+        return session._driver.scheme_client.describe_path(table_path)
+
+    @staticmethod
+    def _list_directory(session: ydb.Session) -> ydb.Directory:
+        return session._driver.scheme_client.list_directory(session._driver._driver_config.database)
+
+    @staticmethod
+    def _prepare(session: ydb.Session, query: str) -> ydb.DataQuery:
+        return session.prepare(query)
+
+    @staticmethod
+    def _execute_in_tx(
+        tx_context: ydb.TxContext, prepared_query: ydb.DataQuery, parameters: Optional[Mapping[str, Any]]
+    ) -> ydb.convert.ResultSets:
+        return tx_context.execute(prepared_query, parameters, commit_tx=False)
+
+    @staticmethod
+    def _execute_in_session(
+        session: ydb.Session,
+        tx_mode: ydb.AbstractTransactionModeBuilder,
+        prepared_query: ydb.DataQuery,
+        parameters: Optional[Mapping[str, Any]],
+    ) -> ydb.convert.ResultSets:
+        return session.transaction(tx_mode).execute(prepared_query, parameters, commit_tx=True)
+
+    def _run_operation_in_tx(self, callee: collections.abc.Callable, *args, **kwargs):
+        return callee(self.tx_context, *args, **kwargs)
+
+    def _run_operation_in_session(self, callee: collections.abc.Callable, *args, **kwargs):
+        return callee(self.tx_context.session, *args, **kwargs)
+
+    def _retry_operation_in_pool(self, callee: collections.abc.Callable, *args, **kwargs):
+        return self.session_pool.retry_operation_sync(callee, None, *args, **kwargs)
+
+    def _rows_iterable(self, chunks_iterable: ydb.convert.ResultSets):
+        try:
+            for chunk in chunks_iterable:
+                self.description = [
+                    (
+                        col.name,
+                        get_column_type(col.type),
+                        None,
+                        None,
+                        None,
+                        None,
+                        None,
+                    )
+                    for col in chunk.columns
+                ]
+                for row in chunk.rows:
+                    # returns tuple to be compatible with SqlAlchemy and because
+                    #  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, original_error=e) from e
+
+    def _ensure_prefetched(self):
+        if self.rows is not None and self._rows_prefetched is None:
+            self._rows_prefetched = list(self.rows)
+            self.rows = iter(self._rows_prefetched)
+        return self._rows_prefetched
+
+    def executemany(self, operation: YdbQuery, seq_of_parameters: Optional[Sequence[Mapping[str, Any]]]):
+        for parameters in seq_of_parameters:
+            self.execute(operation, parameters)
+
+    def executescript(self, script):
+        return self.execute(script)
+
+    def fetchone(self):
+        return next(self.rows or [], None)
+
+    def fetchmany(self, size=None):
+        return list(itertools.islice(self.rows, size or self.arraysize))
+
+    def fetchall(self):
+        return list(self.rows)
+
+    def nextset(self):
+        self.fetchall()
+
+    def setinputsizes(self, sizes):
+        pass
+
+    def setoutputsize(self, column=None):
+        pass
+
+    def close(self):
+        self.rows = None
+        self._rows_prefetched = None
+
+    @property
+    def rowcount(self):
+        return len(self._ensure_prefetched())
+
+
+class AsyncCursor(Cursor):
+    _await = staticmethod(util.await_only)
+
+    @staticmethod
+    async def _describe_table(session: ydb.aio.table.Session, abs_table_path: str) -> ydb.TableDescription:
+        return await session.describe_table(abs_table_path)
+
+    @staticmethod
+    async def _describe_path(session: ydb.aio.table.Session, table_path: str) -> ydb.SchemeEntry:
+        return await session._driver.scheme_client.describe_path(table_path)
+
+    @staticmethod
+    async def _list_directory(session: ydb.aio.table.Session) -> ydb.Directory:
+        return await session._driver.scheme_client.list_directory(session._driver._driver_config.database)
+
+    @staticmethod
+    async def _execute_scheme(session: ydb.aio.table.Session, query: str) -> ydb.convert.ResultSets:
+        return await session.execute_scheme(query)
+
+    @staticmethod
+    async def _prepare(session: ydb.aio.table.Session, query: str) -> ydb.DataQuery:
+        return await session.prepare(query)
+
+    @staticmethod
+    async def _execute_in_tx(
+        tx_context: ydb.aio.table.TxContext, prepared_query: ydb.DataQuery, parameters: Optional[Mapping[str, Any]]
+    ) -> ydb.convert.ResultSets:
+        return await tx_context.execute(prepared_query, parameters, commit_tx=False)
+
+    @staticmethod
+    async def _execute_in_session(
+        session: ydb.aio.table.Session,
+        tx_mode: ydb.AbstractTransactionModeBuilder,
+        prepared_query: ydb.DataQuery,
+        parameters: Optional[Mapping[str, Any]],
+    ) -> ydb.convert.ResultSets:
+        return await session.transaction(tx_mode).execute(prepared_query, parameters, commit_tx=True)
+
+    def _run_operation_in_tx(self, callee: collections.abc.Coroutine, *args, **kwargs):
+        return self._await(callee(self.tx_context, *args, **kwargs))
+
+    def _run_operation_in_session(self, callee: collections.abc.Coroutine, *args, **kwargs):
+        return self._await(callee(self.tx_context.session, *args, **kwargs))
+
+    def _retry_operation_in_pool(self, callee: collections.abc.Coroutine, *args, **kwargs):
+        return self._await(self.session_pool.retry_operation(callee, *args, **kwargs))

ydb_sqlalchemy/dbapi/errors.py

@@ -0,0 +1,103 @@
+from typing import List, Optional
+
+import ydb
+from google.protobuf.message import Message
+
+
+class Warning(Exception):
+    pass
+
+
+class Error(Exception):
+    def __init__(
+        self,
+        message: str,
+        original_error: Optional[ydb.Error] = None,
+    ):
+        super(Error, self).__init__(message)
+
+        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):
+    pass
+
+
+class DatabaseError(Error):
+    pass
+
+
+class DataError(DatabaseError):
+    pass
+
+
+class OperationalError(DatabaseError):
+    pass
+
+
+class IntegrityError(DatabaseError):
+    pass
+
+
+class InternalError(DatabaseError):
+    pass
+
+
+class ProgrammingError(DatabaseError):
+    pass
+
+
+class NotSupportedError(DatabaseError):
+    pass
+
+
+def _pretty_issues(issues: List[Message]) -> str:
+    if issues is None:
+        return None
+
+    children_messages = [_get_messages(issue, root=True) for issue in issues]
+
+    if None in children_messages:
+        return None
+
+    return "\n" + "\n".join(children_messages)
+
+
+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
+
+    margin_str = " " * depth * indent
+    pre_message = ""
+    children = ""
+
+    if issue.issues:
+        collapsed_messages = []
+        while not root and len(issue.issues) == 1:
+            collapsed_messages.append(issue.message)
+            issue = issue.issues[0]
+
+        if collapsed_messages:
+            pre_message = f"{margin_str}{', '.join(collapsed_messages)}\n"
+            depth += 1
+            margin_str = " " * depth * indent
+
+        children_messages = [
+            _get_messages(iss, max_depth=max_depth, indent=indent, depth=depth + 1) for iss in issue.issues
+        ]
+
+        if None in children_messages:
+            return None
+
+        children = "\n".join(children_messages)
+
+    return (
+        f"{pre_message}{margin_str}{issue.message}\n{margin_str}"
+        f"severity level: {issue.severity}\n{margin_str}"
+        f"issue code: {issue.issue_code}\n{children}"
+    )