Refactor entitlements

coder · BrunoQuaresma · Sep 14, 2023 · Sep 11, 2023 · Sep 11, 2023 · Sep 11, 2023
commit d34f4cf8fbe897316af088096eaabb8b53313afd
diff --git a/site/src/api/queries/entitlements.ts b/site/src/api/queries/entitlements.ts
@@ -0,0 +1,46 @@
+import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";
+import * as API from "api/api";
+
+const ENTITLEMENTS_QUERY_KEY = ["entitlements"];
+
+export const useEntitlements = () => {
+  return useQuery({
+    queryKey: ENTITLEMENTS_QUERY_KEY,
+    queryFn: fetchEntitlements,
+  });
+};
+
+export const useRefreshEntitlements = ({
+  onSuccess,
+  onError,
+}: {
+  onSuccess: () => void;
+  onError: (error: unknown) => void;
+}) => {
+  const queryClient = useQueryClient();
+  return useMutation({
+    mutationFn: API.refreshEntitlements,
+    onSuccess: async () => {
+      await queryClient.invalidateQueries({
+        queryKey: ENTITLEMENTS_QUERY_KEY,
+      });
+      onSuccess();
+    },
+    onError,
+  });
+};
+
+const fetchEntitlements = () => {
+  // Entitlements is injected by the Coder server into the HTML document.
+  const entitlements = document.querySelector("meta[property=entitlements]");
+  if (entitlements) {
+    const rawContent = entitlements.getAttribute("content");
+    try {
+      return JSON.parse(rawContent as string);
+    } catch (e) {
+      console.warn("Failed to parse entitlements from document", e);
+    }
+  }
+
+  return API.getEntitlements();
+};
diff --git a/site/src/components/Dashboard/DashboardProvider.tsx b/site/src/components/Dashboard/DashboardProvider.tsx
@@ -1,5 +1,6 @@
 import { useMachine } from "@xstate/react";
 import { useBuildInfo } from "api/queries/buildInfo";
+import { useEntitlements } from "api/queries/entitlements";
 import {
   AppearanceConfig,
   BuildInfoResponse,
@@ -9,7 +10,6 @@ import {
 import { FullScreenLoader } from "components/Loader/FullScreenLoader";
 import { createContext, FC, PropsWithChildren, useContext } from "react";
 import { appearanceMachine } from "xServices/appearance/appearanceXService";
-import { entitlementsMachine } from "xServices/entitlements/entitlementsXService";
 import { experimentsMachine } from "xServices/experiments/experimentsMachine";
 
 interface Appearance {
@@ -32,14 +32,16 @@ export const DashboardProviderContext = createContext<
 
 export const DashboardProvider: FC<PropsWithChildren> = ({ children }) => {
   const buildInfoQuery = useBuildInfo();
-  const [entitlementsState] = useMachine(entitlementsMachine);
+  const entitlementsQuery = useEntitlements();
   const [appearanceState, appearanceSend] = useMachine(appearanceMachine);
   const [experimentsState] = useMachine(experimentsMachine);
-  const { entitlements } = entitlementsState.context;
   const { appearance, preview } = appearanceState.context;
   const { experiments } = experimentsState.context;
   const isLoading =
-    !buildInfoQuery.data || !entitlements || !appearance || !experiments;
+    !buildInfoQuery.data ||
+    !entitlementsQuery.data ||
+    !appearance ||
+    !experiments;
 
   const setAppearancePreview = (config: AppearanceConfig) => {
     appearanceSend({
@@ -63,7 +65,7 @@ export const DashboardProvider: FC<PropsWithChildren> = ({ children }) => {
     <DashboardProviderContext.Provider
       value={{
         buildInfo: buildInfoQuery.data,
-        entitlements,
+        entitlements: entitlementsQuery.data,
         experiments,
         appearance: {
           preview,

diff --git a/site/src/hooks/useFeatureVisibility.ts b/site/src/hooks/useFeatureVisibility.ts
@@ -1,6 +1,6 @@
 import { FeatureName } from "api/typesGenerated";
 import { useDashboard } from "components/Dashboard/DashboardProvider";
-import { selectFeatureVisibility } from "xServices/entitlements/entitlementsSelectors";
+import { selectFeatureVisibility } from "utils/entitlements";
 
 export const useFeatureVisibility = (): Record<FeatureName, boolean> => {
   const { entitlements } = useDashboard();

diff --git a/site/src/pages/DeploySettingsPage/LicensesSettingsPage/LicensesSettingsPage.tsx b/site/src/pages/DeploySettingsPage/LicensesSettingsPage/LicensesSettingsPage.tsx
@@ -1,28 +1,40 @@
 import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";
-import { useMachine } from "@xstate/react";
 import { getLicenses, removeLicense } from "api/api";
 import { displayError, displaySuccess } from "components/GlobalSnackbar/utils";
 import { FC, useEffect } from "react";
 import { Helmet } from "react-helmet-async";
 import { useSearchParams } from "react-router-dom";
 import useToggle from "react-use/lib/useToggle";
 import { pageTitle } from "utils/page";
-import { entitlementsMachine } from "xServices/entitlements/entitlementsXService";
 import LicensesSettingsPageView from "./LicensesSettingsPageView";
 import { getErrorMessage } from "api/errors";
+import {
+  useEntitlements,
+  useRefreshEntitlements,
+} from "api/queries/entitlements";
 
 const LicensesSettingsPage: FC = () => {
   const queryClient = useQueryClient();
-  const [entitlementsState, sendEvent] = useMachine(entitlementsMachine);
-  const { entitlements, getEntitlementsError } = entitlementsState.context;
   const [searchParams, setSearchParams] = useSearchParams();
   const success = searchParams.get("success");
   const [confettiOn, toggleConfettiOn] = useToggle(false);
-  if (getEntitlementsError) {
-    displayError(
-      getErrorMessage(getEntitlementsError, "Failed to fetch entitlements"),
-    );
-  }
+  const entitlements = useEntitlements();
+  const refreshEntitlements = useRefreshEntitlements({
+    onSuccess: () => {
+      displaySuccess("Successfully refreshed licenses");
+    },
+    onError: (error) => {
+      displayError(getErrorMessage(error, "Failed to refresh entitlements"));
+    },
+  });
+
+  useEffect(() => {
+    if (entitlements.error) {
+      displayError(
+        getErrorMessage(entitlements.error, "Failed to fetch entitlements"),
+      );
+    }
+  }, [entitlements.error]);
 
   const { mutate: removeLicenseApi, isLoading: isRemovingLicense } =
     useMutation(removeLicense, {
@@ -59,14 +71,14 @@ const LicensesSettingsPage: FC = () => {
       <LicensesSettingsPageView
         showConfetti={confettiOn}
         isLoading={isLoading}
-        userLimitActual={entitlements?.features.user_limit.actual}
-        userLimitLimit={entitlements?.features.user_limit.limit}
+        isRefreshing={refreshEntitlements.isLoading}
+        userLimitActual={entitlements.data?.features.user_limit.actual}
+        userLimitLimit={entitlements.data?.features.user_limit.limit}
         licenses={licenses}
         isRemovingLicense={isRemovingLicense}
         removeLicense={(licenseId: number) => removeLicenseApi(licenseId)}
         refreshEntitlements={() => {
-          const x = sendEvent("REFRESH");
-          return !x.context.getEntitlementsError;
+          refreshEntitlements.mutate();
         }}
       />
     </>

diff --git a/site/src/pages/DeploySettingsPage/LicensesSettingsPage/LicensesSettingsPageView.tsx b/site/src/pages/DeploySettingsPage/LicensesSettingsPage/LicensesSettingsPageView.tsx
@@ -12,8 +12,8 @@ import Confetti from "react-confetti";
 import { Link } from "react-router-dom";
 import useWindowSize from "react-use/lib/useWindowSize";
 import MuiLink from "@mui/material/Link";
-import { displaySuccess } from "components/GlobalSnackbar/utils";
 import Tooltip from "@mui/material/Tooltip";
+import { LoadingButton } from "components/LoadingButton/LoadingButton";
 
 type Props = {
   showConfetti: boolean;
@@ -22,8 +22,9 @@ type Props = {
   userLimitLimit?: number;
   licenses?: GetLicensesResponse[];
   isRemovingLicense: boolean;
+  isRefreshing: boolean;
   removeLicense: (licenseId: number) => void;
-  refreshEntitlements?: () => boolean;
+  refreshEntitlements: () => void;
 };
 
 const LicensesSettingsPageView: FC<Props> = ({
@@ -33,6 +34,7 @@ const LicensesSettingsPageView: FC<Props> = ({
   userLimitLimit,
   licenses,
   isRemovingLicense,
+  isRefreshing,
   removeLicense,
   refreshEntitlements,
 }) => {
@@ -69,18 +71,13 @@ const LicensesSettingsPageView: FC<Props> = ({
             Add a license
           </Button>
           <Tooltip title="Refresh license entitlements. This is done automatically every 10 minutes.">
-            <Button
-              onClick={() => {
-                if (refreshEntitlements) {
-                  if (refreshEntitlements()) {
-                    displaySuccess("Successfully refreshed licenses");
-                  }
-                }
-              }}
+            <LoadingButton
+              loading={isRefreshing}
+              onClick={refreshEntitlements}
               startIcon={<RefreshIcon />}
             >
               Refresh
-            </Button>
+            </LoadingButton>
           </Tooltip>
         </Stack>
       </Stack>

diff --git a/...ntitlements/entitlementsSelectors.test.ts → site/src/utils/entitlements.test.ts b/...ntitlements/entitlementsSelectors.test.ts → site/src/utils/entitlements.test.ts
@@ -1,4 +1,4 @@
-import { getFeatureVisibility } from "./entitlementsSelectors";
+import { getFeatureVisibility } from "./entitlements";
 
 describe("getFeatureVisibility", () => {
   it("returns empty object if there is no license", () => {

diff --git a/...ces/entitlements/entitlementsSelectors.ts → site/src/utils/entitlements.ts b/...ces/entitlements/entitlementsSelectors.ts → site/src/utils/entitlements.ts