fix: extra query not working

Author: miello

packages/core/src/builder.handler.ts

@@ -7,7 +7,6 @@ import {
   Many,
   One,
   or,
-  sql,
   type TableRelationalConfig,
 } from 'drizzle-orm'
 import type { NodePgQueryResultHKT } from 'drizzle-orm/node-postgres'
@@ -42,33 +41,25 @@ export function createDefaultApiHandlers<
   schema: Record<string, unknown>
   fields: TFields
   tableTsKey: string
-  identiferColumn: string
+  identifierColumn: string
   tableNamesMap: Record<string, string>
   tables: Record<string, TableRelationalConfig>
 }): CollectionAdminApi<TContext, TFields> {
-  const { fields, tableTsKey, tableNamesMap, tables, schema, identiferColumn } = args
+  const { fields, tableTsKey, tableNamesMap, tables, schema, identifierColumn } = args
 
   const tableRelationalConfig = tables[tableTsKey]
   const primaryKeyColumn = getPrimaryColumn(tableRelationalConfig)
-  const identifierKeyColumn = tableRelationalConfig.columns[identiferColumn]
+  const identifierKeyColumn = tableRelationalConfig.columns[identifierColumn]
   const tableName = tableRelationalConfig.tsName
   const tableSchema = getTableFromSchema(schema, tableTsKey)
-  const queryPayload = createDrizzleQuery(fields)
-
-  const extrasQuery = {
-    extras: {
-      __pk: sql<string | number>`${primaryKeyColumn}`.as('__pk'),
-      __id: sql<string | number>`${identifierKeyColumn}`.as('__id'),
-    },
-  }
+  const queryPayload = createDrizzleQuery(fields, identifierColumn)
 
   const findOne: ApiFindOneHandler<TContext, TFields> = async (args) => {
     const db = args.context.db
     const query = db.query[tableName as keyof typeof db.query] as RelationalQueryBuilder<any, any>
 
     const result = await query.findFirst({
       ...queryPayload,
-      ...extrasQuery,
       where: eq(primaryKeyColumn, args.id),
     })
     if (!result) {
@@ -94,7 +85,6 @@ export function createDefaultApiHandlers<
 
     const result = await query.findMany({
       ...queryPayload,
-      ...extrasQuery,
       limit: args.limit,
       offset: args.offset,
       orderBy: orderBy
@@ -136,7 +126,6 @@ export function createDefaultApiHandlers<
 
     const result = await query.findFirst({
       ...queryPayload,
-      ...extrasQuery,
       where: eq(primaryKeyColumn, pk),
     })
 
@@ -165,7 +154,6 @@ export function createDefaultApiHandlers<
       const pk = await apiHandler.update(args.id, tx, args.data)
       const result = await query.findFirst({
         ...queryPayload,
-        ...extrasQuery,
         where: eq(primaryKeyColumn, pk),
       })
 
@@ -574,19 +562,19 @@ function mapResultToFields(fields: Fields<any>, result: Record<string, any>): Re
         const value = result[field._.relationTsName]
         if (!value) return []
 
-        const primaryColumnTsName = field._.primaryColumnTsName
         if (Array.isArray(value)) {
           const values = value.map((v) => ({
             ...mapResultToFields(field.fields, v),
-            __pk: v[primaryColumnTsName],
+            __pk: v['__pk'],
+            __id: v['__id'],
           }))
           return [[field.fieldName, values]]
         }
 
         return [
           [
             field.fieldName,
-            { ...mapResultToFields(field.fields, value), __pk: value[primaryColumnTsName] },
+            { ...mapResultToFields(field.fields, value), __pk: value['__pk'], __id: value['__id'] },
           ],
         ]
       }

packages/core/src/builder.ts

@@ -82,7 +82,7 @@ export class Builder<
     const defaultHandlers = createDefaultApiHandlers({
       schema: this.config.schema,
       fields: config.fields,
-      identiferColumn: config.identifierColumn as string,
+      identifierColumn: config.identifierColumn as string,
       tableTsKey: tableTsName,
       tables: this.tableRelationalConfigByTableTsName,
       tableNamesMap: this.tableTsNameByTableDbName,

packages/core/src/utils.ts

@@ -1,5 +1,5 @@
-import type { Column, TableRelationalConfig } from 'drizzle-orm'
-import { is, Table } from 'drizzle-orm'
+import type { Column, SQL, TableRelationalConfig } from 'drizzle-orm'
+import { is, sql, Table } from 'drizzle-orm'
 import type { IsNever, Simplify, ValueOf } from 'type-fest'
 import type { ZodIssue, ZodObject, ZodOptional, ZodType } from 'zod'
 
@@ -9,7 +9,14 @@ import type {
   ApiRouteHandlerPayloadWithContext,
   ApiRouteSchema,
 } from './endpoint'
-import type { Field, FieldRelation, Fields, FieldsInitial, FieldsWithFieldName } from './field'
+import type {
+  Field,
+  FieldMetadataColumns,
+  FieldRelation,
+  Fields,
+  FieldsInitial,
+  FieldsWithFieldName,
+} from './field'
 
 export function isRelationField(field: Field): field is FieldRelation {
   return field._.source === 'relation'
@@ -74,7 +81,43 @@ export function getTableFromSchema(schema: Record<string, unknown>, tableTsName:
   return schema[tableTsName]
 }
 
-export function createDrizzleQuery(fields: Fields<any>): Record<string, any> {
+const getExtraField = (fields: Fields<any>, identifierColumn?: string) => {
+  const firstColumn = Object.values(fields).find((field) => field._.source === 'column')
+    ?._ as FieldMetadataColumns
+
+  if (!firstColumn) {
+    throw new Error('No column fields found in the provided fields')
+  }
+
+  const extraWith: [string, SQL.Aliased<string | number>][] = []
+
+  const table = firstColumn.column.table
+  const tableColumns = Object.values(table)
+  const primaryKeyColumn = tableColumns.find((col) => col.primary)
+
+  extraWith.push(['__pk', sql`${primaryKeyColumn}`.as('__pk') as SQL.Aliased<string | number>])
+
+  if (identifierColumn) {
+    const identifierKeyColumn = tableColumns.find((col) => col.name === identifierColumn)
+
+    if (!identifierKeyColumn) {
+      throw new Error(`Identifier column ${identifierColumn} not found in table ${table._.name}`)
+    }
+
+    if (!primaryKeyColumn) {
+      throw new Error(`Primary key column not found in table ${table._.name}`)
+    }
+
+    extraWith.push(['__id', sql`${identifierKeyColumn}`.as('__id') as SQL.Aliased<string | number>])
+  }
+
+  return Object.fromEntries(extraWith)
+}
+
+export function createDrizzleQuery(
+  fields: Fields<any>,
+  identifierColumn?: string
+): Record<string, any> {
   const queryColumns = Object.fromEntries(
     Object.values(fields).flatMap((field) => {
       if (field._.source !== 'column') return []
@@ -94,6 +137,7 @@ export function createDrizzleQuery(fields: Fields<any>): Record<string, any> {
   return {
     columns: queryColumns,
     with: queryWith,
+    extras: getExtraField(fields, identifierColumn),
   }
 }