Add option to initialize drizzle client with context

.changeset/warm-trains-occur.md

@@ -0,0 +1,5 @@
+---
+"@pothos/plugin-drizzle": minor
+---
+
+Add option to initialize drizzle client with context

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

@@ -32,6 +32,7 @@ import type {
   TypesForRelation,
   VariantFieldOptions,
 } from './types';
+import { getSchemaConfig } from './utils/config';
 import {
   drizzleCursorConnectionQuery,
   getCursorFormatter,
@@ -195,10 +196,12 @@ export class DrizzleObjectFieldBuilder<
     connectionOptions = {},
     edgeOptions = {},
   ) {
-    const relationField =
-      this.builder.options.drizzle.client._.schema?.[this.table].relations[name as string];
-    const relatedModel =
-      this.builder.options.drizzle.client._.schema?.[relationField?.referencedTableName!]!;
+    const relationField = getSchemaConfig(this.builder).schema?.[this.table].relations[
+      name as string
+    ];
+    const relatedModel = getSchemaConfig(this.builder).schema?.[
+      relationField?.referencedTableName!
+    ]!;
 
     if (!relationField) {
       throw new PothosSchemaError(
@@ -381,8 +384,9 @@ export class DrizzleObjectFieldBuilder<
     >
   ): FieldRef<Types, TypesForRelation<Types, TableConfig['relations'][Field]>, 'Object'> {
     const [options = {} as never] = allArgs;
-    const relationField =
-      this.builder.options.drizzle.client._.schema?.[this.table].relations[name as string];
+    const relationField = getSchemaConfig(this.builder).schema?.[this.table].relations[
+      name as string
+    ];
 
     if (!relationField) {
       throw new PothosSchemaError(

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

@@ -125,10 +125,9 @@ fieldBuilderProto.drizzleConnection = function drizzleConnection<
         context: {},
         info: GraphQLResolveInfo,
       ) => {
-        const drizzleModel =
-          this.builder.options.drizzle.client._.schema?.[
-            typeof type === 'string' ? type : (ref as DrizzleRef<SchemaTypes>).tableName
-          ]!;
+        const drizzleModel = getSchemaConfig(this.builder).schema?.[
+          typeof type === 'string' ? type : (ref as DrizzleRef<SchemaTypes>).tableName
+        ]!;
 
         return resolveDrizzleCursorConnection(
           drizzleModel,

packages/plugin-drizzle/src/model-loader.ts

@@ -1,7 +1,8 @@
 import { type SchemaTypes, createContextCache } from '@pothos/core';
 import { type Column, type TableRelationalConfig, inArray, sql } from 'drizzle-orm';
 import type { GraphQLResolveInfo } from 'graphql';
-import { type PothosDrizzleSchemaConfig, getSchemaConfig } from './utils/config';
+import type { DrizzleClient } from './types';
+import { type PothosDrizzleSchemaConfig, getClient, getSchemaConfig } from './utils/config';
 import { cacheKey, setLoaderMappings } from './utils/loader-map';
 import { selectionStateFromInfo, stateFromInfo } from './utils/map-query';
 import {
@@ -35,6 +36,7 @@ export class ModelLoader {
   table: TableRelationalConfig;
   columns: Column[];
   selectSQL;
+  client: DrizzleClient;
 
   constructor(
     context: object,
@@ -52,6 +54,8 @@ export class ModelLoader {
       this.columns.length > 1
         ? sql`(${sql.join(this.columns, sql`, `)})`
         : this.columns[0].getSQL();
+
+    this.client = getClient(builder, context);
   }
 
   static forModel<Types extends SchemaTypes>(
@@ -166,7 +170,7 @@ export class ModelLoader {
     nextTick.promise
       .then(() => {
         const api = (
-          this.builder.options.drizzle.client.query as Record<
+          this.client.query as Record<
             string,
             { findMany: (...args: unknown[]) => Promise<Record<string, unknown>[]> }
           >

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

@@ -16,6 +16,7 @@ import { ModelLoader } from './model-loader';
 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 { getRefFromModel } from './utils/refs';
 
@@ -35,7 +36,7 @@ schemaBuilderProto.drizzleObject = function drizzleObject(table, { select, field
     extensions: {
       ...options.extensions,
       pothosDrizzleModel: table,
-      pothosDrizzleTable: this.options.drizzle.client._.schema?.[table],
+      pothosDrizzleTable: getSchemaConfig(this).schema?.[table],
       pothosDrizzleSelect: select ?? {},
       pothosDrizzleLoader: ModelLoader.forModel(table, this),
     },
@@ -63,7 +64,7 @@ schemaBuilderProto.drizzleInterface = function drizzleInterface(
     extensions: {
       ...options.extensions,
       pothosDrizzleModel: table,
-      pothosDrizzleTable: this.options.drizzle.client._.schema?.[table],
+      pothosDrizzleTable: getSchemaConfig(this).schema?.[table],
       pothosDrizzleSelect: select ?? {},
       pothosDrizzleLoader: ModelLoader.forModel(table, this),
     },
@@ -88,7 +89,7 @@ schemaBuilderProto.drizzleNode = function drizzleNode(
     ...options
   }: DrizzleNodeOptions<SchemaTypes, keyof SchemaTypes['DrizzleRelationSchema'], {}, {}, []>,
 ) {
-  const tableConfig = this.options.drizzle.client._.schema![table];
+  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?.();
@@ -214,7 +215,7 @@ schemaBuilderProto.drizzleGraphQLOrderBy = function drizzleGraphQLOrderBy(
       drizzleGraphQL: {
         inputType: 'orderBy',
         table,
-        tableConfig: this.options.drizzle.client._.schema?.[table]!,
+        tableConfig: getSchemaConfig(this).schema?.[table]!,
       } satisfies DrizzleGraphQLInputExtensions,
     },
   });
@@ -240,7 +241,7 @@ schemaBuilderProto.drizzleGraphQLFilters = function drizzleGraphQLFilters(
       drizzleGraphQL: {
         inputType: 'filters',
         table,
-        tableConfig: this.options.drizzle.client._.schema?.[table]!,
+        tableConfig: getSchemaConfig(this).schema?.[table]!,
       } satisfies DrizzleGraphQLInputExtensions,
     },
   });
@@ -266,7 +267,7 @@ schemaBuilderProto.drizzleGraphQLInsert = function drizzleGraphQLInsert(
       drizzleGraphQL: {
         inputType: 'insert',
         table,
-        tableConfig: this.options.drizzle.client._.schema?.[table]!,
+        tableConfig: getSchemaConfig(this).schema?.[table]!,
       } satisfies DrizzleGraphQLInputExtensions,
     },
   });
@@ -292,7 +293,7 @@ schemaBuilderProto.drizzleGraphQLUpdate = function drizzleGraphQLUpdate(
       drizzleGraphQL: {
         inputType: 'update',
         table,
-        tableConfig: this.options.drizzle.client._.schema?.[table]!,
+        tableConfig: getSchemaConfig(this).schema?.[table]!,
       } satisfies DrizzleGraphQLInputExtensions,
     },
   });

packages/plugin-drizzle/src/types.ts

@@ -41,11 +41,24 @@ import type { DrizzleRef } from './interface-ref';
 import type { IndirectInclude } from './utils/map-query';
 import type { SelectionMap } from './utils/selections';
 
-export interface DrizzlePluginOptions<_Types extends SchemaTypes> {
-  client: { _: Partial<RelationalSchemaConfig<TablesRelationalConfig>>; query: {} };
+export type DrizzleClient = {
+  _: Partial<RelationalSchemaConfig<TablesRelationalConfig>>;
+  query: {};
+};
+
+export type DrizzlePluginOptions<Types extends SchemaTypes> = {
   maxConnectionSize?: number;
   defaultConnectionSize?: number;
-}
+} & (
+  | {
+      client: DrizzleClient;
+      schema?: Partial<RelationalSchemaConfig<TablesRelationalConfig>>;
+    }
+  | {
+      client: (ctx: Types['Context']) => DrizzleClient;
+      schema: Partial<RelationalSchemaConfig<TablesRelationalConfig>>;
+    }
+);
 
 export const drizzleTableName = Symbol.for('Pothos.drizzleTableName');
 

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

@@ -4,12 +4,17 @@ import type {
   TableRelationalConfig,
   TablesRelationalConfig,
 } from 'drizzle-orm';
+import type { DrizzleClient } from '../types';
 
 export interface PothosDrizzleSchemaConfig extends RelationalSchemaConfig<TablesRelationalConfig> {
   dbToSchema: Record<string, TableRelationalConfig>;
 }
-const cache = createContextCache(
-  (config: RelationalSchemaConfig<TablesRelationalConfig>): PothosDrizzleSchemaConfig => {
+const configCache = createContextCache(
+  (builder: PothosSchemaTypes.SchemaBuilder<SchemaTypes>): PothosDrizzleSchemaConfig => {
+    const config = (builder.options.drizzle.schema ??
+      (builder.options.drizzle.client as DrizzleClient)
+        ._) as RelationalSchemaConfig<TablesRelationalConfig>;
+
     const dbToSchema = Object.values(config.schema).reduce<Record<string, TableRelationalConfig>>(
       (acc, table) => {
         acc[table.dbName] = table;
@@ -25,8 +30,29 @@ const cache = createContextCache(
   },
 );
 
+const clientCache = createContextCache((builder: PothosSchemaTypes.SchemaBuilder<SchemaTypes>) => {
+  const clientConfig = builder.options.drizzle.client;
+  const getClient =
+    typeof clientConfig === 'function'
+      ? createContextCache((ctx) => clientConfig(ctx))
+      : (_ctx: object) => clientConfig;
+
+  return createContextCache((context: object) => {
+    const client = getClient(context);
+
+    return client;
+  });
+});
+
 export function getSchemaConfig<Types extends SchemaTypes>(
   builder: PothosSchemaTypes.SchemaBuilder<Types>,
 ) {
-  return cache(builder.options.drizzle.client._ as RelationalSchemaConfig<TablesRelationalConfig>);
+  return configCache(builder as never);
+}
+
+export function getClient<Types extends SchemaTypes>(
+  builder: PothosSchemaTypes.SchemaBuilder<Types>,
+  context: object,
+) {
+  return clientCache(builder as never)(context);
 }