postgresql: pass through JIT-enabled variant of non-JIT postgres and vice versa

This is useful if your postgresql version is dependant on
`system.stateVersion` and not pinned down manually. Then it's not
necessary to find out which version exactly is in use and define
`package` manually, but just stay with what NixOS provides as default:

    $ nix-instantiate -A postgresql
    /nix/store/82fzmb77mz2b787dgj7mn4a8i4f6l6sn-postgresql-14.7.drv
    $ nix-instantiate -A postgresql_jit
    /nix/store/qsjkb72fcrrfpsszrwbsi9q9wgp39m50-postgresql-14.7.drv
    $ nix-instantiate -A postgresql.withJIT
    /nix/store/qsjkb72fcrrfpsszrwbsi9q9wgp39m50-postgresql-14.7.drv
    $ nix-instantiate -A postgresql.withJIT.withoutJIT
    /nix/store/82fzmb77mz2b787dgj7mn4a8i4f6l6sn-postgresql-14.7.drv

I.e. you can use postgresql with JIT (for complex queries only[1]) like
this:

    services.postgresql = {
      enable = true;
      enableJIT = true;
    };

Performing a new override instead of re-using the `_jit`-variants for
that has the nice property that overlays for the original package apply
to the JIT-enabled variant, i.e.

    with import ./. {
      overlays = [
        (self: super: {
          postgresql = super.postgresql.overrideAttrs (_: { fnord = "snens"; });
        })
      ];
    };
    postgresql.withJIT.fnord

still gives the string `snens` whereas `postgresql_jit` doesn't have the
attribute `fnord` in its derivation.

[1] https://www.postgresql.org/docs/current/runtime-config-query.html#GUC-JIT-ABOVE-COST

2 files changed, 46 insertions, 5 deletions

diff --git a/nixos/modules/services/databases/postgresql.md b/nixos/modules/services/databases/postgresql.md
index 1805bafe3be..4d66ee38be4 100644
--- a/nixos/modules/services/databases/postgresql.md
+++ b/nixos/modules/services/databases/postgresql.md
@@ -171,3 +171,40 @@ self: super: {
   };
 }
 ```
+
+## JIT (Just-In-Time compilation) {#module-services-postgres-jit}
+
+[JIT](https://www.postgresql.org/docs/current/jit-reason.html)-support in the PostgreSQL package
+is disabled by default because of the ~300MiB closure-size increase from the LLVM dependency. It
+can be optionally enabled in PostgreSQL with the following config option:
+
+```nix
+{
+  services.postgresql.enableJIT = true;
+}
+```
+
+This makes sure that the [`jit`](https://www.postgresql.org/docs/current/runtime-config-query.html#GUC-JIT)-setting
+is set to `on` and a PostgreSQL package with JIT enabled is used. Further tweaking of the JIT compiler, e.g. setting a different
+query cost threshold via [`jit_above_cost`](https://www.postgresql.org/docs/current/runtime-config-query.html#GUC-JIT-ABOVE-COST)
+can be done manually via [`services.postgresql.settings`](#opt-services.postgresql.settings).
+
+The attribute-names of JIT-enabled PostgreSQL packages are suffixed with `_jit`, i.e. for each `pkgs.postgresql`
+(and `pkgs.postgresql_<major>`) in `nixpkgs` there's also a `pkgs.postgresql_jit` (and `pkgs.postgresql_<major>_jit`).
+Alternatively, a JIT-enabled variant can be derived from a given `postgresql` package via `postgresql.withJIT`.
+This is also useful if it's not clear which attribute from `nixpkgs` was originally used (e.g. when working with
+[`config.services.postgresql.package`](#opt-services.postgresql.package) or if the package was modified via an
+overlay) since all modifications are propagated to `withJIT`. I.e.
+
+```nix
+with import <nixpkgs> {
+  overlays = [
+    (self: super: {
+      postgresql = super.postgresql.overrideAttrs (_: { pname = "foobar"; });
+    })
+  ];
+};
+postgresql.withJIT.pname
+```
+
+evaluates to `"foobar"`.
diff --git a/nixos/modules/services/databases/postgresql.nix b/nixos/modules/services/databases/postgresql.nix
index 7bbe1ad2259..16f76248b3a 100644
--- a/nixos/modules/services/databases/postgresql.nix
+++ b/nixos/modules/services/databases/postgresql.nix
@@ -42,6 +42,8 @@ in
 
       enable = mkEnableOption (lib.mdDoc "PostgreSQL Server");
 
+      enableJIT = mkEnableOption (lib.mdDoc "JIT support");
+
       package = mkOption {
         type = types.package;
         example = literalExpression "pkgs.postgresql_11";
@@ -435,19 +437,21 @@ in
         log_line_prefix = cfg.logLinePrefix;
         listen_addresses = if cfg.enableTCPIP then "*" else "localhost";
         port = cfg.port;
+        jit = mkDefault (if cfg.enableJIT then "on" else "off");
       };
 
     services.postgresql.package = let
         mkThrow = ver: throw "postgresql_${ver} was removed, please upgrade your postgresql version.";
+        base = if versionAtLeast config.system.stateVersion "22.05" then pkgs.postgresql_14
+            else if versionAtLeast config.system.stateVersion "21.11" then pkgs.postgresql_13
+            else if versionAtLeast config.system.stateVersion "20.03" then pkgs.postgresql_11
+            else if versionAtLeast config.system.stateVersion "17.09" then mkThrow "9_6"
+            else mkThrow "9_5";
     in
       # Note: when changing the default, make it conditional on
       # ‘system.stateVersion’ to maintain compatibility with existing
       # systems!
-      mkDefault (if versionAtLeast config.system.stateVersion "22.05" then pkgs.postgresql_14
-            else if versionAtLeast config.system.stateVersion "21.11" then pkgs.postgresql_13
-            else if versionAtLeast config.system.stateVersion "20.03" then pkgs.postgresql_11
-            else if versionAtLeast config.system.stateVersion "17.09" then mkThrow "9_6"
-            else mkThrow "9_5");
+      mkDefault (if cfg.enableJIT then base.withJIT else base);
 
     services.postgresql.dataDir = mkDefault "/var/lib/postgresql/${cfg.package.psqlSchema}";