Merge: postgresqlPackages.pgvectorscale: init at 0.7.0, buildPgrxExtension: support cargo workspaces (#393689)

pkgs/development/tools/rust/cargo-pgrx/buildPgrxExtension.nix

@@ -53,6 +53,7 @@
   buildType ? "release",
   buildFeatures ? [ ],
   cargoBuildFlags ? [ ],
+  cargoPgrxFlags ? [ ],
   postgresql,
   # cargo-pgrx calls rustfmt on generated bindings, this is not strictly necessary, so we avoid the
   # dependency here. Set to false and provide rustfmt in nativeBuildInputs, if you need it, e.g.
@@ -115,6 +116,8 @@ let
     "usePgTestCheckFeature"
   ];
 
+  cargoPgrxFlags' = lib.escapeShellArgs cargoPgrxFlags;
+
   # so we don't accidentally `(rustPlatform.buildRustPackage argsForBuildRustPackage) // { ... }` because
   # we forgot parentheses
   finalArgs = argsForBuildRustPackage // {
@@ -140,6 +143,7 @@ let
       PGRX_BUILD_FLAGS="--frozen -j $NIX_BUILD_CORES ${builtins.concatStringsSep " " cargoBuildFlags}" \
       ${lib.optionalString stdenv.hostPlatform.isDarwin ''RUSTFLAGS="''${RUSTFLAGS:+''${RUSTFLAGS} }-Clink-args=-Wl,-undefined,dynamic_lookup"''} \
       cargo pgrx package \
+        ${cargoPgrxFlags'} \
         --pg-config ${lib.getDev postgresql}/bin/pg_config \
         ${maybeDebugFlag} \
         --features "${builtins.concatStringsSep " " buildFeatures}" \
@@ -159,7 +163,7 @@ let
 
       ${maybeEnterBuildAndTestSubdir}
 
-      cargo-pgrx pgrx stop all
+      cargo-pgrx pgrx stop all ${cargoPgrxFlags'}
 
       mv $out/${postgresql}/* $out
       rm -rf $out/nix

pkgs/servers/sql/postgresql/ext/pgvectorscale/package.nix

@@ -0,0 +1,89 @@
+{
+  buildPgrxExtension,
+  postgresql,
+  fetchFromGitHub,
+  lib,
+  postgresqlTestExtension,
+}:
+
+let
+  finalPackage = buildPgrxExtension rec {
+    pname = "pgvectorscale";
+    version = "0.7.0";
+
+    src = fetchFromGitHub {
+      owner = "timescale";
+      repo = "pgvectorscale";
+      tag = version;
+      hash = "sha256-dy481k2SvyYXwwcsyLZSl3XlhSk9C5+4LfEfciB1DK4=";
+    };
+
+    doCheck = false;
+
+    useFetchCargoVendor = true;
+    cargoHash = "sha256-CeRyDn9VhxfjWFJ1/Z/XvOUQOSnDoHHZAqgfYTeKU0o=";
+    cargoPatches = [
+      ./add-Cargo.lock.patch
+    ];
+
+    cargoPgrxFlags = [
+      "-p"
+      "vectorscale"
+    ];
+
+    inherit postgresql;
+
+    passthru.tests.extension = postgresqlTestExtension {
+      inherit finalPackage;
+      withPackages = [ "pgvector" ];
+      sql =
+        let
+          genCheck =
+            id: compare: expected:
+            let
+              vecStr = "[${lib.concatMapStringsSep "," toString compare}]";
+            in
+            ''
+              ASSERT (
+                SELECT id
+                FROM document_embedding
+                WHERE ${toString expected} = (embedding <-> '${vecStr}')
+              ) = ${toString id},
+              'Expected vector of row with ID=${toString id} to have a euclidean distance from ${vecStr} of ${toString expected}';
+            '';
+        in
+        ''
+          CREATE EXTENSION vectorscale CASCADE;
+          CREATE TABLE IF NOT EXISTS document_embedding  (
+              id BIGINT PRIMARY KEY GENERATED BY DEFAULT AS IDENTITY,
+              embedding VECTOR(3)
+          );
+
+          INSERT INTO document_embedding (id, embedding) VALUES
+          (1, '[1,2,4]'),
+          (2, '[1,2,5]');
+
+          CREATE INDEX document_embedding_idx ON document_embedding
+          USING diskann (embedding vector_cosine_ops);
+
+          DO $$
+            BEGIN
+            ${genCheck 1 [ 1 2 3 ] 1}
+            ${genCheck 2 [ 1 2 3 ] 2}
+            END;
+          $$
+          LANGUAGE PLPGSQL;
+        '';
+    };
+
+    meta = {
+      homepage = "https://github.com/timescale/pgvectorscale";
+      maintainers = lib.teams.flyingcircus.members;
+      description = "Complement to pgvector for high performance, cost efficient vector search on large workloads";
+      license = lib.licenses.postgresql;
+      platforms = postgresql.meta.platforms;
+      changelog = "https://github.com/timescale/pgvectorscale/releases/tag/${version}";
+    };
+  };
+in
+finalPackage