CORE-8311: Add locust_plugins dependency and parameters to export results to timescale db

Author: zaslavskii

chainbench/main.py

@@ -106,6 +106,22 @@ def cli(ctx: click.Context):
     "You may specify this option multiple times",
     multiple=True,
 )
+@click.option(
+    "--timescale", is_flag=True, help="Export data to PG with timescale extension"
+)
+@click.option(
+    "--pg-host", default=None, help="Hostname of PG instance with timescale extension"
+)
+@click.option(
+    "--pg-port",
+    default=5432,
+    help="Port of PG instance with timescale extension",
+    show_default=True,
+)
+@click.option(
+    "--pg-username", default="postgres", help="PG username", show_default=True
+)
+@click.option("--pg-password", default=None, help="PG password")
 @click.pass_context
 def start(
     ctx: click.Context,
@@ -126,6 +142,11 @@ def start(
     notify: str | None,
     debug_trace_methods: bool,
     exclude_tags: list[str],
+    timescale: bool,
+    pg_host: str | None,
+    pg_port: int,
+    pg_username: str,
+    pg_password: str | None,
 ):
     if notify:
         click.echo(f"Notify when test is finished using topic: {notify}")
@@ -139,6 +160,15 @@ def start(
         click.echo("Target is required when running in headless mode")
         sys.exit(1)
 
+    if timescale and any(
+        pg_arg is None for pg_arg in (pg_host, pg_port, pg_username, pg_password)
+    ):
+        click.echo(
+            "PG connection parameters are required "
+            "when --timescale flag is used: pg_host, pg_port, pg_username, pg_password"
+        )
+        sys.exit(1)
+
     if not profile_dir:
         profile_dir = get_base_path(__file__)
 
@@ -180,6 +210,11 @@ def start(
         headless=headless,
         target=target,
         exclude_tags=custom_exclude_tags,
+        timescale=timescale,
+        pg_host=pg_host,
+        pg_port=pg_port,
+        pg_username=pg_username,
+        pg_password=pg_password,
     )
     if headless:
         click.echo(f"Starting master in headless mode for {profile}")
@@ -203,6 +238,11 @@ def start(
             worker_id=worker_id,
             log_level=log_level,
             exclude_tags=custom_exclude_tags,
+            timescale=timescale,
+            pg_host=pg_host,
+            pg_port=pg_port,
+            pg_username=pg_username,
+            pg_password=pg_password,
         )
         worker_args = shlex.split(worker_command, posix=is_posix)
         worker_process = subprocess.Popen(worker_args)

chainbench/user/base.py

@@ -9,6 +9,10 @@
 from chainbench.util.event import setup_event_listeners
 from chainbench.util.rpc import generate_request
 
+# importing plugins here as all profiles depend on it
+import locust_plugins  # isort: skip  # noqa
+
+
 setup_event_listeners()
 
 

chainbench/util/cli.py

@@ -40,6 +40,18 @@ def ensure_results_dir(
     return results_dir
 
 
+def get_timescale_args(
+    pg_host: str | None,
+    pg_port: int | None,
+    pg_username: str | None,
+    pg_password: str | None,
+) -> str:
+    return (
+        f" --timescale --pghost={pg_host} --pgport={pg_port}"
+        f" --pgpassword={pg_password} --pguser={pg_username}"
+    )
+
+
 def get_master_command(
     profile_path: Path,
     host: str,
@@ -53,6 +65,11 @@ def get_master_command(
     exclude_tags: list[str],
     target: str | None = None,
     headless: bool = False,
+    timescale: bool = False,
+    pg_host: str | None = None,
+    pg_port: int | None = None,
+    pg_username: str | None = None,
+    pg_password: str | None = None,
 ) -> str:
     """Generate master command."""
     command = (
@@ -65,6 +82,9 @@ def get_master_command(
         f"--loglevel {log_level} --expect-workers {workers}"
     )
 
+    if timescale:
+        command += get_timescale_args(pg_host, pg_port, pg_username, pg_password)
+
     if target is not None:
         command += f" --host {target}"
 
@@ -87,13 +107,21 @@ def get_worker_command(
     target: str | None = None,
     headless: bool = False,
     worker_id: int = 0,
+    timescale: bool = False,
+    pg_host: str | None = None,
+    pg_port: int | None = None,
+    pg_username: str | None = None,
+    pg_password: str | None = None,
 ) -> str:
     """Generate worker command."""
     command = (
         f"locust -f {profile_path} --worker --master-host {host} --master-port {port} "
         f"--logfile {results_path}/worker_{worker_id}.log --loglevel {log_level}"
     )
 
+    if timescale:
+        command += get_timescale_args(pg_host, pg_port, pg_username, pg_password)
+
     if target is not None:
         command += f" --host {target}"