Complete 0.5.0: enable available diagnostic extensions by default

Author: O2eg

README.md

@@ -152,7 +152,11 @@ instead of `--passfile`; `PGPASSWORD` is accepted only as an admin-password comp
 `prepare-db --recreate` drops only the selected workload database after terminating its sessions;
 it does not drop the workload role.
 Unless `--extensions` is supplied, it creates both `pg_stat_statements` and
-`pg_buffercache`, matching the default pg_diag metric set used by pg_play.
+`pg_buffercache`, and also installs `pg_stat_kcache` and `pg_wait_sampling` when
+the server provides them. Managed `pg_stand` images include and preload the required
+libraries, and new databases inherit the extensions from `template1`. An explicit
+`--extensions` list replaces both default lists. Optional extension choices are
+included in the machine plan.
 
 ## Targets
 

src/pg_workload/common.py

@@ -32,6 +32,8 @@
 DEFAULT_PGBENCH_CLIENTS = 2
 DEFAULT_PGBENCH_THREADS = 2
 DEFAULT_EXTENSIONS = ("pg_stat_statements", "pg_buffercache")
+# Installed by prepare-db only when the server image provides them.
+DEFAULT_OPTIONAL_EXTENSIONS = ("pg_stat_kcache", "pg_wait_sampling")
 DEFAULT_SCRIPT_PATTERNS = ("sql/[0-9][0-9]_*.sql", "sql/select_*.sql")
 DEFAULT_RESOURCE_DISK_MAX_USED_PCT = 90
 DEFAULT_RESOURCE_MEM_MIN_AVAILABLE_PCT = 10

src/pg_workload/orchestration.py

@@ -8,7 +8,7 @@
 from typing import Any
 
 from pg_workload import __version__
-from pg_workload.common import DEFAULT_EXTENSIONS, csv_list, resolve_relative_path
+from pg_workload.common import DEFAULT_EXTENSIONS, DEFAULT_OPTIONAL_EXTENSIONS, csv_list, resolve_relative_path
 from pg_workload.profiles import Profile
 from pg_workload.state import load_state, state_with_profiles_enabled
 
@@ -140,6 +140,9 @@ def execution_plan(args: Any, profiles: list[Profile], operation: str) -> dict[s
             "workload_superuser": getattr(args, "workload_superuser", False),
             "rotate_workload_password": getattr(args, "rotate_workload_password", False),
             "extensions": sorted(set(csv_list(getattr(args, "extensions", None)) or DEFAULT_EXTENSIONS)),
+            "optional_extensions": (
+                [] if csv_list(getattr(args, "extensions", None)) else sorted(DEFAULT_OPTIONAL_EXTENSIONS)
+            ),
             "preload_libraries": sorted(set(csv_list(getattr(args, "preload_libraries", None)))),
         },
         "execution": {

src/pg_workload/runner.py

@@ -5,6 +5,7 @@
 
 from pg_workload.common import (
     DEFAULT_EXTENSIONS,
+    DEFAULT_OPTIONAL_EXTENSIONS,
     WorkloadError,
     as_string_list,
     csv_list,
@@ -101,6 +102,23 @@ def prepare_database(client: PgClient, config: RuntimeConfig, args: argparse.Nam
             admin_password,
             command=f"CREATE EXTENSION IF NOT EXISTS {safe_extension_name(extension)};",
         )
+    if not csv_list(args.extensions):
+        for extension in DEFAULT_OPTIONAL_EXTENSIONS:
+            available = client.query_scalar(
+                config.dbname,
+                config.admin_user,
+                admin_password,
+                f"SELECT 1 FROM pg_available_extensions WHERE name = {quote_literal(extension)};",
+            )
+            if available:
+                client.run_psql(
+                    config.dbname,
+                    config.admin_user,
+                    admin_password,
+                    command=(f"CREATE EXTENSION IF NOT EXISTS {safe_extension_name(extension)};"),
+                )
+            else:
+                eprint(f"Optional extension {extension} is not available; skipping")
 
     preload_libraries = csv_list(args.preload_libraries) or ["auto_explain", "pg_stat_statements"]
     check_preload_libraries(client, config, admin_password, preload_libraries)

tests/test_workload.py

@@ -912,6 +912,8 @@ def __init__(self):
             def query_scalar(self, dbname, user, password, sql):
                 if "pg_roles" in sql or "pg_database" in sql:
                     return "1"
+                if "pg_available_extensions" in sql:
+                    return "1"
                 return "auto_explain,pg_stat_statements"
 
             def run_psql(self, dbname, user, password, *, command, **kwargs):
@@ -933,6 +935,8 @@ def run_psql(self, dbname, user, password, *, command, **kwargs):
 
             self.assertIn('CREATE EXTENSION IF NOT EXISTS "pg_stat_statements";', client.commands)
             self.assertIn('CREATE EXTENSION IF NOT EXISTS "pg_buffercache";', client.commands)
+            self.assertIn('CREATE EXTENSION IF NOT EXISTS "pg_stat_kcache";', client.commands)
+            self.assertIn('CREATE EXTENSION IF NOT EXISTS "pg_wait_sampling";', client.commands)
 
     def test_job_recover_on_failure_defaults_to_true(self):
         self.assertTrue(self.workload.job_recover_on_failure({}))