WithPgClientResourceStep

with-pgclient-resource.ts

import type { PgClient, PgCodec, WithPgClient } from '@dataplan/pg';
import type { PgTableResource } from '@graphile-contrib/pg-many-to-many';
import {
  ExecutableStep,
  type AccessStep,
  type ExecutionExtra,
  type GrafastResultsList,
  type GrafastValuesList,
  access,
  type SetterStep,
  setter,
  type SetterCapableStep,
} from 'grafast';

type PgResourceAttributes<TResource extends PgTableResource> =
  keyof TResource['codec']['attributes'];

interface ContextValues {
  pgSettings: Record<string, string>;
  withPgClient: WithPgClient;
}

/**
 * Like WithPgClient but with a typed and tracked resource
 */
export class WithTypedResourcePgClientStep<
    TResource extends PgTableResource = PgTableResource,
    TData = any,
    TResult extends Record<PgResourceAttributes<TResource>, any> = Record<
      PgResourceAttributes<TResource>,
      any
    >,
  >
  extends ExecutableStep<TResult>
  implements SetterCapableStep<Record<PgResourceAttributes<TResource>, any>>
{
  static $$export = {
    moduleName: '@litewarp/graphile-nested-mutations',
    exportName: 'WithTypedResourcePgClientStep',
  };

  isSyncAndSafe = false;
  hasSideEffects = true;

  private locked = false;

  /**
   * The resource to be returned by the step
   */
  public readonly resource: TResource;

  /**
   * The id for the PostgreSQL context plan.
   */
  private contextId: number;

  /**
   * The id for the data plan.
   */
  private dataId: number;

  /**
   * Names of the attributes to be selected
   */
  private attributes = new Map<
    PgResourceAttributes<TResource>,
    { name: string; depId: number | null; pgCodec: PgCodec }
  >();

  constructor(
    resource: TResource,
    $data: ExecutableStep<TData>,
    private callback: (
      client: PgClient,
      data: TData,
      selections: {
        attributes: PgResourceAttributes<TResource>[];
        values: [PgResourceAttributes<TResource>, any][];
      },
    ) => Promise<TResult>,
  ) {
    super();
    this.resource = resource;
    this.contextId = this.addDependency(this.resource.executor.context());
    this.dataId = this.addDependency($data);

    // add the primary key to the selection set if exists
    const primaryUnique = this.resource.uniques.find((u) => u.isPrimary);
    if (primaryUnique) {
      primaryUnique.attributes.forEach((attr) => {
        this._setAttribute(attr);
      });
    }
  }

  set(attr: PgResourceAttributes<TResource>, $step: ExecutableStep): void {
    this._setAttribute(attr, this.addDependency($step));
  }

  get(attr: PgResourceAttributes<TResource>): AccessStep<any> {
    this._setAttribute(attr);
    return access(this, attr);
  }

  execute(
    _count: number,
    values: GrafastValuesList<ContextValues | TData | TResult>[],
    _extra: ExecutionExtra,
  ): GrafastResultsList<TResult> {
    const contexts = values[this.contextId] as ContextValues[];
    const datas = values[this.dataId] as TData[];

    const addedValues = [...this.attributes.values()].reduce<
      [PgResourceAttributes<TResource>, any][]
    >((memo, { name, depId }) => {
      if (!depId) {
        return memo;
      }
      const val = values[depId];
      if (Array.isArray(val)) {
        return [...memo, [name, val[0]]];
      }
      return [...memo, [name, val]];
    }, []);

    const selections = {
      attributes: [...this.attributes.keys()],
      values: addedValues,
    };

    return contexts.map(async ({ pgSettings, withPgClient }, i) => {
      const data = datas[i];
      if (!data) {
        throw new Error('Data is undefined');
      }
      return withPgClient(pgSettings, (client) =>
        this.callback(client, data, selections),
      );
    });
  }

  setPlan(): SetterStep<Record<PgResourceAttributes<TResource>, any>, this> {
    if (this.locked) {
      throw new Error(
        `${this}: cannot set values once plan is locked ('setPlan')`,
      );
    }
    return setter(this);
  }

  public finalize(): void {
    if (!this.isFinalized) {
      this.locked = true;

      super.finalize();
    }
  }

  private _setAttribute(
    attr: PgResourceAttributes<TResource>,
    depId?: number | null,
  ): void {
    const attribute = Object.entries(this.resource.codec.attributes).find(
      ([name, _]) => name === attr,
    );

    if (!attribute) {
      throw new Error(
        `${this.resource.name} does not define an attribute named '${String(attr)}'`,
      );
    }

    const { via, codec } = attribute[1];

    if (via) {
      throw new Error(
        `Cannot set or select a 'via' attribute from WithTypedResourcePgClientStep`,
      );
    }

    const exists = this.attributes.has(attr) ? this.attributes.get(attr) : null;

    // if we already have the attribute added, make sure
    // we aren't overwriting a depId related to a step
    if (exists) {
      if (!depId && exists.depId) {
        return;
      }
    }
    this.attributes.set(attr, {
      name: attr.toString(),
      depId: depId || null,
      pgCodec: codec,
    });
  }
}

export function withPgClientResource<
  TResource extends PgTableResource = PgTableResource,
  TData = any,
>(
  resource: TResource,
  $data: ExecutableStep,
  callback: (
    client: PgClient,
    data: TData,
    selections: {
      attributes: PgResourceAttributes<TResource>[];
      values: [PgResourceAttributes<TResource>, any][];
    },
  ) => Promise<Record<PgResourceAttributes<TResource>, any>>,
): WithTypedResourcePgClientStep<TResource> {
  return new WithTypedResourcePgClientStep(
    resource,
    $data,
    (client, data, selections) => callback(client, data, selections),
  );
}