Skip to content

feat: multiple versions for the vector extension #1662

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: develop
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .github/workflows/nix-build.yml
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ jobs:
run: |
cat << 'EOF' | sudo tee /etc/nix/upload-to-cache.sh > /dev/null
#!/usr/bin/env bash
set -eouf
set -euf
export IFS=' '
/nix/var/nix/profiles/default/bin/nix copy --to 's3://nix-postgres-artifacts?secret-key=/etc/nix/nix-secret-key' $OUT_PATHS
EOF
Expand Down
2 changes: 2 additions & 0 deletions flake.nix
Original file line number Diff line number Diff line change
Expand Up @@ -1384,6 +1384,8 @@
devShell = devShells.default;
} // pkgs.lib.optionalAttrs (system == "aarch64-linux") {
inherit (basePackages) postgresql_15_debug postgresql_15_src postgresql_orioledb-17_debug postgresql_orioledb-17_src postgresql_17_debug postgresql_17_src;
} // pkgs.lib.optionalAttrs (system == "x86_64-linux") {
pgvector = import ./nix/ext/tests/pgvector.nix { inherit self; inherit pkgs; };
};

# Apps is a list of names of things that can be executed with 'nix run';
Expand Down
107 changes: 84 additions & 23 deletions nix/ext/pgvector.nix
Original file line number Diff line number Diff line change
@@ -1,31 +1,92 @@
{ lib, stdenv, fetchFromGitHub, postgresql }:
{
pkgs,
lib,
stdenv,
fetchFromGitHub,
postgresql,
}:
let
pname = "vector";

stdenv.mkDerivation rec {
pname = "pgvector";
version = "0.8.0";
# Load version configuration from external file
allVersions = (builtins.fromJSON (builtins.readFile ./versions.json)).${pname};

buildInputs = [ postgresql ];
# Filter versions compatible with current PostgreSQL version
supportedVersions = lib.filterAttrs (
_: value: builtins.elem (lib.versions.major postgresql.version) value.postgresql
) allVersions;

src = fetchFromGitHub {
owner = "pgvector";
repo = pname;
rev = "refs/tags/v${version}";
hash = "sha256-JsZV+I4eRMypXTjGmjCtMBXDVpqTIPHQa28ogXncE/Q=";
};
# Derived version information
versions = lib.naturalSort (lib.attrNames supportedVersions);
latestVersion = lib.last versions;
numberOfVersions = builtins.length versions;
packages = builtins.attrValues (
lib.mapAttrs (name: value: build name value.hash) supportedVersions
);

# Build function for individual versions
build =
version: hash:
stdenv.mkDerivation rec {
inherit pname version;

buildInputs = [ postgresql ];

src = fetchFromGitHub {
owner = "pgvector";
repo = "pgvector";
rev = "refs/tags/v${version}";
inherit hash;
};

installPhase = ''
mkdir -p $out/{lib,share/postgresql/extension}

# Install shared library with version suffix
mv ${pname}${postgresql.dlSuffix} $out/lib/${pname}-${version}${postgresql.dlSuffix}

# Create version-specific control file
sed -e "/^default_version =/d" \
-e "s|^module_pathname = .*|module_pathname = '\$libdir/${pname}'|" \
${pname}.control > $out/share/postgresql/extension/${pname}--${version}.control

# Copy SQL file to install the specific version
cp sql/${pname}.sql $out/share/postgresql/extension/${pname}--${version}.sql

# For the latest version, copy sql upgrade script, default control file and symlink
if [[ "${version}" == "${latestVersion}" ]]; then
cp sql/*.sql $out/share/postgresql/extension
{
echo "default_version = '${latestVersion}'"
cat $out/share/postgresql/extension/${pname}--${latestVersion}.control
} > $out/share/postgresql/extension/${pname}.control
ln -sfn ${pname}-${latestVersion}${postgresql.dlSuffix} $out/lib/${pname}${postgresql.dlSuffix}
fi

installPhase = ''
mkdir -p $out/{lib,share/postgresql/extension}
runHook postInstall
'';

cp *${postgresql.dlSuffix} $out/lib
cp sql/*.sql $out/share/postgresql/extension
cp *.control $out/share/postgresql/extension
'';
meta = with lib; {
description = "Open-source vector similarity search for Postgres";
homepage = "https://github.com/${src.owner}/${src.repo}";
maintainers = with maintainers; [ olirice ];
platforms = postgresql.meta.platforms;
license = licenses.postgresql;
};
};
in
pkgs.buildEnv {
name = pname;
paths = packages;
pathsToLink = [
"/lib"
"/share/postgresql/extension"
];

meta = with lib; {
description = "Open-source vector similarity search for Postgres";
homepage = "https://github.com/${src.owner}/${src.repo}";
maintainers = with maintainers; [ olirice ];
platforms = postgresql.meta.platforms;
license = licenses.postgresql;
passthru = {
inherit versions numberOfVersions;
pname = "${pname}-all";
version =
"multi-" + lib.concatStringsSep "-" (map (v: lib.replaceStrings [ "." ] [ "-" ] v) versions);
};
}
158 changes: 158 additions & 0 deletions nix/ext/tests/pgvector.nix
Original file line number Diff line number Diff line change
@@ -0,0 +1,158 @@
{ self, pkgs }:
let
pname = "vector";
inherit (pkgs) lib;
installedExtension =
postgresMajorVersion: self.packages.${pkgs.system}."psql_${postgresMajorVersion}/exts/${pname}-all";
versions = postgresqlMajorVersion: (installedExtension postgresqlMajorVersion).versions;
postgresqlWithExtension =
postgresql:
let
majorVersion = lib.versions.major postgresql.version;
pkg = pkgs.buildEnv {
name = "postgresql-${majorVersion}-${pname}";
paths = [
postgresql
postgresql.lib
(installedExtension majorVersion)
];
passthru = {
inherit (postgresql) version psqlSchema;
lib = pkg;
withPackages = _: pkg;
};
nativeBuildInputs = [ pkgs.makeWrapper ];
pathsToLink = [
"/"
"/bin"
"/lib"
];
postBuild = ''
wrapProgram $out/bin/postgres --set NIX_PGLIBDIR $out/lib
wrapProgram $out/bin/pg_ctl --set NIX_PGLIBDIR $out/lib
wrapProgram $out/bin/pg_upgrade --set NIX_PGLIBDIR $out/lib
'';
};
in
pkg;
pgsodiumGetKey = lib.getExe (
pkgs.writeShellScriptBin "pgsodium-getkey" ''
echo 0000000000000000000000000000000000000000000000000000000000000000
''
);
in
self.inputs.nixpkgs.lib.nixos.runTest {
name = pname;
hostPkgs = pkgs;
nodes.server =
{ config, ... }:
{
virtualisation = {
forwardPorts = [
{
from = "host";
host.port = 13022;
guest.port = 22;
}
];
};

services.postgresql = {
enable = true;
package = postgresqlWithExtension self.packages.${pkgs.system}.postgresql_15;
settings = {
"shared_preload_libraries" = pname;
"pgsodium.getkey_script" = pgsodiumGetKey;
};
};

specialisation.postgresql17.configuration = {
services.postgresql = {
package = lib.mkForce (postgresqlWithExtension self.packages.${pkgs.system}.postgresql_17);
};

systemd.services.postgresql-migrate = {
serviceConfig = {
Type = "oneshot";
RemainAfterExit = true;
User = "postgres";
Group = "postgres";
StateDirectory = "postgresql";
WorkingDirectory = "${builtins.dirOf config.services.postgresql.dataDir}";
};
script =
let
oldPostgresql = postgresqlWithExtension self.packages.${pkgs.system}.postgresql_15;
newPostgresql = postgresqlWithExtension self.packages.${pkgs.system}.postgresql_17;
oldDataDir = "${builtins.dirOf config.services.postgresql.dataDir}/${oldPostgresql.psqlSchema}";
newDataDir = "${builtins.dirOf config.services.postgresql.dataDir}/${newPostgresql.psqlSchema}";
in
''
if [[ ! -d ${newDataDir} ]]; then
install -d -m 0700 -o postgres -g postgres "${newDataDir}"
${newPostgresql}/bin/initdb -D "${newDataDir}"
echo "shared_preload_libraries = '${pname}'" >> "${newDataDir}/postgresql.conf"
echo "pgsodium.getkey_script = '${pgsodiumGetKey}'" >> "${newDataDir}/postgresql.conf";
${newPostgresql}/bin/pg_upgrade --old-datadir "${oldDataDir}" --new-datadir "${newDataDir}" \
--old-bindir "${oldPostgresql}/bin" --new-bindir "${newPostgresql}/bin"
else
echo "${newDataDir} already exists"
fi
'';
};

systemd.services.postgresql = {
after = [ "postgresql-migrate.service" ];
requires = [ "postgresql-migrate.service" ];
};
};

};
testScript =
{ nodes, ... }:
let
pg17-configuration = "${nodes.server.system.build.toplevel}/specialisation/postgresql17";
in
''
versions = {
"15": [${lib.concatStringsSep ", " (map (s: ''"${s}"'') (versions "15"))}],
"17": [${lib.concatStringsSep ", " (map (s: ''"${s}"'') (versions "17"))}],
}

def run_sql(query):
return server.succeed(f"""sudo -u postgres psql -t -A -F\",\" -c \"{query}\" """).strip()

def check_upgrade_path(pg_version):
with subtest("Check ${pname} upgrade path"):
firstVersion = versions[pg_version][0]
server.succeed("sudo -u postgres psql -c 'DROP EXTENSION IF EXISTS ${pname};'")
run_sql(f"""CREATE EXTENSION ${pname} WITH VERSION '{firstVersion}' CASCADE;""")
installed_version = run_sql(r"""SELECT extversion FROM pg_extension WHERE extname = '${pname}';""")
assert installed_version == firstVersion, f"Expected ${pname} version {firstVersion}, but found {installed_version}"
for version in versions[pg_version][1:]:
run_sql(f"""ALTER EXTENSION ${pname} UPDATE TO '{version}';""")
installed_version = run_sql(r"""SELECT extversion FROM pg_extension WHERE extname = '${pname}';""")
assert installed_version == version, f"Expected ${pname} version {version}, but found {installed_version}"

start_all()

server.wait_for_unit("multi-user.target")
server.wait_for_unit("postgresql.service")

check_upgrade_path("15")

with subtest("Check ${pname} latest extension version"):
server.succeed("sudo -u postgres psql -c 'DROP EXTENSION ${pname};'")
server.succeed("sudo -u postgres psql -c 'CREATE EXTENSION ${pname} CASCADE;'")
installed_extensions=run_sql(r"""SELECT extname, extversion FROM pg_extension;""")
latestVersion = versions["15"][-1]
assert f"${pname},{latestVersion}" in installed_extensions

with subtest("switch to postgresql 17"):
server.succeed(
"${pg17-configuration}/bin/switch-to-configuration test >&2"
)

check_upgrade_path("17")
'';
}
52 changes: 52 additions & 0 deletions nix/ext/versions.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
{
"vector": {
"0.4.0": {
"postgresql": [
"15"
],
"hash": "sha256-bOckX7zvHhgJDDhoAm+VZVIeVIf2hG/3oWZWuTtnZPo="
},
"0.5.1": {
"postgresql": [
"15",
"17"
],
"hash": "sha256-ZNzq+dATZn9LUgeOczsaadr5hwdbt9y/+sAOPIdr77U="
},
"0.6.0": {
"postgresql": [
"15",
"17"
],
"hash": "sha256-hXm+k0BZ9xZP1Tnek14jPoKCPQkA5ovscu9IX2mW7Kc="
},
"0.6.2": {
"postgresql": [
"15",
"17"
],
"hash": "sha256-r+TpFJg6WrMn0L2B7RpmSRvw3XxpHzMRtpFWDCzLvgs="
},
"0.7.0": {
"postgresql": [
"15",
"17"
],
"hash": "sha256-vFn7sNphOYyig6Jl1HILMaC2t9strFQBQ8ywL8Ibx1M="
},
"0.7.4": {
"postgresql": [
"15",
"17"
],
"hash": "sha256-qwPaguQUdDHV8q6GDneLq5MuhVroPizpbqt7f08gKJI="
},
"0.8.0": {
"postgresql": [
"15",
"17"
],
"hash": "sha256-JsZV+I4eRMypXTjGmjCtMBXDVpqTIPHQa28ogXncE/Q="
}
}
}
Loading