@pothos/plugin-drizzle 0.14.0 → 0.16.0

package/src/field-builder.ts

@@ -1,7 +1,8 @@
-import type { FieldRef } from '@pothos/core';
+import type { FieldRef, MaybePromise } from '@pothos/core';
 import { type FieldKind, ObjectRef, RootFieldBuilder, type SchemaTypes } from '@pothos/core';
 import type { TableRelationalConfig } from 'drizzle-orm';
 import type { GraphQLResolveInfo } from 'graphql';
+import { isInterfaceType, isObjectType, Kind } from 'graphql';
 import type { DrizzleRef } from './interface-ref';
 import type { DrizzleConnectionFieldOptions } from './types';
 import { getSchemaConfig } from './utils/config';
@@ -101,6 +102,7 @@ fieldBuilderProto.drizzleConnection = function drizzleConnection<
     maxSize = this.builder.options.drizzle?.maxConnectionSize,
     defaultSize = this.builder.options.drizzle?.defaultConnectionSize,
     resolve,
+    totalCount,
     ...options
   }: DrizzleConnectionFieldOptions<
     SchemaTypes,
@@ -133,6 +135,21 @@ fieldBuilderProto.drizzleConnection = function drizzleConnection<
         context: {},
         info: GraphQLResolveInfo,
       ) => {
+        const returnType = info.returnType;
+        const fields =
+          isObjectType(returnType) || isInterfaceType(returnType) ? returnType.getFields() : {};
+
+        const selections = info.fieldNodes;
+
+        const totalCountOnly = selections.every((selection) =>
+          selection.selectionSet?.selections.every(
+            (s) =>
+              s.kind === Kind.FIELD &&
+              (fields[s.name.value]?.extensions?.pothosDrizzleTotalCount ||
+                s.name.value === '__typename'),
+          ),
+        );
+
         return resolveDrizzleCursorConnection(
           tableName,
           info,
@@ -143,8 +160,13 @@ fieldBuilderProto.drizzleConnection = function drizzleConnection<
             maxSize,
             defaultSize,
             args,
+            totalCount: totalCount && (() => totalCount(parent, args as never, context, info)),
           },
           (q) => {
+            if (totalCountOnly) {
+              return [];
+            }
+
             // return checkIfQueryIsUsed(
             //   this.builder,
             //   query,
@@ -160,6 +182,23 @@ fieldBuilderProto.drizzleConnection = function drizzleConnection<
       ? connectionOptions
       : {
           ...connectionOptions,
+          fields: totalCount
+            ? (
+                t: PothosSchemaTypes.ObjectFieldBuilder<
+                  SchemaTypes,
+                  { totalCount?: () => MaybePromise<number> }
+                >,
+              ) => ({
+                totalCount: t.int({
+                  nullable: false,
+                  extensions: {
+                    pothosDrizzleTotalCount: true,
+                  },
+                  resolve: (parent) => parent.totalCount?.(),
+                }),
+                ...(connectionOptions as { fields?: (t: unknown) => {} }).fields?.(t),
+              })
+            : (connectionOptions as { fields: undefined }).fields,
           extensions: {
             ...(connectionOptions as Record<string, object> | undefined)?.extensions,
           },

package/src/types.ts

@@ -51,6 +51,7 @@ export type DrizzleClient = {
     readonly relations: AnyRelations;
   };
   query: {};
+  $count: (table: Table, filter?: SQL) => SQL<number>;
 };
 
 export type DrizzlePluginOptions<Types extends SchemaTypes> = {
@@ -428,6 +429,18 @@ export type RefForRelation<Types extends SchemaTypes, Rel extends Relation> = Re
   ? ObjectRef<Types, TypesForRelation<Types, Rel>>
   : [ObjectRef<Types, TypesForRelation<Types, Rel>>];
 
+export type RelatedCountOptions<
+  Types extends SchemaTypes,
+  Shape,
+  Args extends InputFieldMap,
+  Where,
+> = Omit<
+  PothosSchemaTypes.ObjectFieldOptions<Types, Shape, 'Int', false, Args, number>,
+  'type' | InferredFieldOptionKeys
+> & {
+  where?: Where | ((args: InputShapeFromFields<Args>, context: Types['Context']) => Where);
+};
+
 export type TypesForRelation<Types extends SchemaTypes, Rel extends Relation> = BuildQueryResult<
   Types['DrizzleRelations'],
   Types['DrizzleRelations'][Rel['targetTable']['_']['name']],
@@ -543,6 +556,12 @@ export type DrizzleConnectionFieldOptions<
         type: Type;
         defaultSize?: number | ((args: ConnectionArgs, ctx: Types['Context']) => number);
         maxSize?: number | ((args: ConnectionArgs, ctx: Types['Context']) => number);
+        totalCount?: (
+          parent: ParentShape,
+          args: ConnectionArgs,
+          context: Types['Context'],
+          info: GraphQLResolveInfo,
+        ) => MaybePromise<number>;
 
         resolve: (
           query: <T = {}>(
@@ -626,7 +645,7 @@ export type RelatedConnectionOptions<
 
         defaultSize?: number | ((args: ConnectionArgs, ctx: Types['Context']) => number);
         maxSize?: number | ((args: ConnectionArgs, ctx: Types['Context']) => number);
-        // totalCount?: NeedsResolve extends false ? boolean : false;
+        totalCount?: boolean;
       }
     : never);
 

package/src/utils/cursors.ts

@@ -462,6 +462,7 @@ export function wrapConnectionResult<T extends {}>(
   cursor: (node: T) => string,
   resolveNode?: (node: Record<string, unknown>) => unknown,
   parent?: unknown,
+  totalCount?: number | (() => MaybePromise<number>) | null,
 ) {
   const gotFullResults = results.length === Math.abs(limit);
   const hasNextPage = args.before ? true : args.last ? false : gotFullResults;
@@ -471,6 +472,7 @@ export function wrapConnectionResult<T extends {}>(
   const connection = {
     parent,
     args,
+    totalCount,
     edges: [] as ({ cursor: string; node: unknown } | null)[],
     pageInfo: {
       startCursor: null as string | null,
@@ -513,14 +515,16 @@ export async function resolveDrizzleCursorConnection<T extends {}>(
   info: GraphQLResolveInfo,
   typeName: string,
   config: PothosDrizzleSchemaConfig,
-  options: Omit<DrizzleCursorConnectionQueryOptions, 'orderBy' | 'config' | 'table'>,
+  options: Omit<DrizzleCursorConnectionQueryOptions, 'orderBy' | 'config' | 'table'> & {
+    totalCount?: () => MaybePromise<number>;
+  },
   resolve: (
     queryFn: (query: QueryForDrizzleConnection<SchemaTypes, TableRelationalConfig>) => SelectionMap,
   ) => MaybePromise<readonly T[]>,
   parent: unknown,
 ) {
   const table = config.relations[tableName];
-  let query: DBQueryConfig<'many'>;
+  let query: DBQueryConfig<'many'> | undefined;
   let formatter: (node: Record<string, unknown>) => string;
   const results = await resolve((q = {}) => {
     const { cursorColumns, ...connectionQuery } = drizzleCursorConnectionQuery({
@@ -562,12 +566,29 @@ export async function resolveDrizzleCursorConnection<T extends {}>(
     return results;
   }
 
+  // Handle totalCountOnly case where resolve returns [] without calling query function
+  if (!query) {
+    return {
+      parent,
+      args: options.args,
+      totalCount: options.totalCount,
+      edges: [],
+      pageInfo: {
+        startCursor: null,
+        endCursor: null,
+        hasPreviousPage: false,
+        hasNextPage: false,
+      },
+    };
+  }
+
   return wrapConnectionResult(
     results,
     options.args,
-    query!.limit as number,
+    query.limit as number,
     formatter!,
     undefined,
     parent,
+    options.totalCount,
   );
 }

package/src/utils/map-query.ts

@@ -401,15 +401,35 @@ function addFieldSelection(
 
   if (fieldSelect && selectionCompatible(state, fieldSelectionMap, true)) {
     mergeSelection(config, state, fieldSelectionMap);
-    state.mappings[selection.alias?.value ?? selection.name.value] = {
-      field: selection.name.value,
-      type: type.name,
-      mappings,
-      indirectPath,
-    };
+
+    state.mappings = mergeMappings(state.mappings, {
+      [selection.alias?.value ?? selection.name.value]: {
+        field: selection.name.value,
+        type: type.name,
+        mappings,
+        indirectPath,
+      },
+    });
   }
 }
 
+function mergeMappings(existing: LoaderMappings, incoming: LoaderMappings): LoaderMappings {
+  const result: LoaderMappings = { ...existing };
+
+  for (const [key, value] of Object.entries(incoming)) {
+    if (result[key]) {
+      result[key] = {
+        ...result[key],
+        mappings: mergeMappings(result[key].mappings, value.mappings),
+      };
+    } else {
+      result[key] = value;
+    }
+  }
+
+  return result;
+}
+
 export interface QueryFromInfoOptions<T extends SelectionMap> {
   config: PothosDrizzleSchemaConfig;
   context: object;