Async driver implementation

.gitignore

@@ -130,3 +130,6 @@ dmypy.json
 
 # PyCharm
 .idea/
+
+# VSCode
+.vscode

README.md

@@ -56,6 +56,7 @@ $ tox -e style
 
 Reformat code:
 ```bash
+$ tox -e isort
 $ tox -e black-format
 ```
 

examples/example.py

@@ -1,11 +1,10 @@
 import datetime
 import logging
-import sqlalchemy as sa
-from sqlalchemy import orm, exc, sql
-from sqlalchemy import Table, Column, Integer, String, Float, TIMESTAMP
 
+import sqlalchemy as sa
 from fill_tables import fill_all_tables, to_days
-from models import Base, Series, Episodes
+from models import Base, Episodes, Series
+from sqlalchemy import TIMESTAMP, Column, Float, Integer, String, Table, exc, orm, sql
 
 
 def describe_table(engine, name):

examples/fill_tables.py

@@ -1,7 +1,6 @@
 import iso8601
-
 import sqlalchemy as sa
-from models import Base, Series, Seasons, Episodes
+from models import Base, Episodes, Seasons, Series
 
 
 def to_days(date):

examples/models.py

@@ -1,7 +1,6 @@
 import sqlalchemy.orm as orm
 from sqlalchemy import Column, Integer, Unicode
 
-
 Base = orm.declarative_base()
 
 

pyproject.toml

@@ -1,2 +1,5 @@
 [tool.black]
 line-length = 120
+
+[tool.isort]
+profile = "black"

setup.cfg

@@ -7,3 +7,5 @@ profile_file=test/profiles.txt
 
 [db]
 default=yql+ydb://localhost:2136/local
+ydb=yql+ydb://localhost:2136/local
+ydb_async=yql+ydb_async://localhost:2136/local

setup.py

@@ -39,6 +39,8 @@
     entry_points={
         "sqlalchemy.dialects": [
             "yql.ydb=ydb_sqlalchemy.sqlalchemy:YqlDialect",
+            "yql.ydb_async=ydb_sqlalchemy.sqlalchemy:AsyncYqlDialect",
+            "ydb_async=ydb_sqlalchemy.sqlalchemy:AsyncYqlDialect",
             "ydb=ydb_sqlalchemy.sqlalchemy:YqlDialect",
             "yql=ydb_sqlalchemy.sqlalchemy:YqlDialect",
         ]

test-requirements.txt

@@ -9,3 +9,5 @@ dockerpty==0.4.1
 flake8==3.9.2
 black==23.3.0
 pytest-cov
+pytest-asyncio
+isort==5.13.2

test/conftest.py

@@ -2,6 +2,8 @@
 from sqlalchemy.dialects import registry
 
 registry.register("yql.ydb", "ydb_sqlalchemy.sqlalchemy", "YqlDialect")
+registry.register("yql.ydb_async", "ydb_sqlalchemy.sqlalchemy", "AsyncYqlDialect")
+registry.register("ydb_async", "ydb_sqlalchemy.sqlalchemy", "AsyncYqlDialect")
 registry.register("ydb", "ydb_sqlalchemy.sqlalchemy", "YqlDialect")
 registry.register("yql", "ydb_sqlalchemy.sqlalchemy", "YqlDialect")
 pytest.register_assert_rewrite("sqlalchemy.testing.assertions")

test/test_core.py

@@ -1,26 +1,27 @@
+import asyncio
 from datetime import date, datetime
 from decimal import Decimal
 from typing import NamedTuple
 
 import pytest
-
 import sqlalchemy as sa
-from sqlalchemy import Table, Column, Integer, Unicode, String
-from sqlalchemy.testing.fixtures import TestBase, TablesTest, config
-
 import ydb
+from sqlalchemy import Column, Integer, String, Table, Unicode
+from sqlalchemy.testing.fixtures import TablesTest, TestBase, config
 from ydb._grpc.v4.protos import ydb_common_pb2
 
-from ydb_sqlalchemy import dbapi, IsolationLevel
-from ydb_sqlalchemy.sqlalchemy import types
+from ydb_sqlalchemy import IsolationLevel, dbapi
 from ydb_sqlalchemy import sqlalchemy as ydb_sa
+from ydb_sqlalchemy.sqlalchemy import types
 
 
 def clear_sql(stm):
     return stm.replace("\n", " ").replace("  ", " ").strip()
 
 
 class TestText(TestBase):
+    __backend__ = True
+
     def test_sa_text(self, connection):
         rs = connection.execute(sa.text("SELECT 1 AS value"))
         assert rs.fetchone() == (1,)
@@ -38,6 +39,8 @@ def test_sa_text(self, connection):
 
 
 class TestCrud(TablesTest):
+    __backend__ = True
+
     @classmethod
     def define_tables(cls, metadata):
         Table(
@@ -82,6 +85,8 @@ def test_sa_crud(self, connection):
 
 
 class TestSimpleSelect(TablesTest):
+    __backend__ = True
+
     @classmethod
     def define_tables(cls, metadata):
         Table(
@@ -174,6 +179,8 @@ def test_sa_select_simple(self, connection):
 
 
 class TestTypes(TablesTest):
+    __backend__ = True
+
     @classmethod
     def define_tables(cls, metadata):
         Table(
@@ -211,6 +218,7 @@ def test_select_types(self, connection):
 
 
 class TestWithClause(TablesTest):
+    __backend__ = True
     run_create_tables = "each"
 
     @staticmethod
@@ -223,10 +231,7 @@ def _create_table_and_get_desc(connection, metadata, **kwargs):
         )
         table.create(connection)
 
-        session: ydb.Session = connection.connection.driver_connection.session_pool.acquire()
-        table_description = session.describe_table("/local/" + table.name)
-        connection.connection.driver_connection.session_pool.release(session)
-        return table_description
+        return connection.connection.driver_connection.describe(table.name)
 
     @pytest.mark.parametrize(
         "auto_partitioning_by_size,res",
@@ -374,6 +379,8 @@ def test_several_keys(self, connection, metadata):
 
 
 class TestTransaction(TablesTest):
+    __backend__ = True
+
     @classmethod
     def define_tables(cls, metadata: sa.MetaData):
         Table(
@@ -462,6 +469,8 @@ def test_not_interactive_transaction(
 
 
 class TestTransactionIsolationLevel(TestBase):
+    __backend__ = True
+
     class IsolationSettings(NamedTuple):
         ydb_mode: ydb.AbstractTransactionModeBuilder
         interactive: bool
@@ -493,7 +502,10 @@ def test_connection_set(self, connection_no_trans: sa.Connection):
 
 
 class TestEngine(TestBase):
-    @pytest.fixture(scope="module")
+    __backend__ = True
+    __only_on__ = "yql+ydb"
+
+    @pytest.fixture(scope="class")
     def ydb_driver(self):
         url = config.db_url
         driver = ydb.Driver(endpoint=f"grpc://{url.host}:{url.port}", database=url.database)
@@ -505,13 +517,14 @@ def ydb_driver(self):
 
         driver.stop()
 
-    @pytest.fixture(scope="module")
+    @pytest.fixture(scope="class")
     def ydb_pool(self, ydb_driver):
         session_pool = ydb.SessionPool(ydb_driver, size=5, workers_threads_count=1)
 
-        yield session_pool
-
-        session_pool.stop()
+        try:
+            yield session_pool
+        finally:
+            session_pool.stop()
 
     def test_sa_queue_pool_with_ydb_shared_session_pool(self, ydb_driver, ydb_pool):
         engine1 = sa.create_engine(config.db_url, poolclass=sa.QueuePool, connect_args={"ydb_session_pool": ydb_pool})
@@ -544,7 +557,36 @@ def test_sa_null_pool_with_ydb_shared_session_pool(self, ydb_driver, ydb_pool):
         assert not ydb_driver._stopped
 
 
+class TestAsyncEngine(TestEngine):
+    __only_on__ = "yql+ydb_async"
+
+    @pytest.fixture(scope="class")
+    def ydb_driver(self):
+        loop = asyncio.get_event_loop()
+        url = config.db_url
+        driver = ydb.aio.Driver(endpoint=f"grpc://{url.host}:{url.port}", database=url.database)
+        try:
+            loop.run_until_complete(driver.wait(timeout=5, fail_fast=True))
+            yield driver
+        finally:
+            loop.run_until_complete(driver.stop())
+
+        loop.run_until_complete(driver.stop())
+
+    @pytest.fixture(scope="class")
+    def ydb_pool(self, ydb_driver):
+        session_pool = ydb.aio.SessionPool(ydb_driver, size=5)
+
+        try:
+            yield session_pool
+        finally:
+            loop = asyncio.get_event_loop()
+            loop.run_until_complete(session_pool.stop())
+
+
 class TestUpsert(TablesTest):
+    __backend__ = True
+
     @classmethod
     def define_tables(cls, metadata):
         Table(
@@ -644,6 +686,8 @@ def test_upsert_from_select(self, connection, metadata):
 
 
 class TestUpsertDoesNotReplaceInsert(TablesTest):
+    __backend__ = True
+
     @classmethod
     def define_tables(cls, metadata):
         Table(

test/test_inspect.py

@@ -1,6 +1,5 @@
 import sqlalchemy as sa
-
-from sqlalchemy import Table, Column, Integer, Unicode, Numeric
+from sqlalchemy import Column, Integer, Numeric, Table, Unicode
 from sqlalchemy.testing.fixtures import TablesTest
 
 

test/test_suite.py

@@ -1,55 +1,86 @@
 import pytest
 import sqlalchemy as sa
 import sqlalchemy.testing.suite.test_types
-
+from sqlalchemy.testing import is_false, is_true
 from sqlalchemy.testing.suite import *  # noqa: F401, F403
-
-from sqlalchemy.testing import is_true, is_false
-from sqlalchemy.testing.suite import eq_, testing, inspect, provide_metadata, config, requirements, fixtures
-from sqlalchemy.testing.suite import func, column, literal_column, select, exists
-from sqlalchemy.testing.suite import MetaData, Column, Table, Integer, String
-
-from sqlalchemy.testing.suite.test_select import (
-    ExistsTest as _ExistsTest,
-    LikeFunctionsTest as _LikeFunctionsTest,
-    CompoundSelectTest as _CompoundSelectTest,
+from sqlalchemy.testing.suite import (
+    Column,
+    Integer,
+    MetaData,
+    String,
+    Table,
+    column,
+    config,
+    eq_,
+    exists,
+    fixtures,
+    func,
+    inspect,
+    literal_column,
+    provide_metadata,
+    requirements,
+    select,
+    testing,
+)
+from sqlalchemy.testing.suite.test_ddl import (
+    LongNameBlowoutTest as _LongNameBlowoutTest,
+)
+from sqlalchemy.testing.suite.test_deprecations import (
+    DeprecatedCompoundSelectTest as _DeprecatedCompoundSelectTest,
+)
+from sqlalchemy.testing.suite.test_dialect import (
+    DifficultParametersTest as _DifficultParametersTest,
+)
+from sqlalchemy.testing.suite.test_dialect import EscapingTest as _EscapingTest
+from sqlalchemy.testing.suite.test_insert import (
+    InsertBehaviorTest as _InsertBehaviorTest,
 )
 from sqlalchemy.testing.suite.test_reflection import (
-    HasTableTest as _HasTableTest,
-    HasIndexTest as _HasIndexTest,
     ComponentReflectionTest as _ComponentReflectionTest,
-    CompositeKeyReflectionTest as _CompositeKeyReflectionTest,
+)
+from sqlalchemy.testing.suite.test_reflection import (
     ComponentReflectionTestExtra as _ComponentReflectionTestExtra,
+)
+from sqlalchemy.testing.suite.test_reflection import (
+    CompositeKeyReflectionTest as _CompositeKeyReflectionTest,
+)
+from sqlalchemy.testing.suite.test_reflection import HasIndexTest as _HasIndexTest
+from sqlalchemy.testing.suite.test_reflection import HasTableTest as _HasTableTest
+from sqlalchemy.testing.suite.test_reflection import (
     QuotedNameArgumentTest as _QuotedNameArgumentTest,
 )
+from sqlalchemy.testing.suite.test_results import RowFetchTest as _RowFetchTest
+from sqlalchemy.testing.suite.test_select import (
+    CompoundSelectTest as _CompoundSelectTest,
+)
+from sqlalchemy.testing.suite.test_select import ExistsTest as _ExistsTest
+from sqlalchemy.testing.suite.test_select import (
+    FetchLimitOffsetTest as _FetchLimitOffsetTest,
+)
+from sqlalchemy.testing.suite.test_select import JoinTest as _JoinTest
+from sqlalchemy.testing.suite.test_select import LikeFunctionsTest as _LikeFunctionsTest
+from sqlalchemy.testing.suite.test_select import OrderByLabelTest as _OrderByLabelTest
+from sqlalchemy.testing.suite.test_types import BinaryTest as _BinaryTest
+from sqlalchemy.testing.suite.test_types import DateTest as _DateTest
 from sqlalchemy.testing.suite.test_types import (
-    IntegerTest as _IntegerTest,
-    NumericTest as _NumericTest,
-    BinaryTest as _BinaryTest,
-    TrueDivTest as _TrueDivTest,
-    TimeTest as _TimeTest,
-    StringTest as _StringTest,
-    NativeUUIDTest as _NativeUUIDTest,
-    TimeMicrosecondsTest as _TimeMicrosecondsTest,
     DateTimeCoercedToDateTimeTest as _DateTimeCoercedToDateTimeTest,
-    DateTest as _DateTest,
+)
+from sqlalchemy.testing.suite.test_types import (
     DateTimeMicrosecondsTest as _DateTimeMicrosecondsTest,
-    DateTimeTest as _DateTimeTest,
-    TimestampMicrosecondsTest as _TimestampMicrosecondsTest,
 )
-from sqlalchemy.testing.suite.test_dialect import (
-    EscapingTest as _EscapingTest,
-    DifficultParametersTest as _DifficultParametersTest,
+from sqlalchemy.testing.suite.test_types import DateTimeTest as _DateTimeTest
+from sqlalchemy.testing.suite.test_types import IntegerTest as _IntegerTest
+from sqlalchemy.testing.suite.test_types import NativeUUIDTest as _NativeUUIDTest
+from sqlalchemy.testing.suite.test_types import NumericTest as _NumericTest
+from sqlalchemy.testing.suite.test_types import StringTest as _StringTest
+from sqlalchemy.testing.suite.test_types import (
+    TimeMicrosecondsTest as _TimeMicrosecondsTest,
 )
-from sqlalchemy.testing.suite.test_select import (
-    JoinTest as _JoinTest,
-    OrderByLabelTest as _OrderByLabelTest,
-    FetchLimitOffsetTest as _FetchLimitOffsetTest,
+from sqlalchemy.testing.suite.test_types import (
+    TimestampMicrosecondsTest as _TimestampMicrosecondsTest,
 )
-from sqlalchemy.testing.suite.test_insert import InsertBehaviorTest as _InsertBehaviorTest
-from sqlalchemy.testing.suite.test_ddl import LongNameBlowoutTest as _LongNameBlowoutTest
-from sqlalchemy.testing.suite.test_results import RowFetchTest as _RowFetchTest
-from sqlalchemy.testing.suite.test_deprecations import DeprecatedCompoundSelectTest as _DeprecatedCompoundSelectTest
+from sqlalchemy.testing.suite.test_types import TimeTest as _TimeTest
+from sqlalchemy.testing.suite.test_types import TrueDivTest as _TrueDivTest
 
 from ydb_sqlalchemy.sqlalchemy import types as ydb_sa_types