feat: add SSH agent authentication

Author: O2eg

README.md

@@ -1,14 +1,17 @@
-# pg_perf_bench
+# pg-perf-bench
 
 ## Overview
 
 Based on [pg_perfbench](https://github.com/TantorLabs/pg_perfbench).
 
-`pg_perf_bench` runs controlled PostgreSQL benchmarks and stores the result
+`pg-perf-bench` runs controlled PostgreSQL benchmarks and stores the result
 together with the facts required to interpret it: the effective workload,
 PostgreSQL configuration, server version, host properties, execution timing,
 raw command output, and collection diagnostics.
 
+The distribution and installed command are `pg-perf-bench`; the import package
+and GitHub repository are named `pg_perf_bench`.
+
 Its benchmark question is maximum TPS for one workload profile and one complete
 environment, including OS and PostgreSQL settings. `pg_workload` schedules and
 runs workload profiles; `pg_perf_bench` sweeps load, resets the dataset for each
@@ -56,7 +59,7 @@ CLI and automation contract
 
 The transport controls PostgreSQL and executes host fact collectors on the
 selected target. The workload commands themselves run on the machine where
-`pg_perf_bench` is invoked:
+`pg-perf-bench` is invoked:
 
 | Transport | Host facts and PostgreSQL lifecycle | `pgbench` / `psql` |
 |---|---|---|
@@ -329,6 +332,35 @@ The workload reaches PostgreSQL through the port published on `--host` and
 --port 55432
 ```
 
+To use an identity already loaded into a local agent, replace `--ssh-key` with
+`--ssh-agent`:
+
+```bash
+eval "$(ssh-agent -s)"
+ssh-add ~/.ssh/id_ed25519
+ssh-add -l
+
+pg-perf-bench collect-all-info \
+  --connection-type ssh \
+  --ssh-host db-host.example \
+  --ssh-user postgres \
+  --ssh-agent \
+  --ssh-known-hosts /secure/path/known_hosts \
+  --remote-pg-host 127.0.0.1 \
+  --remote-pg-port 5432 \
+  --host 127.0.0.1 \
+  --port 55432 \
+  --database appdb \
+  --pg-bin-path /usr/lib/postgresql/18/bin
+```
+
+`--ssh-key` and `--ssh-agent` are mutually exclusive. Agent mode uses the live
+socket inherited through `SSH_AUTH_SOCK`; the socket must remain available
+until collection or benchmarking finishes. The utility does not start an
+agent, run `ssh-add`, or forward the agent to the remote host. Both modes use
+explicit public-key authentication and the selected `known_hosts` policy
+without loading the user's OpenSSH configuration.
+
 For database modes, `--host` and `--port` are the local bind address and
 free port. `--remote-pg-host` and `--remote-pg-port` identify PostgreSQL from
 the SSH server. Host commands run remotely; `asyncpg`, `pgbench`, and `psql`

doc/ssh_mode_usage.md

@@ -30,6 +30,20 @@ Install the public key for the account selected by `--ssh-user`. For benchmark
 mode that account must be able to run `pg_ctl` for the selected cluster; using
 the PostgreSQL service owner is the simplest model.
 
+The key may be referenced directly with `--ssh-key`, or loaded into an
+already-running local agent and selected with `--ssh-agent`:
+
+```bash
+eval "$(ssh-agent -s)"
+ssh-add ~/.ssh/pg_perf_bench
+ssh-add -l
+```
+
+The two options are mutually exclusive. Agent mode requires a live Unix socket
+in the inherited `SSH_AUTH_SOCK` for the complete command. `pg-perf-bench`
+does not start the agent, load identities, or forward the agent to the target.
+Key mode explicitly disables agent fallback.
+
 Host-key verification is enabled by default. Use `--ssh-known-hosts` for a
 dedicated file. Reserve `--ssh-insecure-no-host-key-check` for isolated,
 disposable stands.
@@ -110,8 +124,10 @@ database host through the already authenticated SSH session.
 
 - host-key error: update the selected known-hosts file after independently
   verifying the server key;
-- authentication failure: verify `--ssh-user`, private key permissions, and
-  the installed public key;
+- authentication failure in key mode: verify `--ssh-user`, private key
+  permissions, and the installed public key;
+- authentication failure in agent mode: verify `SSH_AUTH_SOCK`, `ssh-add -l`,
+  the agent lifetime, and the installed public key;
 - local bind failure: choose an unused `--port`;
 - PostgreSQL connection failure with working SSH: verify the remote address,
   PostgreSQL authentication, and listen rules;

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pg-perf-bench"
-version = "0.2.0"
+version = "0.2.1"
 description = "Reproducible PostgreSQL TPS benchmarks with environment evidence and report comparison"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -25,7 +25,7 @@ dependencies = [
   "asyncpg>=0.29,<1",
   "asyncssh>=2.21,<3",
   "docker>=7.1,<8",
-  "pg-diag>=0.10.3,<1",
+  "pg-diag>=0.10.5,<1",
   "aenum>=3.1,<4; python_version < '3.11'",
 ]
 
@@ -47,7 +47,7 @@ dev = [
 pg-perf-bench = "pg_perf_bench.cli:main"
 
 [project.urls]
-Homepage = "https://github.com/O2eg/pg_perf_bench"
+Homepage = "https://o2eg.com/"
 Repository = "https://github.com/O2eg/pg_perf_bench"
 Issues = "https://github.com/O2eg/pg_perf_bench/issues"
 

src/pg_perf_bench/__init__.py

@@ -1,3 +1,3 @@
 """PostgreSQL performance benchmark and evidence collector."""
 
-__version__ = '0.2.0'
+__version__ = '0.2.1'

src/pg_perf_bench/cli.py

@@ -113,7 +113,13 @@ def _add_host_args(parser: argparse.ArgumentParser) -> None:
     parser.add_argument('--ssh-host')
     parser.add_argument('--ssh-port', type=positive_int, default=22)
     parser.add_argument('--ssh-user', default='postgres')
-    parser.add_argument('--ssh-key')
+    ssh_auth = parser.add_mutually_exclusive_group()
+    ssh_auth.add_argument('--ssh-key')
+    ssh_auth.add_argument(
+        '--ssh-agent',
+        action='store_true',
+        help='authenticate through the existing SSH_AUTH_SOCK agent',
+    )
     parser.add_argument('--ssh-known-hosts')
     parser.add_argument('--ssh-insecure-no-host-key-check', action='store_true')
     parser.add_argument('--remote-pg-host')

src/pg_perf_bench/config.py

@@ -3,6 +3,8 @@
 from __future__ import annotations
 
 import math
+import os
+import stat
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any
@@ -22,6 +24,20 @@
 PROTECTED_DATABASES = frozenset({'postgres', 'template0', 'template1'})
 
 
+def resolve_ssh_agent_socket() -> Path:
+    raw_path = os.environ.get('SSH_AUTH_SOCK')
+    if not raw_path:
+        raise ConfigurationError('--ssh-agent requires SSH_AUTH_SOCK to reference a running agent')
+    path = Path(raw_path).expanduser()
+    try:
+        mode = path.stat().st_mode
+    except OSError as exc:
+        raise ConfigurationError(f'SSH agent socket is not available: {path}') from exc
+    if not stat.S_ISSOCK(mode):
+        raise ConfigurationError(f'SSH_AUTH_SOCK is not a socket: {path}')
+    return path
+
+
 def _positive_int(value: Any, option: str) -> int:
     try:
         parsed = int(value)
@@ -100,6 +116,7 @@ class HostConfig:
     ssh_port: int = 22
     ssh_user: str = 'postgres'
     ssh_key: Path | None = None
+    ssh_agent: bool = False
     ssh_known_hosts: Path | None = None
     ssh_insecure_no_host_key_check: bool = False
     remote_pg_host: str | None = None
@@ -131,8 +148,14 @@ def connection_kwargs(
                 'host': self.ssh_host,
                 'port': self.ssh_port,
                 'username': self.ssh_user,
-                'client_keys': str(self.ssh_key) if self.ssh_key else None,
+                'client_keys': [str(self.ssh_key)] if self.ssh_key else [],
+                'agent_path': str(resolve_ssh_agent_socket()) if self.ssh_agent else None,
+                'agent_forwarding': False,
                 'known_hosts': known_hosts,
+                'config': None,
+                'preferred_auth': ['publickey'],
+                'password_auth': False,
+                'kbdint_auth': False,
                 'connect_timeout': min(self.command_timeout, 30.0),
             },
             'env': env,
@@ -267,6 +290,7 @@ def build_runtime_config(args: Any) -> RuntimeConfig:
         ssh_port=_positive_int(values.get('ssh_port') or 22, '--ssh-port'),
         ssh_user=values.get('ssh_user') or 'postgres',
         ssh_key=Path(ssh_key_value).expanduser() if ssh_key_value else None,
+        ssh_agent=bool(values.get('ssh_agent')),
         ssh_known_hosts=(
             Path(values['ssh_known_hosts']).expanduser() if values.get('ssh_known_hosts') else None
         ),
@@ -401,10 +425,12 @@ def _validate_host(host: HostConfig, *, needs_database: bool) -> None:
         return
     _required(host.ssh_host, '--ssh-host')
     _required(host.ssh_user, '--ssh-user')
-    _required(host.ssh_key, '--ssh-key')
-    assert host.ssh_key is not None
-    if not host.ssh_key.is_file():
+    if (host.ssh_key is None) == (not host.ssh_agent):
+        raise ConfigurationError('SSH requires exactly one of --ssh-key or --ssh-agent')
+    if host.ssh_key is not None and not host.ssh_key.is_file():
         raise ConfigurationError(f'SSH private key does not exist: {host.ssh_key}')
+    if host.ssh_agent:
+        resolve_ssh_agent_socket()
     if not host.ssh_insecure_no_host_key_check:
         known_hosts = host.ssh_known_hosts or Path('~/.ssh/known_hosts').expanduser()
         if not known_hosts.is_file():

src/pg_perf_bench/connections/ssh.py

@@ -42,7 +42,7 @@ async def start(self) -> None:
         except asyncssh.PermissionDenied as exc:
             await self.aclose()
             raise PermissionError(
-                'Verify the SSH user and private key for remote server access.'
+                'Verify the SSH user and selected key or agent identity for remote server access.'
             ) from exc
         except (asyncssh.ConnectionLost, OSError, asyncio.TimeoutError) as exc:
             await self.aclose()

src/pg_perf_bench/context/base_context.py

@@ -1,5 +1,6 @@
 from pathlib import Path
 
+from pg_perf_bench.config import resolve_ssh_agent_socket
 from pg_perf_bench.const import ConnectionType
 
 
@@ -43,17 +44,24 @@ def _add_connection_config(self, args):
 
     def _add_ssh_connection_config(self, args):
         """Add SSH connection configuration"""
+        use_agent = bool(getattr(args, 'ssh_agent', False))
         conn_params = {
             'host': args.ssh_host,
             'port': args.ssh_port,
             'username': getattr(args, 'ssh_user', None) or 'postgres',
-            'client_keys': args.ssh_key,
+            'client_keys': [] if use_agent else [args.ssh_key],
+            'agent_path': str(resolve_ssh_agent_socket()) if use_agent else None,
+            'agent_forwarding': False,
             'known_hosts': (
                 None
                 if getattr(args, 'ssh_insecure_no_host_key_check', False)
                 else getattr(args, 'ssh_known_hosts', None)
                 or str(Path('~/.ssh/known_hosts').expanduser())
             ),
+            'config': None,
+            'preferred_auth': ['publickey'],
+            'password_auth': False,
+            'kbdint_auth': False,
             'connect_timeout': float(getattr(args, 'connect_timeout', 5.0)),
         }
 

src/pg_perf_bench/context/benchmark.py

@@ -94,7 +94,6 @@ def filter_none(self, d: dict) -> dict:
         SSHConnectionArgs = [
             'ssh_host',
             'ssh_port',
-            'ssh_key',
             'remote_pg_host',
             'remote_pg_port',
         ]
@@ -122,6 +121,11 @@ def filter_none(self, d: dict) -> dict:
 
         ctype = d.get('connection_type')
         if ctype == ConnectionType.SSH:
+            if bool(d.get('ssh_key')) == bool(d.get('ssh_agent')):
+                raise ValueError(
+                    'Exactly one of "--ssh-key" or "--ssh-agent" must be specified '
+                    f'for connection type "{ConnectionType.SSH}"'
+                )
             for key in SSHConnectionArgs:
                 if d.get(key) is None:
                     raise ValueError(

src/pg_perf_bench/context/collect_info.py

@@ -77,7 +77,7 @@ def __init__(self, args, logger):
         }
 
     def filter_none(self, d: dict) -> dict:
-        SSHConnectionArgs = ['ssh_host', 'ssh_port', 'ssh_key']
+        SSHConnectionArgs = ['ssh_host', 'ssh_port']
         DockerConnectionArgs = [
             'container_name',
         ]
@@ -101,6 +101,11 @@ def filter_none(self, d: dict) -> dict:
 
         ctype = d.get('connection_type')
         if ctype == ConnectionType.SSH:
+            if bool(d.get('ssh_key')) == bool(d.get('ssh_agent')):
+                raise ValueError(
+                    'Exactly one of "--ssh-key" or "--ssh-agent" must be specified '
+                    f'for connection type "{ConnectionType.SSH}"'
+                )
             for key in SSHConnectionArgs:
                 if d.get(key) is None:
                     raise ValueError(