@pothos/plugin-drizzle 0.11.1 → 0.12.0

package/src/schema-builder.ts

@@ -36,8 +36,9 @@ schemaBuilderProto.drizzleObject = function drizzleObject(table, { select, field
     extensions: {
       ...options.extensions,
       pothosDrizzleModel: table,
-      pothosDrizzleTable: getSchemaConfig(this).relations.tablesConfig[table],
-      pothosDrizzleSelect: select ?? true,
+      pothosDrizzleTable: getSchemaConfig(this).relations[table],
+      pothosDrizzleSelect:
+        typeof select === 'object' ? { columns: {}, ...select } : (select ?? true),
       pothosDrizzleLoader: ModelLoader.forModel(table, this),
     },
     name,
@@ -64,8 +65,9 @@ schemaBuilderProto.drizzleInterface = function drizzleInterface(
     extensions: {
       ...options.extensions,
       pothosDrizzleModel: table,
-      pothosDrizzleTable: getSchemaConfig(this).relations.tablesConfig[table],
-      pothosDrizzleSelect: select ?? true,
+      pothosDrizzleTable: getSchemaConfig(this).relations[table],
+      pothosDrizzleSelect:
+        typeof select === 'object' ? { columns: {}, ...select } : (select ?? true),
       pothosDrizzleLoader: ModelLoader.forModel(table, this),
     },
     name,
@@ -97,8 +99,8 @@ schemaBuilderProto.drizzleNode = function drizzleNode(
   >,
 ) {
   const schemaConfig = getSchemaConfig(this);
-  const tableConfig = schemaConfig.relations.tablesConfig[table];
-  const idColumn = typeof column === 'function' ? column(tableConfig.columns) : column;
+  const tableConfig = schemaConfig.relations[table];
+  const idColumn = typeof column === 'function' ? column(tableConfig.table) : column;
   const idColumns = Array.isArray(idColumn) ? idColumn : [idColumn];
   const interfaceRef = this.nodeInterfaceRef?.();
   const resolve = getIDSerializer(idColumns, schemaConfig);
@@ -230,7 +232,7 @@ schemaBuilderProto.drizzleInterfaceFields = function drizzleInterfaceFields(type
 //       drizzleGraphQL: {
 //         inputType: 'orderBy',
 //         table,
-//         tableConfig: getSchemaConfig(this).relations.tablesConfig[table],
+//         tableConfig: getSchemaConfig(this).relations[table],
 //       } satisfies DrizzleGraphQLInputExtensions,
 //     },
 //   });
@@ -256,7 +258,7 @@ schemaBuilderProto.drizzleInterfaceFields = function drizzleInterfaceFields(type
 //       drizzleGraphQL: {
 //         inputType: 'filters',
 //         table,
-//         tableConfig: getSchemaConfig(this).relations.tablesConfig[table],
+//         tableConfig: getSchemaConfig(this).relations[table],
 //       } satisfies DrizzleGraphQLInputExtensions,
 //     },
 //   });
@@ -282,7 +284,7 @@ schemaBuilderProto.drizzleInterfaceFields = function drizzleInterfaceFields(type
 //       drizzleGraphQL: {
 //         inputType: 'insert',
 //         table,
-//         tableConfig: getSchemaConfig(this).relations.tablesConfig[table],
+//         tableConfig: getSchemaConfig(this).relations[table],
 //       } satisfies DrizzleGraphQLInputExtensions,
 //     },
 //   });
@@ -308,7 +310,7 @@ schemaBuilderProto.drizzleInterfaceFields = function drizzleInterfaceFields(type
 //       drizzleGraphQL: {
 //         inputType: 'update',
 //         table,
-//         tableConfig: getSchemaConfig(this).relations.tablesConfig[table],
+//         tableConfig: getSchemaConfig(this).relations[table],
 //       } satisfies DrizzleGraphQLInputExtensions,
 //     },
 //   });

package/src/types.ts

@@ -75,7 +75,7 @@ export type DrizzlePluginOptions<Types extends SchemaTypes> = {
         }[];
         readonly columns: Column[];
       };
-      relations?: Types['DrizzleRelations'];
+      relations: Types['DrizzleRelations'];
     }
 );
 
@@ -93,7 +93,7 @@ type NameOrVariant =
 
 export type DrizzleObjectOptions<
   Types extends SchemaTypes,
-  Table extends keyof Types['DrizzleRelationsConfig'],
+  Table extends keyof Types['DrizzleRelations'],
   Shape,
   Selection,
   Interfaces extends InterfaceParam<Types>[],
@@ -103,7 +103,7 @@ export type DrizzleObjectOptions<
     fields?: (
       t: DrizzleObjectFieldBuilder<
         Types,
-        Types['DrizzleRelationsConfig'][Table],
+        Types['DrizzleRelations'][Table],
         Shape & { [drizzleTableName]?: Table }
       >,
     ) => FieldMap;
@@ -111,7 +111,7 @@ export type DrizzleObjectOptions<
 
 export type DrizzleInterfaceOptions<
   Types extends SchemaTypes,
-  Table extends keyof Types['DrizzleRelationsConfig'],
+  Table extends keyof Types['DrizzleRelations'],
   Shape,
   Selection,
   Interfaces extends InterfaceParam<Types>[],
@@ -121,7 +121,7 @@ export type DrizzleInterfaceOptions<
     fields?: (
       t: DrizzleObjectFieldBuilder<
         Types,
-        Types['DrizzleRelationsConfig'][Table],
+        Types['DrizzleRelations'][Table],
         Shape & { [drizzleTableName]?: Table }
       >,
     ) => FieldMap;
@@ -129,7 +129,7 @@ export type DrizzleInterfaceOptions<
 
 export type DrizzleNodeOptions<
   Types extends SchemaTypes,
-  Table extends keyof Types['DrizzleRelationsConfig'],
+  Table extends keyof Types['DrizzleRelations'],
   Shape,
   Selection,
   Interfaces extends InterfaceParam<Types>[],
@@ -157,14 +157,14 @@ export type DrizzleNodeOptions<
         (
           | Column
           | Column[]
-          | ((columns: Types['DrizzleRelationsConfig'][Table]['columns']) => Column | Column[])
+          | ((columns: Types['DrizzleRelations'][Table]['table']) => Column | Column[])
         );
     };
     select?: Selection;
     fields?: (
       t: DrizzleObjectFieldBuilder<
         Types,
-        Types['DrizzleRelationsConfig'][Table],
+        Types['DrizzleRelations'][Table],
         Shape & { [drizzleTableName]?: Table }
       >,
     ) => FieldMap;
@@ -172,7 +172,7 @@ export type DrizzleNodeOptions<
 
 export type ShapeFromIdColumns<
   Types extends SchemaTypes,
-  Table extends keyof Types['DrizzleRelationsConfig'],
+  Table extends keyof Types['DrizzleRelations'],
   IDColumns,
 > = IDColumns extends Column
   ? IDColumns['_']['data']
@@ -191,7 +191,7 @@ export type ShapeFromIdColumns<
 export type DrizzleFieldOptions<
   Types extends SchemaTypes,
   ParentShape,
-  Table extends keyof Types['DrizzleRelationsConfig'],
+  Table extends keyof Types['DrizzleRelations'],
   Type extends TypeParam<Types>,
   Nullable extends FieldNullability<Type>,
   Args extends InputFieldMap,
@@ -225,7 +225,7 @@ export type DrizzleFieldOptions<
 export type DrizzleFieldWithInputOptions<
   Types extends SchemaTypes,
   ParentShape,
-  Table extends keyof Types['DrizzleRelationsConfig'],
+  Table extends keyof Types['DrizzleRelations'],
   Type extends TypeParam<Types>,
   Nullable extends FieldNullability<Type>,
   Args extends InputFieldMap,
@@ -271,10 +271,13 @@ export type DrizzleObjectFieldOptions<
       unknown extends Select
         ? ParentShape
         : BuildQueryResult<
-            Types['DrizzleRelationsConfig'],
+            Types['DrizzleRelations'],
             ExtractTable<Types, ParentShape>,
-            // biome-ignore lint/suspicious/noExplicitAny: this is fine
-            Record<string, unknown> & (Select extends (...args: any[]) => infer R ? R : Select)
+            Record<string, unknown> &
+              // biome-ignore lint/suspicious/noExplicitAny: this is fine
+              Select extends (...args: any[]) => infer R
+              ? R & { columns: {} }
+              : Select & { columns: {} }
           > &
             ParentShape,
       typeof drizzleTableName
@@ -299,7 +302,7 @@ export type DrizzleObjectFieldOptions<
   > & {
     select?: Select &
       (
-        | DBQueryConfig<'one', Types['DrizzleRelationsConfig'], ExtractTable<Types, ParentShape>>
+        | DBQueryConfig<'one', Types['DrizzleRelations'], ExtractTable<Types, ParentShape>>
         | ((
             args: InputShapeFromFields<Args>,
             ctx: Types['Context'],
@@ -307,11 +310,7 @@ export type DrizzleObjectFieldOptions<
               selection?: Selection,
               path?: string[],
             ) => Selection,
-          ) => DBQueryConfig<
-            'one',
-            Types['DrizzleRelationsConfig'],
-            ExtractTable<Types, ParentShape>
-          >)
+          ) => DBQueryConfig<'one', Types['DrizzleRelations'], ExtractTable<Types, ParentShape>>)
       );
   };
 
@@ -332,9 +331,9 @@ export type DrizzleFieldSelection =
     ) => SelectionMap);
 
 export type ExtractTable<Types extends SchemaTypes, Shape> = Shape extends {
-  [drizzleTableName]?: keyof Types['DrizzleRelationsConfig'];
+  [drizzleTableName]?: keyof Types['DrizzleRelations'];
 }
-  ? Types['DrizzleRelationsConfig'][NonNullable<Shape[typeof drizzleTableName]>]
+  ? Types['DrizzleRelations'][NonNullable<Shape[typeof drizzleTableName]>]
   : never;
 
 export type RelatedFieldOptions<
@@ -368,7 +367,7 @@ export type RelatedFieldOptions<
 
 export type VariantFieldOptions<
   Types extends SchemaTypes,
-  Table extends keyof Types['DrizzleRelationsConfig'],
+  Table extends keyof Types['DrizzleRelations'],
   // biome-ignore lint/suspicious/noExplicitAny: this is fine
   Variant extends DrizzleRef<any, Table> | Table,
   Args extends InputFieldMap,
@@ -397,11 +396,13 @@ export type VariantFieldOptions<
           unknown extends ResolveShape
             ? Shape
             : BuildQueryResult<
-                Types['DrizzleRelationsConfig'],
+                Types['DrizzleRelations'],
                 ExtractTable<Types, Shape>,
                 Record<string, unknown> &
                   // biome-ignore lint/suspicious/noExplicitAny: this is fine
-                  (ResolveShape extends (...args: any[]) => infer R ? R : ResolveShape)
+                  ResolveShape extends (...args: any[]) => infer R
+                  ? R & { columns: {} }
+                  : ResolveShape & { columns: {} }
               > &
                 Shape,
           typeof drizzleTableName
@@ -420,8 +421,8 @@ export type RefForRelation<Types extends SchemaTypes, Rel extends Relation> = Re
   : [ObjectRef<Types, TypesForRelation<Types, Rel>>];
 
 export type TypesForRelation<Types extends SchemaTypes, Rel extends Relation> = BuildQueryResult<
-  Types['DrizzleRelationsConfig'],
-  Types['DrizzleRelationsConfig'][Rel['targetTable']['_']['name']],
+  Types['DrizzleRelations'],
+  Types['DrizzleRelations'][Rel['targetTable']['_']['name']],
   true
 >;
 
@@ -436,16 +437,16 @@ export type QueryForField<
     ? Omit<
         DBQueryConfig<
           'one',
-          Types['DrizzleRelationsConfig'],
-          Types['DrizzleRelationsConfig'][Rel['targetTable']['_']['name']]
+          Types['DrizzleRelations'],
+          Types['DrizzleRelations'][Rel['targetTable']['_']['name']]
         >,
         'columns' | 'extra' | 'with'
       >
     : Omit<
         DBQueryConfig<
           'many',
-          Types['DrizzleRelationsConfig'],
-          Types['DrizzleRelationsConfig'][Rel['targetTable']['_']['name']]
+          Types['DrizzleRelations'],
+          Types['DrizzleRelations'][Rel['targetTable']['_']['name']]
         >,
         'columns' | 'extra' | 'with'
       >
@@ -456,14 +457,9 @@ export type QueryForField<
 export type QueryForDrizzleField<
   Types extends SchemaTypes,
   Param,
-  Table extends keyof Types['DrizzleRelationsConfig'],
+  Table extends keyof Types['DrizzleRelations'],
 > = Omit<
-  DBQueryConfig<
-    'many',
-    Types['DrizzleRelationsConfig'],
-    Types['DrizzleRelationsConfig'][Table],
-    true
-  >,
+  DBQueryConfig<'many', Types['DrizzleRelations'], Types['DrizzleRelations'][Table]>,
   Param extends [unknown] ? never : 'limit'
 >;
 
@@ -472,33 +468,33 @@ export type QueryForRelatedConnection<
   Table extends TableRelationalConfig,
   Args,
 > = Omit<
-  DBQueryConfig<'many', Types['DrizzleRelationsConfig'], Table>,
+  DBQueryConfig<'many', Types['DrizzleRelations'], Table>,
   'limit' | 'offset' | 'columns' | 'extra' | 'with' | 'orderBy'
 > & {
-  orderBy?: ConnectionOrderBy<Table> | ((columns: Table['columns']) => ConnectionOrderBy<Table>);
+  orderBy?: ConnectionOrderBy<Table> | ((table: Table) => ConnectionOrderBy<Table>);
 } extends infer QueryConfig
   ? QueryConfig | ((args: Args, context: Types['Context']) => QueryConfig)
   : never;
 
 export type QueryForDrizzleConnection<
   Types extends SchemaTypes,
-  Table extends TableRelationalConfig,
+  TableConfig extends TableRelationalConfig,
 > = Omit<
-  DBQueryConfig<'many', Types['DrizzleRelationsConfig'], Table>,
+  DBQueryConfig<'many', Types['DrizzleRelations'], TableConfig>,
   'limit' | 'offset' | 'orderBy'
 > & {
-  orderBy?: ConnectionOrderBy<Table> | ((columns: Table['columns']) => ConnectionOrderBy<Table>);
+  orderBy?: ConnectionOrderBy<TableConfig> | ((table: Table) => ConnectionOrderBy<TableConfig>);
 };
 
 export type ListRelation<T extends TableRelationalConfig> = {
-  [K in keyof T['relations']]: T['relations'][K] extends Many<string, string> ? K : never;
+  [K in keyof T['relations']]: T['relations'][K] extends Many<string> ? K : never;
 }[keyof T['relations']];
 
 export type DrizzleConnectionFieldOptions<
   Types extends SchemaTypes,
   ParentShape,
   Type extends // biome-ignore lint/suspicious/noExplicitAny: this is fine
-  DrizzleRef<any, keyof Types['DrizzleRelationsConfig']> | keyof Types['DrizzleRelationsConfig'],
+  DrizzleRef<any, keyof Types['DrizzleRelations']> | keyof Types['DrizzleRelations'],
   TableConfig extends TableRelationalConfig,
   Param extends OutputType<Types>,
   Nullable extends boolean,
@@ -544,7 +540,9 @@ export type DrizzleConnectionFieldOptions<
             selection?: T,
           ) => Omit<T, 'orderBy'> & {
             orderBy: {
-              [K in keyof TableConfig['columns']]?: 'asc' | 'desc' | undefined;
+              [K in TableConfig['table']['_'] extends { columns: infer Columns }
+                ? keyof Columns
+                : never]?: 'asc' | 'desc' | undefined;
             };
           },
           parent: ParentShape,
@@ -567,7 +565,7 @@ export type RelatedConnectionOptions<
   Nullable extends boolean,
   Args extends InputFieldMap,
   NodeTable extends
-    TableRelationalConfig = Types['DrizzleRelationsConfig'][Table['relations'][Field]['targetTable']['_']['name']],
+    TableRelationalConfig = Types['DrizzleRelations'][Table['relations'][Field]['targetTable']['_']['name']],
 > = Omit<
   PothosSchemaTypes.ObjectFieldOptions<
     Types,
@@ -622,7 +620,10 @@ export type ConnectionOrderBy<T extends TableRelationalConfig> =
   | Column
   | Column[]
   | {
-      [K in keyof T['columns']]?: 'asc' | 'desc' | undefined;
+      [K in T['table']['_'] extends { columns: infer Columns } ? keyof Columns : never]?:
+        | 'asc'
+        | 'desc'
+        | undefined;
     };
 
 export type ShapeFromConnection<T> = T extends { shape: unknown } ? T['shape'] : never;

package/src/utils/config.ts

@@ -1,5 +1,5 @@
 import { createContextCache, type SchemaTypes } from '@pothos/core';
-import { type AnyRelations, type Column, getTableName, type Table } from 'drizzle-orm';
+import { type AnyRelations, type Column, getTableColumns, isTable, type Table } from 'drizzle-orm';
 import type { DrizzleClient } from '../types';
 
 export interface PothosDrizzleSchemaConfig {
@@ -17,46 +17,35 @@ const configCache = createContextCache(
       relations = (builder.options.drizzle.client as DrizzleClient)._.relations;
     }
 
-    const dbToSchema = Object.values(relations.tables).reduce<Record<string, Table>>(
-      (acc, table) => {
-        const tableName = getTableName(table);
-        if (tableName) {
-          acc[tableName] = table;
-        }
-        return acc;
-      },
-      {},
-    );
-
-    const columnMappings = Object.values(dbToSchema).reduce<Record<string, Record<string, string>>>(
-      (acc, table) => {
-        acc[getTableName(table)] = Object.entries(table).reduce<Record<string, string>>(
-          (acc, [name, column]) => {
-            acc[(column as Column).name] = name;
-            return acc;
-          },
-          {},
-        );
-        return acc;
-      },
-      {},
-    );
+    const columnNameMappings: Record<string, string> = {};
+
+    Object.values(relations).forEach(({ table }) => {
+      if (isTable(table)) {
+        Object.entries(getTableColumns(table)).forEach(([tsName, col]) => {
+          if (col.uniqueName) {
+            columnNameMappings[col.uniqueName] = tsName;
+          }
+        });
+      }
+    });
 
     return {
       skipDeferredFragments: builder.options.drizzle.skipDeferredFragments ?? true,
       columnToTsName: (column) => {
-        const tableName = getTableName(column.table);
-        const table = columnMappings[tableName];
-        const columnName = table?.[column.name];
+        if (!column.uniqueName) {
+          throw new Error(`Column ${String(column.name)} has no uniqueName`);
+        }
 
-        if (!columnName) {
-          throw new Error(`Could not find column mapping for ${tableName}.${column.name}`);
+        if (!(column.uniqueName in columnNameMappings)) {
+          throw new Error(`Typescript name not found for ${column.uniqueName}`);
         }
 
-        return columnName;
+        return columnNameMappings[column.uniqueName];
       },
       getPrimaryKey: (tableName) => {
-        const tableConfig = builder.options.drizzle.getTableConfig(relations.tables[tableName]);
+        const tableConfig = builder.options.drizzle.getTableConfig(
+          relations[tableName].table as Table,
+        );
 
         const primaryKey = tableConfig.columns.find((column) => column.primary);
 

package/src/utils/cursors.ts

@@ -5,7 +5,14 @@ import {
   PothosValidationError,
   type SchemaTypes,
 } from '@pothos/core';
-import type { Column, DBQueryConfig, SQL, TableConfig, TableRelationalConfig } from 'drizzle-orm';
+import {
+  type Column,
+  type DBQueryConfig,
+  getTableColumns,
+  type SQL,
+  type Table,
+  type TableRelationalConfig,
+} from 'drizzle-orm';
 import type { GraphQLResolveInfo } from 'graphql';
 import type { ConnectionOrderBy, QueryForDrizzleConnection } from '../types';
 import type { PothosDrizzleSchemaConfig } from './config';
@@ -150,18 +157,23 @@ export function parseSerializedIDColumn(id: string, field: Column): unknown {
   }
 
   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}`);
+    if (field.dataType.startsWith('number')) {
+      return Number.parseInt(id, 10);
     }
+
+    if (field.dataType.startsWith('bigint')) {
+      return BigInt(id);
+    }
+
+    if (field.dataType.startsWith('string')) {
+      return id;
+    }
+
+    if (field.dataType === 'object date') {
+      return new Date(id);
+    }
+
+    throw new PothosValidationError(`Unsupported ID type ${field.dataType}`);
   } catch (error: unknown) {
     if (error instanceof PothosValidationError) {
       throw error;
@@ -283,12 +295,12 @@ export interface DrizzleCursorConnectionQueryOptions {
   orderBy: ConnectionOrderBy<TableRelationalConfig>;
   where?: SQL;
   config: PothosDrizzleSchemaConfig;
-  table: TableRelationalConfig | TableConfig;
+  table: TableRelationalConfig;
 }
 
 function parseOrderBy(
   config: PothosDrizzleSchemaConfig,
-  table: TableRelationalConfig | TableConfig,
+  table: TableRelationalConfig,
   orderBy: ConnectionOrderBy<TableRelationalConfig>,
   invert: boolean,
 ) {
@@ -310,13 +322,14 @@ function parseOrderBy(
       normalized.push({ direction: 'asc', column: field });
     }
   } else {
+    const tableColumns = getTableColumns(table.table as Table);
     Object.entries(
       orderBy as {
         [k: string]: 'asc' | 'desc' | undefined;
       },
     ).forEach(([name, direction]) => {
       if (direction) {
-        const column = table.columns[name] as Column;
+        const column = tableColumns[name];
         columns.push(column);
         entries.push([name, direction === 'asc' ? asc : desc]);
         normalized.push({ direction: direction, column });
@@ -506,7 +519,7 @@ export async function resolveDrizzleCursorConnection<T extends {}>(
   ) => MaybePromise<readonly T[]>,
   parent: unknown,
 ) {
-  const table = config.relations.tablesConfig[tableName];
+  const table = config.relations[tableName];
   let query: DBQueryConfig<'many'>;
   let formatter: (node: Record<string, unknown>) => string;
   const results = await resolve((q = {}) => {
@@ -514,8 +527,8 @@ export async function resolveDrizzleCursorConnection<T extends {}>(
       ...options,
       config,
       orderBy:
-        (typeof q.orderBy === 'function' ? q.orderBy(table.columns) : q.orderBy) ??
-        config.getPrimaryKey(table.tsName),
+        (typeof q.orderBy === 'function' ? q.orderBy(table.table as Table) : q.orderBy) ??
+        config.getPrimaryKey(table.name),
       table,
     });
     formatter = getCursorFormatter(cursorColumns, config);

package/src/utils/selections.ts

@@ -1,10 +1,5 @@
 import { PothosValidationError } from '@pothos/core';
-import {
-  type DBQueryConfig,
-  getTableUniqueName,
-  type Table,
-  type TableRelationalConfig,
-} from 'drizzle-orm';
+import type { DBQueryConfig, TableRelationalConfig } from 'drizzle-orm';
 import type { PothosDrizzleSchemaConfig } from './config';
 import { deepEqual } from './deep-equal';
 import type { LoaderMappings } from './loader-map';
@@ -126,18 +121,16 @@ export function mergeSelection(
     extras,
     columns,
     ...query
-  } = typeof selection === 'boolean' ? {} : selection;
+  } = typeof selection === 'boolean' ? (selection === true ? {} : { columns: {} }) : selection;
   if (withSelection) {
     for (const [key, value] of Object.entries(withSelection)) {
       const relation = state.table.relations[key];
 
       if (!relation) {
-        throw new PothosValidationError(`Relation ${key} does not exist on ${state.table.dbName}`);
+        throw new PothosValidationError(`Relation ${key} does not exist on ${state.table.name}`);
       }
 
-      const tableUniqueName = getTableUniqueName(relation.targetTable as Table);
-      const tableName = config.relations.tableNamesMap[tableUniqueName];
-      merge(config.relations.tablesConfig[tableName], key, value as SelectionMap | boolean);
+      merge(config.relations[relation.targetTableName], key, value as SelectionMap | boolean);
     }
   }
 
@@ -155,11 +148,11 @@ export function mergeSelection(
     return;
   }
 
-  if (columns) {
+  if (columns && !state.allColumns) {
     for (const key of Object.keys(columns)) {
       state.columns.add(key);
     }
-  } else if (selection === true) {
+  } else {
     state.allColumns = true;
   }
 
@@ -184,21 +177,18 @@ export function selectionToQuery(
 ): SelectionMap {
   const query: SelectionMap & { extras: Record<string, unknown> } = {
     ...state.query,
+    columns: undefined,
     with: {},
-    columns: {},
     extras: {},
   };
 
-  if (state.allColumns) {
-    for (const key of Object.keys(state.table.columns)) {
-      query.columns![key] = true;
-    }
-  } else {
+  if (!state.allColumns) {
+    query.columns = {};
     for (const key of state.columns) {
       query.columns![key] = true;
     }
 
-    for (const column of config.getPrimaryKey(state.table.tsName)) {
+    for (const column of config.getPrimaryKey(state.table.name)) {
       const tsName = config.columnToTsName(column);
       query.columns![tsName] = true;
     }