PeteDuncanson · October 7, 2021 19:48 · PeteDuncanson · Sep 21, 2021
diff --git a/useDataStoreQuery.tsx b/useDataStoreQuery.tsx
 import {
  DataStore,
  PersistentModel,
  PersistentModelConstructor,
  ProducerModelPredicate,
  SortPredicate,
  ProducerPaginationInput,
 } from "@aws-amplify/datastore";
 import { useAppState } from "providers/AppStateProvider";
 import { useCallback, useEffect, useRef, useState } from "react";
 import { useDebounce } from "react-use";

 /* 
 Ideas/Todo
  - IsDataStoreReady could be a prop, I'm using a Context here to get it, we shouldn't be getting any data until its ready
  - Option to opt out of the subscription updates, so get me the data and thats it
  - Default to no subscriptions if no criteria is provided (ie its a "get everything" query) as we might not want to be grabbing them every time they change. 
  - Work out how the TypeScript signatures work and if we can improve on them as the intellisense isn't as nice as I would like
  - Added debounce as when updating a lot of Models in DataStore the subscriptions that come back can trigger a lot of re-renders, but could do to make the debounce time an optional prop
 */
 export interface QueryReturn<
  TData extends PersistentModel | PersistentModel[]
 > {
  refetch: () => Promise<TData | undefined>;
  data: TData | undefined;
  isLoading: boolean;
  error?: Error;
 }

 export type QueryModelContructorInputParam<T extends PersistentModel> =
  | PersistentModelConstructor<T>
  | null
  | undefined
  | false;

 /**
 * React hook to read from the [Amplify DataStore](https://docs.amplify.aws/lib/datastore/getting-started/q/platform/js).
 * Built-in supports for real-time update using subscriptions.
 *
 * Note that `Predicates.ALL` is not supported as it's use leads to unnecessary re-renders. Simply use `undefined`
 * instead, if you want to query for all data records without any filter.
 *
 * @example
 * // To read from the database, the simplest approach is to query for all records of a given model type.
 * const { data, isLoading, error } = useDataStoreQuery(Post).current)
 *
 * @example
 * // To query for a single item, pass in the ID of the item as the second argument to the query.
 * const { data, isLoading, error } = useDataStoreQuery(Post, "1234567").current)
 *
 * @example
 * // Predicates are supported as well. For example if you wanted a list of all Post Models that have a rating greater than 4:
 * const { data, isLoading, error } = useDataStoreQuery(Post, c => c.rating("gt", 4));
 *
 * @example
 * // Query results can also be sorted by one or more fields.
 * const { data, isLoading, error } = useDataStoreQuery(Post, undefined, useRef({ sort: s => s.rating(SortDirection.ASCENDING) }).current)
 */
 function useDataStoreQuery<T extends PersistentModel>(
  modelConstructor: QueryModelContructorInputParam<T>,
  id: string
 ): QueryReturn<T>;
 function useDataStoreQuery<T extends PersistentModel>(
  modelConstructor: QueryModelContructorInputParam<T>,
  criteria?: ProducerModelPredicate<T>,
  paginationProducer?: ProducerPaginationInput<T>
 ): QueryReturn<T[]>;

 function useDataStoreQuery<T extends PersistentModel>(
  modelConstructor: QueryModelContructorInputParam<T>,
  criteria?: string | ProducerModelPredicate<T>,
  paginationProducer?: ProducerPaginationInput<T>
 ): QueryReturn<any> {
  const shouldFetch = typeof modelConstructor === "function";
  const [isLoading, setIsLoading] = useState(shouldFetch);
  const [error, setError] = useState<Error | undefined>();
  const [dataStoreData, setDataStoreData] = useState<T[] | T | undefined>();
  // This could be a prop, basically we need to know if DataStore has finished syncing before we allow pulling stuff down
  const AppState = useAppState();
  const dataStoreReady = AppState.dataStoreIsReady; 
  const refetchNeeded = useRef(false);

  const [, cancel] = useDebounce(
    () => {
      if (refetchNeeded) {
        fetchAsync().catch(null);
        refetchNeeded.current = false;
      }
    },
    200,
    [refetchNeeded]
  );

  const fetchAsync = useCallback(async () => {
    try {
      setIsLoading(true);
      if (typeof modelConstructor === "function") {
        const data =
          typeof criteria === "string"
            ? await DataStore.query<T>(modelConstructor, criteria)
            : await DataStore.query<T>(
                modelConstructor,
                criteria,
                paginationProducer
              );
        setDataStoreData(data);
        setError(undefined);
      }
    } catch (error) {
      console.error("Error fetching data", { error });
      setError(error as Error);
    } finally {
      setIsLoading(false);
    }
  }, [modelConstructor, criteria, paginationProducer]);

  // Initial data fetch - Go get our data if we haven't already
  useEffect(() => {
    if (dataStoreReady) {
      fetchAsync().catch(null);
    }
  }, [dataStoreReady, fetchAsync]);

  // Invalidate our model to force a refetch if it changes outside of our App (notified via a AppSync subscription)
  useEffect(() => {
    if (dataStoreReady && typeof modelConstructor === "function") {
      const subscription = DataStore.observe<T>(
        modelConstructor,
        criteria
      ).subscribe((msg) => {
        refetchNeeded.current = true;
      });
      // Clean up on unmount
      return () => {
        subscription.unsubscribe();
      };
    }
    return undefined;
  }, [dataStoreReady, modelConstructor, criteria, fetchAsync]);
  return { data: dataStoreData, isLoading, error, refetch: fetchAsync };
 }

 export default useDataStoreQuery;
	import {
	DataStore,
	PersistentModel,
	PersistentModelConstructor,
	ProducerModelPredicate,
	SortPredicate,
	ProducerPaginationInput,
	} from "@aws-amplify/datastore";
	import { useAppState } from "providers/AppStateProvider";
	import { useCallback, useEffect, useRef, useState } from "react";
	import { useDebounce } from "react-use";

	/*
	Ideas/Todo
	- IsDataStoreReady could be a prop, I'm using a Context here to get it, we shouldn't be getting any data until its ready
	- Option to opt out of the subscription updates, so get me the data and thats it
	- Default to no subscriptions if no criteria is provided (ie its a "get everything" query) as we might not want to be grabbing them every time they change.
	- Work out how the TypeScript signatures work and if we can improve on them as the intellisense isn't as nice as I would like
	- Added debounce as when updating a lot of Models in DataStore the subscriptions that come back can trigger a lot of re-renders, but could do to make the debounce time an optional prop
	*/
	export interface QueryReturn<
	TData extends PersistentModel \| PersistentModel[]
	> {
	refetch: () => Promise<TData \| undefined>;
	data: TData \| undefined;
	isLoading: boolean;
	error?: Error;
	}

	export type QueryModelContructorInputParam<T extends PersistentModel> =
	\| PersistentModelConstructor<T>
	\| null
	\| undefined
	\| false;

	/**
	* React hook to read from the [Amplify DataStore](https://docs.amplify.aws/lib/datastore/getting-started/q/platform/js).
	* Built-in supports for real-time update using subscriptions.
	*
	* Note that `Predicates.ALL` is not supported as it's use leads to unnecessary re-renders. Simply use `undefined`
	* instead, if you want to query for all data records without any filter.
	*
	* @example
	* // To read from the database, the simplest approach is to query for all records of a given model type.
	* const { data, isLoading, error } = useDataStoreQuery(Post).current)
	*
	* @example
	* // To query for a single item, pass in the ID of the item as the second argument to the query.
	* const { data, isLoading, error } = useDataStoreQuery(Post, "1234567").current)
	*
	* @example
	* // Predicates are supported as well. For example if you wanted a list of all Post Models that have a rating greater than 4:
	* const { data, isLoading, error } = useDataStoreQuery(Post, c => c.rating("gt", 4));
	*
	* @example
	* // Query results can also be sorted by one or more fields.
	* const { data, isLoading, error } = useDataStoreQuery(Post, undefined, useRef({ sort: s => s.rating(SortDirection.ASCENDING) }).current)
	*/
	function useDataStoreQuery<T extends PersistentModel>(
	modelConstructor: QueryModelContructorInputParam<T>,
	id: string
	): QueryReturn<T>;
	function useDataStoreQuery<T extends PersistentModel>(
	modelConstructor: QueryModelContructorInputParam<T>,
	criteria?: ProducerModelPredicate<T>,
	paginationProducer?: ProducerPaginationInput<T>
	): QueryReturn<T[]>;

	function useDataStoreQuery<T extends PersistentModel>(
	modelConstructor: QueryModelContructorInputParam<T>,
	criteria?: string \| ProducerModelPredicate<T>,
	paginationProducer?: ProducerPaginationInput<T>
	): QueryReturn<any> {
	const shouldFetch = typeof modelConstructor === "function";
	const [isLoading, setIsLoading] = useState(shouldFetch);
	const [error, setError] = useState<Error \| undefined>();
	const [dataStoreData, setDataStoreData] = useState<T[] \| T \| undefined>();
	// This could be a prop, basically we need to know if DataStore has finished syncing before we allow pulling stuff down
	const AppState = useAppState();
	const dataStoreReady = AppState.dataStoreIsReady;
	const refetchNeeded = useRef(false);

	const [, cancel] = useDebounce(
	() => {
	if (refetchNeeded) {
	fetchAsync().catch(null);
	refetchNeeded.current = false;
	}
	},
	200,
	[refetchNeeded]
	);

	const fetchAsync = useCallback(async () => {
	try {
	setIsLoading(true);
	if (typeof modelConstructor === "function") {
	const data =
	typeof criteria === "string"
	? await DataStore.query<T>(modelConstructor, criteria)
	: await DataStore.query<T>(
	modelConstructor,
	criteria,
	paginationProducer
	);
	setDataStoreData(data);
	setError(undefined);
	}
	} catch (error) {
	console.error("Error fetching data", { error });
	setError(error as Error);
	} finally {
	setIsLoading(false);
	}
	}, [modelConstructor, criteria, paginationProducer]);

	// Initial data fetch - Go get our data if we haven't already
	useEffect(() => {
	if (dataStoreReady) {
	fetchAsync().catch(null);
	}
	}, [dataStoreReady, fetchAsync]);

	// Invalidate our model to force a refetch if it changes outside of our App (notified via a AppSync subscription)
	useEffect(() => {
	if (dataStoreReady && typeof modelConstructor === "function") {
	const subscription = DataStore.observe<T>(
	modelConstructor,
	criteria
	).subscribe((msg) => {
	refetchNeeded.current = true;
	});
	// Clean up on unmount
	return () => {
	subscription.unsubscribe();
	};
	}
	return undefined;
	}, [dataStoreReady, modelConstructor, criteria, fetchAsync]);
	return { data: dataStoreData, isLoading, error, refetch: fetchAsync };
	}

	export default useDataStoreQuery;