[v14] Database Automatic User Provisioning support for Redshift (#34126)

* Database Automatic User Provisioning support for Redshift (#33307)

* Database Automatic User Provisioning support for Redshift

* capitalize sql script

* DeleteUser to fallback to deactivate.

* add TPxxx code to RAISE messages

* fix missing space in scripts

* feat(postgres): support auto-provisioned user deletion Redshift (#34006)

---------

Co-authored-by: Gabriel Corado <gabriel.oliveira@goteleport.com>

api/types/database.go

@@ -311,7 +311,7 @@ func (d *DatabaseV3) SupportsAutoUsers() bool {
 	switch d.GetProtocol() {
 	case DatabaseProtocolPostgreSQL:
 		switch d.GetType() {
-		case DatabaseTypeSelfHosted, DatabaseTypeRDS:
+		case DatabaseTypeSelfHosted, DatabaseTypeRDS, DatabaseTypeRedshift:
 			return true
 		}
 	case DatabaseProtocolMySQL:

lib/srv/db/postgres/sql/redshift-activate-user.sql

@@ -0,0 +1,40 @@
+CREATE OR REPLACE PROCEDURE teleport_activate_user(username varchar, roles text)
+LANGUAGE plpgsql
+AS $$
+DECLARE
+    roles_length integer;
+    cur_roles_length integer;
+BEGIN
+    roles_length := JSON_ARRAY_LENGTH(roles);
+
+    -- If the user already exists and was provisioned by Teleport, reactivate
+    -- it, otherwise provision a new one.
+    IF EXISTS (SELECT user_id FROM svv_user_grants WHERE user_name = username AND admin_option = false AND role_name = 'teleport-auto-user') THEN
+        -- If the user has active connections, make sure the provided roles
+        -- match what the user currently has.
+        IF EXISTS (SELECT user_name FROM stv_sessions WHERE user_name = CONCAT('IAM:', username)) THEN
+          SELECT INTO cur_roles_length COUNT(role_name) FROM svv_user_grants WHERE user_name = username AND admin_option=false AND role_name != 'teleport-auto-user';
+          IF roles_length != cur_roles_length THEN
+            RAISE EXCEPTION 'TP002: User has active connections and roles have changed';
+          END IF;
+          FOR i IN 0..roles_length-1 LOOP
+            IF NOT EXISTS (SELECT role_name FROM svv_user_grants WHERE user_name = username AND admin_option=false AND role_name = JSON_EXTRACT_ARRAY_ELEMENT_TEXT(roles,i)) THEN
+                RAISE EXCEPTION 'TP002: User has active connections and roles have changed';
+            END IF;
+          END LOOP;
+          RETURN;
+        END IF;
+        -- Otherwise reactivate the user, but first strip it of all roles to
+        -- account for scenarios with left-over roles if database agent crashed
+        -- and failed to cleanup upon session termination.
+        CALL teleport_deactivate_user(username);
+        EXECUTE 'ALTER USER ' || QUOTE_IDENT(username) || ' CONNECTION LIMIT UNLIMITED';
+    ELSE
+        EXECUTE 'CREATE USER ' || QUOTE_IDENT(username) || ' WITH PASSWORD DISABLE';
+        EXECUTE 'GRANT ROLE "teleport-auto-user" TO ' || QUOTE_IDENT(username);
+    END IF;
+    -- Assign all roles to the created/activated user.
+    FOR i in 0..roles_length-1 LOOP
+        EXECUTE 'GRANT ROLE ' || QUOTE_IDENT(JSON_EXTRACT_ARRAY_ELEMENT_TEXT(roles,i)) || ' TO ' || QUOTE_IDENT(username);
+    END LOOP;
+END;$$;

lib/srv/db/postgres/sql/redshift-deactivate-user.sql

@@ -0,0 +1,20 @@
+CREATE OR REPLACE PROCEDURE teleport_deactivate_user(username varchar)
+LANGUAGE plpgsql
+AS $$
+DECLARE
+    rec record;
+BEGIN
+    -- Only deactivate if the user doesn't have other active sessions.
+    -- Update to pg_stat_activity is delayed for a few hundred ms. Use
+    -- stv_sessions instead.
+    IF EXISTS (SELECT user_name FROM stv_sessions WHERE user_name = CONCAT('IAM:', username)) THEN
+        RAISE EXCEPTION 'TP000: User has active connections';
+    ELSE
+        -- Revoke all role memberships except teleport-auto-user.
+        FOR rec IN select role_name FROM svv_user_grants WHERE user_name = username AND admin_option = false AND role_name != 'teleport-auto-user' LOOP
+             EXECUTE 'REVOKE ROLE ' || QUOTE_IDENT(rec.role_name) || ' FROM ' || QUOTE_IDENT(username);
+        END LOOP;
+        -- Disable ability to login for the user.
+        EXECUTE 'ALTER USER ' || QUOTE_IDENT(username) || ' WITH CONNECTION LIMIT 0';
+    END IF;
+END;$$;

lib/srv/db/postgres/sql/redshift-delete-user.sql

@@ -0,0 +1,20 @@
+CREATE OR REPLACE PROCEDURE teleport_delete_user(username varchar)
+LANGUAGE plpgsql
+AS $$
+BEGIN
+    -- Only drop if the user doesn't have other active sessions.
+    IF EXISTS (SELECT usename FROM pg_stat_activity WHERE usename = username) THEN
+        RAISE NOTICE 'User has active connections';
+    ELSE
+        BEGIN
+            EXECUTE 'DROP USER ' || QUOTE_IDENT(username);
+        EXCEPTION WHEN OTHERS THEN
+            -- Redshift only support OTHERS as exception condition, so we handle
+            -- any error that might happen.
+
+            -- Drop user/role will fail if user has dependent objects.
+            -- In this scenario, fallback into disabling the user.
+            CALL teleport_deactivate_user(username);
+        END;
+    END IF;
+END;$$;

lib/srv/db/postgres/users.go

@@ -19,12 +19,14 @@ package postgres
 import (
 	"context"
 	_ "embed"
+	"encoding/json"
 	"fmt"
 	"strings"
 
 	"github.com/gravitational/trace"
 	"github.com/jackc/pgx/v4"
 
+	"github.com/gravitational/teleport/api/types"
 	"github.com/gravitational/teleport/lib/srv/db/common"
 )
 
@@ -43,27 +45,25 @@ func (e *Engine) ActivateUser(ctx context.Context, sessionCtx *common.Session) e
 	// We could call this once when the database is being initialized but
 	// doing it here has a nice "self-healing" property in case the Teleport
 	// bookkeeping group or stored procedures get deleted or changed offband.
-	err = e.initAutoUsers(ctx, conn)
+	err = e.initAutoUsers(ctx, sessionCtx, conn)
 	if err != nil {
 		return trace.Wrap(err)
 	}
 
-	roles := sessionCtx.DatabaseRoles
-	if sessionCtx.Database.IsRDS() {
-		roles = append(roles, "rds_iam")
+	roles, err := prepareRoles(sessionCtx)
+	if err != nil {
+		return trace.Wrap(err)
 	}
 
 	e.Log.Infof("Activating PostgreSQL user %q with roles %v.", sessionCtx.DatabaseUser, roles)
 
 	_, err = conn.Exec(ctx, activateQuery, sessionCtx.DatabaseUser, roles)
 	if err != nil {
-		if strings.Contains(err.Error(), "already exists") {
-			return trace.AlreadyExists("user %q already exists in this PostgreSQL database and is not managed by Teleport", sessionCtx.DatabaseUser)
-		}
-		return trace.Wrap(err)
+		e.Log.Debugf("Call teleport_activate_user failed: %v", err)
+		return trace.Wrap(convertActivateError(sessionCtx, err))
 	}
-
 	return nil
+
 }
 
 // DeactivateUser disables the database user.
@@ -103,26 +103,55 @@ func (e *Engine) DeleteUser(ctx context.Context, sessionCtx *common.Session) err
 	e.Log.Infof("Deleting PostgreSQL user %q.", sessionCtx.DatabaseUser)
 
 	var state string
-	err = conn.QueryRow(ctx, deleteQuery, sessionCtx.DatabaseUser).Scan(&state)
+	switch {
+	case sessionCtx.Database.IsRedshift():
+		err = e.deleteUserRedshift(ctx, sessionCtx, conn, &state)
+	default:
+		err = conn.QueryRow(ctx, deleteQuery, sessionCtx.DatabaseUser).Scan(&state)
+	}
 	if err != nil {
 		return trace.Wrap(err)
 	}
 
 	switch state {
 	case common.SQLStateUserDropped:
-		e.Log.Debug("User %q deleted successfully.", sessionCtx.DatabaseUser)
+		e.Log.Debugf("User %q deleted successfully.", sessionCtx.DatabaseUser)
 	case common.SQLStateUserDeactivated:
-		e.Log.Infof("Unable to delete user %q, it was disabled instead", sessionCtx.DatabaseUser)
+		e.Log.Infof("Unable to delete user %q, it was disabled instead.", sessionCtx.DatabaseUser)
 	default:
 		e.Log.Warnf("Unable to determine user %q deletion state.", sessionCtx.DatabaseUser)
 	}
 
 	return nil
 }
 
+// deleteUserRedshift deletes the Redshift database user.
+//
+// Failures inside Redshift default procedures are always rethrown exceptions if
+// the exception handler completes successfully. Given this, we need to assert
+// into the returned error instead of doing this on state returned (like regular
+// PostgreSQL).
+func (e *Engine) deleteUserRedshift(ctx context.Context, sessionCtx *common.Session, conn *pgx.Conn, state *string) error {
+	_, err := conn.Exec(ctx, deleteQuery, sessionCtx.DatabaseUser)
+	if err == nil {
+		*state = common.SQLStateUserDropped
+		return nil
+	}
+
+	// Redshift returns SQLSTATE 55006 (object_in_use) when DROP USER fails due
+	// to user owning resources.
+	// https://docs.aws.amazon.com/redshift/latest/dg/r_DROP_USER.html#r_DROP_USER-notes
+	if strings.Contains(err.Error(), "55006") {
+		*state = common.SQLStateUserDeactivated
+		return nil
+	}
+
+	return trace.Wrap(err)
+}
+
 // initAutoUsers installs procedures for activating and deactivating users and
 // creates the bookkeeping role for auto-provisioned users.
-func (e *Engine) initAutoUsers(ctx context.Context, conn *pgx.Conn) error {
+func (e *Engine) initAutoUsers(ctx context.Context, sessionCtx *common.Session, conn *pgx.Conn) error {
 	// Create a role/group which all auto-created users will be a part of.
 	_, err := conn.Exec(ctx, fmt.Sprintf("create role %q", teleportAutoUserRole))
 	if err != nil {
@@ -133,8 +162,9 @@ func (e *Engine) initAutoUsers(ctx context.Context, conn *pgx.Conn) error {
 	} else {
 		e.Log.Debugf("Created PostgreSQL role %q.", teleportAutoUserRole)
 	}
+
 	// Install stored procedures for creating and disabling database users.
-	for name, sql := range procs {
+	for name, sql := range pickProcedures(sessionCtx) {
 		_, err := conn.Exec(ctx, sql)
 		if err != nil {
 			return trace.Wrap(err)
@@ -159,6 +189,44 @@ func (e *Engine) pgxConnect(ctx context.Context, sessionCtx *common.Session) (*p
 	return pgx.ConnectConfig(ctx, pgxConf)
 }
 
+func prepareRoles(sessionCtx *common.Session) (any, error) {
+	switch sessionCtx.Database.GetType() {
+	case types.DatabaseTypeRDS:
+		return append(sessionCtx.DatabaseRoles, "rds_iam"), nil
+
+	case types.DatabaseTypeRedshift:
+		// Redshift does not support array. Encode roles in JSON (type text).
+		rolesJSON, err := json.Marshal(sessionCtx.DatabaseRoles)
+		if err != nil {
+			return nil, trace.Wrap(err)
+		}
+		return string(rolesJSON), nil
+
+	default:
+		return sessionCtx.DatabaseRoles, nil
+	}
+}
+
+func convertActivateError(sessionCtx *common.Session, err error) error {
+	switch {
+	case strings.Contains(err.Error(), "already exists"):
+		return trace.AlreadyExists("user %q already exists in this PostgreSQL database and is not managed by Teleport", sessionCtx.DatabaseUser)
+
+	case strings.Contains(err.Error(), "TP002: User has active connections and roles have changed"):
+		return trace.CompareFailed("roles for user %q has changed. Please quit all active connections and try again.", sessionCtx.DatabaseUser)
+
+	default:
+		return trace.Wrap(err)
+	}
+}
+
+func pickProcedures(sessionCtx *common.Session) map[string]string {
+	if sessionCtx.Database.IsRedshift() {
+		return redshiftProcs
+	}
+	return procs
+}
+
 const (
 	// activateProcName is the name of the stored procedure Teleport will use
 	// to automatically provision/activate database users.
@@ -176,24 +244,36 @@ const (
 )
 
 var (
-	//go:embed activate-user.sql
+	//go:embed sql/activate-user.sql
 	activateProc string
 	// activateQuery is the query for calling user activation procedure.
 	activateQuery = fmt.Sprintf(`call %v($1, $2)`, activateProcName)
 
-	//go:embed deactivate-user.sql
+	//go:embed sql/deactivate-user.sql
 	deactivateProc string
 	// deactivateQuery is the query for calling user deactivation procedure.
 	deactivateQuery = fmt.Sprintf(`call %v($1)`, deactivateProcName)
 
-	//go:embed delete-user.sql
+	//go:embed sql/delete-user.sql
 	deleteProc string
 	// deleteQuery is the query for calling user deletion procedure.
 	deleteQuery = fmt.Sprintf(`call %v($1)`, deleteProcName)
 
+	//go:embed sql/redshift-activate-user.sql
+	redshiftActivateProc string
+	//go:embed sql/redshift-deactivate-user.sql
+	redshiftDeactivateProc string
+	//go:embed sql/redshift-delete-user.sql
+	redshiftDeleteProc string
+
 	procs = map[string]string{
 		activateProcName:   activateProc,
 		deactivateProcName: deactivateProc,
 		deleteProcName:     deleteProc,
 	}
+	redshiftProcs = map[string]string{
+		activateProcName:   redshiftActivateProc,
+		deactivateProcName: redshiftDeactivateProc,
+		deleteProcName:     redshiftDeleteProc,
+	}
 )

lib/srv/db/postgres/users_test.go

@@ -0,0 +1,84 @@
+/*
+Copyright 2023 Gravitational, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package postgres
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/gravitational/teleport/api/types"
+	"github.com/gravitational/teleport/lib/defaults"
+	"github.com/gravitational/teleport/lib/srv/db/common"
+)
+
+func Test_prepareRoles(t *testing.T) {
+	selfHostedDatabase, err := types.NewDatabaseV3(types.Metadata{
+		Name: "self-hosted",
+	}, types.DatabaseSpecV3{
+		Protocol: defaults.ProtocolPostgres,
+		URI:      "localhost:5432",
+	})
+	require.NoError(t, err)
+
+	rdsDatabase, err := types.NewDatabaseV3(types.Metadata{
+		Name: "rds",
+	}, types.DatabaseSpecV3{
+		Protocol: defaults.ProtocolPostgres,
+		URI:      "aurora-instance-1.abcdefghijklmnop.us-west-1.rds.amazonaws.com:5432",
+	})
+	require.NoError(t, err)
+
+	redshiftDatabase, err := types.NewDatabaseV3(types.Metadata{
+		Name: "redshift",
+	}, types.DatabaseSpecV3{
+		Protocol: defaults.ProtocolPostgres,
+		URI:      "redshift-cluster-1.abcdefghijklmnop.us-east-1.redshift.amazonaws.com:5439",
+	})
+	require.NoError(t, err)
+
+	tests := []struct {
+		inputDatabase types.Database
+		expectRoles   any
+	}{
+		{
+			inputDatabase: selfHostedDatabase,
+			expectRoles:   []string{"role1", "role2"},
+		},
+		{
+			inputDatabase: rdsDatabase,
+			expectRoles:   []string{"role1", "role2", "rds_iam"},
+		},
+		{
+			inputDatabase: redshiftDatabase,
+			expectRoles:   `["role1","role2"]`,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.inputDatabase.GetName(), func(t *testing.T) {
+			sessionCtx := &common.Session{
+				Database:      test.inputDatabase,
+				DatabaseRoles: []string{"role1", "role2"},
+			}
+
+			actualRoles, err := prepareRoles(sessionCtx)
+			require.NoError(t, err)
+			require.Equal(t, test.expectRoles, actualRoles)
+		})
+	}
+}