test_runner: pass password to pgbench via PGPASSWORD (#2468)

test_runner/fixtures/log_helper.py

@@ -1,6 +1,5 @@
 import logging
 import logging.config
-import re
 
 """
 This file configures logging to use in python tests.
@@ -30,17 +29,6 @@ LOGGING = {
 }
 
 
-class PasswordFilter(logging.Filter):
-    """Filter out password from logs."""
-
-    # Good enough to filter our passwords produced by PgProtocol.connstr
-    FILTER = re.compile(r"(\s*)password=[^\s]+(\s*)")
-
-    def filter(self, record: logging.LogRecord) -> bool:
-        record.msg = self.FILTER.sub(r"\1password=<hidden>\2", str(record.msg))
-        return True
-
-
 def getLogger(name="root") -> logging.Logger:
     """Method to get logger for tests.
 
@@ -50,6 +38,5 @@ def getLogger(name="root") -> logging.Logger:
 
 # default logger for tests
 log = getLogger()
-log.addFilter(PasswordFilter())
 
 logging.config.dictConfig(LOGGING)

test_runner/fixtures/neon_fixtures.py

@@ -283,10 +283,15 @@ class PgProtocol:
         return str(make_dsn(**self.conn_options(**kwargs)))
 
     def conn_options(self, **kwargs):
+        """
+        Construct a dictionary of connection options from default values and extra parameters.
+        An option can be dropped from the returning dictionary by None-valued extra parameter.
+        """
         result = self.default_options.copy()
         if "dsn" in kwargs:
             result.update(parse_dsn(kwargs["dsn"]))
         result.update(kwargs)
+        result = {k: v for k, v in result.items() if v is not None}
 
         # Individual statement timeout in seconds. 2 minutes should be
         # enough for our tests, but if you need a longer, you can

test_runner/performance/test_perf_pgbench.py

@@ -4,7 +4,7 @@ import os
 import timeit
 from datetime import datetime
 from pathlib import Path
-from typing import List
+from typing import Dict, List
 
 import pytest
 from fixtures.benchmark_fixture import MetricReport, PgBenchInitResult, PgBenchRunResult
@@ -24,14 +24,18 @@ def utc_now_timestamp() -> int:
     return calendar.timegm(datetime.utcnow().utctimetuple())
 
 
-def init_pgbench(env: PgCompare, cmdline):
+def init_pgbench(env: PgCompare, cmdline, password: None):
+    environ: Dict[str, str] = {}
+    if password is not None:
+        environ["PGPASSWORD"] = password
+
     # calculate timestamps and durations separately
     # timestamp is intended to be used for linking to grafana and logs
     # duration is actually a metric and uses float instead of int for timestamp
     start_timestamp = utc_now_timestamp()
     t0 = timeit.default_timer()
     with env.record_pageserver_writes("init.pageserver_writes"):
-        out = env.pg_bin.run_capture(cmdline)
+        out = env.pg_bin.run_capture(cmdline, env=environ)
         env.flush()
 
     duration = timeit.default_timer() - t0
@@ -48,13 +52,15 @@ def init_pgbench(env: PgCompare, cmdline):
     env.zenbenchmark.record_pg_bench_init_result("init", res)
 
 
-def run_pgbench(env: PgCompare, prefix: str, cmdline):
+def run_pgbench(env: PgCompare, prefix: str, cmdline, password: None):
+    environ: Dict[str, str] = {}
+    if password is not None:
+        environ["PGPASSWORD"] = password
+
     with env.record_pageserver_writes(f"{prefix}.pageserver_writes"):
         run_start_timestamp = utc_now_timestamp()
         t0 = timeit.default_timer()
-        out = env.pg_bin.run_capture(
-            cmdline,
-        )
+        out = env.pg_bin.run_capture(cmdline, env=environ)
         run_duration = timeit.default_timer() - t0
         run_end_timestamp = utc_now_timestamp()
         env.flush()
@@ -82,10 +88,14 @@ def run_pgbench(env: PgCompare, prefix: str, cmdline):
 def run_test_pgbench(env: PgCompare, scale: int, duration: int, workload_type: PgBenchLoadType):
     env.zenbenchmark.record("scale", scale, "", MetricReport.TEST_PARAM)
 
+    password = env.pg.default_options.get("password", None)
+    options = "-cstatement_timeout=1h " + env.pg.default_options.get("options", "")
+    # drop password from the connection string by passing password=None and set password separately
+    connstr = env.pg.connstr(password=None, options=options)
+
     if workload_type == PgBenchLoadType.INIT:
         # Run initialize
-        options = "-cstatement_timeout=1h " + env.pg.default_options.get("options", "")
-        init_pgbench(env, ["pgbench", f"-s{scale}", "-i", env.pg.connstr(options=options)])
+        init_pgbench(env, ["pgbench", f"-s{scale}", "-i", connstr], password=password)
 
     if workload_type == PgBenchLoadType.SIMPLE_UPDATE:
         # Run simple-update workload
@@ -99,8 +109,9 @@ def run_test_pgbench(env: PgCompare, scale: int, duration: int, workload_type: P
                 f"-T{duration}",
                 "-P2",
                 "--progress-timestamp",
-                env.pg.connstr(),
+                connstr,
             ],
+            password=password,
         )
 
     if workload_type == PgBenchLoadType.SELECT_ONLY:
@@ -115,8 +126,9 @@ def run_test_pgbench(env: PgCompare, scale: int, duration: int, workload_type: P
                 f"-T{duration}",
                 "-P2",
                 "--progress-timestamp",
-                env.pg.connstr(),
+                connstr,
             ],
+            password=password,
         )
 
     env.report_size()