nixos/postgresql/citus: fix syscall filter and add test

nixos/modules/services/databases/postgresql.nix

@@ -727,10 +727,16 @@ in
           RestrictRealtime = true;
           RestrictSUIDSGID = true;
           SystemCallArchitectures = "native";
-          SystemCallFilter = [
-            "@system-service"
-            "~@privileged @resources"
-          ] ++ lib.optionals (any extensionInstalled [ "plv8" ]) [ "@pkey" ];
+          SystemCallFilter =
+            [
+              "@system-service"
+              "~@privileged @resources"
+            ]
+            ++ lib.optionals (any extensionInstalled [ "plv8" ]) [ "@pkey" ]
+            ++ lib.optionals (any extensionInstalled [ "citus" ]) [
+              "getpriority"
+              "setpriority"
+            ];
           UMask = if groupAccessAvailable then "0027" else "0077";
         }
         (mkIf (cfg.dataDir != "/var/lib/postgresql/${cfg.package.psqlSchema}") {

nixos/tests/postgresql/citus.nix

@@ -0,0 +1,165 @@
+{
+  pkgs,
+  makeTest,
+}:
+
+let
+  inherit (pkgs) lib;
+
+  # commands taken straight from the README at https://github.com/citusdata/citus?tab=readme-ov-file#using-citus
+  test-sql = pkgs.writeText "postgresql-test" ''
+    CREATE EXTENSION citus;
+
+    CREATE TABLE events (
+      device_id bigint,
+      event_id bigserial,
+      event_time timestamptz default now(),
+      data jsonb not null,
+      PRIMARY KEY (device_id, event_id)
+    );
+
+    -- distribute the events table across shards placed locally or on the worker nodes
+    SELECT create_distributed_table('events', 'device_id');
+
+    -- insert some events
+    INSERT INTO events (device_id, data)
+    SELECT s % 100, ('{"measurement":'||random()||'}')::jsonb FROM generate_series(1,1000) s;
+
+    -- get the last 3 events for device 1, routed to a single node
+    SELECT * FROM events WHERE device_id = 1 ORDER BY event_time DESC, event_id DESC LIMIT 3;
+
+    CREATE TABLE devices (
+      device_id bigint primary key,
+      device_name text,
+      device_type_id int
+    );
+    CREATE INDEX ON devices (device_type_id);
+
+    -- co-locate the devices table with the events table
+    SELECT create_distributed_table('devices', 'device_id', colocate_with := 'events');
+
+    -- insert device metadata
+    INSERT INTO devices (device_id, device_name, device_type_id)
+    SELECT s, 'device-'||s, 55 FROM generate_series(0, 99) s;
+
+    -- optionally: make sure the application can only insert events for a known device
+    ALTER TABLE events ADD CONSTRAINT device_id_fk
+    FOREIGN KEY (device_id) REFERENCES devices (device_id);
+
+    -- get the average measurement across all devices of type 55, parallelized across shards
+    SELECT avg((data->>'measurement')::double precision)
+    FROM events JOIN devices USING (device_id)
+    WHERE device_type_id = 55;
+
+    CREATE TABLE device_logs (
+      device_id bigint primary key,
+      log text
+    );
+
+    -- insert device logs
+    INSERT INTO device_logs (device_id, log)
+    SELECT s, 'device log:'||s FROM generate_series(0, 99) s;
+
+    -- convert device_logs into a distributed table without interrupting the application
+    SELECT create_distributed_table_concurrently('device_logs', 'device_id', colocate_with := 'devices');
+
+
+    -- get the count of the logs, parallelized across shards
+    SELECT count(*) FROM device_logs;
+
+    CREATE TABLE device_types (
+      device_type_id int primary key,
+      device_type_name text not null unique
+    );
+
+    -- replicate the table across all nodes to enable foreign keys and joins on any column
+    SELECT create_reference_table('device_types');
+
+    -- insert a device type
+    INSERT INTO device_types (device_type_id, device_type_name) VALUES (55, 'laptop');
+
+    -- optionally: make sure the application can only insert devices with known types
+    ALTER TABLE devices ADD CONSTRAINT device_type_fk
+    FOREIGN KEY (device_type_id) REFERENCES device_types (device_type_id);
+
+    -- get the last 3 events for devices whose type name starts with laptop, parallelized across shards
+    SELECT device_id, event_time, data->>'measurement' AS value, device_name, device_type_name
+    FROM events JOIN devices USING (device_id) JOIN device_types USING (device_type_id)
+    WHERE device_type_name LIKE 'laptop%' ORDER BY event_time DESC LIMIT 3;
+
+    CREATE TABLE events_columnar (
+      device_id bigint,
+      event_id bigserial,
+      event_time timestamptz default now(),
+      data jsonb not null
+    )
+    USING columnar;
+
+    -- insert some data
+    INSERT INTO events_columnar (device_id, data)
+    SELECT d, '{"hello":"columnar"}' FROM generate_series(1,1000) d;
+
+    -- create a row-based table to compare
+    CREATE TABLE events_row AS SELECT * FROM events_columnar;
+  '';
+
+  makeTestFor =
+    package:
+    makeTest {
+      name = "citus-${package.name}";
+      meta = with lib.maintainers; {
+        maintainers = [ typetetris ];
+      };
+
+      nodes.machine =
+        { ... }:
+        {
+          services.postgresql = {
+            inherit package;
+            enable = true;
+            enableJIT = lib.hasInfix "-jit-" package.name;
+            extensions =
+              ps: with ps; [
+                citus
+              ];
+            settings = {
+              shared_preload_libraries = "citus";
+            };
+          };
+        };
+
+      testScript = ''
+        def check_count(statement, lines):
+            return 'test $(sudo -u postgres psql postgres -tAc "{}") -eq {}'.format(
+                statement, lines
+            )
+
+
+        machine.start()
+        machine.wait_for_unit("postgresql")
+
+        with subtest("Postgresql with extension citus is available just after unit start"):
+            machine.succeed(
+                "sudo -u postgres psql -f ${test-sql}"
+            )
+
+        machine.succeed(check_count("SELECT count(*) FROM events;", 1000))
+        machine.succeed(check_count("SELECT count(*) FROM devices;", 100))
+        machine.succeed(check_count("SELECT count(*) FROM events_columnar;", 1000))
+        machine.succeed(check_count("SELECT count(*) FROM events_row;", 1000))
+
+        machine.shutdown()
+      '';
+    };
+in
+# Not run by default, because this requires allowUnfree.
+# To run these tests:
+#   NIXPKGS_ALLOW_UNFREE=1 nix-build -A nixosTests.postgresql.citus
+lib.dontRecurseIntoAttrs (
+  lib.concatMapAttrs (n: p: { ${n} = makeTestFor p; }) (
+    lib.filterAttrs (_: p: !p.pkgs.citus.meta.broken) pkgs.postgresqlVersions
+  )
+  // {
+    passthru.override = p: makeTestFor p;
+  }
+)

nixos/tests/postgresql/default.nix

@@ -18,6 +18,7 @@ in
 
   # extensions
   anonymizer = importWithArgs ./anonymizer.nix;
+  citus = importWithArgs ./citus.nix;
   pgjwt = importWithArgs ./pgjwt.nix;
   pgvecto-rs = importWithArgs ./pgvecto-rs.nix;
   timescaledb = importWithArgs ./timescaledb.nix;