Timescale listener: Fix forwarding of pghost from command line. Add -…

…-pgport and --pgdatabase.
SvenskaSpel · Mar 9, 2022 · 1425b11 · 1425b11
1 parent 93493db
commit 1425b11
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 2 deletions.
diff --git a/locust_plugins/__init__.py b/locust_plugins/__init__.py
@@ -65,6 +65,13 @@ def add_checks_arguments(parser: configargparse.ArgumentParser):
         env_var="PGHOST",
         default="",
     )
+    locust_dashboards.add_argument(
+        "--pgport",
+        type=str,
+        help="",
+        env_var="PGPORT",
+        default="",
+    )
     locust_dashboards.add_argument(
         "--pgpassword",
         type=str,
@@ -79,6 +86,13 @@ def add_checks_arguments(parser: configargparse.ArgumentParser):
         env_var="PGUSER",
         default="",
     )
+    locust_dashboards.add_argument(
+        "--pgdatabase",
+        type=str,
+        help="",
+        env_var="PGDATABASE",
+        default="",
+    )
     run_info = parser.add_argument_group(
         "locust-plugins - Run info",
         "Extra run info for listeners",

diff --git a/locust_plugins/listeners.py b/locust_plugins/listeners.py
@@ -129,8 +129,10 @@ def on_start(self, environment: locust.env.Environment):
     def _dbconn(self) -> psycopg2.extensions.connection:
         try:
             conn = psycopg2.connect(
-                host=os.environ.get("PGHOST", "localhost"),
-                password=self.env.parsed_options.pgpassword,  # type: ignore[union-attr]
+                host=self.env.parsed_options.pghost,
+                password=self.env.parsed_options.pgpassword,
+                database=self.env.parsed_options.pgdatabase,
+                port=self.env.parsed_options.pgport,
                 keepalives_idle=120,
                 keepalives_interval=20,
                 keepalives_count=6,