From 4bd241af271d1e8a6f89ef35c665424011d60a3b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kevin=20Gr=C3=BCneberg?= <k.grueneberg1994@gmail.com>
Date: Wed, 13 Sep 2023 00:21:53 +0200
Subject: [PATCH] feat: over-usage banners for orgs

---
 .../BillingBreakdown/BillingBreakdown.tsx     | 13 ++++-
 .../interfaces/Organization/Usage/Usage.tsx   | 40 ++++++++++++++-
 .../LayoutHeader/LayoutHeader.tsx             | 49 ++++++++++++++++---
 .../ui/OveragesBanner/OveragesBanner.utils.ts | 17 ++++++-
 4 files changed, 108 insertions(+), 11 deletions(-)

diff --git a/studio/components/interfaces/Organization/BillingSettingsV2/BillingBreakdown/BillingBreakdown.tsx b/studio/components/interfaces/Organization/BillingSettingsV2/BillingBreakdown/BillingBreakdown.tsx
index 44049b6780a..c806533ad1b 100644
--- a/studio/components/interfaces/Organization/BillingSettingsV2/BillingBreakdown/BillingBreakdown.tsx
+++ b/studio/components/interfaces/Organization/BillingSettingsV2/BillingBreakdown/BillingBreakdown.tsx
@@ -14,8 +14,10 @@ import { Alert, Button } from 'ui'
 import { BILLING_BREAKDOWN_METRICS } from './BillingBreakdown.constants'
 import BillingMetric from './BillingMetric'
 import UpcomingInvoice from './UpcomingInvoice'
+import { useOrgSettingsPageStateSnapshot } from 'state/organization-settings'
 
 const BillingBreakdown = () => {
+  const snap = useOrgSettingsPageStateSnapshot()
   const { slug: orgSlug } = useParams()
   const {
     data: usage,
@@ -93,7 +95,16 @@ const BillingBreakdown = () => {
                 variant="danger"
                 title="Your organization's usage has exceeded its included quota"
                 actions={[
-                  <Button key="upgrade-button" type="default" className="ml-8" onClick={() => {}}>
+                  <Button
+                    key="upgrade-button"
+                    type="default"
+                    className="ml-8"
+                    onClick={() =>
+                      snap.setPanelKey(
+                        currentPlan?.id === 'free' ? 'subscriptionPlan' : 'costControl'
+                      )
+                    }
+                  >
                     {currentPlan?.id === 'free' ? 'Upgrade plan' : 'Change spend cap'}
                   </Button>,
                 ]}
diff --git a/studio/components/interfaces/Organization/Usage/Usage.tsx b/studio/components/interfaces/Organization/Usage/Usage.tsx
index 66157592ef6..34942cc2873 100644
--- a/studio/components/interfaces/Organization/Usage/Usage.tsx
+++ b/studio/components/interfaces/Organization/Usage/Usage.tsx
@@ -11,12 +11,13 @@ import { useProjectsQuery } from 'data/projects/projects-query'
 import { useOrgSubscriptionQuery } from 'data/subscriptions/org-subscription-query'
 import { useSelectedOrganization } from 'hooks'
 import { TIME_PERIODS_BILLING, TIME_PERIODS_REPORTS } from 'lib/constants'
-import { Button, IconExternalLink, IconInfo, Listbox } from 'ui'
+import { Alert, Button, IconExternalLink, IconInfo, Listbox } from 'ui'
 import Activity from './Activity'
 import Bandwidth from './Bandwidth'
 import SizeAndCounts from './SizeAndCounts'
 import InformationBox from 'components/ui/InformationBox'
 import Link from 'next/link'
+import { useOrgUsageQuery } from 'data/usage/org-usage-query'
 
 const Usage = () => {
   const { slug, projectRef } = useParams()
@@ -32,6 +33,9 @@ const Usage = () => {
     isError: isErrorSubscription,
     isSuccess: isSuccessSubscription,
   } = useOrgSubscriptionQuery({ orgSlug: slug })
+
+  const { data: usage } = useOrgUsageQuery({ orgSlug: slug })
+
   const orgProjects = projects?.filter((project) => project.organization_id === organization?.id)
 
   useEffect(() => {
@@ -88,6 +92,13 @@ const Usage = () => {
     ? orgProjects?.find((it) => it.ref === selectedProjectRef)
     : undefined
 
+  const hasExceededAnyLimits = Boolean(
+    usage?.usages.find(
+      (metric) =>
+        !metric.unlimited && metric.capped && metric.usage > (metric?.pricing_free_units ?? 0)
+    )
+  )
+
   return (
     <>
       <ScaffoldContainer className="sticky top-0 border-b bg-scale-200 z-10 overflow-hidden">
@@ -156,6 +167,33 @@ const Usage = () => {
         </div>
       </ScaffoldContainer>
 
+      {!selectedProject && subscription && hasExceededAnyLimits && (
+        <ScaffoldContainer className="mt-5">
+          <Alert
+            withIcon
+            variant="danger"
+            title="Your organization's usage has exceeded its included quota"
+            actions={[
+              <Link
+                key="upgrade-button"
+                href={`/org/${slug}/billing?panel=${
+                  subscription.plan.id === 'free' ? 'subscriptionPlan' : 'costControl'
+                }`}
+              >
+                <Button asChild key="upgrade-button" type="default" className="ml-8">
+                  <a>{subscription.plan.id === 'free' ? 'Upgrade plan' : 'Change spend cap'}</a>
+                </Button>
+              </Link>,
+            ]}
+          >
+            Your projects can become unresponsive or enter read only mode.{' '}
+            {subscription.plan.id === 'free'
+              ? 'Please upgrade to the Pro plan to ensure that your projects remain available.'
+              : 'Please disable spend cap to ensure that your projects remain available.'}
+          </Alert>
+        </ScaffoldContainer>
+      )}
+
       {selectedProjectRef && (
         <ScaffoldContainer className="mt-5">
           <InformationBox
diff --git a/studio/components/layouts/ProjectLayout/LayoutHeader/LayoutHeader.tsx b/studio/components/layouts/ProjectLayout/LayoutHeader/LayoutHeader.tsx
index 05088113ee6..3d9ae98cdea 100644
--- a/studio/components/layouts/ProjectLayout/LayoutHeader/LayoutHeader.tsx
+++ b/studio/components/layouts/ProjectLayout/LayoutHeader/LayoutHeader.tsx
@@ -2,7 +2,10 @@ import { useParams } from 'common'
 import Link from 'next/link'
 import { Badge } from 'ui'
 
-import { getResourcesExceededLimits } from 'components/ui/OveragesBanner/OveragesBanner.utils'
+import {
+  getResourcesExceededLimits,
+  getResourcesExceededLimitsOrg,
+} from 'components/ui/OveragesBanner/OveragesBanner.utils'
 import { useProjectReadOnlyQuery } from 'data/config/project-read-only-query'
 import { useProjectUsageQuery } from 'data/usage/project-usage-query'
 import { useFlag, useSelectedOrganization, useSelectedProject } from 'hooks'
@@ -14,6 +17,8 @@ import NotificationsPopover from './NotificationsPopover'
 import OrgDropdown from './OrgDropdown'
 import ProjectDropdown from './ProjectDropdown'
 import { useProjectSubscriptionV2Query } from 'data/subscriptions/project-subscription-v2-query'
+import { useOrgUsageQuery } from 'data/usage/org-usage-query'
+import { useMemo } from 'react'
 
 const LayoutHeader = ({ customHeaderComponents, breadcrumbs = [], headerBorder = true }: any) => {
   const selectedOrganization = useSelectedOrganization()
@@ -26,26 +31,48 @@ const LayoutHeader = ({ customHeaderComponents, breadcrumbs = [], headerBorder =
   })
 
   // Skip with org-based-billing, as quota is for the entire org
-  const { data: usage } = useProjectUsageQuery(
+  const { data: projectUsage } = useProjectUsageQuery(
     { projectRef },
-    { enabled: selectedOrganization && !selectedOrganization.subscription_id }
+    { enabled: Boolean(selectedOrganization && !selectedOrganization.subscription_id) }
   )
-  const resourcesExceededLimits = getResourcesExceededLimits(usage)
+  const { data: orgUsage } = useOrgUsageQuery(
+    { orgSlug: selectedOrganization?.slug },
+    { enabled: Boolean(selectedOrganization && selectedOrganization.subscription_id) }
+  )
+
+  const exceedingLimits = useMemo(() => {
+    if (orgUsage) {
+      return getResourcesExceededLimitsOrg(orgUsage?.usages || []).length > 0
+    } else if (projectUsage) {
+      return getResourcesExceededLimits(projectUsage).length > 0
+    } else {
+      return false
+    }
+  }, [projectUsage, orgUsage])
 
   // Skip with org-based-billing, as quota is for the entire org
-  const { data: subscription } = useProjectSubscriptionV2Query(
+  const { data: projectSubscription } = useProjectSubscriptionV2Query(
     { projectRef },
     { enabled: selectedOrganization && !selectedOrganization.subscription_id }
   )
 
+  const { data: orgSubscription } = useProjectSubscriptionV2Query(
+    { projectRef },
+    { enabled: Boolean(selectedOrganization && selectedOrganization.subscription_id) }
+  )
+
+  const subscription = useMemo(() => {
+    return projectSubscription || orgSubscription
+  }, [projectSubscription, orgSubscription])
+
   const projectHasNoLimits = subscription?.usage_billing_enabled === true
 
   const showOverUsageBadge =
     useFlag('overusageBadge') &&
     subscription !== undefined &&
-    (subscription.plan.id === 'free' || subscription?.plan.id === 'pro') &&
+    (subscription.plan.id === 'free' || subscription.plan.id === 'pro') &&
     !projectHasNoLimits &&
-    resourcesExceededLimits.length > 0
+    exceedingLimits
 
   return (
     <div
@@ -95,7 +122,13 @@ const LayoutHeader = ({ customHeaderComponents, breadcrumbs = [], headerBorder =
 
                 {showOverUsageBadge && (
                   <div className="ml-2">
-                    <Link href={`/project/${projectRef}/settings/billing/usage`}>
+                    <Link
+                      href={
+                        selectedOrganization?.subscription_id
+                          ? `/org/${selectedOrganization.slug}/usage`
+                          : `/project/${projectRef}/settings/billing/usage`
+                      }
+                    >
                       <a>
                         <Badge color="red">Exceeding usage limits</Badge>
                       </a>
diff --git a/studio/components/ui/OveragesBanner/OveragesBanner.utils.ts b/studio/components/ui/OveragesBanner/OveragesBanner.utils.ts
index 6c66916241f..69438ce4b20 100644
--- a/studio/components/ui/OveragesBanner/OveragesBanner.utils.ts
+++ b/studio/components/ui/OveragesBanner/OveragesBanner.utils.ts
@@ -1,5 +1,6 @@
 import { compact } from 'lodash'
 import { USAGE_APPROACHING_THRESHOLD } from 'lib/constants'
+import { OrgUsageResponse, UsageMetric } from 'data/usage/org-usage-query'
 
 export const getResourcesApproachingLimits = (usages: any) => {
   if (!usages) return []
@@ -22,7 +23,21 @@ export const getResourcesExceededLimits = (usages: any) => {
       .filter((resourceName) => usages[resourceName] !== null)
       .map((resourceName) => {
         const resource = usages[resourceName]
-        if (resource.limit > 0 && (resource.usage / resource.limit) > 1) return resourceName
+        if (resource.limit > 0 && resource.usage / resource.limit > 1) return resourceName
       })
   )
 }
+
+export const getResourcesExceededLimitsOrg = (usageMetrics: UsageMetric[]): string[] => {
+  if (!usageMetrics.length) return []
+
+  return usageMetrics
+    .filter((usageMetric) => {
+      if (!usageMetric.capped || !usageMetric.available_in_plan || usageMetric.unlimited) return false
+
+      const freeUnits = usageMetric.pricing_free_units || 0
+
+      return usageMetric.usage > freeUnits
+    })
+    .map((it) => it.metric)
+}