Merge pull request #22 from ydb-platform/fix_pending_query

Fix pending query issue

Author: vgvoleg

tests/test_connections.py

@@ -29,9 +29,11 @@ def _test_isolation_level_read_only(
         cursor = connection.cursor()
         with suppress(dbapi.DatabaseError):
             maybe_await(cursor.execute_scheme("DROP TABLE foo"))
-        maybe_await(cursor.execute_scheme(
-            "CREATE TABLE foo(id Int64 NOT NULL, PRIMARY KEY (id))"
-        ))
+        maybe_await(
+            cursor.execute_scheme(
+                "CREATE TABLE foo(id Int64 NOT NULL, PRIMARY KEY (id))"
+            )
+        )
 
         connection.set_isolation_level(isolation_level)
         cursor = connection.cursor()
@@ -60,9 +62,11 @@ def _test_connection(self, connection: dbapi.Connection) -> None:
         with pytest.raises(dbapi.ProgrammingError):
             maybe_await(connection.describe("/local/foo"))
 
-        maybe_await(cur.execute_scheme(
-            "CREATE TABLE foo(id Int64 NOT NULL, PRIMARY KEY (id))"
-        ))
+        maybe_await(
+            cur.execute_scheme(
+                "CREATE TABLE foo(id Int64 NOT NULL, PRIMARY KEY (id))"
+            )
+        )
 
         assert maybe_await(connection.check_exists("/local/foo"))
 
@@ -84,26 +88,28 @@ def _test_cursor_raw_query(self, connection: dbapi.Connection) -> None:
             "CREATE TABLE test(id Int64 NOT NULL, text Utf8, PRIMARY KEY (id))"
         ))
 
-        maybe_await(cur.execute(
-            """
+        maybe_await(
+            cur.execute(
+                """
             DECLARE $data AS List<Struct<id:Int64, text: Utf8>>;
 
             INSERT INTO test SELECT id, text FROM AS_TABLE($data);
             """,
-            {
-                "$data": ydb.TypedValue(
-                    [
-                        {"id": 17, "text": "seventeen"},
-                        {"id": 21, "text": "twenty one"},
-                    ],
-                    ydb.ListType(
-                        ydb.StructType()
-                        .add_member("id", ydb.PrimitiveType.Int64)
-                        .add_member("text", ydb.PrimitiveType.Utf8)
-                    ),
-                )
-            },
-        ))
+                {
+                    "$data": ydb.TypedValue(
+                        [
+                            {"id": 17, "text": "seventeen"},
+                            {"id": 21, "text": "twenty one"},
+                        ],
+                        ydb.ListType(
+                            ydb.StructType()
+                            .add_member("id", ydb.PrimitiveType.Int64)
+                            .add_member("text", ydb.PrimitiveType.Utf8)
+                        ),
+                    )
+                },
+            )
+        )
 
         maybe_await(cur.execute_scheme("DROP TABLE test"))
 
@@ -112,13 +118,15 @@ def _test_cursor_raw_query(self, connection: dbapi.Connection) -> None:
     def _test_errors(
         self,
         connection: dbapi.Connection,
-        connect_method: callable = dbapi.connect
+        connect_method: callable = dbapi.connect,
     ) -> None:
         with pytest.raises(dbapi.InterfaceError):
-            maybe_await(connect_method(
-                "localhost:2136",  # type: ignore
-                database="/local666",  # type: ignore
-            ))
+            maybe_await(
+                connect_method(
+                    "localhost:2136",  # type: ignore
+                    database="/local666",  # type: ignore
+                )
+            )
 
         cur = connection.cursor()
 
@@ -137,9 +145,9 @@ def _test_errors(
         with pytest.raises(dbapi.ProgrammingError):
             maybe_await(cur.execute("SELECT * FROM test"))
 
-        maybe_await(cur.execute_scheme(
-            "CREATE TABLE test(id Int64, PRIMARY KEY (id))"
-        ))
+        maybe_await(
+            cur.execute_scheme("CREATE TABLE test(id Int64, PRIMARY KEY (id))")
+        )
 
         maybe_await(cur.execute("INSERT INTO test(id) VALUES(1)"))
 
@@ -154,8 +162,9 @@ def _test_bulk_upsert(self, connection: dbapi.Connection) -> None:
         with suppress(dbapi.DatabaseError):
             maybe_await(cursor.execute_scheme("DROP TABLE pet"))
 
-        maybe_await(cursor.execute_scheme(
-            """
+        maybe_await(
+            cursor.execute_scheme(
+                """
             CREATE TABLE pet (
             pet_id INT,
             name TEXT NOT NULL,
@@ -165,7 +174,8 @@ def _test_bulk_upsert(self, connection: dbapi.Connection) -> None:
             PRIMARY KEY (pet_id)
             );
             """
-        ))
+            )
+        )
 
         column_types = (
             ydb.BulkUpsertColumns()
@@ -182,14 +192,14 @@ def _test_bulk_upsert(self, connection: dbapi.Connection) -> None:
                 "name": "Lester",
                 "pet_type": "Hamster",
                 "birth_date": "2020-06-23",
-                "owner": "Lily"
+                "owner": "Lily",
             },
             {
                 "pet_id": 4,
                 "name": "Quincy",
                 "pet_type": "Parrot",
                 "birth_date": "2013-08-11",
-                "owner": "Anne"
+                "owner": "Anne",
             },
         ]
 
@@ -204,18 +214,19 @@ def _test_error_with_interactive_tx(
         self,
         connection: dbapi.Connection,
     ) -> None:
-
         cur = connection.cursor()
-        maybe_await(cur.execute_scheme(
-            """
+        maybe_await(
+            cur.execute_scheme(
+                """
             DROP TABLE IF EXISTS test;
             CREATE TABLE test (
             id Int64 NOT NULL,
             val Int64,
             PRIMARY KEY(id)
             )
             """
-        ))
+            )
+        )
 
         connection.set_isolation_level(dbapi.IsolationLevel.SERIALIZABLE)
         maybe_await(connection.begin())
@@ -274,8 +285,8 @@ def test_bulk_upsert(self, connection: dbapi.Connection) -> None:
         self._test_bulk_upsert(connection)
 
     def test_errors_with_interactive_tx(
-            self, connection: dbapi.Connection
-        ) -> None:
+        self, connection: dbapi.Connection
+    ) -> None:
         self._test_error_with_interactive_tx(connection)
 
 
@@ -291,8 +302,10 @@ def connect() -> dbapi.AsyncConnection:
         try:
             yield conn
         finally:
+
             def close() -> None:
                 maybe_await(conn.close())
+
             await greenlet_spawn(close)
 
     @pytest.mark.asyncio
@@ -315,7 +328,9 @@ async def test_isolation_level_read_only(
     ) -> None:
         await greenlet_spawn(
             self._test_isolation_level_read_only,
-            connection, isolation_level, read_only
+            connection,
+            isolation_level,
+            read_only,
         )
 
     @pytest.mark.asyncio

tests/test_cursors.py

@@ -161,7 +161,6 @@ class TestCursor(BaseCursorTestSuit):
     def sync_cursor(
         self, session_pool_sync: ydb.QuerySessionPool
     ) -> Generator[Cursor]:
-
         cursor = Cursor(
             FakeSyncConnection(),
             session_pool_sync,
@@ -195,9 +194,7 @@ def test_cursor_fetch_all_multiple_result_sets(
     ) -> None:
         self._test_cursor_fetch_all_multiple_result_sets(sync_cursor)
 
-    def test_cursor_state_after_error(
-        self, sync_cursor: Cursor
-    ) -> None:
+    def test_cursor_state_after_error(self, sync_cursor: Cursor) -> None:
         self._test_cursor_state_after_error(sync_cursor)
 
 
@@ -255,6 +252,4 @@ async def test_cursor_fetch_all_multiple_result_sets(
     async def test_cursor_state_after_error(
         self, async_cursor: AsyncCursor
     ) -> None:
-        await greenlet_spawn(
-            self._test_cursor_state_after_error, async_cursor
-        )
+        await greenlet_spawn(self._test_cursor_state_after_error, async_cursor)

ydb_dbapi/connections.py

@@ -102,7 +102,9 @@ def __init__(
                 database=self.database,
                 credentials=self.credentials,
                 query_client_settings=self._get_client_settings(),
-                root_certificates=ydb.load_ydb_root_certificate(root_certificates_path),
+                root_certificates=ydb.load_ydb_root_certificate(
+                    root_certificates_path
+                ),
             )
             self._driver = self._driver_cls(driver_config)
             self._session_pool = self._pool_cls(self._driver, size=5)