Add bulk upsert to connection

Author: vgvoleg

tests/test_connections.py

@@ -149,6 +149,57 @@ def _test_errors(
         maybe_await(cur.execute_scheme("DROP TABLE test"))
         maybe_await(cur.close())
 
+    def _test_bulk_upsert(self, connection: dbapi.Connection) -> None:
+        cursor = connection.cursor()
+        with suppress(dbapi.DatabaseError):
+            maybe_await(cursor.execute_scheme("DROP TABLE pet"))
+
+        maybe_await(cursor.execute_scheme(
+            """
+            CREATE TABLE pet (
+            pet_id INT,
+            name TEXT NOT NULL,
+            pet_type TEXT NOT NULL,
+            birth_date TEXT NOT NULL,
+            owner TEXT NOT NULL,
+            PRIMARY KEY (pet_id)
+            );
+            """
+        ))
+
+        column_types = (
+            ydb.BulkUpsertColumns()
+            .add_column("pet_id", ydb.OptionalType(ydb.PrimitiveType.Int32))
+            .add_column("name", ydb.PrimitiveType.Utf8)
+            .add_column("pet_type", ydb.PrimitiveType.Utf8)
+            .add_column("birth_date", ydb.PrimitiveType.Utf8)
+            .add_column("owner", ydb.PrimitiveType.Utf8)
+        )
+
+        rows = [
+            {
+                "pet_id": 3,
+                "name": "Lester",
+                "pet_type": "Hamster",
+                "birth_date": "2020-06-23",
+                "owner": "Lily"
+            },
+            {
+                "pet_id": 4,
+                "name": "Quincy",
+                "pet_type": "Parrot",
+                "birth_date": "2013-08-11",
+                "owner": "Anne"
+            },
+        ]
+
+        maybe_await(connection.bulk_upsert("pet", rows, column_types))
+
+        maybe_await(cursor.execute("SELECT * FROM pet"))
+        assert cursor.rowcount == 2
+
+        maybe_await(cursor.execute_scheme("DROP TABLE pet"))
+
 
 class TestConnection(BaseDBApiTestSuit):
     @pytest.fixture
@@ -191,6 +242,9 @@ def test_cursor_raw_query(self, connection: dbapi.Connection) -> None:
     def test_errors(self, connection: dbapi.Connection) -> None:
         self._test_errors(connection)
 
+    def test_bulk_upsert(self, connection: dbapi.Connection) -> None:
+        self._test_bulk_upsert(connection)
+
 
 class TestAsyncConnection(BaseDBApiTestSuit):
     @pytest_asyncio.fixture
@@ -244,3 +298,9 @@ async def test_cursor_raw_query(
     @pytest.mark.asyncio
     async def test_errors(self, connection: dbapi.AsyncConnection) -> None:
         await greenlet_spawn(self._test_errors, connection)
+
+    @pytest.mark.asyncio
+    async def test_bulk_upsert(
+        self, connection: dbapi.AsyncConnection
+    ) -> None:
+        await greenlet_spawn(self._test_bulk_upsert, connection)

ydb_dbapi/connections.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import posixpath
+from collections.abc import Sequence
 from enum import Enum
 from typing import NamedTuple
 
@@ -301,6 +302,25 @@ def callee() -> ydb.Directory:
                 result.extend(self._get_table_names(child_abs_path))
         return result
 
+    @handle_ydb_errors
+    def bulk_upsert(
+        self,
+        table_name: str,
+        rows: Sequence,
+        column_types: ydb.BulkUpsertColumns,
+    ) -> None:
+        settings = self._get_request_settings()
+        abs_table_path = posixpath.join(
+            self.database, self.table_path_prefix, table_name
+        )
+
+        self._driver.table_client.bulk_upsert(
+            abs_table_path,
+            rows=rows,
+            column_types=column_types,
+            settings=settings,
+        )
+
 
 class AsyncConnection(BaseConnection):
     _driver_cls = ydb.aio.Driver
@@ -446,6 +466,25 @@ async def callee() -> ydb.Directory:
                 result.extend(await self._get_table_names(child_abs_path))
         return result
 
+    @handle_ydb_errors
+    async def bulk_upsert(
+        self,
+        table_name: str,
+        rows: Sequence,
+        column_types: ydb.BulkUpsertColumns,
+    ) -> None:
+        settings = self._get_request_settings()
+        abs_table_path = posixpath.join(
+            self.database, self.table_path_prefix, table_name
+        )
+
+        await self._driver.table_client.bulk_upsert(
+            abs_table_path,
+            rows=rows,
+            column_types=column_types,
+            settings=settings,
+        )
+
 
 def connect(*args: tuple, **kwargs: dict) -> Connection:
     conn = Connection(*args, **kwargs)  # type: ignore