refactor connection

Author: vgvoleg

pyproject.toml

@@ -40,7 +40,40 @@ target-version = "py39"
 src = ["ydb_dbapi", "tests"]
 
 [tool.ruff.lint]
-select = ["F", "E"]
+ignore = [
+    "D", # Allow not to have docstring for each method
+    "ANN101", # Allow not to specify `Self` type for `self` param
+    "ANN102", # Allow not to specify `Class` type for `cls` param
+    "ANN401", # Allow to use Any type
+    "FIX002", # Allow to use TODO in code
+    "TD001", # Allow to specify FIXME (still trigger FIX001)
+    "TD002", # Allow not to specify author of todo
+    "TD003", # Allow not to specify ticket of todo (it doesn't work)
+    "PERF401", # Allow not to use list comprehension each time
+    "FBT", # Allow boolean positional arguments
+    "TCH002", # Allow not to use `if TYPE_CHECKING` for all imports, which don't used explicitly
+    "TCH003", # Allow not to use `if TYPE_CHECKING` for std lib imports
+    "ISC001", # Conflicts with formatter
+    "COM812", # Conflicts with formatter,
+    # Ignores below could be deleted
+    "EM101", # Allow to use string literals in exceptions
+    "TRY003", # Allow specifying long messages outside the exception class
+]
+select = ["ALL"]
+
+# Allow fix for all enabled rules (when `--fix`) is provided.
+fixable = ["ALL"]
+unfixable = ["B"]
+
+# Allow unused variables when underscore-prefixed.
+dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
+
+[tool.ruff.lint.isort]
+force-single-line = true
+
+[tool.ruff.lint.per-file-ignores]
+"**/test_*.py" = ["S", "SLF", "ANN201", "ARG", "PLR2004"]
+"__init__.py" = ["F401", "F403"]
 
 [tool.pytest.ini_options]
 asyncio_mode = "auto"

tests/conftest.py

@@ -1,13 +1,11 @@
-import pytest
-import ydb
+from __future__ import annotations
 
-from typing import (
-    Any,
-    Callable,
-    Generator,
-    Optional,
-)
+from collections.abc import Generator
+from typing import Any
+from typing import Callable
 
+import pytest
+import ydb
 from testcontainers.core.generic import DbContainer
 from testcontainers.core.generic import wait_container_is_ready
 from testcontainers.core.utils import setup_logger
@@ -18,7 +16,7 @@
 class YDBContainer(DbContainer):
     def __init__(
         self,
-        name: Optional[str] = None,
+        name: str | None = None,
         port: str = "2135",
         image: str = "ydbplatform/local-ydb:trunk",
         **kwargs: Any,
@@ -116,7 +114,7 @@ def connection_kwargs(ydb_container: YDBContainer) -> dict:
     }
 
 
-@pytest.fixture()
+@pytest.fixture
 async def driver(ydb_container, event_loop):
     driver = ydb.aio.Driver(
         connection_string=ydb_container.get_connection_string()
@@ -147,3 +145,11 @@ async def session_pool(driver: ydb.aio.Driver):
         )
 
         yield session_pool
+
+@pytest.fixture
+async def session(session_pool: ydb.aio.QuerySessionPool):
+    session = await session_pool.acquire()
+
+    yield session
+
+    await session_pool.release(session)

tests/test_connection.py

@@ -3,7 +3,6 @@
 import pytest
 import pytest_asyncio
 import ydb
-
 import ydb_dbapi as dbapi
 
 
@@ -25,14 +24,13 @@ async def _test_isolation_level_read_only(
 
         connection.set_isolation_level(isolation_level)
 
-        await connection.begin()
-
         async with connection.cursor() as cursor:
             query = "UPSERT INTO foo(id) VALUES (1)"
             if read_only:
                 with pytest.raises(dbapi.DatabaseError):
                     await cursor.execute(query)
                     await cursor.finish_query()
+
             else:
                 await cursor.execute(query)
 

tests/test_cursor.py

@@ -3,8 +3,8 @@
 
 
 @pytest.mark.asyncio
-async def test_cursor_ddl(session_pool):
-    cursor = ydb_dbapi.Cursor(session_pool=session_pool)
+async def test_cursor_ddl(session):
+    cursor = ydb_dbapi.Cursor(session=session)
 
     yql = """
         CREATE TABLE table (
@@ -27,8 +27,8 @@ async def test_cursor_ddl(session_pool):
 
 
 @pytest.mark.asyncio
-async def test_cursor_dml(session_pool):
-    cursor = ydb_dbapi.Cursor(session_pool=session_pool)
+async def test_cursor_dml(session):
+    cursor = ydb_dbapi.Cursor(session=session)
     yql_text = """
     INSERT INTO table (id, val) VALUES
     (1, 1),
@@ -39,7 +39,7 @@ async def test_cursor_dml(session_pool):
     await cursor.execute(query=yql_text)
     assert await cursor.fetchone() is None
 
-    cursor = ydb_dbapi.Cursor(session_pool=session_pool)
+    cursor = ydb_dbapi.Cursor(session=session)
 
     yql_text = """
     SELECT COUNT(*) FROM table as sum
@@ -53,8 +53,8 @@ async def test_cursor_dml(session_pool):
 
 
 @pytest.mark.asyncio
-async def test_cursor_fetch_one(session_pool):
-    cursor = ydb_dbapi.Cursor(session_pool=session_pool)
+async def test_cursor_fetch_one(session):
+    cursor = ydb_dbapi.Cursor(session=session)
     yql_text = """
     INSERT INTO table (id, val) VALUES
     (1, 1),
@@ -64,7 +64,7 @@ async def test_cursor_fetch_one(session_pool):
     await cursor.execute(query=yql_text)
     assert await cursor.fetchone() is None
 
-    cursor = ydb_dbapi.Cursor(session_pool=session_pool)
+    cursor = ydb_dbapi.Cursor(session=session)
 
     yql_text = """
     SELECT id, val FROM table
@@ -82,8 +82,8 @@ async def test_cursor_fetch_one(session_pool):
 
 
 @pytest.mark.asyncio
-async def test_cursor_fetch_many(session_pool):
-    cursor = ydb_dbapi.Cursor(session_pool=session_pool)
+async def test_cursor_fetch_many(session):
+    cursor = ydb_dbapi.Cursor(session=session)
     yql_text = """
     INSERT INTO table (id, val) VALUES
     (1, 1),
@@ -95,7 +95,7 @@ async def test_cursor_fetch_many(session_pool):
     await cursor.execute(query=yql_text)
     assert await cursor.fetchone() is None
 
-    cursor = ydb_dbapi.Cursor(session_pool=session_pool)
+    cursor = ydb_dbapi.Cursor(session=session)
 
     yql_text = """
     SELECT id, val FROM table
@@ -120,8 +120,8 @@ async def test_cursor_fetch_many(session_pool):
 
 
 @pytest.mark.asyncio
-async def test_cursor_fetch_all(session_pool):
-    cursor = ydb_dbapi.Cursor(session_pool=session_pool)
+async def test_cursor_fetch_all(session):
+    cursor = ydb_dbapi.Cursor(session=session)
     yql_text = """
     INSERT INTO table (id, val) VALUES
     (1, 1),
@@ -132,7 +132,7 @@ async def test_cursor_fetch_all(session_pool):
     await cursor.execute(query=yql_text)
     assert await cursor.fetchone() is None
 
-    cursor = ydb_dbapi.Cursor(session_pool=session_pool)
+    cursor = ydb_dbapi.Cursor(session=session)
 
     yql_text = """
     SELECT id, val FROM table
@@ -152,8 +152,8 @@ async def test_cursor_fetch_all(session_pool):
 
 
 @pytest.mark.asyncio
-async def test_cursor_next_set(session_pool):
-    cursor = ydb_dbapi.Cursor(session_pool=session_pool)
+async def test_cursor_next_set(session):
+    cursor = ydb_dbapi.Cursor(session=session)
     yql_text = """SELECT 1 as val; SELECT 2 as val;"""
 
     await cursor.execute(query=yql_text)

ydb_dbapi/__init__.py

@@ -1,3 +1,5 @@
-from .errors import *  # noqa
-from .connection import Connection, IsolationLevel, connect  # noqa
-from .cursors import Cursor  # noqa
+from .connection import Connection
+from .connection import IsolationLevel
+from .connection import connect
+from .cursors import Cursor
+from .errors import *