diff --git a/controlplane/src/core/controllers/auth.ts b/controlplane/src/core/controllers/auth.ts index c49688d4fe..448e482410 100644 --- a/controlplane/src/core/controllers/auth.ts +++ b/controlplane/src/core/controllers/auth.ts @@ -1,7 +1,7 @@ import { FastifyPluginCallback } from 'fastify'; import fp from 'fastify-plugin'; import { PostgresJsDatabase } from 'drizzle-orm/postgres-js'; -import { eq } from 'drizzle-orm'; +import { eq, sql } from 'drizzle-orm'; import { lru } from 'tiny-lru'; import { uid } from 'uid'; import { PlatformEventName } from '@wundergraph/cosmo-connect/dist/notifications/events_pb'; @@ -255,64 +255,97 @@ const plugin: FastifyPluginCallback = function Auth(fasti return insertedSessions[0]; }); - const orgs = await opts.organizationRepository.memberships({ - userId, - }); + const orgs = await opts.db.transaction(async (tx) => { + const advisoryLockRows = await tx.execute( + sql`select pg_try_advisory_xact_lock(hashtext(${userId})) as acquired`, + ); + + if (!advisoryLockRows?.[0]?.acquired) { + // We need to identify when we failed to acquire the lock because another request already acquired it + return -1; + } - if (orgs.length === 0) { + const orgRepo = new OrganizationRepository(req.log, tx, opts.defaultBillingPlanId); + + // Check if the user is already a member of at least one organization + const existingMemberships = await tx + .select({ one: sql`1`.as('one') }) + .from(organizationsMembers) + .where(eq(organizationsMembers.userId, userId)) + .limit(1) + .execute(); + + if (existingMemberships.length > 0) { + return existingMemberships.length; + } + + // Authenticate on Keycloak and create the organization group await opts.keycloakClient.authenticateClient(); const organizationSlug = uid(8); - const [kcRootGroupId, kcCreatedGroups] = await opts.keycloakClient.seedGroup({ userID: userId, organizationSlug, realm: opts.keycloakRealm, }); - await opts.db.transaction(async (tx) => { - const orgRepo = new OrganizationRepository(req.log, tx, opts.defaultBillingPlanId); - const orgGroupRepo = new OrganizationGroupRepository(tx); - - const insertedOrg = await orgRepo.createOrganization({ - organizationName: userEmail.split('@')[0], - organizationSlug, - ownerID: userId, - kcGroupId: kcRootGroupId, - }); + // Create the new organization and add the user as a member of the organization + const insertedOrg = await orgRepo.createOrganization({ + organizationName: userEmail.split('@')[0], + organizationSlug, + ownerID: userId, + kcGroupId: kcRootGroupId, + }); - const orgMember = await orgRepo.addOrganizationMember({ - organizationID: insertedOrg.id, - userID: userId, - }); + const orgMember = await orgRepo.addOrganizationMember({ + organizationID: insertedOrg.id, + userID: userId, + }); - await orgGroupRepo.importKeycloakGroups({ - organizationId: insertedOrg.id, - kcGroups: kcCreatedGroups, - }); + // Create the organization groups + const orgGroupRepo = new OrganizationGroupRepository(tx); - const orgAdminGroup = await orgGroupRepo.byName({ - organizationId: insertedOrg.id, - name: 'admin', - }); + await orgGroupRepo.importKeycloakGroups({ + organizationId: insertedOrg.id, + kcGroups: kcCreatedGroups, + }); - if (orgAdminGroup) { - await orgGroupRepo.addUserToGroup({ - organizationMemberId: orgMember.id, - groupId: orgAdminGroup.groupId, - }); - } + const orgAdminGroup = await orgGroupRepo.byName({ + organizationId: insertedOrg.id, + name: 'admin', + }); - const namespaceRepo = new NamespaceRepository(tx, insertedOrg.id); - const ns = await namespaceRepo.create({ - name: DefaultNamespace, - createdBy: userId, + if (orgAdminGroup) { + await orgGroupRepo.addUserToGroup({ + organizationMemberId: orgMember.id, + groupId: orgAdminGroup.groupId, }); - if (!ns) { - throw new Error(`Could not create ${DefaultNamespace} namespace`); - } + } + + // Create the default namespace for the organization + const namespaceRepo = new NamespaceRepository(tx, insertedOrg.id); + const ns = await namespaceRepo.create({ + name: DefaultNamespace, + createdBy: userId, }); + if (!ns) { + throw new Error(`Could not create ${DefaultNamespace} namespace`); + } + + // We return an empty even when we just created the organization, that way we can still send the + // user registered webhook and prompt the user to migrate + return 0; + }); + + if (orgs === -1) { + // We failed to acquire the lock, so we need to retry the request + await res.code(429).send('Slow down'); + return; + } + + if (orgs === 0) { + // Send a notification to the platform that a new user has been created opts.platformWebhooks.send(PlatformEventName.USER_REGISTER_SUCCESS, { user_id: userId, user_email: userEmail, @@ -343,7 +376,7 @@ const plugin: FastifyPluginCallback = function Auth(fasti } else { res.redirect(opts.webBaseUrl); } - } else if (orgs.length === 0) { + } else if (orgs === 0) { res.redirect(opts.webBaseUrl + '?migrate=true'); } else { res.redirect(opts.webBaseUrl); diff --git a/controlplane/src/core/repositories/OrganizationRepository.ts b/controlplane/src/core/repositories/OrganizationRepository.ts index 014ce4d045..0327f194c3 100644 --- a/controlplane/src/core/repositories/OrganizationRepository.ts +++ b/controlplane/src/core/repositories/OrganizationRepository.ts @@ -95,12 +95,13 @@ export class OrganizationRepository { return org; } - public async updateOrganization(input: { id: string; slug?: string; name?: string }) { + public async updateOrganization(input: { id: string; slug?: string; name?: string; kcGroupId?: string }) { await this.db .update(organizations) .set({ name: input.name, slug: input.slug, + kcGroupId: input.kcGroupId, }) .where(eq(organizations.id, input.id)) .execute(); diff --git a/controlplane/test/utils.test.ts b/controlplane/test/utils.test.ts index aa6f5a6a19..a9cb2cb584 100644 --- a/controlplane/test/utils.test.ts +++ b/controlplane/test/utils.test.ts @@ -1,5 +1,10 @@ import { describe, expect, test } from 'vitest'; -import { isValidLabelMatchers, mergeUrls, normalizeLabelMatchers, isGoogleCloudStorageUrl } from '../src/core/util.js'; +import { + isValidLabelMatchers, + mergeUrls, + normalizeLabelMatchers, + isGoogleCloudStorageUrl, +} from '../src/core/util.js'; describe('Utils', () => { test('isValidLabelMatchers', () => {