From 092313a9363dcc2d603f46c94740e880522d6a04 Mon Sep 17 00:00:00 2001
From: Andrew Farries <andyrb@gmail.com>
Date: Wed, 31 Jul 2024 11:26:01 +0100
Subject: [PATCH] Add `roll.WithSearchPath` option (#380)

Add a `roll.WithSearchPath` option to allow widening of the search path
during migration execution.

By default, the session that executes a migration has its search path
restricted to the schema in which the migration is applied. By using
`WithSearchPath` the search path can include arbitrarily many other
schema.
---
 pkg/roll/execute_test.go | 36 ++++++++++++++++++++++++++++++++++++
 pkg/roll/options.go      | 12 ++++++++++++
 pkg/roll/roll.go         |  4 +++-
 3 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/pkg/roll/execute_test.go b/pkg/roll/execute_test.go
index d35b60b2..f562b8ec 100644
--- a/pkg/roll/execute_test.go
+++ b/pkg/roll/execute_test.go
@@ -759,6 +759,42 @@ func TestSQLTransformerOptionIsUsedWhenCreatingTriggers(t *testing.T) {
 	})
 }
 
+func TestWithSearchPathOptionIsRespected(t *testing.T) {
+	t.Parallel()
+
+	opts := []roll.Option{roll.WithSearchPath("public")}
+
+	testutils.WithMigratorInSchemaAndConnectionToContainerWithOptions(t, "foo", opts, func(mig *roll.Roll, db *sql.DB) {
+		ctx := context.Background()
+
+		// Create a function in the public schema
+		_, err := db.ExecContext(ctx, `CREATE OR REPLACE FUNCTION say_hello()
+      RETURNS TEXT AS $$
+        SELECT 'hello world';
+      $$ LANGUAGE sql;
+    `)
+		require.NoError(t, err)
+
+		// Apply a migration in the foo schema that references the function in the public schema
+		err = mig.Start(ctx, &migrations.Migration{
+			Name: "01_raw_sql",
+			Operations: migrations.Operations{
+				&migrations.OpRawSQL{
+					Up: "SELECT say_hello()",
+				},
+			},
+		})
+		require.NoError(t, err)
+
+		// Complete the migration
+		err = mig.Complete(ctx)
+		require.NoError(t, err)
+
+		// No assertions required as the migration would have failed if the
+		// function reference was not found
+	})
+}
+
 func createTableOp(tableName string) *migrations.OpCreateTable {
 	return &migrations.OpCreateTable{
 		Name: tableName,
diff --git a/pkg/roll/options.go b/pkg/roll/options.go
index 617ae1b1..96212f7c 100644
--- a/pkg/roll/options.go
+++ b/pkg/roll/options.go
@@ -22,6 +22,9 @@ type options struct {
 	// disable creation of version schema for raw SQL migrations
 	noVersionSchemaForRawSQL bool
 
+	// additional entries to add to the search_path during migration execution
+	searchPath []string
+
 	migrationHooks MigrationHooks
 }
 
@@ -86,3 +89,12 @@ func WithSQLTransformer(transformer migrations.SQLTransformer) Option {
 		o.sqlTransformer = transformer
 	}
 }
+
+// WithSearchPath sets the search_path to use during migration execution. The
+// schema in which the migration is run is always included in the search path,
+// regardless of this setting.
+func WithSearchPath(schemas ...string) Option {
+	return func(o *options) {
+		o.searchPath = schemas
+	}
+}
diff --git a/pkg/roll/roll.go b/pkg/roll/roll.go
index 4b86ff34..68e71174 100644
--- a/pkg/roll/roll.go
+++ b/pkg/roll/roll.go
@@ -6,6 +6,7 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"strings"
 
 	"github.com/lib/pq"
 
@@ -78,7 +79,8 @@ func setupConn(ctx context.Context, pgURL, schema string, options options) (*sql
 		dsn = pgURL
 	}
 
-	dsn += " search_path=" + schema
+	searchPath := append([]string{schema}, options.searchPath...)
+	dsn += " search_path=" + strings.Join(searchPath, ",")
 
 	conn, err := sql.Open("postgres", dsn)
 	if err != nil {