diff --git a/.changeset/warm-trains-occur.md b/.changeset/warm-trains-occur.md new file mode 100644 index 000000000..10d352d5c --- /dev/null +++ b/.changeset/warm-trains-occur.md @@ -0,0 +1,5 @@ +--- +"@pothos/plugin-drizzle": minor +--- + +Add option to initialize drizzle client with context diff --git a/packages/plugin-drizzle/src/drizzle-field-builder.ts b/packages/plugin-drizzle/src/drizzle-field-builder.ts index 7fd597812..895010021 100644 --- a/packages/plugin-drizzle/src/drizzle-field-builder.ts +++ b/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, '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( diff --git a/packages/plugin-drizzle/src/field-builder.ts b/packages/plugin-drizzle/src/field-builder.ts index 05b166261..336759a68 100644 --- a/packages/plugin-drizzle/src/field-builder.ts +++ b/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).tableName - ]!; + const drizzleModel = getSchemaConfig(this.builder).schema?.[ + typeof type === 'string' ? type : (ref as DrizzleRef).tableName + ]!; return resolveDrizzleCursorConnection( drizzleModel, diff --git a/packages/plugin-drizzle/src/model-loader.ts b/packages/plugin-drizzle/src/model-loader.ts index 6606a0927..10d16bb4a 100644 --- a/packages/plugin-drizzle/src/model-loader.ts +++ b/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( @@ -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[]> } > diff --git a/packages/plugin-drizzle/src/schema-builder.ts b/packages/plugin-drizzle/src/schema-builder.ts index 464e9b599..ef664f366 100644 --- a/packages/plugin-drizzle/src/schema-builder.ts +++ b/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, ) { - 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, }, }); diff --git a/packages/plugin-drizzle/src/types.ts b/packages/plugin-drizzle/src/types.ts index 5477eb570..944702d1a 100644 --- a/packages/plugin-drizzle/src/types.ts +++ b/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>; query: {} }; +export type DrizzleClient = { + _: Partial>; + query: {}; +}; + +export type DrizzlePluginOptions = { maxConnectionSize?: number; defaultConnectionSize?: number; -} +} & ( + | { + client: DrizzleClient; + schema?: Partial>; + } + | { + client: (ctx: Types['Context']) => DrizzleClient; + schema: Partial>; + } +); export const drizzleTableName = Symbol.for('Pothos.drizzleTableName'); diff --git a/packages/plugin-drizzle/src/utils/config.ts b/packages/plugin-drizzle/src/utils/config.ts index a330b2a9b..a401e1a3d 100644 --- a/packages/plugin-drizzle/src/utils/config.ts +++ b/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 { dbToSchema: Record; } -const cache = createContextCache( - (config: RelationalSchemaConfig): PothosDrizzleSchemaConfig => { +const configCache = createContextCache( + (builder: PothosSchemaTypes.SchemaBuilder): PothosDrizzleSchemaConfig => { + const config = (builder.options.drizzle.schema ?? + (builder.options.drizzle.client as DrizzleClient) + ._) as RelationalSchemaConfig; + const dbToSchema = Object.values(config.schema).reduce>( (acc, table) => { acc[table.dbName] = table; @@ -25,8 +30,29 @@ const cache = createContextCache( }, ); +const clientCache = createContextCache((builder: PothosSchemaTypes.SchemaBuilder) => { + 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( builder: PothosSchemaTypes.SchemaBuilder, ) { - return cache(builder.options.drizzle.client._ as RelationalSchemaConfig); + return configCache(builder as never); +} + +export function getClient( + builder: PothosSchemaTypes.SchemaBuilder, + context: object, +) { + return clientCache(builder as never)(context); }