From 6ecfa62eb70bf60b7eeded3205ddd54684554082 Mon Sep 17 00:00:00 2001
From: Joshen Lim <joshenlimek@gmail.com>
Date: Tue, 10 Jun 2025 19:05:18 +0800
Subject: [PATCH] Manually updated org query data after creating a new org
 (#36301)

---
 .../Organization/NewOrg/NewOrgForm.tsx        |  8 +------
 .../organization-create-mutation.ts           | 21 ++++++++++++++-----
 2 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/apps/studio/components/interfaces/Organization/NewOrg/NewOrgForm.tsx b/apps/studio/components/interfaces/Organization/NewOrg/NewOrgForm.tsx
index 1e7b3555c4b..2f16d232697 100644
--- a/apps/studio/components/interfaces/Organization/NewOrg/NewOrgForm.tsx
+++ b/apps/studio/components/interfaces/Organization/NewOrg/NewOrgForm.tsx
@@ -1,6 +1,5 @@
 import { PaymentElement, useElements, useStripe } from '@stripe/react-stripe-js'
 import type { PaymentMethod } from '@stripe/stripe-js'
-import { useQueryClient } from '@tanstack/react-query'
 import _ from 'lodash'
 import { Edit2, ExternalLink, HelpCircle } from 'lucide-react'
 import Link from 'next/link'
@@ -14,10 +13,7 @@ import { LOCAL_STORAGE_KEYS } from 'common'
 import SpendCapModal from 'components/interfaces/Billing/SpendCapModal'
 import Panel from 'components/ui/Panel'
 import { useOrganizationCreateMutation } from 'data/organizations/organization-create-mutation'
-import {
-  invalidateOrganizationsQuery,
-  useOrganizationsQuery,
-} from 'data/organizations/organizations-query'
+import { useOrganizationsQuery } from 'data/organizations/organizations-query'
 import { useProjectsQuery } from 'data/projects/projects-query'
 import { useLocalStorageQuery } from 'hooks/misc/useLocalStorage'
 import { BASE_PATH, PRICING_TIER_LABELS_ORG } from 'lib/constants'
@@ -94,7 +90,6 @@ const NewOrgForm = ({ onPaymentMethodReset }: NewOrgFormProps) => {
   const { data: projects } = useProjectsQuery()
   const stripe = useStripe()
   const elements = useElements()
-  const queryClient = useQueryClient()
 
   const [lastVisitedOrganization] = useLocalStorageQuery(
     LOCAL_STORAGE_KEYS.LAST_VISITED_ORGANIZATION,
@@ -156,7 +151,6 @@ const NewOrgForm = ({ onPaymentMethodReset }: NewOrgFormProps) => {
 
   const { mutate: createOrganization } = useOrganizationCreateMutation({
     onSuccess: async (org) => {
-      await invalidateOrganizationsQuery(queryClient)
       const prefilledProjectName = user.profile?.username
         ? user.profile.username + `'s Project`
         : 'My Project'
diff --git a/apps/studio/data/organizations/organization-create-mutation.ts b/apps/studio/data/organizations/organization-create-mutation.ts
index acd8c924c12..f3b20c8d821 100644
--- a/apps/studio/data/organizations/organization-create-mutation.ts
+++ b/apps/studio/data/organizations/organization-create-mutation.ts
@@ -1,11 +1,11 @@
 import { useMutation, UseMutationOptions, useQueryClient } from '@tanstack/react-query'
 import { toast } from 'sonner'
 
+import type { components } from 'data/api'
 import { handleError, post } from 'data/fetchers'
 import { permissionKeys } from 'data/permissions/keys'
 import type { ResponseError } from 'types'
 import { organizationKeys } from './keys'
-import type { components } from 'data/api'
 
 export type OrganizationCreateVariables = {
   name: string
@@ -61,10 +61,21 @@ export const useOrganizationCreateMutation = ({
     (vars) => createOrganization(vars),
     {
       async onSuccess(data, variables, context) {
-        await Promise.all([
-          queryClient.invalidateQueries(organizationKeys.list()),
-          queryClient.invalidateQueries(permissionKeys.list()),
-        ])
+        // [Joshen] We're manually updating the query client here as the org's subscription is
+        // created async, and the invalidation will happen too quick where the GET organizations
+        // endpoint will error out with a 500 since the subscription isn't created yet.
+        queryClient.setQueriesData(
+          {
+            queryKey: organizationKeys.list(),
+            exact: true,
+          },
+          (prev: any) => {
+            if (!prev) return prev
+            return [...prev, data]
+          }
+        )
+
+        await queryClient.invalidateQueries(permissionKeys.list())
         await onSuccess?.(data, variables, context)
       },
       async onError(data, variables, context) {