test

Author: ljmnoonan

Dockerfile

@@ -20,17 +20,25 @@ ENV CERTS="{}" \
     HBA_EXTRA_RULES=""
 RUN apk add --no-cache python3 py3-netifaces \
  && if [ "${PG_MAJOR:-0}" -ge 12 ]; then \
-        apk add --no-cache --virtual .pgvector-build build-base linux-headers ca-certificates \
-        && for v in 23 22 21 20 19; do \
-             apk add --no-cache --virtual .pgvector-build clang${v} llvm${v} && break; \
-           done \
+        PG_CLANG="$(pg_config --configure | tr "'" "\n" | sed -n 's/^CLANG=clang-//p')" \
+        && test -n "${PG_CLANG}" \
+        && PG_LLVM="$(pg_config --configure | tr "'" "\n" | sed -n 's#^LLVM_CONFIG=/usr/lib/llvm\([0-9][0-9]*\)/bin/llvm-config#\1#p')" \
+        && test -n "${PG_LLVM}" \
+        && test "${PG_CLANG}" = "${PG_LLVM}" \
+        && apk add --no-cache --virtual .pgvector-build \
+             build-base \
+             linux-headers \
+             ca-certificates \
+             "clang${PG_CLANG}" \
+             "llvm${PG_LLVM}" \
         && wget -qO- "https://github.com/pgvector/pgvector/archive/refs/tags/v${PGVECTOR_VERSION}.tar.gz" \
-          | tar -xz -C /tmp; \
-        cd "/tmp/pgvector-${PGVECTOR_VERSION}" \
-          && make PG_CONFIG=/usr/local/bin/pg_config \
-          && make install PG_CONFIG=/usr/local/bin/pg_config; \
-        cd / && rm -rf "/tmp/pgvector-${PGVECTOR_VERSION}"; \
-        apk del .pgvector-build; \
+          | tar -xz -C /tmp \
+        && cd "/tmp/pgvector-${PGVECTOR_VERSION}" \
+        && make PG_CONFIG=/usr/local/bin/pg_config \
+        && make install PG_CONFIG=/usr/local/bin/pg_config \
+        && cd / \
+        && rm -rf "/tmp/pgvector-${PGVECTOR_VERSION}" \
+        && apk del .pgvector-build; \
     fi \
  && mkdir -p /etc/postgres \
  && chmod a=rwx /etc/postgres

tests/test.py

@@ -58,30 +58,31 @@ def _generate_certs(self):
         """Generate certificates for testing the image."""
         certgen("example.localdomain", "test_user")
 
+    def _exec_local_psql(self, command, dbname="test_db", username="test_user"):
+        """Execute a SQL command inside the Postgres container."""
+        return docker(
+            "container",
+            "exec",
+            self.postgres_container,
+            "psql",
+            "--command",
+            command,
+            "--dbname",
+            dbname,
+            "--no-align",
+            "--tuples-only",
+            "--username",
+            username,
+        )
+
     def _check_local_connection(self):
         """Check that local connection works fine."""
         # The 1st test could fail while postgres boots
         for attempt in range(10):
             try:
                 time.sleep(15)
                 # Test local connections via unix socket work
-                self.assertEqual(
-                    "1\n",
-                    docker(
-                        "container",
-                        "exec",
-                        self.postgres_container,
-                        "psql",
-                        "--command",
-                        "SELECT 1",
-                        "--dbname",
-                        "test_db",
-                        "--no-align",
-                        "--tuples-only",
-                        "--username",
-                        "test_user",
-                    ),
-                )
+                self.assertEqual("1\n", self._exec_local_psql("SELECT 1"))
             except (AssertionError, ProcessExecutionError):
                 if attempt < 9:
                     print("Failure number {}. Retrying...".format(attempt))
@@ -313,6 +314,45 @@ def test_certs_falsy_lan(self):
         with self.assertRaises(ProcessExecutionError):
             self._check_password_auth("example.localdomain")
 
+    def test_pgvector_extension_available(self):
+        """Supported PostgreSQL images expose a working vector extension."""
+        self.postgres_container = docker(
+            "container",
+            "run",
+            *self._platform_args(),
+            "-d",
+            "--network",
+            "lan",
+            "-e",
+            "POSTGRES_DB=test_db",
+            "-e",
+            "POSTGRES_PASSWORD=test_password",
+            "-e",
+            "POSTGRES_USER=test_user",
+            CONF_EXTRA,
+            self.image,
+        ).strip()
+        self._check_local_connection()
+
+        server_version_num = int(
+            self._exec_local_psql("SHOW server_version_num").strip()
+        )
+        if server_version_num < 120000:
+            self.skipTest("pgvector is only built for PostgreSQL 12+")
+
+        vector_version = self._exec_local_psql(
+            "SELECT default_version FROM pg_available_extensions WHERE name = 'vector'"
+        ).strip()
+        self.assertTrue(vector_version)
+
+        self._exec_local_psql("CREATE EXTENSION vector")
+        self.assertEqual(
+            "1\n",
+            self._exec_local_psql(
+                "SELECT ('[1,2,3]'::vector <-> '[1,2,4]'::vector)::int"
+            ),
+        )
+
     def test_hba_extra_rules_added(self):
         """Test that HBA_EXTRA_RULES lines are added to pg_hba.conf."""
         if "9.6" in self.image: