nixos/tests/postgresql/pgvecto-rs.nix

   1 {
   2   pkgs,
   3   makeTest,
   4 }:
   5
   6 let
   7   inherit (pkgs) lib;
   8
   9   # Test cases from https://docs.pgvecto.rs/use-cases/hybrid-search.html
  10   test-sql = pkgs.writeText "postgresql-test" ''
  11     CREATE EXTENSION vectors;
  12
  13     CREATE TABLE items (
  14       id bigserial PRIMARY KEY,
  15       content text NOT NULL,
  16       embedding vectors.vector(3) NOT NULL -- 3 dimensions
  17     );
  18
  19     INSERT INTO items (content, embedding) VALUES
  20       ('a fat cat sat on a mat and ate a fat rat', '[1, 2, 3]'),
  21       ('a fat dog sat on a mat and ate a fat rat', '[4, 5, 6]'),
  22       ('a thin cat sat on a mat and ate a thin rat', '[7, 8, 9]'),
  23       ('a thin dog sat on a mat and ate a thin rat', '[10, 11, 12]');
  24   '';
  25
  26   makeTestFor =
  27     package:
  28     makeTest {
  29       name = "pgvecto-rs-${package.name}";
  30       meta = with lib.maintainers; {
  31         maintainers = [ diogotcorreia ];
  32       };
  33
  34       nodes.machine =
  35         { ... }:
  36         {
  37           services.postgresql = {
  38             inherit package;
  39             enable = true;
  40             enableJIT = lib.hasInfix "-jit-" package.name;
  41             extraPlugins =
  42               ps: with ps; [
  43                 pgvecto-rs
  44               ];
  45             settings.shared_preload_libraries = "vectors";
  46           };
  47         };
  48
  49       testScript =
  50         { nodes, ... }:
  51         let
  52           inherit (nodes.machine.services.postgresql.package.pkgs) pgvecto-rs;
  53         in
  54         ''
  55           def check_count(statement, lines):
  56               return 'test $(sudo -u postgres psql postgres -tAc "{}"|wc -l) -eq {}'.format(
  57                   statement, lines
  58               )
  59
  60
  61           machine.start()
  62           machine.wait_for_unit("postgresql")
  63
  64           with subtest("Postgresql with extension vectors is available just after unit start"):
  65               machine.succeed(check_count("SELECT * FROM pg_available_extensions WHERE name = 'vectors' AND default_version = '${pgvecto-rs.version}';", 1))
  66
  67           machine.succeed("sudo -u postgres psql -f ${test-sql}")
  68
  69           machine.succeed(check_count("SELECT content, embedding FROM items WHERE to_tsvector('english', content) @@ 'cat & rat'::tsquery;", 2))
  70
  71           machine.shutdown()
  72         '';
  73     };
  74 in
  75 lib.recurseIntoAttrs (
  76   lib.concatMapAttrs (n: p: { ${n} = makeTestFor p; }) (
  77     lib.filterAttrs (_: p: !p.pkgs.pgvecto-rs.meta.broken) pkgs.postgresqlVersions
  78   )
  79   // {
  80     passthru.override = p: makeTestFor p;
  81   }
  82 )