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; data: TData | undefined; isLoading: boolean; error?: Error; } export type QueryModelContructorInputParam = | PersistentModelConstructor | 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( modelConstructor: QueryModelContructorInputParam, id: string ): QueryReturn; function useDataStoreQuery( modelConstructor: QueryModelContructorInputParam, criteria?: ProducerModelPredicate, paginationProducer?: ProducerPaginationInput ): QueryReturn; function useDataStoreQuery( modelConstructor: QueryModelContructorInputParam, criteria?: string | ProducerModelPredicate, paginationProducer?: ProducerPaginationInput ): QueryReturn { const shouldFetch = typeof modelConstructor === "function"; const [isLoading, setIsLoading] = useState(shouldFetch); const [error, setError] = useState(); const [dataStoreData, setDataStoreData] = useState(); // 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(modelConstructor, criteria) : await DataStore.query( 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( 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;