postgresql: monitor collation warnings, upgrade automatically if poss…

…ible

Fixes PL-131544

nixos/services/default.nix

@@ -33,7 +33,7 @@ in {
     ./opensearch.nix
     ./opensearch_dashboards.nix
     ./percona.nix
-    ./postgresql.nix
+    ./postgresql
     ./prometheus.nix
     ./rabbitmq.nix
     ./raid

nixos/services/postgresql.nix → nixos/services/postgresql/default.nix

@@ -61,6 +61,9 @@ let
     then { include_dir = "${localConfigPath}"; }
     else {};
 
+  collationVerifierScript = pkgs.writers.writePython3Bin "verify-collations"
+      {} (builtins.readFile ./verify-collations.py);
+
 in {
   options = with lib; {
 
@@ -159,13 +162,12 @@ in {
         fi
       '';
 
-      systemd.services.postgresql.postStart =
-      let
-        psql = "${postgresqlPkg}/bin/psql --port=${toString upstreamCfg.port}";
-      in ''
+      systemd.services.postgresql.postStart = ''
         ln -sfT ${postgresqlPkg} ${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}
+      '');
 
       systemd.services.postgresql.serviceConfig = {
         Restart = "always";
@@ -364,6 +366,12 @@ in {
               '';
             interval = 30;
           };
+        } // lib.optionalAttrs (lib.versionAtLeast cfg.majorVersion "15") {
+            postgresql-collation-warnings = {
+              notification = "PostgreSQL is reporting collation warnings with affected objects. Check /run/postgresql-collation-warnings and PL-131544.";
+              command = "check_file_age -i -c 10 -w 5 ${upstreamCfg.dataDir}/postgresql-collation-warnings";
+              interval = 600;
+            };
         } // lib.optionalAttrs (cfg.autoUpgrade.enable && cfg.autoUpgrade.checkExpectedDatabases) {
             postgresql-autoupgrade-possible = {
               notification = "Unexpected PostgreSQL databases present, autoupgrade will fail!";

nixos/services/postgresql/verify-collations.py

@@ -0,0 +1,84 @@
+import csv
+import io
+import os
+import pathlib
+import subprocess
+import sys
+
+DATADIR = sys.argv[1]
+
+IGNORE = "template0"
+
+GET_AFFECTED_OBJECTS = """
+SELECT pg_describe_object(refclassid, refobjid, refobjsubid) AS "Collation",
+       pg_describe_object(classid, objid, objsubid) AS "Object"
+FROM pg_depend d JOIN pg_collation c
+     ON refclassid = 'pg_collation'::regclass AND refobjid = c.oid
+WHERE c.collversion <> pg_collation_actual_version(c.oid)
+ORDER BY 1, 2;
+"""
+
+GET_COLLATION_VERSIONS = """
+SELECT datname, datcollate AS db_collation,
+       datcollversion,
+       pg_database_collation_actual_version(oid) AS oscollversion
+FROM pg_database
+WHERE datname != 'template0'
+"""
+
+
+def psql(*args):
+    env = os.environ.copy()
+    env["PGCLIENTENCODING"] = "utf-8"
+    output = subprocess.check_output(
+        ("psql", "--csv") + args, env=env, encoding="utf-8"
+    )
+    output = io.StringIO(output)
+    rows = csv.DictReader(output, delimiter=",", quotechar='"')
+    return rows
+
+
+print("Verifying database collation versions ...")
+
+databases = psql("-c", GET_COLLATION_VERSIONS)
+
+warnings = []
+
+for database in databases:
+    db_name = database["datname"]
+    if db_name in IGNORE:
+        continue
+    db_collation_version = database["datcollversion"]
+    os_collation_version = database["oscollversion"]
+    print(db_name)
+    print(f"\tdb collation version: {db_collation_version}")
+    print(f"\tos collation version: {os_collation_version}")
+    if db_collation_version == os_collation_version:
+        print("\tOK")
+        continue
+
+    affected_objects = list(psql("-c", GET_AFFECTED_OBJECTS, db_name))
+    if not affected_objects:
+        print("\tUpdating collation")
+        psql(
+            "-c",
+            f"ALTER DATABASE {db_name} REFRESH COLLATION VERSION",
+            db_name,
+        )
+    else:
+        for row in affected_objects:
+            obj, collation = row["Object"], row["Collation"]
+            print(f"\tobject: {obj} collation: {collation}")
+            warnings.append((db_name, obj, collation))
+        # XXX create a warning for the agent so we can figure out what the
+        # proper next step is See PL-131544 for context. potentially create
+        # a maintenance item here
+
+warning_file = pathlib.Path(DATADIR) / "postgresql-collation-warnings"
+
+if warnings:
+    with open(warning_file, "w") as f:
+        for warning in warnings:
+            f.write(str(warning))
+else:
+    warning_file.unlink(missing_ok=True)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.black]
-line-length = 80
+line-length = 79
 
 [tool.isort]
 profile = "black"
-line_length = 80
+line_length = 79