Support data streams in create-track

esrally/rally.py

@@ -182,9 +182,15 @@ def add_track_source(subparser):
     create_track_parser.add_argument(
         "--indices",
         type=non_empty_list,
-        required=True,
+        required=False,
         help="Comma-separated list of indices to include in the track",
     )
+    create_track_parser.add_argument(
+        "--datastreams",
+        type=non_empty_list,
+        required=False,
+        help="Comma-separated list of data-streams to include in the track",
+    )
     create_track_parser.add_argument(
         "--target-hosts",
         default="",
@@ -987,9 +993,18 @@ def dispatch_sub_command(arg_parser, args, cfg):
             cfg.add(config.Scope.applicationOverride, "generator", "output.path", args.output_path)
             generate(cfg)
         elif sub_command == "create-track":
-            cfg.add(config.Scope.applicationOverride, "generator", "indices", args.indices)
-            cfg.add(config.Scope.applicationOverride, "generator", "output.path", args.output_path)
-            cfg.add(config.Scope.applicationOverride, "track", "track.name", args.track)
+            if args.datastreams is not None and args.indices is None:
+                cfg.add(config.Scope.applicationOverride, "generator", "indices", "*")
+                cfg.add(config.Scope.applicationOverride, "generator", "datastreams", args.datastreams)
+                cfg.add(config.Scope.applicationOverride, "generator", "output.path", args.output_path)
+                cfg.add(config.Scope.applicationOverride, "track", "track.name", args.track)
+            elif args.indices is not None:
+                cfg.add(config.Scope.applicationOverride, "generator", "indices", args.indices)
+                cfg.add(config.Scope.applicationOverride, "generator", "datastreams", args.datastreams)
+                cfg.add(config.Scope.applicationOverride, "generator", "output.path", args.output_path)
+                cfg.add(config.Scope.applicationOverride, "track", "track.name", args.track)
+            else:
+                raise exceptions.SystemSetupError("Following arguments are required: --indices or specify --datastream ", sub_command)
             configure_connection_params(arg_parser, args, cfg)
 
             tracker.create_track(cfg)
@@ -1043,7 +1058,6 @@ def main():
     if args.subcommand is None:
         arg_parser.print_help()
         sys.exit(0)
-
     console.init(quiet=args.quiet)
     console.println(BANNER)
 

esrally/tracker/index.py

@@ -46,10 +46,19 @@ def update_index_setting_parameters(settings):
             settings[s] = param.format(orig=orig_value)
 
 
+allowed_list_hidden_indices = [
+    ".ds-metrics-kubernetes",
+    ".ds-metrics-system",
+    ".ds-metrics-elastic_agent",
+    ".ds-logs-kubernetes",
+    ".ds-logs-elastic_agent",
+]
+
+
 def is_valid(index_name):
     if len(index_name) == 0:
         return False, "Index name is empty"
-    if index_name.startswith("."):
+    if index_name.startswith(".") and not index_name.startswith(tuple(allowed_list_hidden_indices)):
         return False, f"Index [{index_name}] is hidden"
     return True, None
 
@@ -65,7 +74,7 @@ def extract_index_mapping_and_settings(client, index_pattern):
     results = {}
     logger = logging.getLogger(__name__)
     # the response might contain multiple indices if a wildcard was provided
-    response = client.indices.get(index=index_pattern)
+    response = client.indices.get(index=index_pattern, params={"expand_wildcards": "all"})
     for index, details in response.items():
         valid, reason = is_valid(index)
         if valid:
@@ -104,3 +113,20 @@ def extract(client, outdir, index_pattern):
             }
         )
     return results
+
+
+def datastreamextract(client, datastream_pattern):
+    """
+    Calls Elasticsearch client get_data_stream function to retrieve list of indexes
+    :param client: Elasticsearch client
+    :param index_pattern: name of datastream
+    :return: index creation dictionary
+    """
+    results = []
+    # the response might contain multiple indices if a wildcard was provided
+    params_defined = {"expand_wildcards": "all", "filter_path": "data_streams.name"}
+    results_datastreams = client.indices.get_data_stream(name=datastream_pattern, params=params_defined)
+
+    for indices in results_datastreams["data_streams"]:
+        results.append(indices.get("name"))
+    return results

esrally/tracker/tracker.py

@@ -35,6 +35,19 @@ def process_template(templates_path, template_filename, template_vars, output_pa
         f.write(template.render(template_vars))
 
 
+def extract_indexes_from_datastreams(client, datastreams_to_extract):
+    indices = []
+    # first extract index metadata (which is cheap) and defer extracting data to reduce the potential for
+    # errors due to invalid index names late in the process.
+    for datastream_name in datastreams_to_extract:
+        try:
+            indices += index.datastreamextract(client, datastream_name)
+        except ElasticsearchException:
+            logging.getLogger(__name__).exception("Failed to extract indexes from datastream [%s]", datastream_name)
+
+    return indices
+
+
 def extract_mappings_and_corpora(client, output_path, indices_to_extract):
     indices = []
     corpora = []
@@ -63,8 +76,7 @@ def create_track(cfg):
     root_path = cfg.opts("generator", "output.path")
     target_hosts = cfg.opts("client", "hosts")
     client_options = cfg.opts("client", "options")
-
-    logger.info("Creating track [%s] matching indices [%s]", track_name, indices)
+    datastreams = cfg.opts("generator", "datastreams")
 
     client = EsClientFactory(
         hosts=target_hosts.all_hosts[opts.TargetHosts.DEFAULT], client_options=client_options.all_client_options[opts.TargetHosts.DEFAULT]
@@ -76,6 +88,14 @@ def create_track(cfg):
     output_path = os.path.abspath(os.path.join(io.normalize_path(root_path), track_name))
     io.ensure_dir(output_path)
 
+    if datastreams is None:
+        logger.info("Creating track [%s] matching indices [%s]", track_name, indices)
+    else:
+        logger.info("Creating track [%s] matching datastreams [%s]", track_name, datastreams)
+        extracted_indices = extract_indexes_from_datastreams(client, datastreams)
+        indices = extracted_indices
+        console.info(f"##### List datastreams to extract [{indices}]", logger=logger)
+
     indices, corpora = extract_mappings_and_corpora(client, output_path, indices)
     if len(indices) == 0:
         raise RuntimeError("Failed to extract any indices for track!")