Create experimental ability to provision Postgres helper functions

Towards NSL-5018

Author: n-g

library/oss/postgres/prepare/database/functions.sql

@@ -0,0 +1,49 @@
+-- Collection of optional helper functions
+-- To provision these functions add
+--   provision_helper_functions: true
+-- to the database intent.
+
+-- `CREATE TABLE IF NOT EXISTS` and `ALTER TABLE … ADD COLUMN IF NOT EXISTS`
+-- both require exclusive locks with Postgres, even if the table/column already exists.
+-- The functions below provide ensure semantics while only acquiring exclusive locks on mutations.
+
+-- fn_ensure_table is a lock-friendly replacement for `CREATE TABLE IF NOT EXISTS`.
+--
+-- Example usage:
+--
+-- SELECT fn_ensure_table('testtable', $$
+--   UserID TEXT NOT NULL,
+--   PRIMARY KEY(UserID)
+-- $$);
+CREATE OR REPLACE FUNCTION fn_ensure_table(tname TEXT, def TEXT)
+  RETURNS void
+  LANGUAGE plpgsql AS
+$func$
+BEGIN
+  IF NOT EXISTS (
+    SELECT 1 FROM pg_tables
+    WHERE schemaname = 'public' AND tablename = LOWER(tname)
+  ) THEN
+    EXECUTE 'CREATE TABLE IF NOT EXISTS ' || tname || ' (' || def || ');';
+  END IF;
+END
+$func$;
+
+-- fn_ensure_column is a lock-friendly replacement for `ALTER TABLE ... ADD COLUMN IF NOT EXISTS`.
+--
+-- Example usage:
+--
+-- SELECT fn_ensure_column('testtable', 'CreatedAt', 'TIMESTAMP DEFAULT CURRENT_TIMESTAMP');
+CREATE OR REPLACE FUNCTION fn_ensure_column(tname TEXT, cname TEXT, def TEXT)
+  RETURNS void
+  LANGUAGE plpgsql AS
+$func$
+BEGIN
+  IF NOT EXISTS (
+    SELECT 1 FROM information_schema.columns
+    WHERE table_name = LOWER(tname) AND column_name = LOWER(cname)
+  ) THEN
+    EXECUTE 'ALTER TABLE ' || tname || ' ADD COLUMN IF NOT EXISTS ' || cname || ' ' || def;
+  END IF;
+END
+$func$;

library/oss/postgres/prepare/database/main.go

@@ -6,7 +6,9 @@ package main
 
 import (
 	"context"
+	"embed"
 	"fmt"
+	"io/fs"
 	"log"
 	"math/rand"
 	"os"
@@ -27,32 +29,44 @@ const (
 	connIdleTimeout = 15 * time.Minute
 	connTimeout     = 5 * time.Minute
 
-	caCertPath = "/tmp/ca.pem"
+	caCertPath          = "/tmp/ca.pem"
+	helperFunctionsPath = "functions.sql"
+)
+
+var (
+	//go:embed *.sql
+	data embed.FS
 )
 
 func main() {
 	ctx, p := provider.MustPrepare[*postgres.DatabaseIntent]()
 
+	if err := run(ctx, p); err != nil {
+		log.Fatal(err)
+	}
+}
+
+func run(ctx context.Context, p *provider.Provider[*postgres.DatabaseIntent]) error {
 	cluster := &postgresclass.ClusterInstance{}
 	if err := p.Resources.Unmarshal(fmt.Sprintf("%s:cluster", providerPkg), cluster); err != nil {
-		log.Fatalf("unable to read required resource \"cluster\": %v", err)
+		return fmt.Errorf("unable to read required resource \"cluster\": %w", err)
 	}
 
 	// TODO inject file as secret ref and propagate secret ref to server, too.
 	if cluster.CaCert != "" {
 		if err := os.WriteFile(caCertPath, []byte(cluster.CaCert), 0644); err != nil {
-			log.Fatalf("failed to write %q: %v", caCertPath, err)
+			return fmt.Errorf("failed to write %q: %w", caCertPath, err)
 		}
 
 		if err := os.Setenv("PGSSLROOTCERT", caCertPath); err != nil {
-			log.Fatalf("failed to set PGSSLROOTCERT: %v", err)
+			return fmt.Errorf("failed to set PGSSLROOTCERT: %w", err)
 		}
 
 	}
 
 	exists, err := ensureDatabase(ctx, cluster, p.Intent.Name)
 	if err != nil {
-		log.Fatalf("unable to create database %q: %v", p.Intent.Name, err)
+		return fmt.Errorf("unable to create database %q: %w", p.Intent.Name, err)
 	}
 
 	instance := &postgresclass.DatabaseInstance{
@@ -72,29 +86,45 @@ func main() {
 			MaxConnIdleTime: connIdleTimeout,
 		})
 		if err != nil {
-			log.Fatalf("unable to open connection: %v", err)
+			return fmt.Errorf("unable to open connection: %w", err)
 		}
 		defer func() {
 			if err := db.Close(); err != nil {
 				log.Printf("unable to close database connection: %v", err)
 			}
 		}()
 
+		if p.Intent.ProvisionHelperFunctions {
+			content, err := fs.ReadFile(data, helperFunctionsPath)
+			if err != nil {
+				return fmt.Errorf("failed to read %s: %w", helperFunctionsPath, err)
+			}
+
+			if err := applyWithRetry(ctx, db, string(content)); err != nil {
+				return fmt.Errorf("unable to apply helper functions: %w", err)
+			}
+		}
+
 		for _, schema := range p.Intent.Schema {
-			if err := backoff.Retry(func() error {
-				_, err := db.Exec(ctx, string(schema.Contents))
-				return err
-			}, backOff{
-				interval: 100 * time.Millisecond,
-				deadline: time.Now().Add(15 * time.Second),
-				jitter:   100 * time.Millisecond,
-			}); err != nil {
-				log.Fatalf("unable to apply schema %q: %v", schema.Path, err)
+			if err := applyWithRetry(ctx, db, string(schema.Contents)); err != nil {
+				return fmt.Errorf("unable to apply schema %q: %w", schema.Path, err)
 			}
 		}
 	}
 
 	p.EmitResult(instance)
+	return nil
+}
+
+func applyWithRetry(ctx context.Context, db *universepg.DB, sql string) error {
+	return backoff.Retry(func() error {
+		_, err := db.Exec(ctx, sql)
+		return err
+	}, backOff{
+		interval: 100 * time.Millisecond,
+		deadline: time.Now().Add(15 * time.Second),
+		jitter:   100 * time.Millisecond,
+	})
 }
 
 func ensureDatabase(ctx context.Context, cluster *postgresclass.ClusterInstance, name string) (bool, error) {

library/oss/postgres/types.pb.go

@@ -20,7 +20,8 @@ message ClusterIntent {
 
 message DatabaseIntent {
     // The database name is applied as is (e.g. it is case-sensitive).
-    string   name                                  = 1;
-    repeated foundation.schema.FileContents schema = 2;
-    bool skip_schema_initialization_if_exists = 3;
+    string   name                                                                = 1;
+    repeated foundation.schema.FileContents schema                               = 2;
+    bool                                    skip_schema_initialization_if_exists = 3;
+    bool                                    provision_helper_functions           = 4;
 }