Add option to customize admin schema for PostgreSQL auto user provisioning

api/proto/teleport/legacy/types/types.proto

@@ -344,6 +344,11 @@ message DatabaseAdminUser {
   // Depending on the database type, this database may be used to store
   // procedures or data for managing database users.
   string DefaultDatabase = 2 [(gogoproto.jsontag) = "default_database"];
+  // DefaultSchema is the schema that the privileged database user will use.
+  //
+  // Depending on the database type, this schema may be used to store
+  // procedures.
+  string DefaultSchema = 3 [(gogoproto.jsontag) = "default_schema"];
 }
 
 // OracleOptions contains information about privileged database user used

api/types/constants.go

@@ -690,6 +690,10 @@ const (
 	// the admin user logs into by default.
 	DatabaseAdminDefaultDatabaseLabel = TeleportNamespace + "/db-admin-default-database"
 
+	// DatabaseAdminDefaultSchemaLabel is used to identify the schema that the
+	// admin user will use by default.
+	DatabaseAdminDefaultSchemaLabel = TeleportNamespace + "/db-admin-default-schema"
+
 	// cloudKubeClusterNameOverrideLabel is a cloud agnostic label key for
 	// overriding kubernetes cluster name in discovered cloud kube clusters.
 	// It's used for AWS, GCP, and Azure, but not exported to decouple the

api/types/database.go

@@ -304,6 +304,9 @@ func (d *DatabaseV3) GetAdminUser() (ret DatabaseAdminUser) {
 		if ret.DefaultDatabase == "" {
 			ret.DefaultDatabase = d.Metadata.Labels[DatabaseAdminDefaultDatabaseLabel]
 		}
+		if ret.DefaultSchema == "" {
+			ret.DefaultSchema = d.Metadata.Labels[DatabaseAdminDefaultSchemaLabel]
+		}
 	}
 	return
 }

lib/config/configuration.go

@@ -1832,6 +1832,7 @@ func applyDatabasesConfig(fc *FileConfig, cfg *servicecfg.Config) error {
 			AdminUser: servicecfg.DatabaseAdminUser{
 				Name:            database.AdminUser.Name,
 				DefaultDatabase: database.AdminUser.DefaultDatabase,
+				DefaultSchema:   database.AdminUser.DefaultSchema,
 			},
 			Oracle: convOracleOptions(database.Oracle),
 			AWS: servicecfg.DatabaseAWS{

lib/config/fileconf.go

@@ -1815,6 +1815,11 @@ type DatabaseAdminUser struct {
 	// Depending on the database type, this database may be used to store
 	// procedures or data for managing database users.
 	DefaultDatabase string `yaml:"default_database"`
+	// DefaultSchema is the schema that the privileged database user will use.
+	//
+	// Depending on the database type, this schema may be used to store
+	// procedures.
+	DefaultSchema string `yaml:"default_schema"`
 }
 
 // DatabaseAD contains database Active Directory configuration.

lib/service/servicecfg/database.go

@@ -83,6 +83,11 @@ type DatabaseAdminUser struct {
 	// Depending on the database type, this database may be used to store
 	// procedures or data for managing database users.
 	DefaultDatabase string
+	// DefaultSchema is the schema that the privileged database user will use.
+	//
+	// Depending on the database type, this schema may be used to store
+	// procedures.
+	DefaultSchema string
 }
 
 // OracleOptions are additional Oracle options.
@@ -133,6 +138,7 @@ func (d *Database) ToDatabase() (types.Database, error) {
 		AdminUser: &types.DatabaseAdminUser{
 			Name:            d.AdminUser.Name,
 			DefaultDatabase: d.AdminUser.DefaultDatabase,
+			DefaultSchema:   d.AdminUser.DefaultSchema,
 		},
 		Oracle: convOracleOptions(d.Oracle),
 		AWS: types.AWS{

lib/services/database.go

@@ -210,6 +210,10 @@ func ValidateDatabase(db types.Database) error {
 		}
 	}
 
+	if err := validateAdminDefaultSchema(db); err != nil {
+		return trace.Wrap(err)
+	}
+
 	return nil
 }
 
@@ -347,6 +351,28 @@ func ValidateSQLServerURI(uri string) error {
 	return nil
 }
 
+// validateAdminDefaultSchema validates the admin user default schema value.
+func validateAdminDefaultSchema(db types.Database) error {
+	defaultSchema := db.GetAdminUser().DefaultSchema
+	// Only validate if auto user provisioning is configured to the database.
+	if !db.SupportsAutoUsers() || db.GetAdminUser().Name == "" || defaultSchema == "" {
+		return nil
+	}
+
+	if db.GetProtocol() != defaults.ProtocolPostgres || (db.GetType() != types.DatabaseTypeSelfHosted && !db.IsRDS() && !db.IsRedshift()) {
+		return trace.BadParameter("admin user default schema is only supported for PostgreSQL databases")
+	}
+
+	// Currently, PostgreSQL procedures are only created on activate user. If
+	// the schema is set to the temporary the procedures won't be available for
+	// auto users teardown.
+	if defaultSchema == PostgresTempSchema {
+		return trace.BadParameter("admin user default schema cannot be set to %q", PostgresTempSchema)
+	}
+
+	return nil
+}
+
 func isDNSError(err error) bool {
 	if err == nil {
 		return false
@@ -356,6 +382,11 @@ func isDNSError(err error) bool {
 	return errors.As(err, &dnsErr)
 }
 
+const (
+	// PostgresTempSchema is the name of the temporary schema on PostgreSQL.
+	PostgresTempSchema = "pg_temp"
+)
+
 const (
 	// RDSDescribeTypeInstance is used to filter for Instances type of RDS DBs when describing RDS Databases.
 	RDSDescribeTypeInstance = "instance"

lib/services/database_test.go

@@ -492,6 +492,87 @@ func TestValidateDatabase(t *testing.T) {
 			},
 			expectError: false,
 		},
+		{
+			inputName: "other-protocol-default-schema",
+			inputSpec: types.DatabaseSpecV3{
+				Protocol: defaults.ProtocolMySQL,
+				URI:      "localhost:3306",
+				AdminUser: &types.DatabaseAdminUser{
+					Name:          "tp",
+					DefaultSchema: "random-schema",
+				},
+			},
+			expectError: true,
+		},
+		{
+			inputName: "self-hosted-postgres-temp-default-schema-disabled-auto-provisioning",
+			inputSpec: types.DatabaseSpecV3{
+				Protocol: defaults.ProtocolPostgres,
+				URI:      "localhost:5432",
+				AdminUser: &types.DatabaseAdminUser{
+					DefaultSchema: PostgresTempSchema,
+				},
+			},
+			expectError: false,
+		},
+		{
+			inputName: "self-hosted-postgres-temp-default-schema",
+			inputSpec: types.DatabaseSpecV3{
+				Protocol: defaults.ProtocolPostgres,
+				URI:      "localhost:5432",
+				AdminUser: &types.DatabaseAdminUser{
+					Name:          "tp",
+					DefaultSchema: PostgresTempSchema,
+				},
+			},
+			expectError: true,
+		},
+		{
+			inputName: "rds-postgres-temp-default-schema",
+			inputSpec: types.DatabaseSpecV3{
+				Protocol: defaults.ProtocolPostgres,
+				URI:      "localhost:5432",
+				AWS: types.AWS{
+					RDS: types.RDS{
+						InstanceID: "example-rds",
+					},
+				},
+				AdminUser: &types.DatabaseAdminUser{
+					Name:          "tp",
+					DefaultSchema: PostgresTempSchema,
+				},
+			},
+			expectError: true,
+		},
+		{
+			inputName: "rds-postgres-default-schema",
+			inputSpec: types.DatabaseSpecV3{
+				Protocol: defaults.ProtocolPostgres,
+				URI:      "localhost:5432",
+				AWS: types.AWS{
+					RDS: types.RDS{
+						InstanceID: "example-rds",
+					},
+				},
+				AdminUser: &types.DatabaseAdminUser{
+					Name:          "tp",
+					DefaultSchema: "example-schema",
+				},
+			},
+			expectError: false,
+		},
+		{
+			inputName: "rds-postgres-default-schema",
+			inputSpec: types.DatabaseSpecV3{
+				Protocol: defaults.ProtocolPostgres,
+				URI:      "localhost:5432",
+				AdminUser: &types.DatabaseAdminUser{
+					Name:          "tp",
+					DefaultSchema: "example-schema",
+				},
+			},
+			expectError: false,
+		},
 	}
 
 	for _, test := range tests {

lib/srv/db/autousers_test.go

@@ -50,8 +50,10 @@ func TestAutoUsersPostgres(t *testing.T) {
 		importRuleSpec      *dbobjectimportrulev1.DatabaseObjectImportRuleSpec
 
 		adminDefaultDatabase   string
+		adminDefaultSchema     string
 		connectionErrorMessage string
 		expectAdminDatabase    string
+		expectSchema           string
 	}{
 		"activate/deactivate users": {
 			mode:                types.CreateDatabaseUserMode_DB_USER_MODE_KEEP,
@@ -77,6 +79,27 @@ func TestAutoUsersPostgres(t *testing.T) {
 			adminDefaultDatabase: "admin-db",
 			expectAdminDatabase:  "admin-db",
 		},
+		"admin user default schema": {
+			mode:                types.CreateDatabaseUserMode_DB_USER_MODE_KEEP,
+			databaseRoles:       []string{"reader", "writer"},
+			adminDefaultSchema:  "example",
+			expectAdminDatabase: "user-db",
+			expectSchema:        "example",
+		},
+		"admin user default schema and permissions": {
+			mode:               types.CreateDatabaseUserMode_DB_USER_MODE_KEEP,
+			adminDefaultSchema: "example",
+			databasePermissions: types.DatabasePermissions{
+				{
+					Permissions: []string{"SELECT"},
+					Match: map[string]apiutils.Strings{
+						"can_select": []string{"true"},
+					},
+				},
+			},
+			expectAdminDatabase: "user-db",
+			expectSchema:        "example",
+		},
 		"roles and permissions are mutually exclusive": {
 			mode:          types.CreateDatabaseUserMode_DB_USER_MODE_BEST_EFFORT_DROP,
 			databaseRoles: []string{"reader", "writer"},
@@ -138,6 +161,7 @@ func TestAutoUsersPostgres(t *testing.T) {
 				db.Spec.AdminUser = &types.DatabaseAdminUser{
 					Name:            "postgres",
 					DefaultDatabase: tc.adminDefaultDatabase,
+					DefaultSchema:   tc.adminDefaultSchema,
 				}
 			}))
 			go testCtx.startHandlingConnections()
@@ -173,15 +197,15 @@ func TestAutoUsersPostgres(t *testing.T) {
 
 			// Verify incoming connections.
 			// 1. Admin connecting to admin database
-			requirePostgresConnection(t, testCtx.postgres["postgres"].db.ParametersCh(), "postgres", tc.expectAdminDatabase)
+			requirePostgresConnection(t, testCtx.postgres["postgres"].db.ParametersCh(), "postgres", tc.expectAdminDatabase, tc.expectSchema)
 
 			// 2. If there are any database permissions: admin connecting to session database.
 			if len(tc.databasePermissions) > 0 {
-				requirePostgresConnection(t, testCtx.postgres["postgres"].db.ParametersCh(), "postgres", "user-db")
+				requirePostgresConnection(t, testCtx.postgres["postgres"].db.ParametersCh(), "postgres", "user-db", "")
 			}
 
 			// 3. User connecting to session database.
-			requirePostgresConnection(t, testCtx.postgres["postgres"].db.ParametersCh(), "alice", "user-db")
+			requirePostgresConnection(t, testCtx.postgres["postgres"].db.ParametersCh(), "alice", "user-db", "")
 
 			// Verify user was activated.
 			select {
@@ -224,13 +248,14 @@ func TestAutoUsersPostgres(t *testing.T) {
 	}
 }
 
-func requirePostgresConnection(t *testing.T, parametersCh chan map[string]string, expectUser, expectDatabase string) {
+func requirePostgresConnection(t *testing.T, parametersCh chan map[string]string, expectUser, expectDatabase, expectSchema string) {
 	t.Helper()
 	select {
 	case parameters := <-parametersCh:
 		require.NotNil(t, parameters)
 		require.Equal(t, expectUser, parameters["user"])
 		require.Equal(t, expectDatabase, parameters["database"])
+		require.Equal(t, expectSchema, parameters["search_path"])
 	case <-time.After(5 * time.Second):
 		t.Fatal("no connection after 5s")
 	}

lib/srv/db/postgres/users.go

@@ -38,15 +38,23 @@ import (
 )
 
 // connectAsAdmin connect to the database from db route as admin user.
-// If useDefaultDatabase is true and a default database is configured for the admin user, it will be used instead.
-func (e *Engine) connectAsAdmin(ctx context.Context, sessionCtx *common.Session, useDefaultDatabase bool) (*pgx.Conn, error) {
+// If useDefaultValues is true and a default values are configured for the admin user, they will be used instead.
+func (e *Engine) connectAsAdmin(ctx context.Context, sessionCtx *common.Session, useDefaultValues bool) (*pgx.Conn, error) {
+	loginSchema := ""
 	loginDatabase := sessionCtx.DatabaseName
-	if useDefaultDatabase && sessionCtx.Database.GetAdminUser().DefaultDatabase != "" {
-		loginDatabase = sessionCtx.Database.GetAdminUser().DefaultDatabase
-	} else {
-		e.Log.WithField("database", loginDatabase).Info("Connecting to session database")
+	if useDefaultValues {
+		if sessionCtx.Database.GetAdminUser().DefaultDatabase != "" {
+			loginDatabase = sessionCtx.Database.GetAdminUser().DefaultDatabase
+		} else {
+			e.Log.WithField("database", loginDatabase).Info("Connecting to session database")
+		}
+
+		if sessionCtx.Database.GetAdminUser().DefaultSchema != "" {
+			loginSchema = sessionCtx.Database.GetAdminUser().DefaultSchema
+		}
 	}
-	conn, err := e.pgxConnect(ctx, sessionCtx.WithUserAndDatabase(sessionCtx.Database.GetAdminUser().Name, loginDatabase))
+
+	conn, err := e.pgxConnect(ctx, sessionCtx.WithUserAndDatabase(sessionCtx.Database.GetAdminUser().Name, loginDatabase), loginSchema)
 	return conn, trace.Wrap(err)
 }
 
@@ -408,15 +416,25 @@ func (e *Engine) initAutoUsers(ctx context.Context, sessionCtx *common.Session,
 
 // pgxConnect connects to the database using pgx driver which is higher-level
 // than pgconn and is easier to use for executing queries.
-func (e *Engine) pgxConnect(ctx context.Context, sessionCtx *common.Session) (*pgx.Conn, error) {
+func (e *Engine) pgxConnect(ctx context.Context, sessionCtx *common.Session, searchPath string) (*pgx.Conn, error) {
 	config, err := e.getConnectConfig(ctx, sessionCtx)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
+
+	// Always reset runtime params to avoid using user provided flags to admin
+	// connections.
+	config.RuntimeParams = map[string]string{}
+
+	if searchPath != "" {
+		config.RuntimeParams[searchPathParam] = searchPath
+	}
+
 	pgxConf, err := pgx.ParseConfig("")
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
+
 	pgxConf.Config = *config
 	return pgx.ConnectConfig(ctx, pgxConf)
 }
@@ -459,6 +477,13 @@ func pickProcedures(sessionCtx *common.Session) map[string]string {
 	return procs
 }
 
+const (
+	// searchPathParam is the name of the startup parameter used to set the
+	// connection `search_path`.
+	// https://www.postgresql.org/docs/16/ddl-schemas.html#DDL-SCHEMAS-PATH
+	searchPathParam = "search_path"
+)
+
 const (
 	// activateProcName is the name of the stored procedure Teleport will use
 	// to automatically provision/activate database users.