createSocketIOGraphQLClient.ts

import type { Socket as IOSocket } from "socket.io-client";
import {
  makePushPullAsyncIterableIterator,
  withHandlers,
} from "@n1ru4l/push-pull-async-iterable-iterator";

export type ExecutionParameter = {
  operation: string;
  operationName?: string | null;
  variables?: { [key: string]: any };
  extensions?: { [key: string]: any };
};

export type SocketIOGraphQLClient<TExecutionResult = unknown> = {
  execute: (
    opts: ExecutionParameter
  ) => AsyncIterableIterator<TExecutionResult>;
  destroy: () => void;
};

type OperationRecord<TExecutionResult = unknown> = {
  iterator: AsyncIterableIterator<TExecutionResult>;
  publishValue: (value: TExecutionResult) => void;
  execute: () => void;
};

export const createSocketIOGraphQLClient = <TExecutionResult = unknown>(
  socket: IOSocket
): SocketIOGraphQLClient<TExecutionResult> => {
  let currentOperationId = 0;
  const operations = new Map<number, OperationRecord<TExecutionResult>>();
  const onExecutionResult = ({ id, isFinal, ...result }: any) => {
    const record = operations.get(id);
    if (!record) {
      return;
    }
    record.publishValue(result);

    if (isFinal) {
      record.iterator.return?.();
      operations.delete(id);
    }
  };

  let isOffline = false;

  const onDisconnect = () => {
    isOffline = true;
  };
  const onConnect = () => {
    if (isOffline) {
      isOffline = false;
      Array.from(operations.values()).forEach((record) => {
        record.execute();
      });
    }
  };

  socket.on("@graphql/result", onExecutionResult);
  socket.on("connect", onConnect);
  socket.on("disconnect", onDisconnect);

  const destroy = () => {
    socket.off("@graphql/result", onExecutionResult);
    socket.off("connect", onConnect);
    socket.off("disconnect", onDisconnect);
  };

  const execute = ({
    operation,
    variables,
    operationName,
    extensions,
  }: ExecutionParameter): AsyncIterableIterator<TExecutionResult> => {
    const operationId = currentOperationId;
    currentOperationId = currentOperationId + 1;
    const { asyncIterableIterator: source, pushValue: publishValue } =
      makePushPullAsyncIterableIterator<TExecutionResult>();

    const stream = withHandlers(source, () => {
      if (operations.delete(operationId) === false) {
        return;
      }

      socket.emit("@graphql/unsubscribe", {
        id: operationId,
      });
      source.return();
    });

    const record: OperationRecord<TExecutionResult> = {
      execute: () => {
        socket.emit("@graphql/execute", {
          id: operationId,
          operationName,
          operation,
          variables,
          extensions,
        });
      },
      iterator: stream,
      publishValue,
    };

    operations.set(operationId, record);

    if (!isOffline) {
      record.execute();
    }

    return stream;
  };

  return {
    execute,
    destroy,
  };
};