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

test_dbapi/test_dbapi.py

@@ -0,0 +1,196 @@
+from contextlib import suppress
+
+import pytest
+import pytest_asyncio
+import sqlalchemy.util as util
+import ydb
+
+import ydb_sqlalchemy.dbapi as dbapi
+
+
+class BaseDBApiTestSuit:
+    def _test_isolation_level_read_only(self, connection: dbapi.Connection, isolation_level: str, read_only: bool):
+        connection.cursor().execute(
+            dbapi.YdbQuery("CREATE TABLE foo(id Int64 NOT NULL, PRIMARY KEY (id))", is_ddl=True)
+        )
+        connection.set_isolation_level(isolation_level)
+
+        cursor = connection.cursor()
+
+        connection.begin()
+
+        query = dbapi.YdbQuery("UPSERT INTO foo(id) VALUES (1)")
+        if read_only:
+            with pytest.raises(dbapi.DatabaseError):
+                cursor.execute(query)
+        else:
+            cursor.execute(query)
+
+        connection.rollback()
+
+        connection.cursor().execute(dbapi.YdbQuery("DROP TABLE foo", is_ddl=True))
+        connection.cursor().close()
+
+    def _test_connection(self, connection: dbapi.Connection):
+        connection.commit()
+        connection.rollback()
+
+        cur = connection.cursor()
+        with suppress(dbapi.DatabaseError):
+            cur.execute(dbapi.YdbQuery("DROP TABLE foo", is_ddl=True))
+
+        assert not connection.check_exists("/local/foo")
+        with pytest.raises(dbapi.ProgrammingError):
+            connection.describe("/local/foo")
+
+        cur.execute(dbapi.YdbQuery("CREATE TABLE foo(id Int64 NOT NULL, PRIMARY KEY (id))", is_ddl=True))
+
+        assert connection.check_exists("/local/foo")
+
+        col = connection.describe("/local/foo").columns[0]
+        assert col.name == "id"
+        assert col.type == ydb.PrimitiveType.Int64
+
+        cur.execute(dbapi.YdbQuery("DROP TABLE foo", is_ddl=True))
+        cur.close()
+
+    def _test_cursor_raw_query(self, connection: dbapi.Connection):
+        cur = connection.cursor()
+        assert cur
+
+        with suppress(dbapi.DatabaseError):
+            cur.execute(dbapi.YdbQuery("DROP TABLE test", is_ddl=True))
+
+        cur.execute(dbapi.YdbQuery("CREATE TABLE test(id Int64 NOT NULL, text Utf8, PRIMARY KEY (id))", is_ddl=True))
+
+        cur.execute(
+            dbapi.YdbQuery(
+                """
+                DECLARE $data AS List<Struct<id:Int64, text: Utf8>>;
+
+                INSERT INTO test SELECT id, text FROM AS_TABLE($data);
+                """,
+                parameters_types={
+                    "$data": ydb.ListType(
+                        ydb.StructType()
+                        .add_member("id", ydb.PrimitiveType.Int64)
+                        .add_member("text", ydb.PrimitiveType.Utf8)
+                    )
+                },
+            ),
+            {
+                "$data": [
+                    {"id": 17, "text": "seventeen"},
+                    {"id": 21, "text": "twenty one"},
+                ]
+            },
+        )
+
+        cur.execute(dbapi.YdbQuery("DROP TABLE test", is_ddl=True))
+
+        cur.close()
+
+    def _test_errors(self, connection: dbapi.Connection):
+        with pytest.raises(dbapi.InterfaceError):
+            dbapi.YdbDBApi().connect("localhost:2136", database="/local666")
+
+        cur = connection.cursor()
+
+        with suppress(dbapi.DatabaseError):
+            cur.execute(dbapi.YdbQuery("DROP TABLE test", is_ddl=True))
+
+        with pytest.raises(dbapi.DataError):
+            cur.execute(dbapi.YdbQuery("SELECT 18446744073709551616"))
+
+        with pytest.raises(dbapi.DataError):
+            cur.execute(dbapi.YdbQuery("SELECT * FROM 拉屎"))
+
+        with pytest.raises(dbapi.DataError):
+            cur.execute(dbapi.YdbQuery("SELECT floor(5 / 2)"))
+
+        with pytest.raises(dbapi.ProgrammingError):
+            cur.execute(dbapi.YdbQuery("SELECT * FROM test"))
+
+        cur.execute(dbapi.YdbQuery("CREATE TABLE test(id Int64, PRIMARY KEY (id))", is_ddl=True))
+
+        cur.execute(dbapi.YdbQuery("INSERT INTO test(id) VALUES(1)"))
+        with pytest.raises(dbapi.IntegrityError):
+            cur.execute(dbapi.YdbQuery("INSERT INTO test(id) VALUES(1)"))
+
+        cur.execute(dbapi.YdbQuery("DROP TABLE test", is_ddl=True))
+        cur.close()
+
+
+class TestSyncConnection(BaseDBApiTestSuit):
+    @pytest.fixture
+    def sync_connection(self) -> dbapi.Connection:
+        conn = dbapi.YdbDBApi().connect(host="localhost", port="2136", database="/local")
+        try:
+            yield conn
+        finally:
+            conn.close()
+
+    @pytest.mark.parametrize(
+        "isolation_level, read_only",
+        [
+            (dbapi.IsolationLevel.SERIALIZABLE, False),
+            (dbapi.IsolationLevel.AUTOCOMMIT, False),
+            (dbapi.IsolationLevel.ONLINE_READONLY, True),
+            (dbapi.IsolationLevel.ONLINE_READONLY_INCONSISTENT, True),
+            (dbapi.IsolationLevel.STALE_READONLY, True),
+            (dbapi.IsolationLevel.SNAPSHOT_READONLY, True),
+        ],
+    )
+    def test_isolation_level_read_only(self, isolation_level: str, read_only: bool, sync_connection: dbapi.Connection):
+        self._test_isolation_level_read_only(sync_connection, isolation_level, read_only)
+
+    def test_connection(self, sync_connection: dbapi.Connection):
+        self._test_connection(sync_connection)
+
+    def test_cursor_raw_query(self, sync_connection: dbapi.Connection):
+        return self._test_cursor_raw_query(sync_connection)
+
+    def test_errors(self, sync_connection: dbapi.Connection):
+        return self._test_errors(sync_connection)
+
+
+class TestAsyncConnection(BaseDBApiTestSuit):
+    @pytest_asyncio.fixture
+    async def async_connection(self) -> dbapi.AsyncConnection:
+        def connect():
+            return dbapi.YdbDBApi().async_connect(host="localhost", port="2136", database="/local")
+
+        conn = await util.greenlet_spawn(connect)
+        try:
+            yield conn
+        finally:
+            await util.greenlet_spawn(conn.close)
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "isolation_level, read_only",
+        [
+            (dbapi.IsolationLevel.SERIALIZABLE, False),
+            (dbapi.IsolationLevel.AUTOCOMMIT, False),
+            (dbapi.IsolationLevel.ONLINE_READONLY, True),
+            (dbapi.IsolationLevel.ONLINE_READONLY_INCONSISTENT, True),
+            (dbapi.IsolationLevel.STALE_READONLY, True),
+            (dbapi.IsolationLevel.SNAPSHOT_READONLY, True),
+        ],
+    )
+    async def test_isolation_level_read_only(
+        self, isolation_level: str, read_only: bool, async_connection: dbapi.AsyncConnection
+    ):
+        await util.greenlet_spawn(self._test_isolation_level_read_only, async_connection, isolation_level, read_only)
+
+    @pytest.mark.asyncio
+    async def test_connection(self, async_connection: dbapi.AsyncConnection):
+        await util.greenlet_spawn(self._test_connection, async_connection)
+
+    @pytest.mark.asyncio
+    async def test_cursor_raw_query(self, async_connection: dbapi.AsyncConnection):
+        await util.greenlet_spawn(self._test_cursor_raw_query, async_connection)
+
+    @pytest.mark.asyncio
+    async def test_errors(self, async_connection: dbapi.AsyncConnection):
+        await util.greenlet_spawn(self._test_errors, async_connection)

ydb_sqlalchemy/__init__.py

@@ -0,0 +1,3 @@
+from .dbapi import IsolationLevel  # noqa: F401
+from .sqlalchemy import Upsert, types, upsert  # noqa: F401
+from _version import VERSION

ydb_sqlalchemy/_version.py

@@ -0,0 +1 @@
+VERSION = "0.0.1b4"

ydb_sqlalchemy/dbapi/__init__.py

@@ -0,0 +1,43 @@
+from .connection import AsyncConnection, Connection, IsolationLevel  # noqa: F401
+from .cursor import AsyncCursor, Cursor, YdbQuery  # noqa: F401
+from .errors import (
+    DatabaseError,
+    DataError,
+    Error,
+    IntegrityError,
+    InterfaceError,
+    InternalError,
+    NotSupportedError,
+    OperationalError,
+    ProgrammingError,
+    Warning,
+)
+
+
+class YdbDBApi:
+    def __init__(self):
+        self.paramstyle = "pyformat"
+        self.threadsafety = 0
+        self.apilevel = "1.0"
+        self._init_dbapi_attributes()
+
+    def _init_dbapi_attributes(self):
+        for name, value in {
+            "Warning": Warning,
+            "Error": Error,
+            "InterfaceError": InterfaceError,
+            "DatabaseError": DatabaseError,
+            "DataError": DataError,
+            "OperationalError": OperationalError,
+            "IntegrityError": IntegrityError,
+            "InternalError": InternalError,
+            "ProgrammingError": ProgrammingError,
+            "NotSupportedError": NotSupportedError,
+        }.items():
+            setattr(self, name, value)
+
+    def connect(self, *args, **kwargs) -> Connection:
+        return Connection(*args, **kwargs)
+
+    def async_connect(self, *args, **kwargs) -> AsyncConnection:
+        return AsyncConnection(*args, **kwargs)

ydb_sqlalchemy/dbapi/connection.py

@@ -0,0 +1,180 @@
+import collections.abc
+import posixpath
+from typing import Any, List, NamedTuple, Optional
+
+import sqlalchemy.util as util
+import ydb
+
+from .cursor import AsyncCursor, Cursor
+from .errors import InterfaceError, InternalError, NotSupportedError
+
+
+class IsolationLevel:
+    SERIALIZABLE = "SERIALIZABLE"
+    ONLINE_READONLY = "ONLINE READONLY"
+    ONLINE_READONLY_INCONSISTENT = "ONLINE READONLY INCONSISTENT"
+    STALE_READONLY = "STALE READONLY"
+    SNAPSHOT_READONLY = "SNAPSHOT READONLY"
+    AUTOCOMMIT = "AUTOCOMMIT"
+
+
+class Connection:
+    _await = staticmethod(util.await_only)
+
+    _is_async = False
+    _ydb_driver_class = ydb.Driver
+    _ydb_session_pool_class = ydb.SessionPool
+    _ydb_table_client_class = ydb.TableClient
+    _cursor_class = Cursor
+
+    def __init__(
+        self,
+        host: str = "",
+        port: str = "",
+        database: str = "",
+        **conn_kwargs: Any,
+    ):
+        self.endpoint = f"grpc://{host}:{port}"
+        self.database = database
+        self.conn_kwargs = conn_kwargs
+
+        if "ydb_session_pool" in self.conn_kwargs:  # Use session pool managed manually
+            self._shared_session_pool = True
+            self.session_pool: ydb.SessionPool = self.conn_kwargs.pop("ydb_session_pool")
+            self.driver = (
+                self.session_pool._driver
+                if hasattr(self.session_pool, "_driver")
+                else self.session_pool._pool_impl._driver
+            )
+            self.driver.table_client = self._ydb_table_client_class(self.driver, self._get_table_client_settings())
+        else:
+            self._shared_session_pool = False
+            self.driver = self._create_driver()
+            self.session_pool = self._ydb_session_pool_class(self.driver, size=5)
+
+        self.interactive_transaction: bool = False  # AUTOCOMMIT
+        self.tx_mode: ydb.AbstractTransactionModeBuilder = ydb.SerializableReadWrite()
+        self.tx_context: Optional[ydb.TxContext] = None
+
+    def cursor(self):
+        return self._cursor_class(self.session_pool, self.tx_mode, self.tx_context)
+
+    def describe(self, table_path: str) -> ydb.TableDescription:
+        abs_table_path = posixpath.join(self.database, table_path)
+        cursor = self.cursor()
+        return cursor.describe_table(abs_table_path)
+
+    def check_exists(self, table_path: str) -> ydb.SchemeEntry:
+        cursor = self.cursor()
+        return cursor.check_exists(table_path)
+
+    def get_table_names(self) -> List[str]:
+        cursor = self.cursor()
+        return cursor.get_table_names()
+
+    def set_isolation_level(self, isolation_level: str):
+        class IsolationSettings(NamedTuple):
+            ydb_mode: ydb.AbstractTransactionModeBuilder
+            interactive: bool
+
+        ydb_isolation_settings_map = {
+            IsolationLevel.AUTOCOMMIT: IsolationSettings(ydb.SerializableReadWrite(), interactive=False),
+            IsolationLevel.SERIALIZABLE: IsolationSettings(ydb.SerializableReadWrite(), interactive=True),
+            IsolationLevel.ONLINE_READONLY: IsolationSettings(ydb.OnlineReadOnly(), interactive=False),
+            IsolationLevel.ONLINE_READONLY_INCONSISTENT: IsolationSettings(
+                ydb.OnlineReadOnly().with_allow_inconsistent_reads(), interactive=False
+            ),
+            IsolationLevel.STALE_READONLY: IsolationSettings(ydb.StaleReadOnly(), interactive=False),
+            IsolationLevel.SNAPSHOT_READONLY: IsolationSettings(ydb.SnapshotReadOnly(), interactive=True),
+        }
+        ydb_isolation_settings = ydb_isolation_settings_map[isolation_level]
+        if self.tx_context and self.tx_context.tx_id:
+            raise InternalError("Failed to set transaction mode: transaction is already began")
+        self.tx_mode = ydb_isolation_settings.ydb_mode
+        self.interactive_transaction = ydb_isolation_settings.interactive
+
+    def get_isolation_level(self) -> str:
+        if self.tx_mode.name == ydb.SerializableReadWrite().name:
+            if self.interactive_transaction:
+                return IsolationLevel.SERIALIZABLE
+            else:
+                return IsolationLevel.AUTOCOMMIT
+        elif self.tx_mode.name == ydb.OnlineReadOnly().name:
+            if self.tx_mode.settings.allow_inconsistent_reads:
+                return IsolationLevel.ONLINE_READONLY_INCONSISTENT
+            else:
+                return IsolationLevel.ONLINE_READONLY
+        elif self.tx_mode.name == ydb.StaleReadOnly().name:
+            return IsolationLevel.STALE_READONLY
+        elif self.tx_mode.name == ydb.SnapshotReadOnly().name:
+            return IsolationLevel.SNAPSHOT_READONLY
+        else:
+            raise NotSupportedError(f"{self.tx_mode.name} is not supported")
+
+    def begin(self):
+        self.tx_context = None
+        if self.interactive_transaction:
+            session = self._maybe_await(self.session_pool.acquire)
+            self.tx_context = session.transaction(self.tx_mode)
+            self._maybe_await(self.tx_context.begin)
+
+    def commit(self):
+        if self.tx_context and self.tx_context.tx_id:
+            self._maybe_await(self.tx_context.commit)
+            self._maybe_await(self.session_pool.release, self.tx_context.session)
+            self.tx_context = None
+
+    def rollback(self):
+        if self.tx_context and self.tx_context.tx_id:
+            self._maybe_await(self.tx_context.rollback)
+            self._maybe_await(self.session_pool.release, self.tx_context.session)
+            self.tx_context = None
+
+    def close(self):
+        self.rollback()
+        if not self._shared_session_pool:
+            self._maybe_await(self.session_pool.stop)
+            self._stop_driver()
+
+    @classmethod
+    def _maybe_await(cls, callee: collections.abc.Callable, *args, **kwargs) -> Any:
+        if cls._is_async:
+            return cls._await(callee(*args, **kwargs))
+        return callee(*args, **kwargs)
+
+    def _get_table_client_settings(self) -> ydb.TableClientSettings:
+        return (
+            ydb.TableClientSettings()
+            .with_native_date_in_result_sets(True)
+            .with_native_datetime_in_result_sets(True)
+            .with_native_timestamp_in_result_sets(True)
+            .with_native_interval_in_result_sets(True)
+            .with_native_json_in_result_sets(False)
+        )
+
+    def _create_driver(self):
+        driver_config = ydb.DriverConfig(
+            endpoint=self.endpoint,
+            database=self.database,
+            table_client_settings=self._get_table_client_settings(),
+        )
+        driver = self._ydb_driver_class(driver_config)
+        try:
+            self._maybe_await(driver.wait, timeout=5, fail_fast=True)
+        except ydb.Error as 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
+        return driver
+
+    def _stop_driver(self):
+        self._maybe_await(self.driver.stop)
+
+
+class AsyncConnection(Connection):
+    _is_async = True
+    _ydb_driver_class = ydb.aio.Driver
+    _ydb_session_pool_class = ydb.aio.SessionPool
+    _ydb_table_client_class = ydb.aio.table.TableClient
+    _cursor_class = AsyncCursor

ydb_sqlalchemy/dbapi/constants.py

@@ -0,0 +1,218 @@
+YDB_KEYWORDS = {
+    "abort",
+    "action",
+    "add",
+    "after",
+    "all",
+    "alter",
+    "analyze",
+    "and",
+    "ansi",
+    "any",
+    "array",
+    "as",
+    "asc",
+    "assume",
+    "async",
+    "attach",
+    "autoincrement",
+    "before",
+    "begin",
+    "bernoulli",
+    "between",
+    "bitcast",
+    "by",
+    "cascade",
+    "case",
+    "cast",
+    "changefeed",
+    "check",
+    "collate",
+    "column",
+    "columns",
+    "commit",
+    "compact",
+    "conditional",
+    "conflict",
+    "constraint",
+    "consumer",
+    "cover",
+    "create",
+    "cross",
+    "cube",
+    "current",
+    "current_date",
+    "current_time",
+    "current_timestamp",
+    "data",
+    "database",
+    "decimal",
+    "declare",
+    "default",
+    "deferrable",
+    "deferred",
+    "define",
+    "delete",
+    "desc",
+    "detach",
+    "disable",
+    "discard",
+    "distinct",
+    "do",
+    "drop",
+    "each",
+    "else",
+    "empty",
+    "empty_action",
+    "encrypted",
+    "end",
+    "erase",
+    "error",
+    "escape",
+    "evaluate",
+    "except",
+    "exclude",
+    "exclusion",
+    "exclusive",
+    "exists",
+    "explain",
+    "export",
+    "external",
+    "fail",
+    "family",
+    "filter",
+    "flatten",
+    "following",
+    "for",
+    "foreign",
+    "from",
+    "full",
+    "function",
+    "glob",
+    "group",
+    "grouping",
+    "groups",
+    "hash",
+    "having",
+    "hop",
+    "if",
+    "ignore",
+    "ilike",
+    "immediate",
+    "import",
+    "in",
+    "index",
+    "indexed",
+    "inherits",
+    "initially",
+    "inner",
+    "insert",
+    "instead",
+    "intersect",
+    "into",
+    "is",
+    "isnull",
+    "join",
+    "json_exists",
+    "json_query",
+    "json_value",
+    "key",
+    "left",
+    "like",
+    "limit",
+    "local",
+    "match",
+    "natural",
+    "no",
+    "not",
+    "notnull",
+    "null",
+    "nulls",
+    "object",
+    "of",
+    "offset",
+    "on",
+    "only",
+    "or",
+    "order",
+    "others",
+    "outer",
+    "over",
+    "partition",
+    "passing",
+    "password",
+    "plan",
+    "pragma",
+    "preceding",
+    "presort",
+    "primary",
+    "process",
+    "raise",
+    "range",
+    "reduce",
+    "references",
+    "regexp",
+    "reindex",
+    "release",
+    "rename",
+    "replace",
+    "replication",
+    "reset",
+    "respect",
+    "restrict",
+    "result",
+    "return",
+    "returning",
+    "revert",
+    "right",
+    "rlike",
+    "rollback",
+    "rollup",
+    "row",
+    "rows",
+    "sample",
+    "savepoint",
+    "schema",
+    "select",
+    "semi",
+    "sets",
+    "source",
+    "stream",
+    "subquery",
+    "symbols",
+    "sync",
+    "system",
+    "table",
+    "tablesample",
+    "tablestore",
+    "temp",
+    "temporary",
+    "then",
+    "ties",
+    "to",
+    "topic",
+    "transaction",
+    "trigger",
+    "type",
+    "unbounded",
+    "unconditional",
+    "union",
+    "unique",
+    "unknown",
+    "update",
+    "upsert",
+    "use",
+    "user",
+    "using",
+    "vacuum",
+    "values",
+    "view",
+    "virtual",
+    "when",
+    "where",
+    "window",
+    "with",
+    "without",
+    "wrapper",
+    "xor",
+}