make AllUserIDs a fully-fledged query citizen

coder · johnstcn · Sep 15, 2023 · Sep 15, 2023 · Sep 15, 2023 · Sep 15, 2023
commit 11f85bff7436a1006a91bc47a5bcf447c727f159
diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -664,6 +664,15 @@ func (q *querier) ActivityBumpWorkspace(ctx context.Context, arg uuid.UUID) erro
 	return update(q.log, q.auth, fetch, q.db.ActivityBumpWorkspace)(ctx, arg)
 }
 
+func (q *querier) AllUserIDs(ctx context.Context) ([]uuid.UUID, error) {
+	// Although this technically only reads users, only system-related functions should be
+	// allowed to call this.
+	if err := q.authorizeContext(ctx, rbac.ActionRead, rbac.ResourceSystem); err != nil {
+		return nil, err
+	}
+	return q.db.AllUserIDs(ctx)
+}
+
 func (q *querier) CleanTailnetCoordinators(ctx context.Context) error {
 	if err := q.authorizeContext(ctx, rbac.ActionDelete, rbac.ResourceTailnetCoordinator); err != nil {
 		return err

diff --git a/coderd/database/dbfake/dbfake.go b/coderd/database/dbfake/dbfake.go
@@ -812,6 +812,14 @@ func (q *FakeQuerier) ActivityBumpWorkspace(ctx context.Context, workspaceID uui
 	return sql.ErrNoRows
 }
 
+func (q *FakeQuerier) AllUserIDs(ctx context.Context) ([]uuid.UUID, error) {
+	userIDs := make([]uuid.UUID, 0, len(q.users))
+	for idx := range q.users {
+		userIDs[idx] = q.users[idx].ID
+	}
+	return userIDs, nil
+}
+
 func (*FakeQuerier) CleanTailnetCoordinators(_ context.Context) error {
 	return ErrUnimplemented
 }

diff --git a/coderd/database/dbmetrics/dbmetrics.go b/coderd/database/dbmetrics/dbmetrics.go
diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go
diff --git a/coderd/database/querier.go b/coderd/database/querier.go
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/dbcrypt.sql b/coderd/database/queries/dbcrypt.sql
@@ -16,3 +16,4 @@ AND
 INSERT INTO dbcrypt_keys
 	(number, active_key_digest, created_at, test)
 VALUES (@number::int, @active_key_digest::text, CURRENT_TIMESTAMP, @test::text);
+
diff --git a/coderd/database/queries/users.sql b/coderd/database/queries/users.sql
@@ -262,3 +262,8 @@ WHERE
     last_seen_at < @last_seen_after :: timestamp
     AND status = 'active'::user_status
 RETURNING id, email, last_seen_at;
+
+-- AllUserIDs returns all UserIDs regardless of user status or deletion.
+-- name: AllUserIDs :many
+SELECT DISTINCT id FROM USERS;
+
diff --git a/enterprise/dbcrypt/cliutil.go b/enterprise/dbcrypt/cliutil.go
@@ -6,8 +6,6 @@ import (
 
 	"golang.org/x/xerrors"
 
-	"github.com/google/uuid"
-
 	"cdr.dev/slog"
 	"github.com/coder/coder/v2/coderd/database"
 )
@@ -21,7 +19,7 @@ func Rotate(ctx context.Context, log slog.Logger, sqlDB *sql.DB, ciphers []Ciphe
 		return xerrors.Errorf("create cryptdb: %w", err)
 	}
 
-	userIDs, err := allUserIDs(ctx, sqlDB)
+	userIDs, err := db.AllUserIDs(ctx)
 	if err != nil {
 		return xerrors.Errorf("get users: %w", err)
 	}
@@ -105,7 +103,7 @@ func Decrypt(ctx context.Context, log slog.Logger, sqlDB *sql.DB, ciphers []Ciph
 	}
 	cryptDB.primaryCipherDigest = ""
 
-	userIDs, err := allUserIDs(ctx, sqlDB)
+	userIDs, err := db.AllUserIDs(ctx)
 	if err != nil {
 		return xerrors.Errorf("get users: %w", err)
 	}
@@ -214,25 +212,3 @@ func Delete(ctx context.Context, log slog.Logger, sqlDB *sql.DB) error {
 
 	return nil
 }
-
-// allUserIDs returns _all_ user IDs we know about, regardless of status or deletion.
-// We need to encrypt / decrypt tokens regardless of user status or deletion as they
-// may still be valid. While we could check the expiry, we also don't know if the
-// provider is lying about expiry.
-// This function will likely only ever be used here, so keeping it here instead
-// of exposing it in all of our database-related interfaces.
-func allUserIDs(ctx context.Context, sqlDB *sql.DB) ([]uuid.UUID, error) {
-	var id uuid.UUID
-	userIDs := make([]uuid.UUID, 0)
-	rows, err := sqlDB.QueryContext(ctx, `SELECT DISTINCT id FROM users`)
-	if err != nil {
-		return nil, xerrors.Errorf("failed to query all user ids: %w", err)
-	}
-	for rows.Next() {
-		if err := rows.Scan(&id); err != nil {
-			return nil, xerrors.Errorf("failed to scan user_id: %w", err)
-		}
-		userIDs = append(userIDs, id)
-	}
-	return userIDs, nil
-}
Original file line number	Diff line number	Diff line change
Expand Up		@@ -16,3 +16,4 @@ AND
		INSERT INTO dbcrypt_keys
		(number, active_key_digest, created_at, test)
		VALUES (@number::int, @active_key_digest::text, CURRENT_TIMESTAMP, @test::text);