Front-End Web & Mobile

Client-side Caching Strategies for a Next.js app with AWS Amplify

This post builds on the initial posts in this series, Build a Product Roadmap with Next.js and Amplify, where we built an admin page for product managers to login and update the roadmap and then updated the app to add storage of documents.

In this post, we’ll adapt the product management application to include a caching layer to enhance the user experience of the application.  The Amplify GraphQL API library does not include a caching layer and is flexible enough to be used with any caching solution.

TanStack React Query, an open source caching and content invalidation solution for modern web applications, will be featured in this post, despite the existence of various other options like SWR from Vercel, the creators of Next.js.

State Management for Modern Web Applications

Building modern web applications, developers often want a full featured caching layer for optimistic UI updates, auto revalidation based on browser conditions (change in network connectivity or user focusing the application).  They provide a rich end user experience and often rely on local state management to cache data for quick responsiveness to users of the application.  For years, this type of solution was home rolled using Redux, MobX and some combination of localStorage and IndexDB, but it lacked standards and reusability.

In recent years, centralized state management solutions have been abstracted out of real world applications by companies and developers who have solved the issues around state management and caching for modern web applications in a generic and reusable way.

TanStack Query is one of these solutions and describes itself as “Powerful asynchronous state management for TS/JS, React, Solid, Vue and Svelte [that provides] declarative, always-up-to-date auto-managed queries and mutations that directly improve both your developer and user experiences” along with “caching, background updates and stale data out of the box with zero-configuration.”

Install and Configure TanStack React Query

Building on our our previous project, we need to install TanStack Query for React.

npm install @tanstack/react-query

Next, modify pages/_app.js to import QueryClient and QueryClientProvider from @tanstack/react-query create a query client and wrap the React application with a QueryClientProvider.

import "@aws-amplify/ui-react/styles.css";
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import { Amplify } from "aws-amplify";
import React from "react";
import awsExports from "../src/aws-exports";
import "../styles/globals.css";

Amplify.configure({ ...awsExports, ssr: true });
const queryClient = new QueryClient();

function MyApp({ Component, pageProps }) {
  return (
    <QueryClientProvider client={queryClient}>
      <Component {...pageProps} />
    </QueryClientProvider>
  );
}

export default MyApp;

Querying data with Tanstack React Query

Modify the FeaturesTable in components/FeaturesTable.js to query the GraphQL API through @tanstack/react-query.

First import useQuery from @tanstack/react-query.

Next, defined a fetcher function (fetchFeatures) which is an async function fetches a list of features from the GraphQL API using API.graphql method to make an authenticated GraphQL call using Amazon Cognito User Pools. The function runs a query called listFeatures and awaits the result, then returns the items from the result of that query.

// components/FeaturesTable.js
import {
  Button,
  Flex,
  Table,
  TableBody,
  TableCell,
  TableHead,
  TableRow,
  View,
} from "@aws-amplify/ui-react";
import { useQuery } from "@tanstack/react-query";
import { API, graphqlOperation, Storage } from "aws-amplify";
import React, { useEffect } from "react";
import { deleteFeature } from "../src/graphql/mutations";
import { listFeatures } from "../src/graphql/queries";
import {
  onCreateFeature,
  onDeleteFeature,
  onUpdateFeature,
} from "../src/graphql/subscriptions";

const fetchFeatures = async () => {
  const result = await API.graphql({
    authMode: "AMAZON_COGNITO_USER_POOLS",
    query: listFeatures,
  });
  return result.data.listFeatures.items;
};

function FeaturesTable({ initialFeatures = [], setActiveFeature }) {
  // ... component implementation
});

In the FeaturesTable component declare a featuresQueryKey to be used to reference this query in subscription updates, later in the post.

Next, implement the useQuery using the featuresQueryKey, fetchFeatures declared above and finally pass the initialFeatures component prop to preload the query cache from server.

Finally update the no features view using isLoading returned from useQuery.

// components/FeaturesTable.js
// ... imports
const fetchFeatures = async () => {
 // ...
};

function FeaturesTable({ initialFeatures = [], setActiveFeature }) {
  const featuresQueryKey = ["features"];

  const { data: features, isLoading } = useQuery({
    queryKey: featuresQueryKey,
    queryFn: fetchFeatures,
    initialData: initialFeatures,
  });
  
  // ... useEffect and other methods, unmodified
  
  if (isLoading) {
    return <View>No features</View>;
  }
  
  // ... rest of component, unmodified
});

With these updates in place, the features data is managed in the queryClient using the queryKey we provided to useQuery and our existing JSX will continue to reference features to always retrieve the latest from this cache.

Modify subscriptions to update the query cache

Once we have query data managed in the queryClient, we’ll want to ensure the subscriptions are updating the queryClient.

In order to modify the cached query results we need to interact with the queryClient itself.  TanStack React Query provides a react hook, useQueryClient to allow us to do this.

Update the import to include useQueryClient from @tanstack/react-query.

Inside of useEffect, update the subscription methods createSub, updateSub, deleteSub to use queryClient.setQueryData and adjust the logic to modify the data for our query key, featuresQueryKey.

// components/FeaturesTable.js
import {
  // ...
} from "@aws-amplify/ui-react";
import { useQuery, useQueryClient } from "@tanstack/react-query";
// ... additional imports

const fetchFeatures = async () => {
 // ...
};

function FeaturesTable({ initialFeatures = [], setActiveFeature }) {
  const featuresQueryKey = ["features"];

  const queryClient = useQueryClient();

  const { data: features, isLoading } = useQuery({
    queryKey: featuresQueryKey,
    queryFn: fetchFeatures,
    initialData: initialFeatures,
  });

  useEffect(() => {
  const createSub = API.graphql(graphqlOperation(onCreateFeature)).subscribe({
      next: ({ value }) => {
        queryClient.setQueryData(featuresQueryKey, (current) => {
          const toCreateIndex = current.findIndex(
            (item) => item.id === value.data.onCreateFeature.id
          );
          if (toCreateIndex) {
            return current;
          }
          return [...current, value.data.onCreateFeature];
        });
      },
    });

    const updateSub = API.graphql(graphqlOperation(onUpdateFeature)).subscribe({
      next: ({ value }) => {
        queryClient.setQueryData(featuresQueryKey, (current) => {
          const toUpdateIndex = current.findIndex(
            (item) => item.id === value.data.onUpdateFeature.id
          );
          if (toUpdateIndex === -1) {
            return [...current, value.data.onCreateFeature];
          }
          return [
            ...current.slice(0, toUpdateIndex),
            value.data.onUpdateFeature,
            ...current.slice(toUpdateIndex + 1),
          ];
        });
      },
    });

    const deleteSub = API.graphql(graphqlOperation(onDeleteFeature)).subscribe({
      next: ({ value }) => {
        queryClient.setQueryData(featuresQueryKey, (current) => {
          const toDeleteIndex = current.findIndex(
            (item) => item.id === value.data.onDeleteFeature.id
          );
          return [
            ...current.slice(0, toDeleteIndex),
            ...current.slice(toDeleteIndex + 1),
          ];
        });
      },
    });

    return () => {
      createSub.unsubscribe();
      updateSub.unsubscribe();
      deleteSub.unsubscribe();
    };
  }, []);
  
  // rest of component, unmodified

})

For reference, the complete updated version of components/FeaturesTable.js is below.

// components/FeaturesTable.js
import {
  Button,
  Flex,
  Table,
  TableBody,
  TableCell,
  TableHead,
  TableRow,
  View,
} from "@aws-amplify/ui-react";
import { useQuery, useQueryClient } from "@tanstack/react-query";
import { API, graphqlOperation, Storage } from "aws-amplify";
import React, { useEffect } from "react";
import { deleteFeature } from "../src/graphql/mutations";
import { listFeatures } from "../src/graphql/queries";
import {
  onCreateFeature,
  onDeleteFeature,
  onUpdateFeature,
} from "../src/graphql/subscriptions";
const fetchFeatures = async () => {
  const result = await API.graphql({
    authMode: "AMAZON_COGNITO_USER_POOLS",
    query: listFeatures,
  });
  return result.data.listFeatures.items;
};

function FeaturesTable({ initialFeatures = [], setActiveFeature }) {
  const featuresQueryKey = ["features"];

  const queryClient = useQueryClient();

  const { data: features, isLoading } = useQuery({
    queryKey: featuresQueryKey,
    queryFn: fetchFeatures,
    initialData: initialFeatures,
  });

  useEffect(() => {
    const createSub = API.graphql(graphqlOperation(onCreateFeature)).subscribe({
      next: ({ value }) => {
        queryClient.setQueryData(featuresQueryKey, (current) => {
          const toCreateIndex = current.findIndex(
            (item) => item.id === value.data.onCreateFeature.id
          );
          if (toCreateIndex) {
            return current;
          }
          return [...current, value.data.onCreateFeature];
        });
      },
    });

    const updateSub = API.graphql(graphqlOperation(onUpdateFeature)).subscribe({
      next: ({ value }) => {
        queryClient.setQueryData(featuresQueryKey, (current) => {
          const toUpdateIndex = current.findIndex(
            (item) => item.id === value.data.onUpdateFeature.id
          );
          if (toUpdateIndex === -1) {
            return [...current, value.data.onCreateFeature];
          }
          return [
            ...current.slice(0, toUpdateIndex),
            value.data.onUpdateFeature,
            ...current.slice(toUpdateIndex + 1),
          ];
        });
      },
    });

    const deleteSub = API.graphql(graphqlOperation(onDeleteFeature)).subscribe({
      next: ({ value }) => {
        queryClient.setQueryData(featuresQueryKey, (current) => {
          const toDeleteIndex = current.findIndex(
            (item) => item.id === value.data.onDeleteFeature.id
          );
          return [
            ...current.slice(0, toDeleteIndex),
            ...current.slice(toDeleteIndex + 1),
          ];
        });
      },
    });

    return () => {
      createSub.unsubscribe();
      updateSub.unsubscribe();
      deleteSub.unsubscribe();
    };
  }, []);

  function downloadBlob(blob, filename) {
    const url = URL.createObjectURL(blob);
    const a = document.createElement("a");
    a.href = url;
    a.download = filename || "download";
    const clickHandler = () => {
      setTimeout(() => {
        URL.revokeObjectURL(url);
        a.removeEventListener("click", clickHandler);
      }, 150);
    };
    a.addEventListener("click", clickHandler, false);
    a.click();
    return a;
  }

  async function handleDownload(fileKey) {
    const result = await Storage.get(fileKey, { download: true });
    downloadBlob(result.Body, fileKey);
  }

  async function onDeleteInternalDoc(internalDoc) {
    try {
      await Storage.remove(internalDoc);
    } catch ({ errors }) {
      console.error(...errors);
    }
  }

  async function handleDeleteFeature(id) {
    try {
      await API.graphql({
        authMode: "AMAZON_COGNITO_USER_POOLS",
        query: deleteFeature,
        variables: {
          input: {
            id,
          },
        },
      });
    } catch ({ errors }) {
      console.error(...errors);
    }
  }

  if (isLoading) {
    return <View>No features</View>;
  }

  return (
    <Table>
      <TableHead>
        <TableRow>
          <TableCell as="th">Feature</TableCell>
          <TableCell as="th">Released</TableCell>
          <TableCell></TableCell>
        </TableRow>
      </TableHead>
      <TableBody>
        {features.map((feature) => (
          <TableRow key={feature.id}>
            <TableCell>{feature.title}</TableCell>
            <TableCell>{feature.released ? "Yes" : "No"}</TableCell>
            <TableCell>
              <Flex>
                <Button size="small" onClick={() => setActiveFeature(feature)}>
                  Edit
                </Button>
                <Button
                  size="small"
                  onClick={async () =>
                    await Promise.all([
                      // delete the document via Storage
                      onDeleteInternalDoc(feature.internalDoc),
                      handleDeleteFeature(feature.id),
                    ])
                  }
                >
                  Delete
                </Button>
                {feature.internalDoc ? (
                  <Button
                    size="small"
                    onClick={() => handleDownload(feature.internalDoc)}
                  >
                    Download File
                  </Button>
                ) : undefined}
              </Flex>
            </TableCell>
          </TableRow>
        ))}
      </TableBody>
    </Table>
  );
}

export default FeaturesTable;

Rendering new data optimistically

To provide an optimal user experience, our application should instantly update our UI with a newly created Feature.

Optimistic UI updates can be used to achieve this and greatly improve the user experience of web applications. By immediately updating the UI in response to user actions, optimistic updates make the app feel more responsive, reduce waiting and uncertainty about whether an update succeeded, provide context if an error does occur, and create an interaction model that feels more natural to users.

We need to update components/FeatureForm.js to optimistically update the features query cache when a new feature is saved.

First, import useMutation and useQueryClient from @tanstack/react-query, then create a queryClient using useQueryClient.

The useMutation React hook from TanStack React Query makes implementing optimistic UI updates straightforward.

In the onMutate callback, queryClient.getQueryData is used to retrieve data from the cache for the query key provided (in this case “features”).  Then queryClient.setQueryData is called to optimistically update the cache using the same featuresQueryKey used in the FeaturesTable component with newly submitted data, and immediately update the UI.

If the actual mutation succeeds, TanStack React Query will automatically update the cache and UI to match the server data. However, if the mutation fails, queryClient.setQueryData is called in the onError callback, this time setting the cache back to its original state. TanStack React Query will then revert the UI to match.

// components/FeatureForm.js
import {
  Button,
  Flex,
  Heading,
  SwitchField,
  Text,
  TextField,
  View,
} from "@aws-amplify/ui-react";
import { useMutation, useQueryClient } from "@tanstack/react-query";
import { API, Storage } from "aws-amplify";
import React, { useEffect, useState } from "react";
import { v4 } from "uuid";
import { createFeature, updateFeature } from "../src/graphql/mutations";

function FeatureForm({ feature = null, setActiveFeature }) {
  const [id, setId] = useState(undefined);
  const [title, setTitle] = useState("");
  const [description, setDescription] = useState("");
  const [isReleased, setReleased] = useState(false);
  const [internalDoc, setInternalDoc] = useState("");
  const featuresQueryKey = ["features"];

  const queryClient = useQueryClient();

  const saveFeature = useMutation({
    mutationFn: handleSaveFeature,
    onMutate: async (newFeature) => {
      await queryClient.cancelQueries({ queryKey: featuresQueryKey });

      const previousFeatures = queryClient.getQueryData(featuresQueryKey);

      queryClient.setQueryData(featuresQueryKey, (old) => [...old, newFeature]);

      return { previousFeatures };
    },
    onError: (err, newFeature, context) => {
      queryClient.setQueryData(featuresQueryKey, context.previousFeatures);
    },
    onSettled: () => {
      queryClient.invalidateQueries({ queryKey: featuresQueryKey });
    },
  });
  
  // ....
  
})

Next, import v4 from uuid (a module installed with Amplify) to leverage in a centralized createNewFeature function that will create the new record.  Then, modify the handleSaveFeature to work with the newFeature passed to it and specify the authMode of "AMAZON_COGNITO_USER_POOLS" since our project default authorization mode is set to API_KEY and this will be submitted as an authenticated user.

import { ... } from "@aws-amplify/ui-react";
import { useMutation, useQueryClient } from "@tanstack/react-query";
import { API, Storage } from "aws-amplify";
import React, { useEffect, useState } from "react";
import { v4 } from "uuid";
import { createFeature, updateFeature } from "../src/graphql/mutations";

function FeatureForm({ feature = null, setActiveFeature }) {

  // ...
  
  function createNewFeature() {
    const newFeature = {
      id: id || v4(),
      title,
      description,
      released: isReleased,
      internalDoc: internalDoc,
    };

    return newFeature;
  }
  
  async function handleSaveFeature(newFeature) {
    try {
      await API.graphql({
        authMode: "AMAZON_COGNITO_USER_POOLS",
        query: feature ? updateFeature : createFeature,
        variables: {
          input: newFeature,
        },
      });

      feature && setActiveFeature(undefined);
      resetFormFields();
    } catch ({ errors }) {
      console.error(...errors);
      throw new Error(errors[0].message);
    }
  }
  
  // ...
  
})

Finally, modify the onClick for the Save button to call createNewFeature, then pass that to the mutation created earlier, saveFeature.

import { ... } from "@aws-amplify/ui-react";
import { useMutation, useQueryClient } from "@tanstack/react-query";
import { API, Storage } from "aws-amplify";
import React, { useEffect, useState } from "react";
import { v4 } from "uuid";
import { createFeature, updateFeature } from "../src/graphql/mutations";

function FeatureForm({ feature = null, setActiveFeature }) {
  // ...
  
    <View>
      // ... 
          <Button
            onClick={() => {
              const newFeature = createNewFeature();
              saveFeature.mutate(newFeature);
            }}
          >
            Save
          </Button>
        </Flex>
      </Flex>
    </View>
  );
}

export default FeatureForm;

For reference, the complete updated version of components/FeatureForm.js is below.

// components/FeatureForm.js
import {
  Button,
  Flex,
  Heading,
  SwitchField,
  Text,
  TextField,
  View,
} from "@aws-amplify/ui-react";
import { useMutation, useQueryClient } from "@tanstack/react-query";
import { API, Storage } from "aws-amplify";
import React, { useEffect, useState } from "react";
import { v4 } from "uuid";
import { createFeature, updateFeature } from "../src/graphql/mutations";

function FeatureForm({ feature = null, setActiveFeature }) {
  const [id, setId] = useState(undefined);
  const [title, setTitle] = useState("");
  const [description, setDescription] = useState("");
  const [isReleased, setReleased] = useState(false);
  const [internalDoc, setInternalDoc] = useState("");
  const featuresQueryKey = ["features"];

  const queryClient = useQueryClient();

  const saveFeature = useMutation({
    mutationFn: handleSaveFeature,
    onMutate: async (newFeature) => {
      await queryClient.cancelQueries({ queryKey: featuresQueryKey });

      const previousFeatures = queryClient.getQueryData(featuresQueryKey);

      queryClient.setQueryData(featuresQueryKey, (old) => [...old, newFeature]);

      return { previousFeatures };
    },
    onError: (err, newFeature, context) => {
      queryClient.setQueryData(featuresQueryKey, context.previousFeatures);
    },
    onSettled: () => {
      queryClient.invalidateQueries({ queryKey: featuresQueryKey });
    },
  });

  useEffect(() => {
    if (feature) {
      setId(feature.id);
      setTitle(feature.title);
      setDescription(feature.description);
      setReleased(feature.released);
      setInternalDoc(feature.internalDoc);
    }
  }, [feature]);

  async function handleUploadDoc(e) {
    const file = e.target.files[0];
    const fileName = `${Date.now()}-${file.name}`;
    try {
      await Storage.put(fileName, file, {
        contentType: file.type,
      });
      setInternalDoc(fileName);
    } catch (error) {
      console.log("Error uploading file: ", error);
    }
  }

  async function handleRemoveDoc() {
    try {
      await Storage.remove(internalDoc);
      setInternalDoc("");
    } catch (error) {
      console.log("Error removing file: ", error);
    }
  }

  function resetFormFields() {
    setId(undefined);
    setTitle("");
    setDescription("");
    setReleased(false);
    setInternalDoc("");
  }

  function createNewFeature() {
    const newFeature = {
      id: id || v4(),
      title,
      description,
      released: isReleased,
      internalDoc: internalDoc,
    };

    return newFeature;
  }

  async function handleSaveFeature(newFeature) {
    try {
      await API.graphql({
        authMode: "AMAZON_COGNITO_USER_POOLS",
        query: feature ? updateFeature : createFeature,
        variables: {
          input: newFeature,
        },
      });

      feature && setActiveFeature(undefined);
      resetFormFields();
    } catch ({ errors }) {
      console.error(...errors);
      throw new Error(errors[0].message);
    }
  }

  return (
    <View>
      <Heading marginBottom="medium" level={5}>
        {feature ? "Edit" : "New"} Feature
      </Heading>
      <Flex direction={"column"} basis={"max-content"}>
        <TextField
          value={title}
          label="Title"
          errorMessage="There is an error"
          name="title"
          onChange={(e) => setTitle(e.target.value)}
        />

        <TextField
          value={description}
          name="description"
          label="Description"
          errorMessage="There is an error"
          onChange={(e) => setDescription(e.target.value)}
        />

        <SwitchField
          isChecked={isReleased}
          isDisabled={false}
          label="Released?"
          labelPosition="start"
          onChange={() => setReleased(!isReleased)}
        />

        {feature && internalDoc ? (
          <div>
            <Text>Attachment:</Text>
            <Text fontWeight={"bold"}>
              {internalDoc}{" "}
              <Button size="small" onClick={handleRemoveDoc}>
                X
              </Button>
            </Text>
          </div>
        ) : (
          <div>
            <Text>Upload a file:</Text>
            <input type="file" onChange={handleUploadDoc} />
          </div>
        )}

        <Flex marginTop="large">
          <Button
            onClick={() => {
              setActiveFeature(undefined);
              resetFormFields();
            }}
          >
            Cancel
          </Button>
          <Button
            onClick={() => {
              const newFeature = createNewFeature();
              saveFeature.mutate(newFeature);
            }}
          >
            Save
          </Button>
        </Flex>
      </Flex>
    </View>
  );
}

export default FeatureForm;

What we’ve built

In this post, we improved the product management application to include a caching layer using TanStack React Query wrapped around the Amplify GraphQL API to enhance the user experience of the application.  Enhancing our Product Roadmap application with TanStack React Query has enabled capabilities and greater flexibility for users. Optimistic UI updates allow users to instantly see the results after creating a product roadmap item.  If a user loses network connectivity or the application window is in the background, they can continue to use the application and TanStack React Query will automatically revalidate and update it’s query caches based on native browser conditions, such as when network connectivity is restored or the user focuses the application. We’re able to fully leverage the SSR capabilities from Next.js and pre-populate the query cache from data from the server. Finally, we were able to continue to use Amplify subscriptions for realtime updates and adjusted our logic to update the query cache.

For more information on optimistic UI, review our documentation.