Merge pull request #152 from mindsdb/add/pgvector

Add pgvector as default KB store

Author: ZoranPandovski

Dockerfile.pgvector

@@ -0,0 +1,5 @@
+FROM pgvector/pgvector:pg17
+
+COPY init-dbs.sh /docker-entrypoint-initdb.d/init-dbs.sh
+RUN chmod +x /docker-entrypoint-initdb.d/init-dbs.sh
+

Makefile

@@ -1,11 +1,15 @@
 IMAGE?=mindsdb/mindsdb-docker-extension
+PGVECTOR_IMAGE?=mindsdb/mindsdb-docker-extension-pgvector
 
 BUILDER=buildx-multi-arch
 
 INFO_COLOR = \033[0;36m
 NO_COLOR   = \033[m
 
-build-extension: ## Build service image to be deployed as a desktop extension
+build-pgvector: ## Build the pgvector image with init scripts
+	docker build --tag=$(PGVECTOR_IMAGE):latest -f Dockerfile.pgvector .
+
+build-extension: build-pgvector ## Build service image to be deployed as a desktop extension
 	docker build --tag=$(IMAGE):$(TAG) .
 
 install-extension: build-extension ## Install the extension
@@ -20,7 +24,10 @@ prepare-buildx: ## Create buildx builder for multi-arch build, if not exists
 validate-extension: ## Validate extension is ready to be published
 	docker extension validate mindsdb/mindsdb-docker-extension:1.0.0
 
-push-extension: prepare-buildx ## Build & Upload extension image to hub. Do not push if tag already exists: make push-extension tag=0.1
+push-pgvector: prepare-buildx ## Build & Upload pgvector image to hub
+	docker buildx build --push --builder=$(BUILDER) --platform=linux/amd64,linux/arm64 --tag=$(PGVECTOR_IMAGE):latest -f Dockerfile.pgvector .
+
+push-extension: prepare-buildx push-pgvector ## Build & Upload extension image to hub. Do not push if tag already exists: make push-extension tag=0.1
 	docker pull $(IMAGE):$(TAG) && echo "Failure: Tag already exists" || docker buildx build --push --builder=$(BUILDER) --platform=linux/amd64,linux/arm64 --build-arg TAG=$(TAG) --tag=$(IMAGE):$(TAG) .
 
 help: ## Show this help

docker-compose.yaml

@@ -1,18 +1,37 @@
 services:
   mindsdb:
     image: mindsdb/mindsdb:v25.14.0
+    depends_on:
+      pgvector:
+        condition: service_healthy
     ports:
       - "47334:47334"
       - "47335:47335"
-      - "47336:47336"
-      - "47337:47337"
-      - "47338:47338"
     volumes:
       - data:/opt/mindsdb/data
       - config:/opt/mindsdb/config
     container_name: mindsdb_service
     pull_policy: always
+    environment:
+      MINDSDB_DB_CON: "postgresql://mindsdb:mindsdb@pgvector/mindsdb"
+      KB_PGVECTOR_URL: "postgresql://mindsdb:mindsdb@pgvector/kb"
+
+  pgvector:
+    image: mindsdb/mindsdb-docker-extension-pgvector:latest
+    container_name: pgvector_service
+    volumes:
+      - data_pg:/var/lib/postgresql/data
+    environment:
+      POSTGRES_USER: "postgres"
+      POSTGRES_PASSWORD: "postgres"
+      POSTGRES_DB: "postgres"
+    healthcheck:
+      test: ["CMD-SHELL", "pg_isready -U postgres"]
+      interval: 5s
+      timeout: 5s
+      retries: 10
 
 volumes:
   data:
   config:
+  data_pg:

init-dbs.sh

@@ -0,0 +1,74 @@
+#!/bin/bash
+set -e
+
+echo "Initializing databases..."
+
+# Check required environment variables
+if [[ -z "${POSTGRES_USER:-}" || -z "${POSTGRES_DB:-}" ]]; then
+  echo "Error: POSTGRES_USER and POSTGRES_DB environment variables must be set." >&2
+  exit 1
+fi
+
+# Log message for clarity
+echo "Creating additional databases..."
+
+# Function to create a database
+create_database() {
+  local db_name=$1
+  echo "Creating database: $db_name"
+  psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB" <<-EOSQL
+    SELECT 'CREATE DATABASE $db_name'
+    WHERE NOT EXISTS (SELECT FROM pg_database WHERE datname = '$db_name')\gexec
+EOSQL
+  psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB" <<-EOSQL
+    SELECT 'CREATE DATABASE kb'
+    WHERE NOT EXISTS (SELECT FROM pg_database WHERE datname = 'kb')\gexec
+EOSQL
+}
+
+# Function to create a user with permissions
+create_user() {
+  local username=$1
+  local userpassword=$2
+  local db_name=$3
+  echo "Creating user: $username with access to database: $db_name"
+  psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$db_name" <<-EOSQL
+    DO \$\$
+    BEGIN
+      IF NOT EXISTS (SELECT FROM pg_roles WHERE rolname = '$username') THEN
+        CREATE USER $username WITH PASSWORD '$userpassword';
+      ELSE
+        RAISE NOTICE 'User $username already exists. Skipping.';
+      END IF;
+    END
+    \$\$;
+    GRANT ALL PRIVILEGES ON DATABASE $db_name TO $username;
+    GRANT ALL PRIVILEGES ON SCHEMA public TO $username;
+    GRANT ALL ON SCHEMA public TO $username;
+EOSQL
+
+  # Do the same for the kb db
+  psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "kb" <<-EOSQL
+    DO \$\$
+    BEGIN
+      IF NOT EXISTS (SELECT FROM pg_roles WHERE rolname = '$username') THEN
+        CREATE USER $username WITH PASSWORD '$userpassword';
+      ELSE
+        RAISE NOTICE 'User $username already exists. Skipping.';
+      END IF;
+    END
+    \$\$;
+    GRANT ALL PRIVILEGES ON DATABASE kb TO $username;
+    GRANT ALL PRIVILEGES ON SCHEMA public TO $username;
+    GRANT ALL ON SCHEMA public TO $username;
+    CREATE EXTENSION IF NOT EXISTS vector;
+EOSQL
+}
+
+# Create databases
+create_database "mindsdb"
+
+# Create user with permissions in a database
+create_user "mindsdb" "mindsdb" "mindsdb"
+
+echo "Database and user creation completed successfully."