[drizzle] Fix table types for queries on relations to aliased tables

hayes · Sep 10, 2024 · bc73a7b · bc73a7b
1 parent afa9c34
commit bc73a7b
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 6 deletions.
diff --git a/.changeset/light-dragons-deliver.md b/.changeset/light-dragons-deliver.md
@@ -0,0 +1,5 @@
+---
+"@pothos/plugin-drizzle": patch
+---
+
+Fix table types for queries on relations to aliased tables
diff --git a/packages/plugin-drizzle/src/global-types.ts b/packages/plugin-drizzle/src/global-types.ts
@@ -13,6 +13,7 @@ import type {
   TypeParam,
 } from '@pothos/core';
 import type {
+  AnyTable,
   BuildQueryResult,
   Column,
   DBQueryConfig,
@@ -42,6 +43,10 @@ import type { GraphQLInputObjectType } from 'graphql';
 import type { PothosDrizzlePlugin } from '.';
 import type { DrizzleNodeRef } from './node-ref';
 
+type MapToDbName<T> = {
+  [K in keyof T as T[K] extends { dbName: infer K2 extends string } ? K2 : K]: T[K];
+};
+
 declare global {
   export namespace PothosSchemaTypes {
     export interface Plugins<Types extends SchemaTypes> {
@@ -55,11 +60,15 @@ declare global {
     export interface UserSchemaTypes {
       DrizzleSchema: Record<string, unknown>;
       DrizzleRelationSchema: Record<string, TableRelationalConfig>;
+      DrizzleRelationSchemaByDbName: Record<string, TableRelationalConfig>;
     }
 
     export interface ExtendDefaultTypes<PartialTypes extends Partial<UserSchemaTypes>> {
       DrizzleSchema: PartialTypes['DrizzleSchema'] & {};
       DrizzleRelationSchema: ExtractTablesWithRelations<PartialTypes['DrizzleSchema'] & {}>;
+      DrizzleRelationSchemaByDbName: MapToDbName<
+        ExtractTablesWithRelations<PartialTypes['DrizzleSchema'] & {}>
+      >;
     }
 
     export interface SchemaBuilder<Types extends SchemaTypes> {

diff --git a/packages/plugin-drizzle/src/types.ts b/packages/plugin-drizzle/src/types.ts
@@ -393,7 +393,7 @@ export type RefForRelation<Types extends SchemaTypes, Rel extends Relation> = Re
 
 export type TypesForRelation<Types extends SchemaTypes, Rel extends Relation> = BuildQueryResult<
   Types['DrizzleRelationSchema'],
-  Types['DrizzleRelationSchema'][Rel['referencedTable']['_']['name']],
+  Types['DrizzleRelationSchemaByDbName'][Rel['referencedTableName']],
   true
 >;
 
@@ -410,7 +410,7 @@ export type QueryForField<
           'one',
           false,
           Types['DrizzleRelationSchema'],
-          Types['DrizzleRelationSchema'][Rel['referencedTable']['_']['name']]
+          Types['DrizzleRelationSchemaByDbName'][Rel['referencedTableName']]
         >,
         'columns' | 'extra' | 'with'
       >
@@ -419,7 +419,7 @@ export type QueryForField<
           'many',
           false,
           Types['DrizzleRelationSchema'],
-          Types['DrizzleRelationSchema'][Rel['referencedTable']['_']['name']]
+          Types['DrizzleRelationSchemaByDbName'][Rel['referencedTableName']]
         >,
         'columns' | 'extra' | 'with'
       >
@@ -537,7 +537,7 @@ export type RelatedConnectionOptions<
   Nullable extends boolean,
   Args extends InputFieldMap,
   NodeTable extends
-    TableRelationalConfig = Types['DrizzleRelationSchema'][Table['relations'][Field]['referencedTable']['_']['name']],
+    TableRelationalConfig = Types['DrizzleRelationSchemaByDbName'][Table['relations'][Field]['referencedTableName']],
 > = Omit<
   PothosSchemaTypes.ObjectFieldOptions<
     Types,
@@ -580,7 +580,7 @@ export type RelatedConnectionOptions<
         >;
         query?: QueryForRelatedConnection<Types, NodeTable, ConnectionArgs>;
         // biome-ignore lint/suspicious/noExplicitAny: <explanation>
-        type?: DrizzleRef<any, Table['relations'][Field]['referencedTable']['_']['name']>;
+        type?: DrizzleRef<any, Table['relations'][Field]['referencedTableName']>;
 
         defaultSize?: number | ((args: ConnectionArgs, ctx: Types['Context']) => number);
         maxSize?: number | ((args: ConnectionArgs, ctx: Types['Context']) => number);

diff --git a/packages/plugin-drizzle/tests/example/builder.ts b/packages/plugin-drizzle/tests/example/builder.ts
@@ -3,7 +3,6 @@ import AddGraphQL from '@pothos/plugin-add-graphql';
 import ScopeAuthPlugin from '@pothos/plugin-relay';
 import RelayPlugin from '@pothos/plugin-scope-auth';
 import WithInputPlugin from '@pothos/plugin-with-input';
-import { createTableRelationsHelpers, extractTablesRelationalConfig } from 'drizzle-orm';
 import { drizzle } from 'drizzle-orm/libsql';
 import DrizzlePlugin from '../../src';
 import { type DrizzleSchema, client, db, schema } from './db';