From 2bfe9f065b1a8b8178fdbfcb4df2ec4aabbd566b Mon Sep 17 00:00:00 2001
From: Gabriel Corado <gabriel.oliveira@goteleport.com>
Date: Fri, 27 Oct 2023 17:53:06 -0300
Subject: [PATCH] feat(postgres): support auto-provisioned user deletion
 Redshift

---
 .../db/postgres/sql/redshift-delete-user.sql  | 20 +++++++++
 lib/srv/db/postgres/users.go                  | 42 +++++++++++++++----
 2 files changed, 54 insertions(+), 8 deletions(-)
 create mode 100644 lib/srv/db/postgres/sql/redshift-delete-user.sql

diff --git a/lib/srv/db/postgres/sql/redshift-delete-user.sql b/lib/srv/db/postgres/sql/redshift-delete-user.sql
new file mode 100644
index 0000000000000..1121e571f47b8
--- /dev/null
+++ b/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;$$;
diff --git a/lib/srv/db/postgres/users.go b/lib/srv/db/postgres/users.go
index f4c1e15f34d8c..cf2d6413137a0 100644
--- a/lib/srv/db/postgres/users.go
+++ b/lib/srv/db/postgres/users.go
@@ -90,12 +90,6 @@ func (e *Engine) DeactivateUser(ctx context.Context, sessionCtx *common.Session)
 
 // DeleteUser deletes the database user.
 func (e *Engine) DeleteUser(ctx context.Context, sessionCtx *common.Session) error {
-	// TODO support DeleteUser for Redshift
-	if sessionCtx.Database.IsRedshift() {
-		e.Log.Debug("DeleteUser is not supported for Redshift yet, it was disabled instead.")
-		return trace.Wrap(e.DeactivateUser(ctx, sessionCtx))
-	}
-
 	if sessionCtx.Database.GetAdminUser().Name == "" {
 		return trace.BadParameter("Teleport does not have admin user configured for this database")
 	}
@@ -109,14 +103,19 @@ 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)
 	default:
@@ -126,6 +125,30 @@ func (e *Engine) DeleteUser(ctx context.Context, sessionCtx *common.Session) err
 	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, sessionCtx *common.Session, conn *pgx.Conn) error {
@@ -240,6 +263,8 @@ var (
 	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,
@@ -249,5 +274,6 @@ var (
 	redshiftProcs = map[string]string{
 		activateProcName:   redshiftActivateProc,
 		deactivateProcName: redshiftDeactivateProc,
+		deleteProcName:     redshiftDeleteProc,
 	}
 )