@pothos/plugin-drizzle 0.17.5 → 0.18.0

package/src/connection-helpers.ts

@@ -10,13 +10,19 @@ import type { DrizzleRef } from './interface-ref.js';
 import type { QueryForDrizzleConnection } from './types.js';
 import { getSchemaConfig } from './utils/config.js';
 import {
+  type DrizzleCursorConnectionQueryOptions,
   drizzleCursorConnectionQuery,
   getCursorFormatter,
   wrapConnectionResult,
 } from './utils/cursors.js';
 import { queryFromInfo } from './utils/map-query.js';
 import { getRefFromModel } from './utils/refs.js';
-import { createState, mergeSelection, selectionToQuery } from './utils/selections.js';
+import {
+  createState,
+  mergeSelection,
+  omitUndefinedKeys,
+  selectionToQuery,
+} from './utils/selections.js';
 
 export function drizzleConnectionHelpers<
   Types extends SchemaTypes,
@@ -88,8 +94,8 @@ export function drizzleConnectionHelpers<
     ctx: Types['Context'],
     parent?: Parent,
   ) {
-    const { select, cursorColumns } = getQueryArgs(args, ctx);
-    const formatCursor = getCursorFormatter(cursorColumns, config);
+    const { select, cursorFields } = getQueryArgs(args, ctx);
+    const formatCursor = getCursorFormatter(cursorFields, config);
     return wrapConnectionResult(
       list,
       args,
@@ -117,7 +123,7 @@ export function drizzleConnectionHelpers<
       (typeof query === 'function' ? query(args, ctx) : query) ?? {};
     const table = config.relations[tableName];
 
-    const { cursorColumns, columns, ...connectionQuery } = drizzleCursorConnectionQuery({
+    const { cursorFields, columns, ...connectionQuery } = drizzleCursorConnectionQuery({
       ctx,
       maxSize,
       defaultSize: typeof defaultSize === 'function' ? defaultSize(args, ctx) : defaultSize,
@@ -125,6 +131,7 @@ export function drizzleConnectionHelpers<
       orderBy:
         (typeof orderBy === 'function' ? orderBy(table) : orderBy) ??
         getSchemaConfig(builder).getPrimaryKey(tableName),
+      extras: (fieldQuery as { extras?: DrizzleCursorConnectionQueryOptions['extras'] }).extras,
       where,
       config,
       table,
@@ -140,7 +147,7 @@ export function drizzleConnectionHelpers<
         },
         limit: Math.abs(limit ?? connectionQuery.limit),
       },
-      cursorColumns,
+      cursorFields,
     };
   };
 
@@ -181,10 +188,10 @@ export function drizzleConnectionHelpers<
     const baseQuery = typeof query === 'function' ? query(args, ctx) : (query ?? {});
     const queryResult = selectionToQuery(config, selectState);
 
-    return {
+    return omitUndefinedKeys({
       ...baseQuery,
       ...queryResult,
-    } as unknown as Omit<Selection, 'orderBy'> & {
+    }) as unknown as Omit<Selection, 'orderBy'> & {
       orderBy: {
         [K in TableConfig['table']['_'] extends { columns: infer Columns }
           ? keyof Columns

package/src/drizzle-field-builder.ts

@@ -51,12 +51,13 @@ import type {
 } from './types.js';
 import { getClient, getSchemaConfig } from './utils/config.js';
 import {
+  type DrizzleCursorConnectionQueryOptions,
   drizzleCursorConnectionQuery,
   getCursorFormatter,
   wrapConnectionResult,
 } from './utils/cursors.js';
 import { getRefFromModel } from './utils/refs.js';
-import type { SelectionMap } from './utils/selections.js';
+import { omitUndefinedKeys, type SelectionMap } from './utils/selections.js';
 
 // Workaround for FieldKind not being extended on Builder classes
 const RootBuilder: {
@@ -236,9 +237,10 @@ export class DrizzleObjectFieldBuilder<
         orderBy?: unknown;
         where?: SQL;
         columns?: Record<string, boolean>;
+        extras?: DrizzleCursorConnectionQueryOptions['extras'];
       };
 
-      const { cursorColumns, columns, ...connectionQuery } = drizzleCursorConnectionQuery({
+      const { cursorFields, columns, ...connectionQuery } = drizzleCursorConnectionQuery({
         ctx,
         maxSize,
         defaultSize,
@@ -246,6 +248,7 @@ export class DrizzleObjectFieldBuilder<
         orderBy:
           (typeof orderBy === 'function' ? orderBy(relatedTable.table) : orderBy) ??
           getSchemaConfig(this.builder).getPrimaryKey(relationField.targetTableName),
+        extras: fieldQuery.extras,
         where,
         config: schemaConfig,
         table: relatedTable,
@@ -261,7 +264,7 @@ export class DrizzleObjectFieldBuilder<
           },
           limit: Math.abs(limit ?? connectionQuery.limit),
         },
-        cursorColumns,
+        cursorFields,
       };
     };
 
@@ -376,13 +379,13 @@ export class DrizzleObjectFieldBuilder<
             };
           }
 
-          const { select, cursorColumns } = getQuery(args, context);
+          const { select, cursorFields } = getQuery(args, context);
 
           return wrapConnectionResult(
             parentRecord[name] as readonly {}[],
             args,
             select.limit,
-            getCursorFormatter(cursorColumns, schemaConfig),
+            getCursorFormatter(cursorFields, schemaConfig),
             undefined,
             parent,
             countValue,
@@ -526,7 +529,7 @@ export class DrizzleObjectFieldBuilder<
       const relQuery = {
         columns: {},
         with: {
-          [name]: {
+          [name]: omitUndefinedKeys({
             ...nestedQuery(query),
             ...((typeof query === 'function'
               ? (query as (args: {}, context: {}, pathInfo: PathInfo) => {})(
@@ -535,7 +538,7 @@ export class DrizzleObjectFieldBuilder<
                   pathInfo,
                 )
               : query) as {}),
-          },
+          }),
         },
       };
 

package/src/types.ts

@@ -68,6 +68,9 @@ type GetTableConfigFn<TTable = Table> = (table: TTable) => {
   primaryKeys: {
     readonly columns: Column[];
   }[];
+  uniqueConstraints?: {
+    readonly columns: Column[];
+  }[];
   readonly columns: Column[];
 };
 
@@ -677,11 +680,12 @@ export type RelatedConnectionOptions<
 export type ConnectionOrderBy<T extends TableRelationalConfig> =
   | Column
   | Column[]
+  // keys name a column, or an extra the same query selects. `string & {}` keeps
+  // the column names in autocomplete while letting extra names through
   | {
-      [K in T['table']['_'] extends { columns: infer Columns } ? keyof Columns : never]?:
-        | 'asc'
-        | 'desc'
-        | undefined;
+      [K in
+        | (T['table']['_'] extends { columns: infer Columns } ? keyof Columns : never)
+        | (string & {})]?: 'asc' | 'desc' | undefined;
     };
 
 export type ShapeFromConnection<T> = T extends { shape: unknown } ? T['shape'] : never;

package/src/utils/config.ts

@@ -5,7 +5,10 @@ import type { DrizzleClient } from '../types.js';
 export interface PothosDrizzleSchemaConfig {
   skipDeferredFragments: boolean;
   relations: AnyRelations;
+  findPrimaryKey: (tableName: string) => Column[] | null;
+  findTieBreaker: (tableName: string) => Column[] | null;
   getPrimaryKey: (tableName: string) => Column[];
+  getUniqueConstraints: (tableName: string) => Column[][];
   columnToTsName: (column: Column) => string;
 }
 const configCache = createContextCache(
@@ -27,6 +30,113 @@ const configCache = createContextCache(
       }
     });
 
+    const buildTableConfig = (tableName: string) => {
+      const table = relations[tableName].table as Table;
+      const tableConfig = builder.options.drizzle.getTableConfig(table);
+      const tableColumns = Object.values(getColumns(table));
+
+      // The pg dialect builds primaryKeys and uniqueConstraints from
+      // ExtraConfigColumns rather than the table's own columns, so the objects
+      // it hands back are not the ones columnNameMappings knows about, and do
+      // not carry notNull. Look each one back up by name.
+      const toTableColumn = (column: Column) =>
+        tableColumns.find((candidate) => candidate.name === column.name) ?? column;
+
+      return {
+        columns: tableConfig.columns.map(toTableColumn),
+        primaryKeys: tableConfig.primaryKeys.map((key) => key.columns.map(toTableColumn)),
+        uniqueConstraints: (tableConfig.uniqueConstraints ?? []).map((constraint) =>
+          constraint.columns.map(toTableColumn),
+        ),
+      };
+    };
+
+    const tableConfigs = new Map<string, ReturnType<typeof buildTableConfig>>();
+
+    // drizzle rebuilds this from the table's config callback on every call
+    const getTableConfig = (tableName: string) => {
+      let tableConfig = tableConfigs.get(tableName);
+
+      if (!tableConfig) {
+        tableConfig = buildTableConfig(tableName);
+        tableConfigs.set(tableName, tableConfig);
+      }
+
+      return tableConfig;
+    };
+
+    // Every set of columns the database guarantees is unique. Unique indexes
+    // are not included: missing one only costs a redundant order column, while
+    // wrongly reporting a set as unique would break pagination.
+    const getUniqueConstraints = (tableName: string) => {
+      const tableConfig = getTableConfig(tableName);
+
+      return [
+        ...tableConfig.columns
+          .filter((column) => column.primary || column.isUnique)
+          .map((column) => [column]),
+        ...tableConfig.primaryKeys,
+        ...tableConfig.uniqueConstraints,
+      ].filter((columns) => columns.length > 0);
+    };
+
+    // What a cursor can rely on to break ties. Primary key columns are
+    // non-nullable whatever the drizzle schema says, because SQL requires it,
+    // so a composite key declared without notNull() on each column still
+    // counts. A column that is only unique has to say so itself: comparing a
+    // cursor against null matches nothing, and rows sharing a null tie anyway.
+    const findTieBreaker = (tableName: string) => {
+      const tableConfig = getTableConfig(tableName);
+
+      const primaryKey = tableConfig.columns.find((column) => column.primary);
+
+      if (primaryKey) {
+        return [primaryKey];
+      }
+
+      const compositeKey = tableConfig.primaryKeys.find((columns) => columns.length > 0);
+
+      if (compositeKey) {
+        return compositeKey;
+      }
+
+      const uniqueColumn = tableConfig.columns.find((column) => column.isUnique && column.notNull);
+
+      if (uniqueColumn) {
+        return [uniqueColumn];
+      }
+
+      const uniqueConstraint = tableConfig.uniqueConstraints.find(
+        (columns) => columns.length > 0 && columns.every((column) => column.notNull),
+      );
+
+      return uniqueConstraint ?? null;
+    };
+
+    const findPrimaryKey = (tableName: string) => {
+      const tableConfig = getTableConfig(tableName);
+
+      const primaryKey = tableConfig.columns.find((column) => column.primary);
+
+      if (primaryKey) {
+        return [primaryKey];
+      }
+
+      const primaryKeys = tableConfig.primaryKeys.find((columns) => columns.length > 0);
+
+      if (primaryKeys) {
+        return primaryKeys;
+      }
+
+      const uniqueColumn = tableConfig.columns.find((column) => column.isUnique);
+
+      if (uniqueColumn) {
+        return [uniqueColumn];
+      }
+
+      return null;
+    };
+
     return {
       skipDeferredFragments: builder.options.drizzle.skipDeferredFragments ?? true,
       columnToTsName: (column) => {
@@ -38,30 +148,17 @@ const configCache = createContextCache(
 
         return tsName;
       },
+      findPrimaryKey,
+      findTieBreaker,
+      getUniqueConstraints,
       getPrimaryKey: (tableName) => {
-        const tableConfig = builder.options.drizzle.getTableConfig(
-          relations[tableName].table as Table,
-        );
-
-        const primaryKey = tableConfig.columns.find((column) => column.primary);
-
-        if (primaryKey) {
-          return [primaryKey];
-        }
-
-        const primaryKeys = tableConfig.primaryKeys.find((key) => key.columns.length > 0);
-
-        if (primaryKeys) {
-          return primaryKeys.columns;
-        }
-
-        const uniqueColumn = tableConfig.columns.find((column) => column.isUnique);
+        const primaryKey = findPrimaryKey(tableName);
 
-        if (uniqueColumn) {
-          return [uniqueColumn];
+        if (!primaryKey) {
+          throw new Error(`Could not find primary key for table ${tableName}`);
         }
 
-        throw new Error(`Could not find primary key for table ${tableName}`);
+        return primaryKey;
       },
       relations,
     };