@pothos/plugin-drizzle 0.16.2 → 0.17.1

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@pothos/plugin-drizzle",
-  "version": "0.16.2",
+  "version": "0.17.1",
   "description": "A Pothos plugin for drizzle",
   "main": "./lib/index.js",
   "types": "./dts/index.d.ts",
@@ -39,20 +39,20 @@
     "graphql": "^16.10.0"
   },
   "devDependencies": {
-    "@libsql/client": "^0.15.15",
+    "@libsql/client": "^0.17.2",
     "drizzle-graphql": "^0.8.5",
     "drizzle-kit": "1.0.0-beta.10-42c284e",
     "drizzle-orm": "1.0.0-beta.10-42c284e",
     "graphql-scalars": "^1.25.0",
     "graphql-tag": "^2.12.6",
-    "postgres": "^3.4.7",
+    "postgres": "^3.4.9",
+    "@pothos/core": "4.12.0",
     "@pothos/plugin-add-graphql": "4.3.0",
     "@pothos/plugin-errors": "4.9.0",
-    "@pothos/plugin-relay": "4.6.2",
+    "@pothos/plugin-relay": "4.7.0",
     "@pothos/plugin-scope-auth": "4.1.6",
     "@pothos/plugin-with-input": "4.1.4",
-    "@pothos/test-utils": "2.1.3",
-    "@pothos/core": "4.12.0"
+    "@pothos/test-utils": "2.1.3"
   },
   "scripts": {
     "type": "tsc --project tsconfig.type.json",

package/src/drizzle-field-builder.ts

@@ -41,6 +41,7 @@ import type { DrizzleRef } from './interface-ref';
 import type {
   DrizzleConnectionShape,
   ListRelation,
+  PathInfo,
   RelatedConnectionOptions,
   RelatedCountOptions,
   RelatedFieldOptions,
@@ -219,9 +220,23 @@ export class DrizzleObjectFieldBuilder<
       )!;
     };
 
-    const getQuery = (args: PothosSchemaTypes.DefaultConnectionArguments, ctx: {}) => {
-      const { limit, orderBy, where, ...fieldQuery } =
-        (typeof query === 'function' ? query(args, ctx) : query) ?? {};
+    const getQuery = (
+      args: PothosSchemaTypes.DefaultConnectionArguments,
+      ctx: {},
+      pathInfo?: import('./types').PathInfo,
+    ) => {
+      const { limit, orderBy, where, ...fieldQuery } = ((typeof query === 'function'
+        ? (query as (args: {}, ctx: {}, pathInfo?: import('./types').PathInfo) => {})(
+            args,
+            ctx,
+            pathInfo,
+          )
+        : query) ?? {}) as {
+        limit?: number;
+        orderBy?: unknown;
+        where?: SQL;
+        columns?: Record<string, boolean>;
+      };
 
       const { cursorColumns, columns, ...connectionQuery } = drizzleCursorConnectionQuery({
         ctx,
@@ -255,6 +270,7 @@ export class DrizzleObjectFieldBuilder<
       context: object,
       nestedQuery: (query: unknown, path?: unknown) => { select?: object },
       getSelection: (path: string[]) => FieldNode | null,
+      pathInfo: import('./types').PathInfo,
     ) => {
       typeName ??= this.builder.configStore.getTypeConfig(ref).name;
 
@@ -279,7 +295,7 @@ export class DrizzleObjectFieldBuilder<
         };
       }
 
-      const nested = nestedQuery(getQuery(args, context).select, {
+      const nested = nestedQuery(getQuery(args, context, pathInfo).select, {
         getType: () => typeName!,
         paths: [[{ name: 'nodes' }], [{ name: 'edges' }, { name: 'node' }]],
       }) as SelectionMap;
@@ -500,14 +516,24 @@ export class DrizzleObjectFieldBuilder<
 
     const { query = {}, extensions, ...rest } = options;
 
-    const relationSelect = (args: object, context: object, nestedQuery: (query: unknown) => {}) => {
+    const relationSelect = (
+      args: object,
+      context: object,
+      nestedQuery: (query: unknown) => {},
+      _resolveSelection: unknown,
+      pathInfo: PathInfo,
+    ) => {
       const relQuery = {
         columns: {},
         with: {
           [name]: {
             ...nestedQuery(query),
             ...((typeof query === 'function'
-              ? (query as (args: {}, context: {}) => {})(args, context)
+              ? (query as (args: {}, context: {}, pathInfo: PathInfo) => {})(
+                  args,
+                  context,
+                  pathInfo,
+                )
               : query) as {}),
           },
         },

package/src/index.ts

@@ -121,14 +121,23 @@ export class PothosDrizzlePlugin<Types extends SchemaTypes> extends BasePlugin<T
                   args: {},
                   ctx: Types['Context'],
                   nestedQuery: (query: unknown, path?: string[]) => never,
-                ) => ({
-                  columns: {},
-                  ...(select as (args: unknown, ctx: unknown, nestedQuery: unknown) => {})(
-                    args,
-                    ctx,
-                    nestedQuery,
-                  ),
-                })
+                  _resolveSelection: unknown,
+                  pathInfo: import('./types').PathInfo,
+                ) => {
+                  // Attach PathInfo properties to nestedQuery for t.field select callbacks
+                  const nestedQueryWithPath = Object.assign(nestedQuery, {
+                    path: pathInfo.path,
+                    segments: pathInfo.segments,
+                  });
+                  return {
+                    columns: {},
+                    ...(select as (args: unknown, ctx: unknown, nestedQuery: unknown) => {})(
+                      args,
+                      ctx,
+                      nestedQueryWithPath,
+                    ),
+                  };
+                }
               : {
                   columns: {},
                   ...select,

package/src/types.ts

@@ -44,6 +44,18 @@ import type { DrizzleRef } from './interface-ref';
 import type { IndirectInclude } from './utils/map-query';
 import type { SelectionMap } from './utils/selections';
 
+export interface FieldPathInfo {
+  field: string;
+  alias: string;
+  parentType: string;
+  isList: boolean;
+}
+
+export interface PathInfo {
+  path: string[];
+  segments: FieldPathInfo[];
+}
+
 export type DrizzleClient<TCountSource = Table | SQL | SQLWrapper> = {
   readonly _: {
     readonly schema: unknown;
@@ -323,10 +335,11 @@ export type DrizzleObjectFieldOptions<
         | ((
             args: InputShapeFromFields<Args>,
             ctx: Types['Context'],
-            nestedSelection: <Selection extends boolean | {}>(
+            nestedSelection: (<Selection extends boolean | {}>(
               selection?: Selection,
               path?: string[],
-            ) => Selection,
+            ) => Selection) &
+              PathInfo,
           ) => DBQueryConfig<'one', Types['DrizzleRelations'], ExtractTable<Types, ParentShape>>)
       );
   };
@@ -345,6 +358,7 @@ export type DrizzleFieldSelection =
         type?: string,
       ) => DBQueryConfig<'one'> | boolean,
       resolveSelection: (path: string[]) => FieldNode | null,
+      pathInfo: PathInfo,
     ) => SelectionMap);
 
 export type ExtractTable<Types extends SchemaTypes, Shape> = Shape extends {
@@ -480,7 +494,13 @@ export type QueryForField<
         'columns' | 'extra' | 'with'
       >
 ) extends infer QueryConfig
-  ? QueryConfig | ((args: InputShapeFromFields<Args>, context: Types['Context']) => QueryConfig)
+  ?
+      | QueryConfig
+      | ((
+          args: InputShapeFromFields<Args>,
+          context: Types['Context'],
+          pathInfo: PathInfo,
+        ) => QueryConfig)
   : never;
 
 export type QueryForDrizzleField<
@@ -503,7 +523,7 @@ export type QueryForRelatedConnection<
 > & {
   orderBy?: ConnectionOrderBy<Table> | ((table: Table) => ConnectionOrderBy<Table>);
 } extends infer QueryConfig
-  ? QueryConfig | ((args: Args, context: Types['Context']) => QueryConfig)
+  ? QueryConfig | ((args: Args, context: Types['Context'], pathInfo: PathInfo) => QueryConfig)
   : never;
 
 export type QueryForDrizzleConnection<
@@ -674,28 +694,31 @@ export type DrizzleConnectionShape<
   T,
   Parent,
   Args extends InputFieldMap,
-> = Merge<
-  ShapeFromConnection<PothosSchemaTypes.ConnectionShapeHelper<Types, T, false>> extends infer Shape
-    ? Shape & {
-        parent: Parent;
-        args: InputShapeFromFields<Args> & PothosSchemaTypes.DefaultConnectionArguments;
-      }
-    : never
-> extends infer C
-  ? [C] extends [
-      {
-        edges: infer Edges;
-      },
-    ]
-    ? Merge<
-        Omit<C, 'edges'> & {
-          edges: Edges extends Iterable<MaybePromise<infer Edge> | null | undefined>
-            ? Merge<Edge & { connection: C }>[]
-            : never;
+> =
+  Merge<
+    ShapeFromConnection<
+      PothosSchemaTypes.ConnectionShapeHelper<Types, T, false>
+    > extends infer Shape
+      ? Shape & {
+          parent: Parent;
+          args: InputShapeFromFields<Args> & PothosSchemaTypes.DefaultConnectionArguments;
         }
-      >
-    : C
-  : never;
+      : never
+  > extends infer C
+    ? [C] extends [
+        {
+          edges: infer Edges;
+        },
+      ]
+      ? Merge<
+          Omit<C, 'edges'> & {
+            edges: Edges extends Iterable<MaybePromise<infer Edge> | null | undefined>
+              ? Merge<Edge & { connection: C }>[]
+              : never;
+          }
+        >
+      : C
+    : never;
 
 export type WithBrand<T> = T & { [typeBrandKey]: string };
 

package/src/utils/map-query.ts

@@ -15,11 +15,12 @@ import {
   getNamedType,
   type InlineFragmentNode,
   isInterfaceType,
+  isListType,
   isObjectType,
   Kind,
   type SelectionSetNode,
 } from 'graphql';
-import type { DrizzleFieldSelection } from '../types';
+import type { DrizzleFieldSelection, FieldPathInfo, PathInfo } from '../types';
 import type { PothosDrizzleSchemaConfig } from './config';
 import { type LoaderMappings, setLoaderMappings } from './loader-map';
 import {
@@ -47,6 +48,7 @@ function addTypeSelectionsForField(
   selection: FieldNode,
   indirectPath: string[],
   deferred?: boolean,
+  segments: FieldPathInfo[] = [],
 ) {
   if (selection.name.value.startsWith('__')) {
     return;
@@ -78,6 +80,7 @@ function addTypeSelectionsForField(
           field,
           path,
           deferred,
+          segments,
         );
       },
       deferred,
@@ -92,6 +95,7 @@ function addTypeSelectionsForField(
       selection,
       indirectPath,
       deferred,
+      segments,
     );
     return;
   }
@@ -105,7 +109,16 @@ function addTypeSelectionsForField(
   }
 
   if (selection.selectionSet && (!deferred || !state.skipDeferredFragments)) {
-    addNestedSelections(config, type, context, info, state, selection.selectionSet, indirectPath);
+    addNestedSelections(
+      config,
+      type,
+      context,
+      info,
+      state,
+      selection.selectionSet,
+      indirectPath,
+      segments,
+    );
   }
 }
 
@@ -225,12 +238,13 @@ function addNestedSelections(
   state: SelectionState,
   selections: SelectionSetNode,
   indirectPath: string[],
+  segments: FieldPathInfo[] = [],
 ) {
   let parentType = type;
   for (const selection of selections.selections) {
     switch (selection.kind) {
       case Kind.FIELD:
-        addFieldSelection(config, type, context, info, state, selection, indirectPath);
+        addFieldSelection(config, type, context, info, state, selection, indirectPath, segments);
 
         continue;
       case Kind.FRAGMENT_SPREAD:
@@ -257,6 +271,7 @@ function addNestedSelections(
           state,
           info.fragments[selection.name.value].selectionSet,
           indirectPath,
+          segments,
         );
 
         continue;
@@ -285,6 +300,7 @@ function addNestedSelections(
           state,
           selection.selectionSet,
           indirectPath,
+          segments,
         );
 
         continue;
@@ -305,6 +321,7 @@ function addFieldSelection(
   state: SelectionState,
   selection: FieldNode,
   indirectPath: string[],
+  segments: FieldPathInfo[] = [],
 ) {
   if (selection.name.value.startsWith('__') || fieldSkipped(info, selection)) {
     return;
@@ -314,11 +331,26 @@ function addFieldSelection(
     throw new PothosValidationError(`Unknown field ${selection.name.value} on ${type.name}`);
   }
 
+  const allSegments = [
+    ...segments,
+    {
+      field: selection.name.value,
+      alias: selection.alias?.value ?? selection.name.value,
+      parentType: type.name,
+      isList: isListType(field.type) || isListType(getNamedType(field.type)),
+    },
+  ];
+
   const fieldSelect = field.extensions?.pothosDrizzleSelect as DrizzleFieldSelection | undefined;
 
   let fieldSelectionMap: DBQueryConfig<'one'> | undefined;
   let mappings: LoaderMappings = {};
   if (typeof fieldSelect === 'function') {
+    const pathInfo: PathInfo = {
+      path: allSegments.map((s) => `${s.parentType}.${s.field}`),
+      segments: allSegments,
+    };
+
     const args = getMappedArgumentValues(field, selection, context, info) as Record<
       string,
       unknown
@@ -369,11 +401,22 @@ function addFieldSelection(
                 resolvedField,
                 path,
                 deferred,
+                allSegments,
               );
             },
           );
         }
-        addTypeSelectionsForField(config, returnType, context, info, fieldState, selection, []);
+        addTypeSelectionsForField(
+          config,
+          returnType,
+          context,
+          info,
+          fieldState,
+          selection,
+          [],
+          undefined,
+          allSegments,
+        );
         mappings = fieldState.mappings;
         return selectionToQuery(config, fieldState);
       },
@@ -394,6 +437,7 @@ function addFieldSelection(
         );
         return node;
       },
+      pathInfo,
     );
   } else {
     fieldSelectionMap = fieldSelect!;
@@ -522,7 +566,31 @@ export function stateFromInfo<T extends SelectionMap>({
   } else {
     state = createStateForSelection(config, info, type, undefined, initialSelection);
 
-    addTypeSelectionsForField(config, type, context, info, state, info.fieldNodes[0], []);
+    // Create initial segment for the root query field
+    const rootFieldNode = info.fieldNodes[0];
+    const rootField = info.parentType.getFields()[rootFieldNode.name.value];
+    const initialSegments: FieldPathInfo[] = rootField
+      ? [
+          {
+            field: rootFieldNode.name.value,
+            alias: rootFieldNode.alias?.value ?? rootFieldNode.name.value,
+            parentType: info.parentType.name,
+            isList: isListType(rootField.type) || isListType(getNamedType(rootField.type)),
+          },
+        ]
+      : [];
+
+    addTypeSelectionsForField(
+      config,
+      type,
+      context,
+      info,
+      state,
+      info.fieldNodes[0],
+      [],
+      undefined,
+      initialSegments,
+    );
   }
 
   if (!state) {

package/src/utils/selections.ts

@@ -149,8 +149,10 @@ export function mergeSelection(
   }
 
   if (columns && !state.allColumns) {
-    for (const key of Object.keys(columns)) {
-      state.columns.add(key);
+    for (const [key, value] of Object.entries(columns)) {
+      if (value) {
+        state.columns.add(key);
+      }
     }
   } else {
     state.allColumns = true;