Created for https://reflow.io : react hooks for aws amplify datastore that fill @connection links

useDataStore.ts

import { useEffect, useState } from 'react';
import { DataStore } from '@aws-amplify/datastore';
import type {
  PersistentModel,
  PersistentModelConstructor,
  ProducerModelPredicate,
  ProducerPaginationInput,
} from '@aws-amplify/datastore';

export const useDataStore = <T extends PersistentModel>(
  modelConstructor: PersistentModelConstructor<T>,
  nestedKeys: {
    [K in keyof T]?: {
      model: PersistentModelConstructor<any>;
      criteria: (modelInstance: T) => ProducerModelPredicate<T[K][0]>;
      paginationProducer?: ProducerPaginationInput<T[K][0]>;
    };
  } = {},
  criteria: ProducerModelPredicate<T> | undefined = undefined,
  paginationProducer?: ProducerPaginationInput<T>
): { loaded: boolean; items: T[] } => {
  const [state, setState] = useState<{ loaded: boolean; items: T[] }>({
    loaded: false,
    items: [],
  });

  const getAll = (): Promise<T[]> => {
    return DataStore.query(modelConstructor, criteria, paginationProducer).then((modelValue) => {
      return Promise.all(
        modelValue.map(async (m) => {
          let j: T = {
            ...m,
          };
          for (const key of Object.keys(nestedKeys)) {
            const { model, criteria, paginationProducer } = nestedKeys[key]!;
            const keyItems = await DataStore.query(model, criteria(m) as any, paginationProducer as any);
            j = {
              ...j,
              [key]: keyItems,
            };
          }
          return j;
        })
      );
    });
  };

  useEffect(() => {
    if (!state.loaded) {
      getAll().then((items) => {
        setState({ items, loaded: true });
      });
    }

    const models = [modelConstructor, ...Object.values(nestedKeys).map((key) => key.model)];

    const subscriptions = models.map((model) =>
      DataStore.observe(model).subscribe((msg) => {
        getAll()
          .then((items) => {
            setState({ items: items, loaded: true });
          })
          .catch((e) => {
            console.error(e);
          });
      })
    );

    return () => {
      subscriptions.forEach((subscription) => subscription.unsubscribe());
    };
  });

  return {
    loaded: state.loaded,
    items: state.items,
  };
};

export const useDataStoreItemWithId = <T extends PersistentModel>(
  modelConstructor: PersistentModelConstructor<T>,
  id: string | undefined,
  nestedKeys: {
    [K in keyof T]?: {
      model: PersistentModelConstructor<any>;
      criteria: (modelInstance: T) => string | ProducerModelPredicate<T[K] extends Array<any> ? T[K][0] : T[K]>;
      paginationProducer?: ProducerPaginationInput<T[K][0]>;
      singleton?: boolean;
    };
  } = {}
): { loaded: boolean; item: T | undefined; error: Error | undefined } => {
  const [state, setState] = useState<{
    loaded: boolean;
    error: Error | undefined;
    id: string | undefined;
    item: T | undefined;
  }>({
    loaded: false,
    item: undefined,
    id: undefined,
    error: undefined,
  });

  const get = async (): Promise<void> => {
    try {
      if (id === undefined) {
        return;
      }
      const modelValue = await DataStore.query(modelConstructor, id);
      if (!modelValue) {
        const err = new Error(`no value with id ${id}`);
        setState({ item: undefined, id, loaded: true, error: err });
        return;
      }
      let j: T = modelValue;
      for (const key of Object.keys(nestedKeys)) {
        const { model, criteria, paginationProducer, singleton } = nestedKeys[key]!;
        // key -- just lookup by this key
        let retValue: any;
        if (typeof criteria === 'string') {
          retValue = await DataStore.query(model, criteria);
        } else {
          retValue = await DataStore.query(model, criteria(j) as any, paginationProducer as any);
        }
        if (singleton) {
          j = {
            ...j,
            [key]: retValue?.length > 0 ? retValue[0] : undefined,
          };
        } else {
          j = {
            ...j,
            [key]: retValue || [],
          };
        }
      }
      if (modelValue) {
        setState({ item: j, id, loaded: true, error: undefined });
      } else {
        setState({ item: undefined, id, loaded: true, error: new Error(`no value with id ${id}`) });
      }
    } catch (e) {
      setState({ item: undefined, id, loaded: true, error: e });
      return;
    }
  };

  useEffect(() => {
    if (id !== state.id) {
      setState({ item: undefined, id, loaded: false, error: undefined });
      get();
      return;
    }
    if (id === undefined) {
      return;
    }
    if (!state.loaded) {
      get();
      return;
    }
    const models = [
      { model: modelConstructor, criteria: id },
      ...Object.values(nestedKeys)
        .filter(() => state.item)
        .map((key) => ({ model: key.model, criteria: key.criteria(state.item!) })),
    ];
    const subscriptions = models.map(({ model, criteria }) => {
      try {
        const observer = DataStore.observe(model, criteria);
        const subscription = observer.subscribe(
          (msg) => {
            get();
          },
          (err) => {
            console.error('subscribe', err);
          },
          () => {
            console.log('complete');
          }
        );
        return subscription;
      } catch (e) {
        console.error(e);
        return undefined;
      }
    });

    return () => {
      subscriptions.forEach((subscription) => subscription?.unsubscribe());
    };
  });

  return state;
};