@pothos/plugin-drizzle 0.19.2 → 0.20.1

package/src/utils/map-query.ts

@@ -1,39 +1,60 @@
-import { isThenable } from '@pothos/core';
+import { isThenable, type MaybePromise } from '@pothos/core';
 import { type PathSegment, Plan } from '@pothos/selection-mapper';
+import type { DBQueryConfig } from 'drizzle-orm';
 import type { GraphQLResolveInfo } from 'graphql';
 import { type DrizzlePlan, drizzleAdapter } from './adapter.js';
+import { checkAwaitSelections } from './await-selections.js';
 import type { PothosDrizzleSchemaConfig } from './config.js';
 import type { SelectionMap } from './selections.js';
 
-export interface QueryFromInfoOptions<T extends SelectionMap> {
+export interface QueryFromInfoOptions<Await extends boolean = false> {
   config: PothosDrizzleSchemaConfig;
   context: object;
   info: GraphQLResolveInfo;
   typeName?: string;
-  select?: T;
   path?: PathSegment[];
   paths?: PathSegment[][];
+  /** Whether the caller will await the query. Async selections otherwise throw. */
+  awaitSelections?: Await;
 }
 
-export function queryFromInfo<T extends SelectionMap>({
-  config,
-  select,
-  ...options
-}: QueryFromInfoOptions<T>): T {
-  const plan = Plan.fromInfo(drizzleAdapter(config), {
-    ...options,
-    // A `select` without `columns` merges as "no columns yet", not "every column".
-    initial: select ? { columns: {}, ...select } : undefined,
-  });
+/** Only explicitly selected columns are guaranteed; planning metadata is not part of the query. */
+export type QueryFromInfoResult<T> = Omit<T, keyof QueryFromInfoOptions<boolean> | 'columns'> & {
+  columns: T extends { columns: infer Columns extends {} } ? Columns : {};
+};
 
-  if (!plan) {
-    // Nothing is planned under the paths: the caller gets its own selection back untouched.
-    return (select ?? {}) as T;
-  }
+/** A boolean async option must account for promises, just like Prisma's helper. */
+export type QueryFromInfoReturn<T, Await extends boolean> = [Await] extends [false]
+  ? QueryFromInfoResult<T>
+  : MaybePromise<QueryFromInfoResult<T>>;
+
+/** Merge query options with the GraphQL selection and return a query ready for Drizzle. */
+export function queryFromInfo<const T extends object, Await extends boolean = false>({
+  config,
+  context,
+  info,
+  typeName,
+  path,
+  paths,
+  awaitSelections,
+  ...selection
+}: QueryFromInfoOptions<Await> &
+  T &
+  Record<
+    Exclude<keyof T, keyof QueryFromInfoOptions<Await> | keyof DBQueryConfig<'many'>>,
+    never
+  >): QueryFromInfoReturn<T, Await> {
+  const plan = planFromInfo({ config, context, info, typeName, path, paths });
+  const query = isThenable(plan)
+    ? plan.then((settled) => queryFromPlan(settled as DrizzlePlan | undefined, selection))
+    : queryFromPlan(plan, selection);
 
-  return (
-    isThenable(plan) ? plan.then((settled) => (settled as DrizzlePlan).query()) : plan.query()
-  ) as T;
+  return checkAwaitSelections(
+    query,
+    awaitSelections,
+    'queryFromInfo',
+    `${info.parentType.name}.${info.fieldName}`,
+  ) as QueryFromInfoReturn<T, Await>;
 }
 
 /**
@@ -45,7 +66,7 @@ export function queryFromInfo<T extends SelectionMap>({
 export function planFromInfo({
   config,
   ...options
-}: Omit<QueryFromInfoOptions<SelectionMap>, 'select'>): DrizzlePlan | undefined {
+}: Omit<QueryFromInfoOptions, 'awaitSelections'>): DrizzlePlan | undefined {
   return Plan.fromInfo(drizzleAdapter(config), options);
 }