From 98db4468af65c48941e561922c4426c966d153d8 Mon Sep 17 00:00:00 2001
From: vlad-gogov <vlad-gogov@ydb.tech>
Date: Fri, 29 Nov 2024 14:54:08 +0300
Subject: [PATCH] test

---
 .../scenario/helpers/scenario_tests_helper.py | 20 ++++++
 ydb/tests/olap/scenario/test_insert.py        | 61 +++++++++++++++++++
 ydb/tests/olap/scenario/ya.make               |  1 +
 3 files changed, 82 insertions(+)
 create mode 100644 ydb/tests/olap/scenario/test_insert.py

diff --git a/ydb/tests/olap/scenario/helpers/scenario_tests_helper.py b/ydb/tests/olap/scenario/helpers/scenario_tests_helper.py
index 675231784011..b8395ad17af1 100644
--- a/ydb/tests/olap/scenario/helpers/scenario_tests_helper.py
+++ b/ydb/tests/olap/scenario/helpers/scenario_tests_helper.py
@@ -395,6 +395,26 @@ def execute_scan_query(
         allure.attach(json.dumps(rows), 'result', allure.attachment_type.JSON)
         return ret
 
+    @allure.step('Execute query')
+    def execute_query(
+        self, yql: str, expected_status: ydb.StatusCode | Set[ydb.StatusCode] = ydb.StatusCode.SUCCESS
+    ):
+        """Run a query on the tested database.
+
+        Args:
+            yql: Query text.
+            expected_status: Expected status or set of database response statuses. If the response status is not in the expected set, an exception is thrown.
+
+        Example:
+            tablename = 'testTable'
+            sth = ScenarioTestHelper(ctx)
+            sth.execute_query(f'INSERT INTO `{sth.get_full_path("tablename") }` (key, c) values(1, 100)')
+        """
+
+        allure.attach(yql, 'request', allure.attachment_type.TEXT)
+        with ydb.QuerySessionPool(YdbCluster.get_ydb_driver()) as pool:
+            it = self._run_with_expected_status(lambda: pool.execute_with_retries(yql), expected_status)
+
     def drop_if_exist(self, names: List[str], operation) -> None:
         """Erase entities in the tested database, if it exists.
 
diff --git a/ydb/tests/olap/scenario/test_insert.py b/ydb/tests/olap/scenario/test_insert.py
new file mode 100644
index 000000000000..26b2381ad688
--- /dev/null
+++ b/ydb/tests/olap/scenario/test_insert.py
@@ -0,0 +1,61 @@
+from conftest import BaseTestSet
+from ydb.tests.olap.scenario.helpers import (
+    ScenarioTestHelper,
+    TestContext,
+    CreateTable,
+    DropTable,
+)
+
+from ydb import PrimitiveType
+from typing import List, Dict, Any
+import threading
+
+class TestInsert(BaseTestSet):
+    schema_cnt = (
+        ScenarioTestHelper.Schema()
+        .with_column(name='key', type=PrimitiveType.Int32, not_null=True)
+        .with_column(name='c', type=PrimitiveType.Int64)
+        .with_key_columns('key')
+    )
+
+    schema_log = (
+        ScenarioTestHelper.Schema()
+        .with_column(name='key', type=PrimitiveType.Int32, not_null=True)
+        .with_key_columns('key')
+    )
+
+    def _loop_upsert(self, ctx: TestContext, data: list):
+        sth = ScenarioTestHelper(ctx)
+        for batch in data:
+            sth.bulk_upsert_data("log", self.schema_log, batch)
+
+    def _loop_insert(self, ctx: TestContext):
+        sth = ScenarioTestHelper(ctx)
+        for i in range(100):
+            sth.execute_query(f"$cnt = SELECT CAST(COUNT(*) AS INT64) from `{sth.get_full_path("log")}`; INSERT INTO `{sth.get_full_path("cnt") }` (key, c) values({i}, $cnt)")
+
+    def scenario_read_data_during_bulk_upsert(self, ctx: TestContext):
+        sth = ScenarioTestHelper(ctx)
+        cnt_table_name: str = "cnt"
+        log_table_name: str = "log"
+        sth.execute_scheme_query(CreateTable(cnt_table_name).with_schema(self.schema_cnt))
+        sth.execute_scheme_query(CreateTable(log_table_name).with_schema(self.schema_log))
+        
+        data: list = []
+        for i in range(100):
+            batch: List[Dict[str, Any]] = []
+            for j in range(i, 1000):
+                batch.append({'key' : j})
+            data.append(batch)
+
+        thread1 = threading.Thread(target=self._loop_upsert, args=[ctx, data])
+        thread2 = threading.Thread(target=self._loop_insert, args=[ctx])
+
+        thread1.start()
+        thread2.start()
+
+        thread2.join()
+        thread1.join()
+
+        sth.execute_scheme_query(DropTable(cnt_table_name))
+        sth.execute_scheme_query(DropTable(log_table_name))
diff --git a/ydb/tests/olap/scenario/ya.make b/ydb/tests/olap/scenario/ya.make
index f9eac7689943..2c884561531e 100644
--- a/ydb/tests/olap/scenario/ya.make
+++ b/ydb/tests/olap/scenario/ya.make
@@ -12,6 +12,7 @@ PY3TEST()
         test_simple.py
         test_scheme_load.py
         test_alter_tiering.py
+        test_insert.py
     )
 
     PEERDIR(