Add simple stability test for olap queries (#7543)

4 files changed, 104 insertions, 0 deletions

diff --git a/ydb/tools/cfg/static.py b/ydb/tools/cfg/static.py
index 345701921e..faae12e228 100644
--- a/ydb/tools/cfg/static.py
+++ b/ydb/tools/cfg/static.py
@@ -274,6 +274,10 @@ class StaticConfigGenerator(object):
     @property
     def table_service_config(self):
         return self.__cluster_details.get_service("table_service_config")
+    
+    @property
+    def column_shard_config(self):
+        return self.__cluster_details.get_service("column_shard_config")
 
     @property
     def hive_config(self):
@@ -386,6 +390,9 @@ class StaticConfigGenerator(object):
         if self.table_service_config:
             normalized_config["table_service_config"] = self.table_service_config
 
+        if self.column_shard_config:
+            normalized_config["column_shard_config"] = self.column_shard_config
+
         if self.__cluster_details.blob_storage_config is not None:
             normalized_config["blob_storage_config"] = self.__cluster_details.blob_storage_config
         else:
diff --git a/ydb/tools/olap_workload/__main__.py b/ydb/tools/olap_workload/__main__.py
new file mode 100644
index 0000000000..f49232da9f
--- /dev/null
+++ b/ydb/tools/olap_workload/__main__.py
@@ -0,0 +1,84 @@
+# -*- coding: utf-8 -*-
+import argparse
+import ydb
+import time
+import os
+
+ydb.interceptor.monkey_patch_event_handler()
+
+
+def timestamp():
+    return int(1000 * time.time())
+
+
+def table_name_with_timestamp():
+    return os.path.join("column_table_" + str(timestamp()))
+
+
+class Workload(object):
+    def __init__(self, endpoint, database, duration):
+        self.database = database
+        self.driver = ydb.Driver(ydb.DriverConfig(endpoint, database))
+        self.pool = ydb.SessionPool(self.driver, size=200)
+        self.duration = duration
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.pool.stop()
+        self.driver.stop()
+
+    def create_table(self, table_name):
+        with self.pool.checkout() as s:
+            try:
+                s.execute_scheme(
+                    """
+                    CREATE TABLE %s (
+                    id Int64 NOT NULL,
+                    i64Val Int64,
+                    PRIMARY KEY(id)
+                    )
+                    PARTITION BY HASH(id)
+                    WITH (
+                        STORE = COLUMN
+                    )
+                """
+                    % table_name
+                )
+
+                print("Table %s created" % table_name)
+            except ydb.SchemeError as e:
+                print(e)
+
+    def drop_table(self, table_name):
+        with self.pool.checkout() as s:
+            try:
+                s.drop_table(self.database + "/" + table_name)
+
+                print("Table %s dropped" % table_name)
+            except ydb.SchemeError as e:
+                print(e)
+
+    def run(self):
+        started_at = time.time()
+
+        while time.time() - started_at < self.duration:
+            table_name = table_name_with_timestamp()
+            self.create_table(table_name)
+
+            time.sleep(5)
+
+            self.drop_table(table_name)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(
+        description="olap stability workload", formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+    parser.add_argument('--endpoint', default='localhost:2135', help="An endpoint to be used")
+    parser.add_argument('--database', default=None, required=True, help='A database to connect')
+    parser.add_argument('--duration', default=10**9, type=lambda x: int(x), help='A duration of workload in seconds.')
+    args = parser.parse_args()
+    with Workload(args.endpoint, args.database, args.duration) as workload:
+        workload.run()
diff --git a/ydb/tools/olap_workload/ya.make b/ydb/tools/olap_workload/ya.make
new file mode 100644
index 0000000000..939ecf1af9
--- /dev/null
+++ b/ydb/tools/olap_workload/ya.make
@@ -0,0 +1,12 @@
+PY3_PROGRAM(olap_workload)
+
+PY_SRCS(
+    __main__.py
+)
+
+PEERDIR(
+    ydb/public/sdk/python
+    library/python/monlib
+)
+
+END()
diff --git a/ydb/tools/ya.make b/ydb/tools/ya.make
index fc5e094eb2..375abee364 100644
--- a/ydb/tools/ya.make
+++ b/ydb/tools/ya.make
@@ -4,6 +4,7 @@ RECURSE(
     query_replay
     query_replay_yt
     simple_queue
+    olap_workload
     tsserver
     tstool
     ydbd_slice