@pothos/plugin-drizzle 0.0.1

package/src/index.ts

@@ -0,0 +1,175 @@
+import './global-types';
+import './field-builder';
+import './schema-builder';
+import { GraphQLFieldResolver } from 'graphql';
+import SchemaBuilder, {
+  BasePlugin,
+  createInputValueMapper,
+  mapInputFields,
+  PartialResolveInfo,
+  PothosOutputFieldConfig,
+  SchemaTypes,
+} from '@pothos/core';
+import { ModelLoader } from './model-loader';
+import { DrizzleGraphQLInputExtensions } from './types';
+import {
+  extractFilters,
+  extractOrderBy,
+  remapFromGraphQLSingleInput,
+} from './utils/drizzle-graphql';
+import { getLoaderMapping, setLoaderMappings } from './utils/loader-map';
+
+export * from './types';
+
+const pluginName = 'drizzle';
+
+export default pluginName;
+
+export class PothosDrizzlePlugin<Types extends SchemaTypes> extends BasePlugin<Types> {
+  override onOutputFieldConfig(
+    fieldConfig: PothosOutputFieldConfig<Types>,
+  ): PothosOutputFieldConfig<Types> | null {
+    const argMappings = mapInputFields<Types, DrizzleGraphQLInputExtensions>(
+      fieldConfig.args,
+      this.buildCache,
+      (inputField) => {
+        if (inputField.type.kind === 'InputObject') {
+          const config = this.buildCache.getTypeConfig(inputField.type.ref);
+
+          return (config.extensions?.drizzleGraphQL as DrizzleGraphQLInputExtensions) ?? null;
+        }
+
+        return null;
+      },
+    );
+
+    const argMapper = argMappings
+      ? createInputValueMapper(
+          argMappings,
+          (input, mappings, ctx: Types['Context'], info: PartialResolveInfo) => {
+            if (!mappings.value) {
+              return input;
+            }
+
+            const { table, tableConfig, inputType } = mappings.value;
+
+            switch (inputType) {
+              case 'orderBy':
+                return extractOrderBy(tableConfig, input as never);
+              case 'filters':
+                return extractFilters(tableConfig, table, input as never);
+              case 'insert':
+                return remapFromGraphQLSingleInput(input as never, tableConfig);
+              case 'update':
+                return remapFromGraphQLSingleInput(input as never, tableConfig);
+
+              default:
+                throw new Error(`Unknown drizzle input type: ${inputType}`);
+            }
+          },
+        )
+      : null;
+
+    const argMappers: typeof fieldConfig.argMappers = argMapper
+      ? [
+          ...(fieldConfig.argMappers ?? []),
+          (args, context, info) => argMapper(args, undefined, context, info),
+        ]
+      : fieldConfig.argMappers;
+
+    if (fieldConfig.kind === 'DrizzleObject' && fieldConfig.pothosOptions.select) {
+      const { select } = fieldConfig.pothosOptions;
+      return {
+        ...fieldConfig,
+        argMappers,
+        extensions: {
+          ...fieldConfig.extensions,
+          pothosDrizzleSelect:
+            typeof select === 'function'
+              ? (
+                  args: {},
+                  ctx: Types['Context'],
+                  nestedQuery: (query: unknown, path?: string[]) => never,
+                ) => ({
+                  select: (select as (args: unknown, ctx: unknown, nestedQuery: unknown) => {})(
+                    args,
+                    ctx,
+                    nestedQuery,
+                  ),
+                })
+              : select,
+        },
+      };
+    }
+
+    return {
+      ...fieldConfig,
+      argMappers,
+    };
+  }
+
+  override wrapResolve(
+    resolver: GraphQLFieldResolver<unknown, Types['Context'], object, unknown>,
+    fieldConfig: PothosOutputFieldConfig<Types>,
+  ): GraphQLFieldResolver<unknown, Types['Context'], object> {
+    if (fieldConfig.kind !== 'DrizzleObject' || !fieldConfig.extensions?.pothosDrizzleSelect) {
+      return resolver;
+    }
+
+    const parentConfig = this.buildCache.getTypeConfig(fieldConfig.parentType);
+
+    const parentTypes = new Set([fieldConfig.parentType]);
+
+    if (parentConfig.kind === 'Interface' || parentConfig.kind === 'Object') {
+      parentConfig.interfaces.forEach((iface) => {
+        const interfaceConfig = this.buildCache.getTypeConfig(iface, 'Interface');
+        if (interfaceConfig.extensions?.pothosDrizzleModel) {
+          parentTypes.add(interfaceConfig.name);
+        }
+      });
+    }
+
+    const modelLoader = parentConfig.extensions?.pothosDrizzleLoader as ReturnType<
+      typeof ModelLoader.forModel
+    >;
+
+    if (!modelLoader) {
+      throw new Error(`ModelLoader not found for type ${parentConfig.name}`);
+    }
+
+    return (parent, args, context, info) => {
+      let mapping = getLoaderMapping(context, info.path, info.parentType.name);
+
+      if (!mapping) {
+        for (const parentType of parentTypes) {
+          mapping = getLoaderMapping(context, info.path, parentType);
+          if (mapping) {
+            break;
+          }
+        }
+      }
+
+      if (mapping) {
+        setLoaderMappings(context, info, mapping);
+
+        return resolver(parent, args, context, info);
+      }
+
+      const primaryKeys = (
+        parentConfig.extensions!.pothosDrizzleTable as { primaryKey?: { name: string }[] }
+      ).primaryKey;
+
+      if (!primaryKeys || primaryKeys.length === 0) {
+        throw new Error(
+          `Field ${fieldConfig.name} not resolved in initial query and no primary key found for type ${parentConfig.name}`,
+        );
+      }
+
+      return modelLoader(context)
+        .loadSelection(info, parent as object)
+        .then((result) => resolver(result, args, context, info));
+    };
+  }
+}
+
+SchemaBuilder.registerPlugin(pluginName, PothosDrizzlePlugin, {});

package/src/interface-ref.ts

@@ -0,0 +1,24 @@
+import { InterfaceRef, SchemaTypes } from '@pothos/core';
+import { DrizzleObjectRef, drizzleTableKey } from './object-ref';
+
+export type DrizzleRef<
+  Types extends SchemaTypes,
+  Table extends keyof Types['DrizzleRelationSchema'] = keyof Types['DrizzleRelationSchema'],
+  T = {},
+> = DrizzleInterfaceRef<Types, Table, T> | DrizzleObjectRef<Types, Table, T>;
+
+export class DrizzleInterfaceRef<
+  Types extends SchemaTypes,
+  Table extends keyof Types['DrizzleRelationSchema'] = keyof Types['DrizzleRelationSchema'],
+  T = {},
+> extends InterfaceRef<Types, T> {
+  [drizzleTableKey]!: Table;
+
+  tableName: string;
+
+  constructor(name: string, tableName: string) {
+    super(name);
+
+    this.tableName = tableName;
+  }
+}

package/src/model-loader.ts

@@ -0,0 +1,235 @@
+import { GraphQLResolveInfo } from 'graphql';
+import {
+  Column,
+  inArray,
+  RelationalSchemaConfig,
+  sql,
+  TableRelationalConfig,
+  TablesRelationalConfig,
+} from 'drizzle-orm';
+import { createContextCache, SchemaTypes } from '@pothos/core';
+import { cacheKey, setLoaderMappings } from './utils/loader-map';
+import { selectionStateFromInfo, stateFromInfo } from './utils/map-query';
+import {
+  mergeSelection,
+  selectionCompatible,
+  SelectionMap,
+  SelectionState,
+  selectionToQuery,
+} from './utils/selections';
+import { getSchemaConfig, PothosDrizzleSchemaConfig } from './utils/config';
+
+interface ResolvablePromise<T> {
+  promise: Promise<T>;
+  resolve: (value: T) => void;
+  reject: (err: unknown) => void;
+}
+export class ModelLoader {
+  context: object;
+
+  builder: PothosSchemaTypes.SchemaBuilder<never>;
+
+  modelName: string;
+
+  queryCache = new Map<string, { selection: SelectionState; query: SelectionMap }>();
+
+  staged = new Set<{
+    state: SelectionState;
+    models: Map<object, ResolvablePromise<Record<string, unknown> | null>>;
+  }>();
+
+  config: PothosDrizzleSchemaConfig;
+  table: TableRelationalConfig;
+  columns: Column[];
+  selectSQL;
+
+  constructor(
+    context: object,
+    builder: PothosSchemaTypes.SchemaBuilder<never>,
+    modelName: string,
+    columns?: Column[],
+  ) {
+    this.context = context;
+    this.builder = builder;
+    this.modelName = modelName;
+    this.config = getSchemaConfig(builder);
+    this.table = this.config.schema[modelName];
+    this.columns = columns ?? this.table.primaryKey;
+    this.selectSQL =
+      this.columns.length > 1
+        ? sql`(${sql.join(this.columns, sql`, `)})`
+        : this.columns[0].getSQL();
+  }
+
+  static forModel<Types extends SchemaTypes>(
+    modelName: string,
+    builder: PothosSchemaTypes.SchemaBuilder<Types>,
+    columns?: Column[],
+  ) {
+    return createContextCache(
+      (model) => new ModelLoader(model, builder as never, modelName, columns),
+    );
+  }
+
+  sqlForModel(model: object) {
+    const values = this.columns.map((key) => {
+      const columnName = key.name as keyof typeof model;
+      if (columnName in model) {
+        return model[columnName] as string | number;
+      }
+
+      throw new Error(`Primary key column ${columnName} not found on ${this.modelName}`);
+    });
+
+    return this.columns.length > 1 ? values : values[0];
+  }
+
+  getSelection(info: GraphQLResolveInfo) {
+    const key = cacheKey(info.parentType.name, info.path);
+    if (!this.queryCache.has(key)) {
+      const selection = selectionStateFromInfo(this.config, this.context, info);
+      this.queryCache.set(key, {
+        selection,
+        query: selectionToQuery(selection),
+      });
+    }
+
+    return this.queryCache.get(key)!;
+  }
+
+  getSelectionForField(info: GraphQLResolveInfo, typeName: string) {
+    const key = cacheKey(typeName, info.path);
+    if (!this.queryCache.has(key)) {
+      const selection = stateFromInfo({
+        config: this.config,
+        context: this.context,
+        info,
+        typeName,
+      });
+
+      this.queryCache.set(key, {
+        selection,
+        query: selectionToQuery(selection),
+      });
+    }
+
+    return this.queryCache.get(key)!;
+  }
+
+  async loadSelection(info: GraphQLResolveInfo, model: object) {
+    const { selection, query } = this.getSelection(info);
+
+    const result = await this.stageQuery(selection, query, model);
+
+    if (result) {
+      const mappings = selection.mappings[info.path.key];
+
+      if (mappings) {
+        setLoaderMappings(this.context, info, mappings.mappings);
+      }
+    }
+
+    return result;
+  }
+
+  async loadSelectionForField(info: GraphQLResolveInfo, model: object, returnType: string) {
+    const { selection, query } = this.getSelectionForField(info, returnType);
+
+    const result = await this.stageQuery(selection, query, model);
+
+    if (result) {
+      setLoaderMappings(this.context, info, selection.mappings);
+    }
+
+    return result;
+  }
+
+  async stageQuery(selection: SelectionState, query: SelectionMap, model: object) {
+    for (const entry of this.staged) {
+      if (selectionCompatible(entry.state, query)) {
+        mergeSelection(this.config, entry.state, query);
+
+        if (!entry.models.has(model)) {
+          entry.models.set(model, createResolvablePromise<Record<string, unknown> | null>());
+        }
+
+        return entry.models.get(model)!.promise;
+      }
+    }
+
+    return this.initLoad(selection, model);
+  }
+
+  initLoad(selection: SelectionState, model: object) {
+    const promise = createResolvablePromise<Record<string, unknown> | null>();
+    const entry = {
+      state: selection,
+      models: new Map([[model, promise]]),
+    };
+    this.staged.add(entry);
+
+    const nextTick = createResolvablePromise<void>();
+
+    nextTick.promise
+      .then(() => {
+        const api = (
+          this.builder.options.drizzle.client.query as Record<
+            string,
+            { findMany: (...args: unknown[]) => Promise<Record<string, unknown>[]> }
+          >
+        )[this.modelName];
+
+        const query = api.findMany({
+          ...selectionToQuery(selection),
+          where: inArray(
+            this.selectSQL,
+            [entry.models.keys()].map(([model]) => this.sqlForModel(model)),
+          ),
+        });
+
+        query.then(
+          (results) => {
+            for (const [model, promise] of entry.models.entries()) {
+              const result = results.find((row) =>
+                this.table.primaryKey.every(
+                  (key) => row[key.name] === (model as Record<string, unknown>)[key.name],
+                ),
+              );
+
+              if (result) {
+                promise.resolve(result ?? null);
+              } else
+                promise.reject(
+                  new Error(`Model ${this.modelName}(${this.sqlForModel(model)}) not found`),
+                );
+            }
+          },
+          (err) => {
+            for (const promise of entry.models.values()) {
+              promise.reject(err);
+            }
+          },
+        );
+      })
+      .catch((err) => {
+        for (const promise of entry.models.values()) {
+          promise.reject(err);
+        }
+      });
+
+    setTimeout(() => void nextTick.resolve(), 0);
+
+    return promise.promise;
+  }
+}
+
+function createResolvablePromise<T = unknown>(): ResolvablePromise<T> {
+  let resolveFn!: (value: T) => void;
+  let rejectFn!: (reason?: unknown) => void;
+  const promise = new Promise<T>((resolve, reject) => {
+    resolveFn = resolve;
+    rejectFn = reject;
+  });
+
+  return { promise, resolve: resolveFn, reject: rejectFn };
+}

package/src/node-ref.ts

@@ -0,0 +1,8 @@
+import { SchemaTypes } from '@pothos/core';
+import { DrizzleObjectRef } from './object-ref';
+
+export class DrizzleNodeRef<
+  Types extends SchemaTypes,
+  Table extends keyof Types['DrizzleRelationSchema'] = keyof Types['DrizzleRelationSchema'],
+  T = {},
+> extends DrizzleObjectRef<Types, Table, T> {}

package/src/object-ref.ts

@@ -0,0 +1,51 @@
+import {
+  abstractReturnShapeKey,
+  brandWithType,
+  ObjectRef,
+  SchemaTypes,
+  typeBrandKey,
+} from '@pothos/core';
+import type { WithBrand } from './types';
+
+export const drizzleTableKey = Symbol.for('Pothos.drizzleTableKey');
+
+export class DrizzleObjectRef<
+  Types extends SchemaTypes,
+  Table extends keyof Types['DrizzleRelationSchema'] = keyof Types['DrizzleRelationSchema'],
+  T = {},
+> extends ObjectRef<Types, T> {
+  [drizzleTableKey]!: Types['DrizzleRelationSchema'][Table];
+
+  [abstractReturnShapeKey]!: WithBrand<T>;
+
+  tableName: string;
+
+  constructor(name: string, tableName: string) {
+    super(name);
+
+    this.tableName = tableName;
+  }
+
+  addBrand<V extends T | T[]>(
+    value: V,
+  ): V extends T[] ? { [K in keyof V]: WithBrand<V[K]> } : WithBrand<V> {
+    if (Array.isArray(value)) {
+      value.forEach((val) => void brandWithType(val, this.name as never));
+
+      return value as never;
+    }
+
+    brandWithType(value, this.name as never);
+
+    return value as never;
+  }
+
+  hasBrand(value: unknown) {
+    return (
+      typeof value === 'object' &&
+      value !== null &&
+      typeBrandKey in value &&
+      (value as { [typeBrandKey]?: unknown })[typeBrandKey] === this.name
+    );
+  }
+}