perf: optimize prebuilds membership reconciliation to check orgs not presets (coder#20493) (coder#20555)

Related to PR: coder#20493

(cherry picked from commit 7e8fcb4)

Author: ssncferreira

coderd/database/dbauthz/dbauthz.go

@@ -2649,6 +2649,13 @@ func (q *querier) GetOrganizationsByUserID(ctx context.Context, userID database.
 	return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.GetOrganizationsByUserID)(ctx, userID)
 }
 
+func (q *querier) GetOrganizationsWithPrebuildStatus(ctx context.Context, arg database.GetOrganizationsWithPrebuildStatusParams) ([]database.GetOrganizationsWithPrebuildStatusRow, error) {
+	if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceOrganization.All()); err != nil {
+		return nil, err
+	}
+	return q.db.GetOrganizationsWithPrebuildStatus(ctx, arg)
+}
+
 func (q *querier) GetParameterSchemasByJobID(ctx context.Context, jobID uuid.UUID) ([]database.ParameterSchema, error) {
 	version, err := q.db.GetTemplateVersionByJobID(ctx, jobID)
 	if err != nil {

coderd/database/dbauthz/dbauthz_test.go

@@ -3759,6 +3759,14 @@ func (s *MethodTestSuite) TestPrebuilds() {
 		dbm.EXPECT().GetPrebuildMetrics(gomock.Any()).Return([]database.GetPrebuildMetricsRow{}, nil).AnyTimes()
 		check.Args().Asserts(rbac.ResourceWorkspace.All(), policy.ActionRead)
 	}))
+	s.Run("GetOrganizationsWithPrebuildStatus", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
+		arg := database.GetOrganizationsWithPrebuildStatusParams{
+			UserID:    uuid.New(),
+			GroupName: "test",
+		}
+		dbm.EXPECT().GetOrganizationsWithPrebuildStatus(gomock.Any(), arg).Return([]database.GetOrganizationsWithPrebuildStatusRow{}, nil).AnyTimes()
+		check.Args(arg).Asserts(rbac.ResourceOrganization.All(), policy.ActionRead)
+	}))
 	s.Run("GetPrebuildsSettings", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
 		dbm.EXPECT().GetPrebuildsSettings(gomock.Any()).Return("{}", nil).AnyTimes()
 		check.Args().Asserts()

coderd/database/dbmetrics/querymetrics.go

@@ -327,3 +327,47 @@ SET
 FROM jobs_to_cancel
 WHERE provisioner_jobs.id = jobs_to_cancel.id
 RETURNING jobs_to_cancel.id, jobs_to_cancel.workspace_id, jobs_to_cancel.template_id, jobs_to_cancel.template_version_preset_id;
+
+-- name: GetOrganizationsWithPrebuildStatus :many
+-- GetOrganizationsWithPrebuildStatus returns organizations with prebuilds configured and their
+-- membership status for the prebuilds system user (org membership, group existence, group membership).
+WITH orgs_with_prebuilds AS (
+	-- Get unique organizations that have presets with prebuilds configured
+	SELECT DISTINCT o.id, o.name
+	FROM organizations o
+	INNER JOIN templates t ON t.organization_id = o.id
+	INNER JOIN template_versions tv ON tv.template_id = t.id
+	INNER JOIN template_version_presets tvp ON tvp.template_version_id = tv.id
+	WHERE tvp.desired_instances IS NOT NULL
+),
+prebuild_user_membership AS (
+	-- Check if the user is a member of the organizations
+	SELECT om.organization_id
+	FROM organization_members om
+	INNER JOIN orgs_with_prebuilds owp ON owp.id = om.organization_id
+	WHERE om.user_id = @user_id::uuid
+),
+prebuild_groups AS (
+	-- Check if the organizations have the prebuilds group
+	SELECT g.organization_id, g.id as group_id
+	FROM groups g
+	INNER JOIN orgs_with_prebuilds owp ON owp.id = g.organization_id
+	WHERE g.name = @group_name::text
+),
+prebuild_group_membership AS (
+	-- Check if the user is in the prebuilds group
+	SELECT pg.organization_id
+	FROM prebuild_groups pg
+	INNER JOIN group_members gm ON gm.group_id = pg.group_id
+	WHERE gm.user_id = @user_id::uuid
+)
+SELECT
+	owp.id AS organization_id,
+	owp.name AS organization_name,
+	(pum.organization_id IS NOT NULL)::boolean AS has_prebuild_user,
+	pg.group_id AS prebuilds_group_id,
+	(pgm.organization_id IS NOT NULL)::boolean AS has_prebuild_user_in_group
+FROM orgs_with_prebuilds owp
+LEFT JOIN prebuild_groups pg ON pg.organization_id = owp.id
+LEFT JOIN prebuild_user_membership pum ON pum.organization_id = owp.id
+LEFT JOIN prebuild_group_membership pgm ON pgm.organization_id = owp.id;