review fixes

vgvoleg · vgvoleg · commit af51b1fe1b99 · 2024-10-30T13:10:29.000+03:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -61,6 +61,7 @@ ignore = [
     "TRY003", # Allow specifying long messages outside the exception class
     "SLF001", # Allow access private member,
     "PGH003", # Allow not to specify rule codes
+    "PLR0913", # Allow to have many arguments in function definition
 ]
 select = ["ALL"]
 
diff --git a/ydb_dbapi/connections.py b/ydb_dbapi/connections.py
@@ -1,12 +1,14 @@
 from __future__ import annotations
 
 import posixpath
+from enum import Enum
 from typing import NamedTuple
-from typing import TypedDict
 
 import ydb
-from typing_extensions import NotRequired
-from typing_extensions import Unpack
+from ydb import QuerySessionPool as SessionPool
+from ydb import QueryTxContext as TxContext
+from ydb.aio import QuerySessionPool as AsyncSessionPool
+from ydb.aio import QueryTxContext as AsyncTxContext
 from ydb.retries import retry_operation_async
 from ydb.retries import retry_operation_sync
 
@@ -18,7 +20,7 @@
 from .utils import handle_ydb_errors
 
 
-class IsolationLevel:
+class IsolationLevel(str, Enum):
     SERIALIZABLE = "SERIALIZABLE"
     ONLINE_READONLY = "ONLINE READONLY"
     ONLINE_READONLY_INCONSISTENT = "ONLINE READONLY INCONSISTENT"
@@ -27,49 +29,63 @@ class IsolationLevel:
     AUTOCOMMIT = "AUTOCOMMIT"
 
 
-class ConnectionKwargs(TypedDict):
-    credentials: NotRequired[ydb.AbstractCredentials]
-    ydb_table_path_prefix: NotRequired[str]
-    ydb_session_pool: NotRequired[
-        ydb.QuerySessionPool | ydb.aio.QuerySessionPool
-    ]
+class _IsolationSettings(NamedTuple):
+    ydb_mode: ydb.BaseQueryTxMode
+    interactive: bool
+
+
+_ydb_isolation_settings_map = {
+    IsolationLevel.AUTOCOMMIT: _IsolationSettings(
+        ydb.QuerySerializableReadWrite(), interactive=False
+    ),
+    IsolationLevel.SERIALIZABLE: _IsolationSettings(
+        ydb.QuerySerializableReadWrite(), interactive=True
+    ),
+    IsolationLevel.ONLINE_READONLY: _IsolationSettings(
+        ydb.QueryOnlineReadOnly(), interactive=True
+    ),
+    IsolationLevel.ONLINE_READONLY_INCONSISTENT: _IsolationSettings(
+        ydb.QueryOnlineReadOnly().with_allow_inconsistent_reads(),
+        interactive=True,
+    ),
+    IsolationLevel.STALE_READONLY: _IsolationSettings(
+        ydb.QueryStaleReadOnly(), interactive=True
+    ),
+    IsolationLevel.SNAPSHOT_READONLY: _IsolationSettings(
+        ydb.QuerySnapshotReadOnly(), interactive=True
+    ),
+}
 
 
 class BaseConnection:
-    _tx_mode: ydb.BaseQueryTxMode = ydb.QuerySerializableReadWrite()
-    _tx_context: ydb.QueryTxContext | ydb.aio.QueryTxContext | None = None
-    interactive_transaction: bool = False
-    _shared_session_pool: bool = False
-
     _driver_cls = ydb.Driver
     _pool_cls = ydb.QuerySessionPool
-    _cursor_cls: type[Cursor | AsyncCursor] = Cursor
-
-    _driver: ydb.Driver | ydb.aio.Driver
-    _pool: ydb.QuerySessionPool | ydb.aio.QuerySessionPool
-
-    _current_cursor: AsyncCursor | Cursor | None = None
 
     def __init__(
         self,
         host: str = "",
         port: str = "",
         database: str = "",
-        **conn_kwargs: Unpack[ConnectionKwargs],
+        ydb_table_path_prefix: str = "",
+        credentials: ydb.AbstractCredentials | None = None,
+        ydb_session_pool: SessionPool | AsyncSessionPool | None = None,
+        **kwargs: dict,
     ) -> None:
         self.endpoint = f"grpc://{host}:{port}"
         self.database = database
-        self.conn_kwargs = conn_kwargs
-        self.credentials = self.conn_kwargs.pop("credentials", None)
-        self.table_path_prefix = self.conn_kwargs.pop(
-            "ydb_table_path_prefix", ""
-        )
+        self.credentials = credentials
+        self.table_path_prefix = ydb_table_path_prefix
 
-        if (
-            "ydb_session_pool" in self.conn_kwargs
-        ):  # Use session pool managed manually
+        self.connection_kwargs: dict = kwargs
+
+        self._tx_mode: ydb.BaseQueryTxMode = ydb.QuerySerializableReadWrite()
+        self._tx_context: TxContext | AsyncTxContext | None = None
+        self.interactive_transaction: bool = False
+        self._shared_session_pool: bool = False
+
+        if ydb_session_pool is not None:
             self._shared_session_pool = True
-            self._session_pool = self.conn_kwargs.pop("ydb_session_pool")
+            self._session_pool = ydb_session_pool
             self._driver = self._session_pool._driver
         else:
             driver_config = ydb.DriverConfig(
@@ -82,33 +98,8 @@ def __init__(
 
         self._session: ydb.QuerySession | ydb.aio.QuerySession | None = None
 
-    def set_isolation_level(self, isolation_level: str) -> None:
-        class IsolationSettings(NamedTuple):
-            ydb_mode: ydb.BaseQueryTxMode
-            interactive: bool
-
-        ydb_isolation_settings_map = {
-            IsolationLevel.AUTOCOMMIT: IsolationSettings(
-                ydb.QuerySerializableReadWrite(), interactive=False
-            ),
-            IsolationLevel.SERIALIZABLE: IsolationSettings(
-                ydb.QuerySerializableReadWrite(), interactive=True
-            ),
-            IsolationLevel.ONLINE_READONLY: IsolationSettings(
-                ydb.QueryOnlineReadOnly(), interactive=True
-            ),
-            IsolationLevel.ONLINE_READONLY_INCONSISTENT: IsolationSettings(
-                ydb.QueryOnlineReadOnly().with_allow_inconsistent_reads(),
-                interactive=True,
-            ),
-            IsolationLevel.STALE_READONLY: IsolationSettings(
-                ydb.QueryStaleReadOnly(), interactive=True
-            ),
-            IsolationLevel.SNAPSHOT_READONLY: IsolationSettings(
-                ydb.QuerySnapshotReadOnly(), interactive=True
-            ),
-        }
-        ydb_isolation_settings = ydb_isolation_settings_map[isolation_level]
+    def set_isolation_level(self, isolation_level: IsolationLevel) -> None:
+        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"
@@ -132,7 +123,34 @@ def get_isolation_level(self) -> str:
         msg = f"{self._tx_mode.name} is not supported"
         raise NotSupportedError(msg)
 
-    def cursor(self) -> Cursor | AsyncCursor:
+
+class Connection(BaseConnection):
+    _driver_cls = ydb.Driver
+    _pool_cls = ydb.QuerySessionPool
+    _cursor_cls = Cursor
+
+    def __init__(
+        self,
+        host: str = "",
+        port: str = "",
+        database: str = "",
+        ydb_table_path_prefix: str = "",
+        credentials: ydb.AbstractCredentials | None = None,
+        ydb_session_pool: SessionPool | AsyncSessionPool | None = None,
+        **kwargs: dict,
+    ) -> None:
+        super().__init__(
+            host=host,
+            port=port,
+            database=database,
+            ydb_table_path_prefix=ydb_table_path_prefix,
+            credentials=credentials,
+            ydb_session_pool=ydb_session_pool,
+            **kwargs,
+        )
+        self._current_cursor: Cursor | None = None
+
+    def cursor(self) -> Cursor:
         if self._session is None:
             raise RuntimeError("Connection is not ready, use wait_ready.")
 
@@ -148,16 +166,6 @@ def cursor(self) -> Cursor | AsyncCursor:
         )
         return self._current_cursor
 
-
-class Connection(BaseConnection):
-    _driver_cls = ydb.Driver
-    _pool_cls = ydb.QuerySessionPool
-    _cursor_cls = Cursor
-
-    _driver: ydb.Driver
-    _pool: ydb.QuerySessionPool
-    _current_cursor: Cursor | None = None
-
     def wait_ready(self, timeout: int = 10) -> None:
         try:
             self._driver.wait(timeout, fail_fast=True)
@@ -248,9 +256,42 @@ class AsyncConnection(BaseConnection):
     _pool_cls = ydb.aio.QuerySessionPool
     _cursor_cls = AsyncCursor
 
-    _driver: ydb.aio.Driver
-    _pool: ydb.aio.QuerySessionPool
-    _current_cursor: AsyncCursor | None = None
+    def __init__(
+        self,
+        host: str = "",
+        port: str = "",
+        database: str = "",
+        ydb_table_path_prefix: str = "",
+        credentials: ydb.AbstractCredentials | None = None,
+        ydb_session_pool: SessionPool | AsyncSessionPool | None = None,
+        **kwargs: dict,
+    ) -> None:
+        super().__init__(
+            host=host,
+            port=port,
+            database=database,
+            ydb_table_path_prefix=ydb_table_path_prefix,
+            credentials=credentials,
+            ydb_session_pool=ydb_session_pool,
+            **kwargs,
+        )
+        self._current_cursor: AsyncCursor | None = None
+
+    def cursor(self) -> AsyncCursor:
+        if self._session is None:
+            raise RuntimeError("Connection is not ready, use wait_ready.")
+
+        if self.interactive_transaction:
+            self._tx_context = self._session.transaction(self._tx_mode)
+        else:
+            self._tx_context = None
+
+        self._current_cursor = self._cursor_cls(
+            session=self._session,
+            tx_context=self._tx_context,
+            autocommit=(not self.interactive_transaction),
+        )
+        return self._current_cursor
 
     async def wait_ready(self, timeout: int = 10) -> None:
         try:
diff --git a/ydb_dbapi/cursors.py b/ydb_dbapi/cursors.py
@@ -11,7 +11,6 @@
 from typing_extensions import Self
 
 from .errors import DatabaseError
-from .errors import Error
 from .errors import InterfaceError
 from .errors import ProgrammingError
 from .utils import CursorStatus
@@ -32,11 +31,12 @@ def _get_column_type(type_obj: Any) -> str:
 
 
 class BufferedCursor:
-    arraysize: int = 1
-    _rows: Iterator | None = None
-    _rows_count: int = -1
-    _description: list[tuple] | None = None
-    _state: CursorStatus = CursorStatus.ready
+    def __init__(self) -> None:
+        self.arraysize: int = 1
+        self._rows: Iterator | None = None
+        self._rows_count: int = -1
+        self._description: list[tuple] | None = None
+        self._state: CursorStatus = CursorStatus.ready
 
     @property
     def description(self) -> list[tuple] | None:
@@ -142,6 +142,7 @@ def __init__(
         table_path_prefix: str = "",
         autocommit: bool = True,
     ) -> None:
+        super().__init__()
         self._session = session
         self._tx_context = tx_context
         self._table_path_prefix = table_path_prefix
@@ -167,13 +168,12 @@ def _execute_generic_query(
 
     @handle_ydb_errors
     def _execute_transactional_query(
-        self, query: str, parameters: ParametersType | None = None
+        self,
+        tx_context: ydb.QueryTxContext,
+        query: str,
+        parameters: ParametersType | None = None,
     ) -> Iterator[ydb.convert.ResultSet]:
-        if self._tx_context is None:
-            raise Error(
-                "Unable to execute tx based queries without transaction."
-            )
-        return self._tx_context.execute(
+        return tx_context.execute(
             query=query,
             parameters=parameters,
             commit_tx=self._autocommit,
@@ -188,16 +188,13 @@ def execute(
         self._raise_if_running()
         if self._tx_context is not None:
             self._stream = self._execute_transactional_query(
-                query=query, parameters=parameters
+                tx_context=self._tx_context, query=query, parameters=parameters
             )
         else:
             self._stream = self._execute_generic_query(
                 query=query, parameters=parameters
             )
 
-        if self._stream is None:
-            return
-
         self._begin_query()
 
         self._scroll_stream(replace_current=False)
@@ -256,6 +253,7 @@ def __init__(
         table_path_prefix: str = "",
         autocommit: bool = True,
     ) -> None:
+        super().__init__()
         self._session = session
         self._tx_context = tx_context
         self._table_path_prefix = table_path_prefix
@@ -281,13 +279,12 @@ async def _execute_generic_query(
 
     @handle_ydb_errors
     async def _execute_transactional_query(
-        self, query: str, parameters: ParametersType | None = None
+        self,
+        tx_context: ydb.aio.QueryTxContext,
+        query: str,
+        parameters: ParametersType | None = None,
     ) -> AsyncIterator[ydb.convert.ResultSet]:
-        if self._tx_context is None:
-            raise Error(
-                "Unable to execute tx based queries without transaction."
-            )
-        return await self._tx_context.execute(
+        return await tx_context.execute(
             query=query,
             parameters=parameters,
             commit_tx=self._autocommit,
@@ -302,16 +299,13 @@ async def execute(
         self._raise_if_running()
         if self._tx_context is not None:
             self._stream = await self._execute_transactional_query(
-                query=query, parameters=parameters
+                tx_context=self._tx_context, query=query, parameters=parameters
             )
         else:
             self._stream = await self._execute_generic_query(
                 query=query, parameters=parameters
             )
 
-        if self._stream is None:
-            return
-
         self._begin_query()
 
         await self._scroll_stream(replace_current=False)

Original file line number	Diff line number	Diff line change
`@@ -61,6 +61,7 @@ ignore = [`
`61`	`61`	`"TRY003", # Allow specifying long messages outside the exception class`
`62`	`62`	`"SLF001", # Allow access private member,`
`63`	`63`	`"PGH003", # Allow not to specify rule codes`
	`64`	`+ "PLR0913", # Allow to have many arguments in function definition`
`64`	`65`	`]`
`65`	`66`	`select = ["ALL"]`
`66`	`67`