K8SPG-699: Add pgvector support (#991)

egegunes · hors · web-flow · commit 463bf42e0b0c · 2025-01-02T15:16:45.000+02:00
* K8SPG-699: Add pgvector support

* update extension

* fix golangci-lint

---------

Co-authored-by: Viacheslav Sarzhan &lt;slava.sarzhan@percona.com&gt;
diff --git a/build/crd/crunchy/generated/postgres-operator.crunchydata.com_postgresclusters.yaml b/build/crd/crunchy/generated/postgres-operator.crunchydata.com_postgresclusters.yaml
@@ -8352,6 +8352,8 @@ spec:
                     type: boolean
                   pgStatStatements:
                     type: boolean
+                  pgvector:
+                    type: boolean
                 type: object
               image:
                 description: |-
diff --git a/build/crd/percona/generated/pgv2.percona.com_perconapgclusters.yaml b/build/crd/percona/generated/pgv2.percona.com_perconapgclusters.yaml
@@ -8064,6 +8064,8 @@ spec:
                         type: boolean
                       pg_stat_monitor:
                         type: boolean
+                      pgvector:
+                        type: boolean
                     type: object
                   custom:
                     items:
diff --git a/config/crd/bases/pgv2.percona.com_perconapgclusters.yaml b/config/crd/bases/pgv2.percona.com_perconapgclusters.yaml
@@ -8470,6 +8470,8 @@ spec:
                         type: boolean
                       pg_stat_monitor:
                         type: boolean
+                      pgvector:
+                        type: boolean
                     type: object
                   custom:
                     items:
diff --git a/config/crd/bases/postgres-operator.crunchydata.com_postgresclusters.yaml b/config/crd/bases/postgres-operator.crunchydata.com_postgresclusters.yaml
@@ -8312,6 +8312,8 @@ spec:
                     type: boolean
                   pgStatStatements:
                     type: boolean
+                  pgvector:
+                    type: boolean
                 type: object
               image:
                 description: |-
diff --git a/deploy/bundle.yaml b/deploy/bundle.yaml
@@ -8763,6 +8763,8 @@ spec:
                         type: boolean
                       pg_stat_monitor:
                         type: boolean
+                      pgvector:
+                        type: boolean
                     type: object
                   custom:
                     items:
@@ -33985,6 +33987,8 @@ spec:
                     type: boolean
                   pgStatStatements:
                     type: boolean
+                  pgvector:
+                    type: boolean
                 type: object
               image:
                 description: |-
diff --git a/deploy/cr.yaml b/deploy/cr.yaml
@@ -495,6 +495,7 @@ spec:
 #    builtin:
 #      pg_stat_monitor: true
 #      pg_audit: true
+#      pgvector: false
 #    custom:
 #    - name: pg_cron
 #      version: 1.6.1
diff --git a/deploy/crd.yaml b/deploy/crd.yaml
@@ -8763,6 +8763,8 @@ spec:
                         type: boolean
                       pg_stat_monitor:
                         type: boolean
+                      pgvector:
+                        type: boolean
                     type: object
                   custom:
                     items:
@@ -33985,6 +33987,8 @@ spec:
                     type: boolean
                   pgStatStatements:
                     type: boolean
+                  pgvector:
+                    type: boolean
                 type: object
               image:
                 description: |-
diff --git a/deploy/cw-bundle.yaml b/deploy/cw-bundle.yaml
@@ -8763,6 +8763,8 @@ spec:
                         type: boolean
                       pg_stat_monitor:
                         type: boolean
+                      pgvector:
+                        type: boolean
                     type: object
                   custom:
                     items:
@@ -33985,6 +33987,8 @@ spec:
                     type: boolean
                   pgStatStatements:
                     type: boolean
+                  pgvector:
+                    type: boolean
                 type: object
               image:
                 description: |-
diff --git a/internal/controller/postgrescluster/postgres.go b/internal/controller/postgrescluster/postgres.go
@@ -31,6 +31,7 @@ import (
 	"github.com/percona/percona-postgresql-operator/internal/pgaudit"
 	"github.com/percona/percona-postgresql-operator/internal/pgstatmonitor"
 	"github.com/percona/percona-postgresql-operator/internal/pgstatstatements"
+	"github.com/percona/percona-postgresql-operator/internal/pgvector"
 	"github.com/percona/percona-postgresql-operator/internal/postgis"
 	"github.com/percona/percona-postgresql-operator/internal/postgres"
 	pgpassword "github.com/percona/percona-postgresql-operator/internal/postgres/password"
@@ -234,7 +235,7 @@ func (r *Reconciler) reconcilePostgresDatabases(
 	// Calculate a hash of the SQL that should be executed in PostgreSQL.
 
 	// K8SPG-375, K8SPG-577
-	var pgAuditOK, pgStatMonitorOK, pgStatStatementsOK, postgisInstallOK bool
+	var pgAuditOK, pgStatMonitorOK, pgStatStatementsOK, pgvectorOK, postgisInstallOK bool
 	create := func(ctx context.Context, exec postgres.Executor) error {
 		if cluster.Spec.Extensions.PGStatMonitor {
 			if pgStatMonitorOK = pgstatmonitor.EnableInPostgreSQL(ctx, exec) == nil; !pgStatMonitorOK {
@@ -280,6 +281,18 @@ func (r *Reconciler) reconcilePostgresDatabases(
 			}
 		}
 
+		if cluster.Spec.Extensions.PGVector {
+			if pgvectorOK = pgvector.EnableInPostgreSQL(ctx, exec) == nil; !pgvectorOK {
+				r.Recorder.Event(cluster, corev1.EventTypeWarning, "pgvectorDisabled",
+					"Unable to install pgvector")
+			}
+		} else {
+			if pgvectorOK = pgvector.DisableInPostgreSQL(ctx, exec) == nil; !pgvectorOK {
+				r.Recorder.Event(cluster, corev1.EventTypeWarning, "pgvectorEnabled",
+					"Unable to disable pgvector")
+			}
+		}
+
 		// Enabling PostGIS extensions is a one-way operation
 		// e.g., you can take a PostgresCluster and turn it into a PostGISCluster,
 		// but you cannot reverse the process, as that would potentially remove an extension
diff --git a/internal/pgvector/postgres.go b/internal/pgvector/postgres.go
@@ -0,0 +1,48 @@
+package pgvector
+
+import (
+	"context"
+
+	"github.com/percona/percona-postgresql-operator/internal/logging"
+	"github.com/percona/percona-postgresql-operator/internal/postgres"
+)
+
+// EnableInPostgreSQL installs pgvector triggers into every database.
+func EnableInPostgreSQL(ctx context.Context, exec postgres.Executor) error {
+	log := logging.FromContext(ctx)
+
+	stdout, stderr, err := exec.ExecInAllDatabases(ctx,
+		// Quiet the NOTICE from IF EXISTS, and install the pgAudit event triggers.
+		// - https://www.postgresql.org/docs/current/runtime-config-client.html
+		// - https://github.com/pgaudit/pgaudit#settings
+		`SET client_min_messages = WARNING; CREATE EXTENSION IF NOT EXISTS vector; ALTER EXTENSION vector UPDATE;`,
+		map[string]string{
+			"ON_ERROR_STOP": "on", // Abort when any one command fails.
+			"QUIET":         "on", // Do not print successful commands to stdout.
+		})
+
+	log.V(1).Info("enabled pgvector", "stdout", stdout, "stderr", stderr)
+
+	return err
+}
+
+func DisableInPostgreSQL(ctx context.Context, exec postgres.Executor) error {
+	log := logging.FromContext(ctx)
+
+	stdout, stderr, err := exec.ExecInAllDatabases(ctx,
+		// Quiet the NOTICE from IF EXISTS, and install the pgAudit event triggers.
+		// - https://www.postgresql.org/docs/current/runtime-config-client.html
+		// - https://github.com/pgaudit/pgaudit#settings
+		`SET client_min_messages = WARNING; DROP EXTENSION IF EXISTS vector;`,
+		map[string]string{
+			"ON_ERROR_STOP": "on", // Abort when any one command fails.
+			"QUIET":         "on", // Do not print successful commands to stdout.
+		})
+
+	log.V(1).Info("disabled pgvector", "stdout", stdout, "stderr", stderr)
+
+	return err
+}
+
+// PostgreSQLParameters sets the parameters required by pgAudit.
+func PostgreSQLParameters(outParameters *postgres.Parameters) {}
diff --git a/internal/pgvector/postgres_test.go b/internal/pgvector/postgres_test.go
@@ -0,0 +1,61 @@
+package pgvector
+
+import (
+	"context"
+	"errors"
+	"io"
+	"strings"
+	"testing"
+
+	"gotest.tools/v3/assert"
+)
+
+func TestEnableInPostgreSQL(t *testing.T) {
+	expected := errors.New("whoops")
+	exec := func(
+		_ context.Context, stdin io.Reader, stdout, stderr io.Writer, command ...string,
+	) error {
+		assert.Assert(t, stdout != nil, "should capture stdout")
+		assert.Assert(t, stderr != nil, "should capture stderr")
+
+		assert.Assert(t, strings.Contains(strings.Join(command, "\n"),
+			`SELECT datname FROM pg_catalog.pg_database`,
+		), "expected all databases and templates")
+
+		b, err := io.ReadAll(stdin)
+		assert.NilError(t, err)
+		assert.Equal(t, string(b), strings.Trim(`
+SET client_min_messages = WARNING; CREATE EXTENSION IF NOT EXISTS vector; ALTER EXTENSION vector UPDATE;
+		`, "\t\n"))
+
+		return expected
+	}
+
+	ctx := context.Background()
+	assert.Equal(t, expected, EnableInPostgreSQL(ctx, exec))
+}
+
+func TestDisableInPostgreSQL(t *testing.T) {
+	expected := errors.New("whoops")
+	exec := func(
+		_ context.Context, stdin io.Reader, stdout, stderr io.Writer, command ...string,
+	) error {
+		assert.Assert(t, stdout != nil, "should capture stdout")
+		assert.Assert(t, stderr != nil, "should capture stderr")
+
+		assert.Assert(t, strings.Contains(strings.Join(command, "\n"),
+			`SELECT datname FROM pg_catalog.pg_database`,
+		), "expected all databases and templates")
+
+		b, err := io.ReadAll(stdin)
+		assert.NilError(t, err)
+		assert.Equal(t, string(b), strings.Trim(`
+SET client_min_messages = WARNING; DROP EXTENSION IF EXISTS vector;
+		`, "\t\n"))
+
+		return expected
+	}
+
+	ctx := context.Background()
+	assert.Equal(t, expected, DisableInPostgreSQL(ctx, exec))
+}
diff --git a/pkg/apis/pgv2.percona.com/v2/perconapgcluster_types.go b/pkg/apis/pgv2.percona.com/v2/perconapgcluster_types.go
@@ -206,6 +206,7 @@ func (cr *PerconaPGCluster) Default() {
 	cr.Spec.Proxy.PGBouncer.Metadata.Labels[LabelOperatorVersion] = cr.Spec.CRVersion
 
 	t := true
+	f := false
 
 	if cr.Spec.Backups.TrackLatestRestorableTime == nil {
 		cr.Spec.Backups.TrackLatestRestorableTime = &t
@@ -224,6 +225,9 @@ func (cr *PerconaPGCluster) Default() {
 	if cr.Spec.Extensions.BuiltIn.PGAudit == nil {
 		cr.Spec.Extensions.BuiltIn.PGAudit = &t
 	}
+	if cr.Spec.Extensions.BuiltIn.PGVector == nil {
+		cr.Spec.Extensions.BuiltIn.PGVector = &f
+	}
 
 	if cr.CompareVersion("2.6.0") >= 0 && cr.Spec.AutoCreateUserSchema == nil {
 		cr.Spec.AutoCreateUserSchema = &t
@@ -344,6 +348,7 @@ func (cr *PerconaPGCluster) ToCrunchy(ctx context.Context, postgresCluster *crun
 
 	postgresCluster.Spec.Extensions.PGStatMonitor = *cr.Spec.Extensions.BuiltIn.PGStatMonitor
 	postgresCluster.Spec.Extensions.PGAudit = *cr.Spec.Extensions.BuiltIn.PGAudit
+	postgresCluster.Spec.Extensions.PGVector = *cr.Spec.Extensions.BuiltIn.PGVector
 
 	return postgresCluster, nil
 }
@@ -567,6 +572,7 @@ type CustomExtensionsStorageSpec struct {
 type BuiltInExtensionsSpec struct {
 	PGStatMonitor *bool `json:"pg_stat_monitor,omitempty"`
 	PGAudit       *bool `json:"pg_audit,omitempty"`
+	PGVector      *bool `json:"pgvector,omitempty"`
 }
 
 type ExtensionsSpec struct {
diff --git a/pkg/apis/pgv2.percona.com/v2/zz_generated.deepcopy.go b/pkg/apis/pgv2.percona.com/v2/zz_generated.deepcopy.go
diff --git a/pkg/apis/postgres-operator.crunchydata.com/v1beta1/postgrescluster_types.go b/pkg/apis/postgres-operator.crunchydata.com/v1beta1/postgrescluster_types.go
@@ -186,6 +186,7 @@ type ExtensionsSpec struct {
 	PGStatMonitor    bool `json:"pgStatMonitor,omitempty"`
 	PGAudit          bool `json:"pgAudit,omitempty"`
 	PGStatStatements bool `json:"pgStatStatements,omitempty"`
+	PGVector         bool `json:"pgvector,omitempty"`
 }
 
 // DataSource defines data sources for a new PostgresCluster.
