get elasticsearch ids by key (#1167)

Author: ebattat

benchmark_runner/common/elasticsearch/elasticsearch_operations.py

@@ -364,16 +364,24 @@ def get_uuid_for_ids(self, index: str, ids: list[str]) -> dict:
 
     @typechecked()
     @logger_time_stamp
-    def get_index_ids_between_dates(self, index: str, start_datetime: datetime = None, end_datetime: datetime = None):
+    def get_index_ids_between_dates(
+            self,
+            index: str,
+            start_datetime: datetime = None,
+            end_datetime: datetime = None,
+            key: str = None,
+    ):
         """
-        This method returns list of index ids between dates
-        @param index:
-        @param start_datetime:
-        @param end_datetime:
-        @return: list of ids
+        This method returns list of index ids between dates.
+        Optionally filters documents by a specific key.
         """
-        es_data = self.get_index_data_between_dates(index=index, start_datetime=start_datetime, end_datetime=end_datetime)
-        ids = [hit["_id"] for hit in es_data]
+        es_data = self.get_index_data_between_dates(index=index, start_datetime=start_datetime,
+                                                    end_datetime=end_datetime)
+        if key is None:
+            ids = [hit["_id"] for hit in es_data]
+        else:
+            ids = [hit["_id"] for hit in es_data if key in hit["_source"]]
+
         return ids
 
     @typechecked()

benchmark_runner/workloads/winmssql_vm.py

@@ -74,7 +74,7 @@ def run(self):
             self._data_dict['run_artifacts_url'] = os.path.join(self._run_artifacts_url,
                                                                 f'{self._get_run_artifacts_hierarchy(workload_name=self._workload_name, is_file=True)}-{self._time_stamp_format}.tar.gz')
             self.wait_for_windows_hammerdb_finished()
-            ids = self._get_index_ids_between_dates(index=self._es_index)
+            ids = self._get_index_ids_between_dates(index=self._es_index, key='status')
             # Adding data_updated=True to stamp that this data is already updated and enrich with new product versions fields
             for id in ids:
                 self._update_elasticsearch_index(index=self._es_index, id=id, kind='vm', status='Succeeded', run_artifacts_url=self._data_dict['run_artifacts_url'], database='mssql', vm_name=f'{self._workload_name}-{self._trunc_uuid}', data_updated=True)

benchmark_runner/workloads/workloads_operations.py

@@ -454,10 +454,10 @@ def __get_metadata(self, kind: str = None, status: str = None, database: str = N
 
         return metadata
 
-    def _get_index_ids_between_dates(self, index: str):
+    def _get_index_ids_between_dates(self, index: str, key:str):
         end_datetime = datetime.now(timezone.utc)
         start_datetime = end_datetime - timedelta(hours=4)
-        return self._es_operations.get_index_ids_between_dates(index=index, start_datetime=start_datetime, end_datetime=end_datetime)
+        return self._es_operations.get_index_ids_between_dates(index=index, key=key, start_datetime=start_datetime, end_datetime=end_datetime)
 
     def _get_latest_resource_with_key(self, index: str, key: str) -> dict:
         """