From 324e26716dd151bdec2184a54e7377feecc3298f Mon Sep 17 00:00:00 2001
From: Philipp Spiess <hello@philippspiess.com>
Date: Fri, 5 May 2023 19:35:02 +0200
Subject: [PATCH] Use second rate limit bucket for code completions (#51522)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Based on #51514

This only exists to make @eseliger happy

(Pointed to #51514 so the diff is readable, GitHub stacked diffs when 😢
)

## Test plan

https://user-images.githubusercontent.com/458591/236499559-336d2808-769c-4f39-a04d-f537ae28c13c.mov

<!-- All pull requests REQUIRE a test plan:
https://docs.sourcegraph.com/dev/background-information/testing_principles
-->

(cherry picked from commit 55f68f0ca0f36c0f9e8723407ab1772478208956)
---
 .../settings/quota/UserQuotaProfilePage.tsx   | 71 ++++++++++++++++++-
 client/web/src/user/settings/quota/backend.ts | 11 +++
 2 files changed, 80 insertions(+), 2 deletions(-)

diff --git a/client/web/src/user/settings/quota/UserQuotaProfilePage.tsx b/client/web/src/user/settings/quota/UserQuotaProfilePage.tsx
index a06cbb16f6b3..1fe43ca36b78 100644
--- a/client/web/src/user/settings/quota/UserQuotaProfilePage.tsx
+++ b/client/web/src/user/settings/quota/UserQuotaProfilePage.tsx
@@ -9,7 +9,7 @@ import { LoaderButton } from '../../../components/LoaderButton'
 import { PageTitle } from '../../../components/PageTitle'
 import { Scalars } from '../../../graphql-operations'
 
-import { SET_USER_COMPLETIONS_QUOTA, USER_REQUEST_QUOTAS } from './backend'
+import { SET_USER_CODE_COMPLETIONS_QUOTA, SET_USER_COMPLETIONS_QUOTA, USER_REQUEST_QUOTAS } from './backend'
 
 interface Props {
     user: {
@@ -22,6 +22,8 @@ export const UserQuotaProfilePage: React.FunctionComponent<React.PropsWithChildr
 }) => {
     const { data, loading, error } = useQuery(USER_REQUEST_QUOTAS, { variables: { userID } })
     const [quota, setQuota] = useState<string>('')
+    const [codeCompletionsQuota, setCodeCompletionsQuota] = useState<string>('')
+
     const [
         setUserCompletionsQuota,
         {
@@ -31,6 +33,15 @@ export const UserQuotaProfilePage: React.FunctionComponent<React.PropsWithChildr
         },
     ] = useMutation(SET_USER_COMPLETIONS_QUOTA)
 
+    const [
+        setUserCodeCompletionsQuota,
+        {
+            data: setCodeCompletionsQuotaResponse,
+            loading: setUserCodeCompletionsQuotaLoading,
+            error: setUserCodeCompletionsQuotaError,
+        },
+    ] = useMutation(SET_USER_CODE_COMPLETIONS_QUOTA)
+
     useEffect(() => {
         if (data?.node?.__typename === 'User' && data.node.completionsQuotaOverride !== null) {
             setQuota(data.node.completionsQuotaOverride)
@@ -38,6 +49,12 @@ export const UserQuotaProfilePage: React.FunctionComponent<React.PropsWithChildr
             // No overridden limit.
             setQuota('')
         }
+        if (data?.node?.__typename === 'User' && data.node.codeCompletionsQuotaOverride !== null) {
+            setCodeCompletionsQuota(data.node.codeCompletionsQuotaOverride)
+        } else {
+            // No overridden limit.
+            setCodeCompletionsQuota('')
+        }
     }, [data])
 
     useEffect(() => {
@@ -51,6 +68,17 @@ export const UserQuotaProfilePage: React.FunctionComponent<React.PropsWithChildr
         }
     }, [setCompletionsQuotaResponse])
 
+    useEffect(() => {
+        if (setCodeCompletionsQuotaResponse) {
+            if (setCodeCompletionsQuotaResponse.codeCompletionsQuotaOverride !== null) {
+                setCodeCompletionsQuota(setCodeCompletionsQuotaResponse.codeCompletionsQuotaOverride)
+            } else {
+                // No overridden limit.
+                setCodeCompletionsQuota('')
+            }
+        }
+    }, [setCodeCompletionsQuotaResponse])
+
     const storeCompletionsQuota = useCallback(() => {
         setUserCompletionsQuota({ variables: { userID, quota: quota === '' ? null : parseInt(quota, 10) } }).catch(
             error => {
@@ -59,6 +87,14 @@ export const UserQuotaProfilePage: React.FunctionComponent<React.PropsWithChildr
         )
     }, [quota, userID, setUserCompletionsQuota])
 
+    const storeCodeCompletionsQuota = useCallback(() => {
+        setUserCodeCompletionsQuota({
+            variables: { userID, quota: codeCompletionsQuota === '' ? null : parseInt(codeCompletionsQuota, 10) },
+        }).catch(error => {
+            logger.error(error)
+        })
+    }, [codeCompletionsQuota, userID, setUserCodeCompletionsQuota])
+
     if (loading) {
         return <LoadingSpinner />
     }
@@ -84,7 +120,7 @@ export const UserQuotaProfilePage: React.FunctionComponent<React.PropsWithChildr
             <Container className="mb-3">
                 <H3>Completions</H3>
                 <Text>Number of requests per day allowed against the completions APIs.</Text>
-                <div className="d-flex justify-content-between align-items-end">
+                <div className="d-flex justify-content-between align-items-end mb-5">
                     <Input
                         id="completions-quota"
                         name="completions-quota"
@@ -112,6 +148,37 @@ export const UserQuotaProfilePage: React.FunctionComponent<React.PropsWithChildr
                     />
                 </div>
                 {setUserCompletionsQuotaError && <ErrorAlert error={setUserCompletionsQuotaError} className="mb-0" />}
+                <Text>Number of requests per day allowed against the code completions APIs.</Text>
+                <div className="d-flex justify-content-between align-items-end">
+                    <Input
+                        id="code-completions-quota"
+                        name="code-completions-quota"
+                        type="number"
+                        value={codeCompletionsQuota}
+                        onChange={event => setCodeCompletionsQuota(event.currentTarget.value)}
+                        spellCheck={false}
+                        min={1}
+                        disabled={setUserCodeCompletionsQuotaLoading}
+                        placeholder={`Global limit: ${
+                            data?.site.perUserCodeCompletionsQuota === null
+                                ? 'infinite'
+                                : data?.site.perUserCodeCompletionsQuota
+                        }`}
+                        label="Custom code completions quota"
+                        className="flex-grow-1 mb-0"
+                    />
+                    <LoaderButton
+                        loading={setUserCodeCompletionsQuotaLoading}
+                        label="Save"
+                        onClick={storeCodeCompletionsQuota}
+                        disabled={setUserCodeCompletionsQuotaLoading}
+                        variant="primary"
+                        className="ml-2"
+                    />
+                </div>
+                {setUserCodeCompletionsQuotaError && (
+                    <ErrorAlert error={setUserCodeCompletionsQuotaError} className="mb-0" />
+                )}
             </Container>
         </>
     )
diff --git a/client/web/src/user/settings/quota/backend.ts b/client/web/src/user/settings/quota/backend.ts
index b0e356a81036..adf9b36e442b 100644
--- a/client/web/src/user/settings/quota/backend.ts
+++ b/client/web/src/user/settings/quota/backend.ts
@@ -4,11 +4,13 @@ export const USER_REQUEST_QUOTAS = gql`
     query UserRequestQuotas($userID: ID!) {
         site {
             perUserCompletionsQuota
+            perUserCodeCompletionsQuota
         }
         node(id: $userID) {
             __typename
             ... on User {
                 completionsQuotaOverride
+                codeCompletionsQuotaOverride
             }
         }
     }
@@ -22,3 +24,12 @@ export const SET_USER_COMPLETIONS_QUOTA = gql`
         }
     }
 `
+
+export const SET_USER_CODE_COMPLETIONS_QUOTA = gql`
+    mutation SetUserCodeCompletionsQuota($userID: ID!, $quota: Int) {
+        setUserCodeCompletionsQuota(user: $userID, quota: $quota) {
+            id
+            codeCompletionsQuotaOverride
+        }
+    }
+`