7 files changed, 118 insertions, 9 deletions
diff --git a/nixos/doc/manual/release-notes/rl-2305.section.md b/nixos/doc/manual/release-notes/rl-2305.section.md
index 85fce98b88c0c..e5242847b6fbb 100644
--- a/nixos/doc/manual/release-notes/rl-2305.section.md
+++ b/nixos/doc/manual/release-notes/rl-2305.section.md
@@ -273,6 +273,16 @@ In addition to numerous new and upgraded packages, this release has the followin
 
 - `services.chronyd` is now started with additional systemd sandbox/hardening options for better security.
 
+- PostgreSQL has opt-in support for [JIT compilation](https://www.postgresql.org/docs/current/jit-reason.html). It can be enabled like this:
+  ```nix
+  {
+    services.postgresql = {
+      enable = true;
+      enableJIT = true;
+    };
+  }
+  ```
+
 - `services.dhcpcd` service now don't solicit or accept IPv6 Router Advertisements on interfaces that use static IPv6 addresses.
 
 - The module `services.headscale` was refactored to be compliant with [RFC 0042](https://github.com/NixOS/rfcs/blob/master/rfcs/0042-config-option.md). To be precise, this means that the following things have changed:
diff --git a/nixos/modules/services/databases/postgresql.md b/nixos/modules/services/databases/postgresql.md
index 1805bafe3be38..4d66ee38be426 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 7bbe1ad225955..3d55995aba055 100644
--- a/nixos/modules/services/databases/postgresql.nix
+++ b/nixos/modules/services/databases/postgresql.nix
@@ -7,9 +7,18 @@ let
   cfg = config.services.postgresql;
 
   postgresql =
+    let
+      # ensure that
+      #   services.postgresql = {
+      #     enableJIT = true;
+      #     package = pkgs.postgresql_<major>;
+      #   };
+      # works.
+      base = if cfg.enableJIT && !cfg.package.jitSupport then cfg.package.withJIT else cfg.package;
+    in
     if cfg.extraPlugins == []
-      then cfg.package
-      else cfg.package.withPackages (_: cfg.extraPlugins);
+      then base
+      else base.withPackages (_: cfg.extraPlugins);
 
   toStr = value:
     if true == value then "yes"
@@ -42,6 +51,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 +446,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}";
 
diff --git a/nixos/tests/all-tests.nix b/nixos/tests/all-tests.nix
index cdd9e00d25886..3167a865fc9db 100644
--- a/nixos/tests/all-tests.nix
+++ b/nixos/tests/all-tests.nix
@@ -565,6 +565,7 @@ in {
   postfixadmin = handleTest ./postfixadmin.nix {};
   postgis = handleTest ./postgis.nix {};
   postgresql = handleTest ./postgresql.nix {};
+  postgresql-jit = handleTest ./postgresql-jit.nix {};
   postgresql-wal-receiver = handleTest ./postgresql-wal-receiver.nix {};
   powerdns = handleTest ./powerdns.nix {};
   powerdns-admin = handleTest ./powerdns-admin.nix {};
diff --git a/nixos/tests/postgresql-jit.nix b/nixos/tests/postgresql-jit.nix
new file mode 100644
index 0000000000000..baf26b8da2b39
--- /dev/null
+++ b/nixos/tests/postgresql-jit.nix
@@ -0,0 +1,48 @@
+{ system ? builtins.currentSystem
+, config ? {}
+, pkgs ? import ../.. { inherit system config; }
+}:
+
+with import ../lib/testing-python.nix { inherit system pkgs; };
+
+let
+  inherit (pkgs) lib;
+  packages = builtins.attrNames (import ../../pkgs/servers/sql/postgresql pkgs);
+
+  mkJitTest = packageName: makeTest {
+    name = "${packageName}";
+    meta.maintainers = with lib.maintainers; [ ma27 ];
+    nodes.machine = { pkgs, lib, ... }: {
+      services.postgresql = {
+        enable = true;
+        enableJIT = true;
+        package = pkgs.${packageName};
+        initialScript = pkgs.writeText "init.sql" ''
+          create table demo (id int);
+          insert into demo (id) select generate_series(1, 5);
+        '';
+      };
+    };
+    testScript = ''
+      machine.start()
+      machine.wait_for_unit("postgresql.service")
+
+      with subtest("JIT is enabled"):
+          machine.succeed("sudo -u postgres psql <<<'show jit;' | grep 'on'")
+
+      with subtest("Test JIT works fine"):
+          output = machine.succeed(
+              "cat ${pkgs.writeText "test.sql" ''
+                set jit_above_cost = 1;
+                EXPLAIN ANALYZE SELECT CONCAT('jit result = ', SUM(id)) FROM demo;
+                SELECT CONCAT('jit result = ', SUM(id)) from demo;
+              ''} | sudo -u postgres psql"
+          )
+          assert "JIT:" in output
+          assert "jit result = 15" in output
+
+      machine.shutdown()
+    '';
+  };
+in
+lib.genAttrs packages mkJitTest
diff --git a/nixos/tests/postgresql-wal-receiver.nix b/nixos/tests/postgresql-wal-receiver.nix
index ae2708546f5db..b0bd7711dbcd9 100644
--- a/nixos/tests/postgresql-wal-receiver.nix
+++ b/nixos/tests/postgresql-wal-receiver.nix
@@ -116,4 +116,4 @@ let
     };
 
 # Maps the generic function over all attributes of PostgreSQL packages
-in builtins.listToAttrs (map makePostgresqlWalReceiverTest (builtins.attrNames (import ../../pkgs/servers/sql/postgresql { })))
+in builtins.listToAttrs (map makePostgresqlWalReceiverTest (builtins.attrNames (import ../../pkgs/servers/sql/postgresql pkgs)))
diff --git a/nixos/tests/postgresql.nix b/nixos/tests/postgresql.nix
index 7e0a82c388288..b44849e0a14e5 100644
--- a/nixos/tests/postgresql.nix
+++ b/nixos/tests/postgresql.nix
@@ -137,7 +137,7 @@ let
       maintainers = [ zagy ];
     };
 
-    machine = {...}:
+    nodes.machine = {...}:
       {
         services.postgresql = {
           enable = true;