Repository: ariga/atlas
Pull request #2776 (2024-05-12 16:02:44)

sql/postgres: unset default search_path
When operating on realm-scope we expect all objects to be qualified. However, the default configuration of pg sets the search_path to `$user, public` causes all references to public objects to return from pg_catalog without qualifiers. 


Patch:
diff --git a/sql/postgres/inspect.go b/sql/postgres/inspect.go
index cddd175b..b0d88d80 100644
--- a/sql/postgres/inspect.go
+++ b/sql/postgres/inspect.go
@@ -29,7 +29,12 @@ type inspect struct{ *conn }
 var _ schema.Inspector = (*inspect)(nil)
 
 // InspectRealm returns schema descriptions of all resources in the given realm.
-func (i *inspect) InspectRealm(ctx context.Context, opts *schema.InspectRealmOption) (*schema.Realm, error) {
+func (i *inspect) InspectRealm(ctx context.Context, opts *schema.InspectRealmOption) (_ *schema.Realm, rerr error) {
+	undo, err := i.noSearchPath(ctx)
+	if err != nil {
+		return nil, err
+	}
+	defer func() { rerr = errors.Join(rerr, undo()) }()
 	schemas, err := i.schemas(ctx, opts)
 	if err != nil {
 		return nil, err
@@ -86,6 +91,34 @@ func (i *inspect) InspectRealm(ctx context.Context, opts *schema.InspectRealmOpt
 	return schema.ExcludeRealm(r, opts.Exclude)
 }
 
+// noSearchPath ensures the session search_path is clean when inspecting realms to ensures all
+// referenced objects in the public schema (or any other default search_path) are returned
+// qualified in the inspection.
+func (i *inspect) noSearchPath(ctx context.Context) (func() error, error) {
+	if i.crdb {
+		// Skip logic for CockroachDB.
+		return func() error { return nil }, nil
+	}
+	rows, err := i.QueryContext(ctx, "SELECT current_setting('search_path'), set_config('search_path', '', false)")
+	if err != nil {
+		return nil, err
+	}
+	var prev sql.NullString
+	if err := sqlx.ScanOne(rows, &prev, &sql.NullString{}); err != nil {
+		return nil, err
+	}
+	return func() error {
+		if sqlx.ValidString(prev) {
+			rows, err := i.QueryContext(ctx, "SELECT set_config('search_path', $1, false)", prev.String)
+			if err != nil {
+				return err
+			}
+			return rows.Close()
+		}
+		return nil
+	}, nil
+}
+
 // InspectSchema returns schema descriptions of the tables in the given schema.
 // If the schema name is empty, the result will be the attached schema.
 func (i *inspect) InspectSchema(ctx context.Context, name string, opts *schema.InspectOptions) (s *schema.Schema, err error) {


Test patch:
diff --git a/internal/integration/postgres_test.go b/internal/integration/postgres_test.go
index ace498e1..e6340e59 100644
--- a/internal/integration/postgres_test.go
+++ b/internal/integration/postgres_test.go
@@ -813,7 +813,7 @@ table "users" {
   schema = schema.public
   column "id" { type = serial }
   column "e"  { type = enum.status }
-  column "ae" { type = sql("status[]") }
+  column "ae" { type = sql("public.status[]") }
 }
 
 schema "other" {}
@@ -865,7 +865,7 @@ table "users" {
   }
   column "ae" {
     null = false
-    type = sql("status[]")
+    type = sql("public.status[]")
   }
 }
 enum "status" {
@@ -1420,9 +1420,9 @@ create table atlas_types_sanity
 				},
 				{
 					Name: "tUserDefined",
-					Type: &schema.ColumnType{Type: &postgres.UserDefinedType{T: "address"}, Raw: "USER-DEFINED", Null: true},
+					Type: &schema.ColumnType{Type: &postgres.UserDefinedType{T: "public.address"}, Raw: "USER-DEFINED", Null: true},
 					Default: &schema.RawExpr{
-						X: "'(ab,cd)'::address",
+						X: "'(ab,cd)'::public.address",
 					},
 				},
 			},
diff --git a/sql/postgres/inspect_test.go b/sql/postgres/inspect_test.go
index 2e0b6541..ebe41182 100644
--- a/sql/postgres/inspect_test.go
+++ b/sql/postgres/inspect_test.go
@@ -543,6 +543,13 @@ func TestDriver_Realm(t *testing.T) {
 	mk.version("130000")
 	drv, err := Open(db)
 	require.NoError(t, err)
+	// Set search_path to ''.
+	mk.ExpectQuery(sqltest.Escape("SELECT current_setting('search_path'), set_config('search_path', '', false)")).
+		WillReturnRows(sqltest.Rows(`
+ current_setting | set_config
+-----------------+------------
+       public    |
+`))
 	mk.ExpectQuery(sqltest.Escape(schemasQuery)).
 		WillReturnRows(sqltest.Rows(`
  schema_name | comment 
@@ -555,6 +562,10 @@ func TestDriver_Realm(t *testing.T) {
 	m.ExpectQuery(sqltest.Escape(fmt.Sprintf(tablesQuery, "$1, $2"))).
 		WithArgs("test", "public").
 		WillReturnRows(sqlmock.NewRows([]string{"table_schema", "table_name", "comment", "partition_attrs", "partition_strategy", "partition_exprs"}))
+	// Reset search_path to 'public'.
+	mk.ExpectQuery(sqltest.Escape("SELECT set_config('search_path', $1, false)")).
+		WithArgs("public").
+		WillReturnRows(sqlmock.NewRows(nil))
 	realm, err := drv.InspectRealm(context.Background(), &schema.InspectRealmOption{
 		Mode: schema.InspectSchemas | schema.InspectTables,
 	})
@@ -575,6 +586,13 @@ func TestDriver_Realm(t *testing.T) {
 		return r
 	}(), realm)
 
+	// No need to reset, if the search_path was not set.
+	mk.ExpectQuery(sqltest.Escape("SELECT current_setting('search_path'), set_config('search_path', '', false)")).
+		WillReturnRows(sqltest.Rows(`
+ current_setting | set_config
+-----------------+------------
+                 |
+`))
 	mk.ExpectQuery(sqltest.Escape(fmt.Sprintf(schemasQueryArgs, "IN ($1, $2)"))).
 		WithArgs("test", "public").
 		WillReturnRows(sqltest.Rows(`
@@ -609,6 +627,13 @@ func TestDriver_Realm(t *testing.T) {
 		return r
 	}(), realm)
 
+	// No need to reset, if the search_path was not set.
+	mk.ExpectQuery(sqltest.Escape("SELECT current_setting('search_path'), set_config('search_path', '', false)")).
+		WillReturnRows(sqltest.Rows(`
+ current_setting | set_config
+-----------------+------------
+                 |
+`))
 	mk.ExpectQuery(sqltest.Escape(fmt.Sprintf(schemasQueryArgs, "= $1"))).
 		WithArgs("test").
 		WillReturnRows(sqltest.Rows(`
@@ -643,6 +668,12 @@ func TestInspectMode_InspectRealm(t *testing.T) {
 	require.NoError(t, err)
 	mk := mock{m}
 	mk.version("130000")
+	mk.ExpectQuery(sqltest.Escape("SELECT current_setting('search_path'), set_config('search_path', '', false)")).
+		WillReturnRows(sqltest.Rows(`
+ current_setting | set_config
+-----------------+------------
+                 |
+`))
 	mk.ExpectQuery(sqltest.Escape(schemasQuery)).
 		WillReturnRows(sqltest.Rows(`
  schema_name | comment