fc-postgresql: support upgrades with custom extensions

FC-42163

Upgrading postgresql to a new major with fc-postgresql fails when using
custom extensions, defined via `services.postgresql.extraPlugins`.

This patch modifies the relevant agent code to use a postgresql package
WITH all needed plugins for upgrades:

* for automatic upgrades, these extensions are discovered automatically.
* for manual upgrades, the user must specify those manually with
  `--extension-names`.

The original use-case of a customer was to upgrade postgresql to a newer
version with `postgresqlPackages.anonymizer` installed (FC-42089). This
is a little trickier since it requires `anon` to be listed in
`shared_preload_libraries`. For this, all `shared_preload_libraries` are
now extracted from `postgresql.conf` and passed to `pg_upgrade` via CLI.
This code is based on the assumption that the old `postgresql.conf` is
generated by Nix and thus ignores all edge-cases, the postgresql config
format may have.

changelog.d/20241209_114200_mb_FC_42163_fc_postgresql_extension_support.md

@@ -0,0 +1,22 @@
+<!--
+
+A new changelog entry.
+
+Delete placeholder items that do not apply. Empty sections will be removed
+automatically during release.
+
+Leave the XX.XX as is: this is a placeholder and will be automatically filled
+correctly during the release and helps when backporting over multiple platform
+branches.
+
+-->
+
+### Impact
+
+- None.
+
+### NixOS XX.XX platform
+
+- The `fc-postgresql` command now supports upgrades of databases with preinstalled extensions:
+  - When upgrading manually with `fc-postgresql`, add `--extension-names ext1 --extension-names ext2` to the command line. `ext1`/`ext2` must be the package names of the extensions without the `postgresqlPackages.`-prefix. Usually it's the packages in [`services.postgresql.extraPlugins`](https://search.flyingcircus.io/search/options?q=services.postgresql.extraPlugins&channel=fc-24.05-dev#services.postgresql.extraPlugins).
+  - When using automatic upgrades ([`flyingcircus.services.postgresql.autoUpgrade.enable `](https://search.flyingcircus.io/search/options?q=+flyingcircus.services.postgresql.autoUpgrade.enable+&channel=fc-24.11-dev&page=1#flyingcircus.services.postgresql.autoUpgrade.enable)), existing extensions will be discovered automatically. You don't have to do anything in this case.

doc/src/postgresql.md

@@ -152,6 +152,21 @@ Note that this is done while the old role is still active. It's safe to run
 the command while PostgreSQL is running as it does not have an impact on the
 current cluster and downtime is not required.
 
+If custom extensions are enabled with `services.postgresql.extraPlugins`, make sure
+to add those to the `fc-postgresql` invocation. I.e. for
+
+```nix
+{
+  services.postgresql.extraPlugins = plugins: [ plugins.anonymizer plugins.pgvector ];
+}
+```
+
+you'll need
+
+```
+sudo -u postgres fc-postgresql upgrade --new-version 14 --expected mydb --expected otherdb --extension-names anonymizer --extension-names pgvector
+```
+
 The command should automatically find the old data directory for 13, create
 the new data directory for 14, set it up, and succeed if no problems with the
 old cluster were found. Problems may occur if the old cluster has been
@@ -170,6 +185,9 @@ To actually run the upgrade, use:
 sudo -u postgres fc-postgresql upgrade --new-version 14 --expected mydb --expected otherdb --upgrade-now
 ```
 
+Please note that you'll also need the `--extension-names` parameters as described
+above.
+
 This will stop the postgresql service, prevent it from starting during the
 upgrade, migrate data and mark the old data directory as migrated. This data
 directory cannot be used by the postgresql service anymore after this point.

nixos/services/postgresql/default.nix

@@ -64,6 +64,16 @@ let
   collationVerifierScript = pkgs.writers.writePython3Bin "verify-collations"
       {} (builtins.readFile ./verify-collations.py);
 
+  mkExtensionNamesForAutoUpgrade = postgresqlMajor:
+    let
+      availableExtensions = lib.mapAttrs'
+        (attrName: package: lib.nameValuePair (lib.getName package) attrName)
+        (builtins.removeAttrs pkgs."postgresql${toString postgresqlMajor}Packages" [ "recurseForDerivations" ]);
+      pluginByPname = package: lib.optionals (availableExtensions?${lib.getName package}) [
+        availableExtensions.${lib.getName package}
+      ];
+    in
+    lib.concatMap pluginByPname (config.services.postgresql.extraPlugins config.services.postgresql.package.pkgs);
 in {
   options = with lib; {
 
@@ -159,7 +169,7 @@ in {
       '';
 
       systemd.services.postgresql.postStart = ''
-        ln -sfT ${postgresqlPkg} ${upstreamCfg.dataDir}/package
+        ln -sfT ${config.services.postgresql.package.withPackages config.services.postgresql.extraPlugins} ${upstreamCfg.dataDir}/package
         ln -sfT ${upstreamCfg.dataDir}/package /nix/var/nix/gcroots/per-user/postgres/package_${cfg.majorVersion}
       '' + (lib.optionalString (lib.versionAtLeast cfg.majorVersion "15") ''
         ${collationVerifierScript}/bin/verify-collations ${upstreamCfg.dataDir}
@@ -421,12 +431,13 @@ in {
             "${config.flyingcircus.agent.package}/bin/fc-postgresql upgrade"
             "--new-version ${cfg.majorVersion}"
             "--new-data-dir ${upstreamCfg.dataDir}"
-            "--new-bin-dir ${upstreamCfg.package}/bin"
+            "--new-bin-dir ${upstreamCfg.package.withPackages (ps: attrValues (lib.getAttrs (mkExtensionNamesForAutoUpgrade cfg.majorVersion) ps))}/bin"
             "--no-stop"
             "--nothing-to-do-is-ok"
             "--upgrade-now"
-          ] ++ lib.optional cfg.autoUpgrade.checkExpectedDatabases
-            "--existing-db-check ${expectedDatabaseStr}";
+          ] ++ lib.optionals cfg.autoUpgrade.checkExpectedDatabases [
+            "--existing-db-check ${expectedDatabaseStr}"
+          ] ++ map (extName: "--extension-names ${extName}") (mkExtensionNamesForAutoUpgrade cfg.majorVersion);
         in
           concatStringsSep " \\\n  " upgradeCmd;
 

pkgs/fc/agent/fc/manage/postgresql.py

@@ -150,6 +150,10 @@ def upgrade(
         ),
         default=None,
     ),
+    extension_names: Optional[List[str]] = Option(
+        default=[],
+        help="Extensions that are used in the databases. Must be the package name from https://search.nixos.org without the `postgresqlXXPackages.`-prefix.",
+    ),
     stop: Optional[bool] = Option(default=None),
     nothing_to_do_is_ok: Optional[bool] = False,
 ):
@@ -166,7 +170,7 @@ def upgrade(
             raise Exit(2)
 
         new_bin_dir = fc.util.postgresql.build_new_bin_dir(
-            log, context.pg_data_root, new_version
+            log, context.pg_data_root, new_version, extension_names
         )
         new_data_dir = context.pg_data_root / new_version.value
 
@@ -345,6 +349,10 @@ def prepare_autoupgrade(
         show_choices=True,
         help="PostgreSQL version to upgrade to.",
     ),
+    extension_names: Optional[List[str]] = Option(
+        default=[],
+        help="Extensions that are used in the databases. Must be the package name from https://search.nixos.org without the `postgresqlXXPackages.`-prefix.",
+    ),
     nothing_to_do_is_ok: Optional[bool] = False,
 ):
     with open(config) as f:
@@ -354,7 +362,7 @@ def prepare_autoupgrade(
     log = structlog.get_logger()
 
     new_bin_dir = fc.util.postgresql.build_new_bin_dir(
-        log, context.pg_data_root, new_version
+        log, context.pg_data_root, new_version, extension_names
     )
     new_data_dir = context.pg_data_root / new_version.value
 

pkgs/fc/agent/fc/manage/tests/test_postgresql.py

@@ -1,5 +1,6 @@
 import json
 import traceback
+from typing import List
 from unittest.mock import Mock, patch
 
 import fc.manage.postgresql

pkgs/fc/agent/fc/util/postgresql.py

@@ -7,6 +7,7 @@
 from enum import Enum
 from pathlib import Path
 from subprocess import CalledProcessError, run
+from typing import List
 
 MIGRATED_TO_TEMPLATE = """\
 WARNING: This data directory should not be used anymore!
@@ -79,12 +80,17 @@ def is_service_running():
     return proc.returncode == 0
 
 
-def build_new_bin_dir(log, pg_data_root: Path, new_version: PGVersion):
+def build_new_bin_dir(
+    log,
+    pg_data_root: Path,
+    new_version: PGVersion,
+    extension_names: List[str] = [],
+):
+    extension_names_str = " ".join(extension_names)
     nix_build_new_pg_cmd = [
         "nix-build",
-        "<fc>",
-        "-A",
-        "postgresql_" + new_version.value,
+        "-E",
+        f"with import <fc> {{}}; postgresql_{new_version.value}.withPackages (ps: with ps; [ {extension_names_str} ])",
         "--out-link",
         pg_data_root / "pg_upgrade-package",
     ]
@@ -472,6 +478,22 @@ def run_pg_upgrade_check(
         raise
 
 
+def retrieve_shared_preload_libraries_setting(
+    old_data_dir: Path,
+) -> str | None:
+    """
+    The setting `shared_preload_libraries` must be correctly set for upgrades with
+    e.g. `pg_anon` installed. This function assumes that `old_data_dir` has a `postgresql.conf`
+    managed by NixOS, i.e. a trivial key-value text with quoted strings.
+    """
+    postgresql_conf = old_data_dir / "postgresql.conf"
+    if postgresql_conf.exists():
+        with open(postgresql_conf) as f:
+            for line in f.readlines():
+                if line.startswith("shared_preload_libraries"):
+                    return line.split("=")[1].strip()
+
+
 def run_pg_upgrade(
     log,
     new_bin_dir: Path,
@@ -490,6 +512,17 @@ def run_pg_upgrade(
         new_bin_dir,
     ]
 
+    if shared_preload_libraries := retrieve_shared_preload_libraries_setting(
+        old_data_dir
+    ):
+        setting = f"-c shared_preload_libraries={shared_preload_libraries}"
+        upgrade_cmd += [
+            "-o",
+            setting,
+            "-O",
+            setting,
+        ]
+
     if pg_upgrade_clone_available(
         log,
         new_bin_dir=new_bin_dir,

pkgs/fc/agent/fc/util/tests/test_postgresql.py

@@ -1,4 +1,5 @@
 import json
+from typing import List
 from unittest.mock import Mock
 
 import fc.manage.postgresql

tests/default.nix

@@ -99,6 +99,7 @@ in {
   postgresql15 = callTest ./postgresql { version = "15"; };
   postgresql16 = callTest ./postgresql { version = "16"; };
   postgresql-autoupgrade = callSubTests ./postgresql/upgrade.nix {};
+  postgresql-autoupgrade-exts = callSubTests ./postgresql/upgrade-with-extension.nix {};
   prometheus = callTest ./prometheus.nix {};
   rabbitmq = callTest ./rabbitmq.nix {};
   redis = callTest ./redis.nix {};