nixos/tests/pgvecto-rs: init

3 files changed, 84 insertions, 1 deletions

diff --git a/nixos/tests/all-tests.nix b/nixos/tests/all-tests.nix
index 7376cd40b910a..6c188593a97a4 100644
--- a/nixos/tests/all-tests.nix
+++ b/nixos/tests/all-tests.nix
@@ -686,6 +686,7 @@ in {
   pgbouncer = handleTest ./pgbouncer.nix {};
   pgjwt = handleTest ./pgjwt.nix {};
   pgmanage = handleTest ./pgmanage.nix {};
+  pgvecto-rs = handleTest ./pgvecto-rs.nix {};
   phosh = handleTest ./phosh.nix {};
   photoprism = handleTest ./photoprism.nix {};
   php = handleTest ./php {};
diff --git a/nixos/tests/pgvecto-rs.nix b/nixos/tests/pgvecto-rs.nix
new file mode 100644
index 0000000000000..cd871dab6a0f1
--- /dev/null
+++ b/nixos/tests/pgvecto-rs.nix
@@ -0,0 +1,76 @@
+# mostly copied from ./timescaledb.nix which was copied from ./postgresql.nix
+# as it seemed unapproriate to test additional extensions for postgresql there.
+
+{ system ? builtins.currentSystem
+, config ? { }
+, pkgs ? import ../.. { inherit system config; }
+}:
+
+with import ../lib/testing-python.nix { inherit system pkgs; };
+with pkgs.lib;
+
+let
+  postgresql-versions = import ../../pkgs/servers/sql/postgresql pkgs;
+  # Test cases from https://docs.pgvecto.rs/use-cases/hybrid-search.html
+  test-sql = pkgs.writeText "postgresql-test" ''
+    CREATE EXTENSION vectors;
+
+    CREATE TABLE items (
+      id bigserial PRIMARY KEY,
+      content text NOT NULL,
+      embedding vectors.vector(3) NOT NULL -- 3 dimensions
+    );
+
+    INSERT INTO items (content, embedding) VALUES
+      ('a fat cat sat on a mat and ate a fat rat', '[1, 2, 3]'),
+      ('a fat dog sat on a mat and ate a fat rat', '[4, 5, 6]'),
+      ('a thin cat sat on a mat and ate a thin rat', '[7, 8, 9]'),
+      ('a thin dog sat on a mat and ate a thin rat', '[10, 11, 12]');
+  '';
+  make-postgresql-test = postgresql-name: postgresql-package: makeTest {
+    name = postgresql-name;
+    meta = with pkgs.lib.maintainers; {
+      maintainers = [ diogotcorreia ];
+    };
+
+    nodes.machine = { ... }:
+      {
+        services.postgresql = {
+          enable = true;
+          package = postgresql-package;
+          extraPlugins = ps: with ps; [
+            pgvecto-rs
+          ];
+          settings.shared_preload_libraries = "vectors";
+        };
+      };
+
+    testScript = ''
+      def check_count(statement, lines):
+          return 'test $(sudo -u postgres psql postgres -tAc "{}"|wc -l) -eq {}'.format(
+              statement, lines
+          )
+
+
+      machine.start()
+      machine.wait_for_unit("postgresql")
+
+      with subtest("Postgresql with extension vectors is available just after unit start"):
+          machine.succeed(check_count("SELECT * FROM pg_available_extensions WHERE name = 'vectors' AND default_version = '${postgresql-package.pkgs.pgvecto-rs.version}';", 1))
+
+      machine.succeed("sudo -u postgres psql -f ${test-sql}")
+
+      machine.succeed(check_count("SELECT content, embedding FROM items WHERE to_tsvector('english', content) @@ 'cat & rat'::tsquery;", 2))
+
+      machine.shutdown()
+    '';
+
+  };
+  applicablePostgresqlVersions = filterAttrs (_: value: versionAtLeast value.version "12") postgresql-versions;
+in
+mapAttrs'
+  (name: package: {
+    inherit name;
+    value = make-postgresql-test name package;
+  })
+  applicablePostgresqlVersions
diff --git a/pkgs/servers/sql/postgresql/ext/pgvecto-rs/default.nix b/pkgs/servers/sql/postgresql/ext/pgvecto-rs/default.nix
index d348bcf87b1ca..23d4499029ea2 100644
--- a/pkgs/servers/sql/postgresql/ext/pgvecto-rs/default.nix
+++ b/pkgs/servers/sql/postgresql/ext/pgvecto-rs/default.nix
@@ -5,6 +5,7 @@
 , fetchCrate
 , fetchFromGitHub
 , nix-update-script
+, nixosTests
 , openssl
 , pkg-config
 , postgresql
@@ -98,7 +99,12 @@ in
     RUSTC_BOOTSTRAP = 1;
   };
 
-  passthru.updateScript = nix-update-script { };
+  passthru = {
+    updateScript = nix-update-script { };
+    tests = {
+      pgvecto-rs = nixosTests.pgvecto-rs;
+    };
+  };
 
   meta = with lib; {
     # The pgrx 0.11.2 dependency is broken in aarch64-linux: https://github.com/pgcentralfoundation/pgrx/issues/1429