Merge pull request #604 from supabase/avallete/fix-add-returns-to-builder

fix(types): returns type casting

Author: avallete

jest.config.js

@@ -1,4 +1,5 @@
 module.exports = {
   preset: 'ts-jest',
   testEnvironment: 'node',
+  collectCoverageFrom: ['src/**/*'],
 }

src/PostgrestBuilder.ts

@@ -1,7 +1,14 @@
 // @ts-ignore
 import nodeFetch from '@supabase/node-fetch'
 
-import type { Fetch, PostgrestSingleResponse, PostgrestResponseSuccess } from './types'
+import type {
+  Fetch,
+  PostgrestSingleResponse,
+  PostgrestResponseSuccess,
+  CheckMatchingArrayTypes,
+  MergePartialResult,
+  IsValidResultOverride,
+} from './types'
 import PostgrestError from './PostgrestError'
 
 export default abstract class PostgrestBuilder<Result, ThrowOnError extends boolean = false>
@@ -209,4 +216,57 @@ export default abstract class PostgrestBuilder<Result, ThrowOnError extends bool
 
     return res.then(onfulfilled, onrejected)
   }
+
+  /**
+   * Override the type of the returned `data`.
+   *
+   * @typeParam NewResult - The new result type to override with
+   * @deprecated Use overrideTypes<yourType, { merge: false }>() method at the end of your call chain instead
+   */
+  returns<NewResult>(): PostgrestBuilder<CheckMatchingArrayTypes<Result, NewResult>, ThrowOnError> {
+    /* istanbul ignore next */
+    return this as unknown as PostgrestBuilder<
+      CheckMatchingArrayTypes<Result, NewResult>,
+      ThrowOnError
+    >
+  }
+
+  /**
+   * Override the type of the returned `data` field in the response.
+   *
+   * @typeParam NewResult - The new type to cast the response data to
+   * @typeParam Options - Optional type configuration (defaults to { merge: true })
+   * @typeParam Options.merge - When true, merges the new type with existing return type. When false, replaces the existing types entirely (defaults to true)
+   * @example
+   * ```typescript
+   * // Merge with existing types (default behavior)
+   * const query = supabase
+   *   .from('users')
+   *   .select()
+   *   .overrideTypes<{ custom_field: string }>()
+   *
+   * // Replace existing types completely
+   * const replaceQuery = supabase
+   *   .from('users')
+   *   .select()
+   *   .overrideTypes<{ id: number; name: string }, { merge: false }>()
+   * ```
+   * @returns A PostgrestBuilder instance with the new type
+   */
+  overrideTypes<
+    NewResult,
+    Options extends { merge?: boolean } = { merge: true }
+  >(): PostgrestBuilder<
+    IsValidResultOverride<Result, NewResult, true, false, false> extends true
+      ? MergePartialResult<NewResult, Result, Options>
+      : CheckMatchingArrayTypes<Result, NewResult>,
+    ThrowOnError
+  > {
+    return this as unknown as PostgrestBuilder<
+      IsValidResultOverride<Result, NewResult, true, false, false> extends true
+        ? MergePartialResult<NewResult, Result, Options>
+        : CheckMatchingArrayTypes<Result, NewResult>,
+      ThrowOnError
+    >
+  }
 }

src/PostgrestTransformBuilder.ts

@@ -1,6 +1,6 @@
 import PostgrestBuilder from './PostgrestBuilder'
 import { GetResult } from './select-query-parser/result'
-import { GenericSchema } from './types'
+import { GenericSchema, CheckMatchingArrayTypes } from './types'
 
 export default class PostgrestTransformBuilder<
   Schema extends GenericSchema,
@@ -307,18 +307,19 @@ export default class PostgrestTransformBuilder<
    * Override the type of the returned `data`.
    *
    * @typeParam NewResult - The new result type to override with
+   * @deprecated Use overrideTypes<yourType, { merge: false }>() method at the end of your call chain instead
    */
   returns<NewResult>(): PostgrestTransformBuilder<
     Schema,
     Row,
-    NewResult,
+    CheckMatchingArrayTypes<Result, NewResult>,
     RelationName,
     Relationships
   > {
     return this as unknown as PostgrestTransformBuilder<
       Schema,
       Row,
-      NewResult,
+      CheckMatchingArrayTypes<Result, NewResult>,
       RelationName,
       Relationships
     >

src/types.ts

@@ -1,4 +1,6 @@
 import PostgrestError from './PostgrestError'
+import { ContainsNull } from './select-query-parser/types'
+import { SelectQueryError } from './select-query-parser/utils'
 
 export type Fetch = typeof fetch
 
@@ -89,3 +91,66 @@ type ConditionalSimplifyDeep<
 type NonRecursiveType = BuiltIns | Function | (new (...arguments_: any[]) => unknown)
 type BuiltIns = Primitive | void | Date | RegExp
 type Primitive = null | undefined | string | number | boolean | symbol | bigint
+
+export type IsValidResultOverride<Result, NewResult, Ok, ErrorResult, ErrorNewResult> =
+  Result extends any[]
+    ? NewResult extends any[]
+      ? // Both are arrays - valid
+        Ok
+      : ErrorResult
+    : NewResult extends any[]
+    ? ErrorNewResult
+    : // Neither are arrays - valid
+    // Preserve the optionality of the result if the overriden type is an object (case of chaining with `maybeSingle`)
+    ContainsNull<Result> extends true
+    ? Ok | null
+    : Ok
+
+/**
+ * Utility type to check if array types match between Result and NewResult.
+ * Returns either the valid NewResult type or an error message type.
+ */
+export type CheckMatchingArrayTypes<Result, NewResult> =
+  // If the result is a QueryError we allow the user to override anyway
+  Result extends SelectQueryError<string>
+    ? NewResult
+    : IsValidResultOverride<
+        Result,
+        NewResult,
+        NewResult,
+        {
+          Error: 'Type mismatch: Cannot cast array result to a single object. Use .returns<Array<YourType>> for array results or .single() to convert the result to a single object'
+        },
+        {
+          Error: 'Type mismatch: Cannot cast single object to array type. Remove Array wrapper from return type or make sure you are not using .single() up in the calling chain'
+        }
+      >
+
+type Simplify<T> = T extends object ? { [K in keyof T]: T[K] } : T
+
+type MergeDeep<New, Row> = {
+  [K in keyof New | keyof Row]: K extends keyof Row
+    ? K extends keyof New
+      ? IsPlainObject<New[K]> extends true
+        ? IsPlainObject<Row[K]> extends true
+          ? MergeDeep<New[K], Row[K]>
+          : Row[K]
+        : Row[K]
+      : Row[K]
+    : K extends keyof New
+    ? New[K]
+    : never
+}
+
+// Helper to check if a type is a plain object (not an array)
+type IsPlainObject<T> = T extends any[] ? false : T extends object ? true : false
+
+// Merge the new result with the original (Result) when merge option is true.
+// If NewResult is an array, merge each element.
+export type MergePartialResult<NewResult, Result, Options> = Options extends { merge: true }
+  ? Result extends any[]
+    ? NewResult extends any[]
+      ? Array<Simplify<MergeDeep<NewResult[number], Result[number]>>>
+      : never
+    : Simplify<MergeDeep<NewResult, Result>>
+  : NewResult