[drizzle] improve node ids and fix query being required

.changeset/eighty-kangaroos-buy.md

@@ -0,0 +1,5 @@
+---
+"@pothos/plugin-drizzle": minor
+---
+
+Fix query being required, and improve node IDs

packages/plugin-drizzle/src/drizzle-field-builder.ts

@@ -91,9 +91,9 @@ export class DrizzleObjectFieldBuilder<
         EdgeInterfaces extends InterfaceParam<Types>[] = [],
       >(
         field: Field,
-        options: RelatedConnectionOptions<Types, Shape, TableConfig, Field, Nullable, Args>,
         ...args: NormalizeArgs<
           [
+            options: RelatedConnectionOptions<Types, Shape, TableConfig, Field, Nullable, Args>,
             connectionOptions:
               | ObjectRef<
                   Types,
@@ -182,17 +182,17 @@ export class DrizzleObjectFieldBuilder<
       type?: ObjectRef<Types, unknown, unknown>;
       maxSize?: number | ((args: {}, ctx: {}) => number);
       defaultSize?: number | ((args: {}, ctx: {}) => number);
-      extensions: {};
+      extensions?: {};
       description?: string;
-      query: ((args: {}, ctx: {}) => {}) | {};
-      resolve: (
+      query?: ((args: {}, ctx: {}) => {}) | {};
+      resolve?: (
         query: {},
         parent: unknown,
         args: {},
         ctx: {},
         info: {},
       ) => MaybePromise<readonly {}[]>;
-    },
+    } = {},
     connectionOptions = {},
     edgeOptions = {},
   ) {

packages/plugin-drizzle/src/global-types.ts

@@ -14,6 +14,7 @@ import type {
 } from '@pothos/core';
 import type {
   BuildQueryResult,
+  Column,
   DBQueryConfig,
   ExtractTablesWithRelations,
   SQL,
@@ -116,15 +117,23 @@ declare global {
                   Types['DrizzleRelationSchema'][Table]
                 >
               | true,
+            IDColumn extends Column,
             Shape = BuildQueryResult<
               Types['DrizzleRelationSchema'],
               Types['DrizzleRelationSchema'][Table],
               Selection
             >,
           >(
             table: Table,
-            options: DrizzleNodeOptions<Types, Table, Shape, Selection, Interfaces>,
-          ) => DrizzleNodeRef<Types, Table, Shape>
+            options: DrizzleNodeOptions<Types, Table, Shape, Selection, Interfaces, IDColumn>,
+          ) => DrizzleNodeRef<
+            Types,
+            Table,
+            Shape,
+            {
+              [K in IDColumn['_']['name']]: Extract<IDColumn, { _: { name: K } }>['_']['data'];
+            }
+          >
         : '@pothos/plugin-relay is required to use this method';
 
       drizzleObjectField: <

packages/plugin-drizzle/src/node-ref.ts

@@ -1,8 +1,25 @@
 import type { SchemaTypes } from '@pothos/core';
 import { DrizzleObjectRef } from './object-ref';
 
+export const relayIDShapeKey = Symbol.for('Pothos.relayIDShapeKey');
+
 export class DrizzleNodeRef<
   Types extends SchemaTypes,
   Table extends keyof Types['DrizzleRelationSchema'] = keyof Types['DrizzleRelationSchema'],
   T = {},
-> extends DrizzleObjectRef<Types, Table, T> {}
+  IDShape = string,
+> extends DrizzleObjectRef<Types, Table, T> {
+  [relayIDShapeKey]!: IDShape;
+  parseId: ((id: string, ctx: object) => IDShape) | undefined;
+
+  constructor(
+    name: string,
+    tableName: string,
+    options: {
+      parseId?: (id: string, ctx: object) => IDShape;
+    },
+  ) {
+    super(name, tableName);
+    this.parseId = options.parseId;
+  }
+}

packages/plugin-drizzle/src/schema-builder.ts

@@ -9,6 +9,7 @@ import SchemaBuilder, {
   type FieldRef,
   type InputObjectRef,
 } from '@pothos/core';
+import type { Column } from 'drizzle-orm';
 import type { GraphQLInputObjectType, GraphQLResolveInfo } from 'graphql';
 import { DrizzleObjectFieldBuilder } from './drizzle-field-builder';
 import { DrizzleInterfaceRef } from './interface-ref';
@@ -17,7 +18,7 @@ import { DrizzleNodeRef } from './node-ref';
 import { DrizzleObjectRef } from './object-ref';
 import type { DrizzleGraphQLInputExtensions, DrizzleNodeOptions } from './types';
 import { getSchemaConfig } from './utils/config';
-import { getColumnParser, getColumnSerializer } from './utils/cursors';
+import { getIDParser, getIDSerializer } from './utils/cursors';
 import { getRefFromModel } from './utils/refs';
 
 const schemaBuilderProto = SchemaBuilder.prototype as PothosSchemaTypes.SchemaBuilder<SchemaTypes>;
@@ -87,16 +88,25 @@ schemaBuilderProto.drizzleNode = function drizzleNode(
     name,
     variant,
     ...options
-  }: DrizzleNodeOptions<SchemaTypes, keyof SchemaTypes['DrizzleRelationSchema'], {}, {}, []>,
+  }: DrizzleNodeOptions<
+    SchemaTypes,
+    keyof SchemaTypes['DrizzleRelationSchema'],
+    {},
+    {},
+    [],
+    Column
+  >,
 ) {
   const tableConfig = getSchemaConfig(this).schema![table];
   const idColumn = typeof column === 'function' ? column(tableConfig.columns) : column;
   const idColumns = Array.isArray(idColumn) ? idColumn : [idColumn];
   const interfaceRef = this.nodeInterfaceRef?.();
-  const resolve = getColumnSerializer(idColumns);
-  const idParser = getColumnParser(idColumns);
+  const resolve = getIDSerializer(idColumns);
+  const idParser = getIDParser(idColumns);
   const typeName = variant ?? name ?? table;
-  const nodeRef = new DrizzleNodeRef(typeName, table);
+  const nodeRef = new DrizzleNodeRef(typeName, table, {
+    parseId: idParser,
+  });
   const modelLoader = ModelLoader.forModel(table, this, idColumns);
 
   if (!interfaceRef) {

packages/plugin-drizzle/src/types.ts

@@ -119,6 +119,7 @@ export type DrizzleNodeOptions<
   Shape,
   Selection,
   Interfaces extends InterfaceParam<Types>[],
+  IDColumn extends Column,
 > = NameOrVariant &
   Omit<
     | PothosSchemaTypes.ObjectTypeOptions<Types, Shape>
@@ -139,9 +140,9 @@ export type DrizzleNodeOptions<
       'args' | 'nullable' | 'type' | InferredFieldOptionKeys
     > & {
       column:
-        | Column
-        | Column[]
-        | ((columns: Types['DrizzleRelationSchema'][Table]['columns']) => Column | Column[]);
+        | IDColumn
+        | IDColumn[]
+        | ((columns: Types['DrizzleRelationSchema'][Table]['columns']) => IDColumn | IDColumn[]);
     };
     name: string;
     select?: Selection;
@@ -588,7 +589,7 @@ export type RelatedConnectionOptions<
             Nullable
           >
         >;
-        query: QueryForRelatedConnection<Types, NodeTable, ConnectionArgs>;
+        query?: QueryForRelatedConnection<Types, NodeTable, ConnectionArgs>;
         // biome-ignore lint/suspicious/noExplicitAny: <explanation>
         type?: DrizzleRef<any, Table['relations'][Field]['referencedTable']['_']['name']>;
 

packages/plugin-drizzle/src/utils/cursors.ts

@@ -43,6 +43,35 @@ export function formatDrizzleCursor(record: Record<string, unknown>, fields: Col
   return getCursorFormatter(fields)(record);
 }
 
+export function formatIDChunk(value: unknown) {
+  if (value instanceof Date) {
+    return `${String(Number(value))}`;
+  }
+
+  switch (typeof value) {
+    case 'number':
+    case 'string':
+    case 'bigint':
+      return `${value}`;
+    default:
+      throw new PothosValidationError(`Unsupported ID type ${typeof value}`);
+  }
+}
+
+export function getIDSerializer(fields: Column[]) {
+  if (fields.length === 0) {
+    throw new PothosValidationError('Column serializer must have at least one field');
+  }
+
+  return (value: Record<string, unknown>) => {
+    if (fields.length > 1) {
+      return `${JSON.stringify(fields.map((col) => value[col.name]))}`;
+    }
+
+    return `${formatIDChunk(value[fields[0].name])}`;
+  };
+}
+
 export function getColumnSerializer(fields: Column[]) {
   if (fields.length === 0) {
     throw new PothosValidationError('Column serializer must have at least one field');
@@ -113,32 +142,73 @@ export function parseSerializedDrizzleColumn(value: unknown) {
   }
 }
 
-export function parseID(id: string, dataType: string): unknown {
+export function parseSerializedIDColumn(id: string, field: Column): unknown {
   if (!id) {
     return id;
   }
 
-  switch (dataType) {
-    case 'String':
-      return id;
-    case 'Int':
-      return Number.parseInt(id, 10);
-    case 'BigInt':
-      return BigInt(id);
-    case 'Boolean':
-      return id !== 'false';
-    case 'Float':
-    case 'Decimal':
-      return Number.parseFloat(id);
-    case 'DateTime':
-      return new Date(id);
-    case 'Json':
-      return JSON.parse(id) as unknown;
-    case 'Byte':
-      return Buffer.from(id, 'base64');
-    default:
-      return id;
+  try {
+    switch (field.dataType) {
+      case 'date':
+        return new Date(id);
+      case 'string':
+        return id;
+      case 'number':
+        return Number.parseInt(id, 10);
+      case 'bigint':
+        return BigInt(id);
+      default:
+        throw new PothosValidationError(`Unsupported ID type ${field.dataType}`);
+    }
+  } catch (error: unknown) {
+    if (error instanceof PothosValidationError) {
+      throw error;
+    }
+
+    throw new PothosValidationError(`Invalid serialized ID: ${id}`);
+  }
+}
+
+export function getIDParser(fields: readonly Column[]) {
+  if (fields.length === 0) {
+    throw new PothosValidationError('Column parser must have at least one field');
   }
+
+  return (value: string) => {
+    if (fields.length === 1) {
+      return { [fields[0].name]: parseSerializedIDColumn(value, fields[0]) };
+    }
+
+    try {
+      const parsed = JSON.parse(value) as unknown[];
+
+      if (!Array.isArray(parsed)) {
+        throw new PothosValidationError(
+          `Expected compound ID to contain an array, but got ${value}`,
+        );
+      }
+
+      if (parsed.length !== fields.length) {
+        throw new PothosValidationError(
+          `Expected compound ID to contain ${fields.length} elements, but got ${parsed.length}`,
+        );
+      }
+
+      const record: Record<string, unknown> = {};
+
+      fields.forEach((field, i) => {
+        record[field.name] = parsed[i];
+      });
+
+      return record;
+    } catch (error: unknown) {
+      if (error instanceof PothosValidationError) {
+        throw error;
+      }
+
+      throw new PothosValidationError(`Invalid serialized ID: ${value}`);
+    }
+  };
 }
 
 export function getColumnParser(fields: readonly Column[]) {
@@ -159,6 +229,12 @@ export function getColumnParser(fields: readonly Column[]) {
       );
     }
 
+    if (parsed.length !== fields.length) {
+      throw new PothosValidationError(
+        `Expected compound cursor to contain ${fields.length} elements, but got ${parsed.length}`,
+      );
+    }
+
     const record: Record<string, unknown> = {};
 
     fields.forEach((field, i) => {
@@ -404,7 +480,7 @@ export async function resolveDrizzleCursorConnection<T extends {}>(
 ) {
   let query: DBQueryConfig<'many', false>;
   let formatter: (node: Record<string, unknown>) => string;
-  const results = await resolve((q) => {
+  const results = await resolve((q = {}) => {
     const { cursorColumns, ...connectionQuery } = drizzleCursorConnectionQuery({
       ...options,
       orderBy: q.orderBy

packages/plugin-drizzle/tests/example/schema.graphql

@@ -52,6 +52,7 @@ type Query {
   nodes(ids: [ID!]!): [Node]!
   post(id: ID!): Post
   posts(after: String, before: String, category: String, first: Int, last: Int): QueryPostsConnection
+  user(id: ID!): User
 }
 
 type QueryPostsConnection {

packages/plugin-drizzle/tests/example/schema/query.ts

@@ -1,6 +1,6 @@
 import { builder } from '../builder';
 import { db } from '../db';
-import { Viewer } from './user';
+import { User, Viewer } from './user';
 
 builder.queryType({
   fields: (t) => ({
@@ -13,5 +13,18 @@ builder.queryType({
           }),
         ),
     }),
+    user: t.drizzleField({
+      type: 'users',
+      args: {
+        id: t.arg.globalID({ required: true, for: User }),
+      },
+      resolve: (query, _root, { id }) => {
+        return db.query.users.findFirst(
+          query({
+            where: (user, { eq }) => eq(user.id, id.id.id),
+          }),
+        );
+      },
+    }),
   }),
 });

packages/plugin-drizzle/tests/example/schema/user.ts

@@ -44,7 +44,7 @@ export const Viewer = builder.drizzleObject('users', {
   }),
 });
 
-builder.drizzleNode('users', {
+export const User = builder.drizzleNode('users', {
   name: 'User',
   id: {
     column: (user) => user.id,