Implement Async driver + tests

Author: LuckySting

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,4 @@ dockerpty==0.4.1
 flake8==3.9.2
 black==23.3.0
 pytest-cov
+pytest-asyncio

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,3 +1,4 @@
+import asyncio
 from datetime import date, datetime
 from decimal import Decimal
 from typing import NamedTuple
@@ -21,6 +22,8 @@ def clear_sql(stm):
 
 
 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 +41,8 @@ def test_sa_text(self, connection):
 
 
 class TestCrud(TablesTest):
+    __backend__ = True
+
     @classmethod
     def define_tables(cls, metadata):
         Table(
@@ -82,6 +87,8 @@ def test_sa_crud(self, connection):
 
 
 class TestSimpleSelect(TablesTest):
+    __backend__ = True
+
     @classmethod
     def define_tables(cls, metadata):
         Table(
@@ -174,6 +181,8 @@ def test_sa_select_simple(self, connection):
 
 
 class TestTypes(TablesTest):
+    __backend__ = True
+
     @classmethod
     def define_tables(cls, metadata):
         Table(
@@ -226,6 +235,7 @@ def test_integer_types(self, connection):
 
 
 class TestWithClause(TablesTest):
+    __backend__ = True
     run_create_tables = "each"
 
     @staticmethod
@@ -238,10 +248,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",
@@ -389,6 +396,8 @@ def test_several_keys(self, connection, metadata):
 
 
 class TestTransaction(TablesTest):
+    __backend__ = True
+
     @classmethod
     def define_tables(cls, metadata: sa.MetaData):
         Table(
@@ -477,6 +486,8 @@ def test_not_interactive_transaction(
 
 
 class TestTransactionIsolationLevel(TestBase):
+    __backend__ = True
+
     class IsolationSettings(NamedTuple):
         ydb_mode: ydb.AbstractTransactionModeBuilder
         interactive: bool
@@ -508,7 +519,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)
@@ -520,13 +534,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})
@@ -559,7 +574,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(
@@ -659,6 +703,8 @@ def test_upsert_from_select(self, connection, metadata):
 
 
 class TestUpsertDoesNotReplaceInsert(TablesTest):
+    __backend__ = True
+
     @classmethod
     def define_tables(cls, metadata):
         Table(