From 08db823fee598d5672017ee0537a60a841bdf5a0 Mon Sep 17 00:00:00 2001
From: Martin Gaievski <gaievski@amazon.com>
Date: Tue, 23 Jul 2024 00:54:10 +0000
Subject: [PATCH] Adding operation for enabling concurrent segment search

Signed-off-by: Martin Gaievski <gaievski@amazon.com>
---
 osbenchmark/worker_coordinator/runner.py | 14 ++++++++++++++
 osbenchmark/workload/workload.py         |  3 +++
 tests/worker_coordinator/runner_test.py  | 20 ++++++++++++++++++++
 3 files changed, 37 insertions(+)

diff --git a/osbenchmark/worker_coordinator/runner.py b/osbenchmark/worker_coordinator/runner.py
index d280ddffd..3d9e1306b 100644
--- a/osbenchmark/worker_coordinator/runner.py
+++ b/osbenchmark/worker_coordinator/runner.py
@@ -2696,3 +2696,17 @@ async def __call__(self, opensearch, params):
 
     def __repr__(self, *args, **kwargs):
         return "deploy-ml-model"
+
+class EnableConcurrentSegmentSearch(Runner):
+    @time_func
+    async def __call__(self, opensearch, params):
+        enable_setting = params.get("enable", "false")
+        body = {
+            "transient": {
+                "search.concurrent_segment_search.enabled": enable_setting
+              } 
+            }
+        await opensearch.cluster.put_settings(body=body)
+
+    def __repr__(self, *args, **kwargs):
+        return "enable-concurrent-segment-search"
\ No newline at end of file
diff --git a/osbenchmark/workload/workload.py b/osbenchmark/workload/workload.py
index 95545fb59..c58b1e817 100644
--- a/osbenchmark/workload/workload.py
+++ b/osbenchmark/workload/workload.py
@@ -635,6 +635,7 @@ class OperationType(Enum):
     DeleteMlModel = 1041
     RegisterMlModel = 1042
     DeployMlModel = 1043
+    EnableConcurrentSegmentSearch = 1044
 
     @property
     def admin_op(self):
@@ -752,6 +753,8 @@ def from_hyphenated_string(cls, v):
             return OperationType.TrainKnnModel
         elif v == "delete-knn-model":
             return OperationType.DeleteKnnModel
+        elif v == "enable-concurrent-segment-search":
+            return OperationType.EnableConcurrentSegmentSearch
         else:
             raise KeyError(f"No enum value for [{v}]")
 
diff --git a/tests/worker_coordinator/runner_test.py b/tests/worker_coordinator/runner_test.py
index 4f4b303b2..d0b83d461 100644
--- a/tests/worker_coordinator/runner_test.py
+++ b/tests/worker_coordinator/runner_test.py
@@ -6785,3 +6785,23 @@ async def test_param_id_mandatory(self, opensearch, on_client_request_start, on_
             await r(opensearch, params)
 
         self.assertEqual(0, opensearch.transport.perform_request.call_count)
+
+class EnableConcurrentSegmentSearchTests(TestCase):
+    @mock.patch('osbenchmark.client.RequestContextHolder.on_client_request_end')
+    @mock.patch('osbenchmark.client.RequestContextHolder.on_client_request_start')
+    @mock.patch("opensearchpy.OpenSearch")
+    @run_async
+    async def test_enable_concurrent_segment_search(self, opensearch, on_client_request_start, on_client_request_end):
+        opensearch.cluster.put_settings.return_value = as_future()
+        params = {
+            "enable": "true"
+        }
+
+        r = runner.EnableConcurrentSegmentSearch()
+        await r(opensearch, params)
+
+        opensearch.cluster.put_settings.assert_called_once_with(body={
+            "transient": {
+                "search.concurrent_segment_search.enabled": "true"
+            } 
+        })
\ No newline at end of file