Skip to content

Commit

Permalink
chore(base): Add comments and docstrings to initial react-query code …
Browse files Browse the repository at this point in the history
…MAAASENG-4252 (#5577)

- Added comments and docstrings to helper functions surrounding react-query implementation

Resolves [MAAASENG-4252](https://warthogs.atlassian.net/browse/MAAASENG-4252)
  • Loading branch information
ndv99 authored Jan 15, 2025
1 parent ef18791 commit a04eefb
Show file tree
Hide file tree
Showing 5 changed files with 61 additions and 0 deletions.
12 changes: 12 additions & 0 deletions src/app/api/base.ts
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,21 @@ type ApiEndpoint = typeof API_ENDPOINTS;
export type ApiEndpointKey = keyof ApiEndpoint;
type ApiUrl = `${typeof SERVICE_API}${ApiEndpoint[ApiEndpointKey]}`;

/**
* Constructs a complete API URL from a given endpoint key.
* @param endpoint The endpoint key
* @returns A full API URL (string)
*/
export const getFullApiUrl = (endpoint: ApiEndpointKey): ApiUrl =>
`${SERVICE_API}${API_ENDPOINTS[endpoint]}`;

/**
* Runs a given fetch request with the appropriate headers and CSRF token.
*
* @param url The URL to fetch
* @param options RequestInit options
* @returns The JSON response from the fetch
*/
export const fetchWithAuth = async (url: string, options: RequestInit = {}) => {
const csrftoken = getCookie("csrftoken");
const headers = {
Expand Down
5 changes: 5 additions & 0 deletions src/app/api/endpoints.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,10 @@
import { fetchWithAuth, getFullApiUrl } from "@/app/api/base";
import type { Zone } from "@/app/store/zone/types";

/**
* Fetches a list of zones.
*
* @returns The list of zones
*/
export const fetchZones = (): Promise<Zone[]> =>
fetchWithAuth(getFullApiUrl("zones"));
6 changes: 6 additions & 0 deletions src/app/api/query-client.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import { QueryClient } from "@tanstack/react-query";

// Different query keys for different methods.
export const queryKeys = {
zones: {
list: ["zones"],
Expand All @@ -10,23 +11,28 @@ type QueryKeys = typeof queryKeys;
type QueryKeyCategories = keyof QueryKeys;
type QueryKeySubcategories<T extends QueryKeyCategories> = keyof QueryKeys[T];

// This basically exists to get us the query key as it appears in react query, i.e. a string in an array.
export type QueryKey =
QueryKeys[QueryKeyCategories][QueryKeySubcategories<QueryKeyCategories>];

// first element of the queryKeys array
// Idk what Peter meant by the above comment, but QueryModel is basically QueryKey but not in an array from what I understand.
export type QueryModel = QueryKey[number];

// 5 minutes feels rather long for default stale time.
export const defaultQueryOptions = {
staleTime: 5 * 60 * 1000, // 5 minutes
cacheTime: 15 * 60 * 1000, // 15 minutes
refetchOnWindowFocus: true,
} as const;

// 0 is far too quick lol but at least we're not using this.
export const realTimeQueryOptions = {
staleTime: 0,
cacheTime: 60 * 1000, // 1 minute
} as const;

// This just creates a query client and provides the options specified above.
export const createQueryClient = () =>
new QueryClient({
defaultOptions: {
Expand Down
24 changes: 24 additions & 0 deletions src/app/api/query/base.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,22 +11,32 @@ import statusSelectors from "@/app/store/status/selectors";
import type { WebSocketEndpointModel } from "@/websocket-client";
import { WebSocketMessageType } from "@/websocket-client";

/**
* Provides a hook to subscribe to NOTIFY messages from the websocket.
*
* @returns An object with a subscribe function that takes a callback to run when a NOTIFY message is received.
*/
export const useWebSocket = () => {
const websocketClient = useContext(WebSocketContext);

if (!websocketClient) {
throw new Error("useWebSocket must be used within a WebSocketProvider");
}

// Listen for NOTIFY messages and run a callback when received
const subscribe = useCallback(
(callback: (msg: any) => void) => {
if (!websocketClient.rws) return;

const messageHandler = (messageEvent: MessageEvent) => {
const data = JSON.parse(messageEvent.data);
// if we get a NOTIFY, run the provided callback
if (data.type === WebSocketMessageType.NOTIFY) callback(data);
};
// add an event listener for NOTIFY messages
websocketClient.rws.addEventListener("message", messageHandler);

// this is a function to remove that event listener, it gets called in a cleanup effect down below.
return () =>
websocketClient.rws?.removeEventListener("message", messageHandler);
},
Expand All @@ -40,6 +50,16 @@ const wsToQueryKeyMapping: Partial<Record<WebSocketEndpointModel, string>> = {
zone: "zones",
// Add more mappings as needed
} as const;

/**
* A function to run a query which invalidates the query cache when a
* websocket message is received, or when the websocket reconnects.
*
* @param queryKey The query key to use
* @param queryFn The query function to run
* @param options Options for useQuery
* @returns The return value of useQuery
*/
export function useWebsocketAwareQuery<
TQueryFnData = unknown,
TError = unknown,
Expand All @@ -60,13 +80,17 @@ export function useWebsocketAwareQuery<
const previousConnectedCount = usePrevious(connectedCount);

useEffect(() => {
// connectedCount will change if the websocket reconnects - if this happens, we should invalidate the query
if (connectedCount !== previousConnectedCount) {
queryClient.invalidateQueries({ queryKey });
}
}, [connectedCount, previousConnectedCount, queryClient, queryKey]);

useEffect(() => {
// subscribe returns a function to remove the event listener for NOTIFY messages;
// This function will be used as the cleanup function for the effect.
return subscribe(
// This callback function will be called when a NOTIFY message is received
({ name: model }: { action: string; name: WebSocketEndpointModel }) => {
const mappedKey = wsToQueryKeyMapping[model];
const modelQueryKey = queryKey[0];
Expand Down
14 changes: 14 additions & 0 deletions src/app/api/query/zones.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,15 +8,29 @@ const zoneKeys = {
list: ["zones"] as const,
};

/**
* Fetches a list of zones.
*
* @returns The zones list as a UseQueryResult
*/
export const useZones = () => {
return useWebsocketAwareQuery(zoneKeys.list, fetchZones);
};

/**
* Fetches the number of zones.
* @returns The number of zones as a UseQueryResult
*/
export const useZoneCount = () =>
useWebsocketAwareQuery<Zone[], Zone[], number>(["zones"], fetchZones, {
select: (data) => data?.length ?? 0,
});

/**
* Get a zone by its ID from the list result.
* @param id The zone's ID
* @returns A Zone, or null as a UseQueryResult
*/
export const useZoneById = (id?: ZonePK | null) =>
useWebsocketAwareQuery(zoneKeys.list, fetchZones, {
select: selectById<Zone>(id ?? null),
Expand Down

0 comments on commit a04eefb

Please sign in to comment.