From c25d51a749a758ebe9886370124fa0bfbd2afe03 Mon Sep 17 00:00:00 2001
From: Romeo Van Snick <romeo@romeovansnick.be>
Date: Thu, 11 Jul 2024 09:52:32 +0200
Subject: [PATCH] feat(codegen): Allow query reponse types to be overridden
 through SanityQueries (#858)

---
 src/SanityClient.ts           | 89 +++++++++++++++++++++++------------
 src/types.ts                  |  9 ++++
 test/codegen/client.test-d.ts | 65 +++++++++++++++++++++++++
 3 files changed, 134 insertions(+), 29 deletions(-)
 create mode 100644 test/codegen/client.test-d.ts

diff --git a/src/SanityClient.ts b/src/SanityClient.ts
index a6e229e3..a8136e65 100644
--- a/src/SanityClient.ts
+++ b/src/SanityClient.ts
@@ -17,6 +17,7 @@ import type {
   BaseActionOptions,
   BaseMutationOptions,
   ClientConfig,
+  ClientReturn,
   FilteredResponseQueryOptions,
   FirstDocumentIdMutationOptions,
   FirstDocumentMutationOptions,
@@ -142,10 +143,11 @@ export class ObservableSanityClient {
    *
    * @param query - GROQ-query to perform
    */
-  fetch<R = Any, Q extends QueryWithoutParams = QueryWithoutParams>(
-    query: string,
-    params?: Q | QueryWithoutParams,
-  ): Observable<R>
+  fetch<
+    R = Any,
+    Q extends QueryWithoutParams = QueryWithoutParams,
+    const G extends string = string,
+  >(query: G, params?: Q | QueryWithoutParams): Observable<ClientReturn<G, R>>
   /**
    * Perform a GROQ-query against the configured dataset.
    *
@@ -153,11 +155,15 @@ export class ObservableSanityClient {
    * @param params - Optional query parameters
    * @param options - Optional request options
    */
-  fetch<R = Any, Q extends QueryWithoutParams | QueryParams = QueryParams>(
-    query: string,
+  fetch<
+    R = Any,
+    Q extends QueryWithoutParams | QueryParams = QueryParams,
+    const G extends string = string,
+  >(
+    query: G,
     params: Q extends QueryWithoutParams ? QueryWithoutParams : Q,
     options?: FilteredResponseQueryOptions,
-  ): Observable<R>
+  ): Observable<ClientReturn<G, R>>
   /**
    * Perform a GROQ-query against the configured dataset.
    *
@@ -165,11 +171,15 @@ export class ObservableSanityClient {
    * @param params - Optional query parameters
    * @param options - Request options
    */
-  fetch<R = Any, Q extends QueryWithoutParams | QueryParams = QueryParams>(
+  fetch<
+    R = Any,
+    Q extends QueryWithoutParams | QueryParams = QueryParams,
+    const G extends string = string,
+  >(
     query: string,
     params: Q extends QueryWithoutParams ? QueryWithoutParams : Q,
     options: UnfilteredResponseQueryOptions,
-  ): Observable<RawQueryResponse<R>>
+  ): Observable<RawQueryResponse<ClientReturn<G, R>>>
   /**
    * Perform a GROQ-query against the configured dataset.
    *
@@ -177,13 +187,17 @@ export class ObservableSanityClient {
    * @param params - Optional query parameters
    * @param options - Request options
    */
-  fetch<R = Any, Q extends QueryWithoutParams | QueryParams = QueryParams>(
-    query: string,
+  fetch<
+    R = Any,
+    Q extends QueryWithoutParams | QueryParams = QueryParams,
+    const G extends string = string,
+  >(
+    query: G,
     params: Q extends QueryWithoutParams ? QueryWithoutParams : Q,
     options: UnfilteredResponseWithoutQuery,
-  ): Observable<RawQuerylessQueryResponse<R>>
-  fetch<R, Q>(
-    query: string,
+  ): Observable<RawQuerylessQueryResponse<ClientReturn<G, R>>>
+  fetch<R, Q, const G extends string>(
+    query: G,
     params?: Q,
     options?: QueryOptions,
   ): Observable<RawQueryResponse<R> | R> {
@@ -808,10 +822,11 @@ export class SanityClient {
    *
    * @param query - GROQ-query to perform
    */
-  fetch<R = Any, Q extends QueryWithoutParams = QueryWithoutParams>(
-    query: string,
-    params?: Q | QueryWithoutParams,
-  ): Promise<R>
+  fetch<
+    R = Any,
+    Q extends QueryWithoutParams = QueryWithoutParams,
+    const G extends string = string,
+  >(query: G, params?: Q | QueryWithoutParams): Promise<ClientReturn<G, R>>
   /**
    * Perform a GROQ-query against the configured dataset.
    *
@@ -819,11 +834,15 @@ export class SanityClient {
    * @param params - Optional query parameters
    * @param options - Optional request options
    */
-  fetch<R = Any, Q extends QueryWithoutParams | QueryParams = QueryParams>(
-    query: string,
+  fetch<
+    R = Any,
+    Q extends QueryWithoutParams | QueryParams = QueryParams,
+    const G extends string = string,
+  >(
+    query: G,
     params: Q extends QueryWithoutParams ? QueryWithoutParams : Q,
     options?: FilteredResponseQueryOptions,
-  ): Promise<R>
+  ): Promise<ClientReturn<G, R>>
   /**
    * Perform a GROQ-query against the configured dataset.
    *
@@ -831,11 +850,15 @@ export class SanityClient {
    * @param params - Optional query parameters
    * @param options - Request options
    */
-  fetch<R = Any, Q extends QueryWithoutParams | QueryParams = QueryParams>(
-    query: string,
+  fetch<
+    R = Any,
+    Q extends QueryWithoutParams | QueryParams = QueryParams,
+    const G extends string = string,
+  >(
+    query: G,
     params: Q extends QueryWithoutParams ? QueryWithoutParams : Q,
     options: UnfilteredResponseQueryOptions,
-  ): Promise<RawQueryResponse<R>>
+  ): Promise<RawQueryResponse<ClientReturn<G, R>>>
   /**
    * Perform a GROQ-query against the configured dataset.
    *
@@ -843,14 +866,22 @@ export class SanityClient {
    * @param params - Optional query parameters
    * @param options - Request options
    */
-  fetch<R = Any, Q extends QueryWithoutParams | QueryParams = QueryParams>(
-    query: string,
+  fetch<
+    R = Any,
+    Q extends QueryWithoutParams | QueryParams = QueryParams,
+    const G extends string = string,
+  >(
+    query: G,
     params: Q extends QueryWithoutParams ? QueryWithoutParams : Q,
     options: UnfilteredResponseWithoutQuery,
-  ): Promise<RawQuerylessQueryResponse<R>>
-  fetch<R, Q>(query: string, params?: Q, options?: QueryOptions): Promise<RawQueryResponse<R> | R> {
+  ): Promise<RawQuerylessQueryResponse<ClientReturn<G, R>>>
+  fetch<R, Q, const G extends string>(
+    query: G,
+    params?: Q,
+    options?: QueryOptions,
+  ): Promise<RawQueryResponse<ClientReturn<G, R>> | ClientReturn<G, R>> {
     return lastValueFrom(
-      dataMethods._fetch<R, Q>(
+      dataMethods._fetch<ClientReturn<G, R>, Q>(
         this,
         this.#httpRequest,
         this.#clientConfig.stega,
diff --git a/src/types.ts b/src/types.ts
index 2a77726f..4b4d1f01 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -1236,6 +1236,15 @@ export interface LiveEventMessage {
   tags: SyncTag[]
 }
 
+/** @public */
+export interface SanityQueries {}
+
+/** @public */
+export type ClientReturn<
+  GroqString extends string,
+  Fallback = Any,
+> = GroqString extends keyof SanityQueries ? SanityQueries[GroqString] : Fallback
+
 export type {
   ContentSourceMapParsedPath,
   ContentSourceMapParsedPathKeyedSegment,
diff --git a/test/codegen/client.test-d.ts b/test/codegen/client.test-d.ts
new file mode 100644
index 00000000..95ef95c3
--- /dev/null
+++ b/test/codegen/client.test-d.ts
@@ -0,0 +1,65 @@
+import {createClient, type RawQueryResponse} from '@sanity/client'
+import {describe, expectTypeOf, test} from 'vitest'
+
+type FooResult = {
+  bar: number
+}
+
+// This would normally be generated by @sanity/codegen
+declare module '@sanity/client' {
+  export interface SanityQueries {
+    "*[_type == 'foo']": FooResult
+  }
+}
+
+describe('client.fetch', () => {
+  const client = createClient({})
+
+  describe('without params', () => {
+    test('known query type', async () => {
+      const resp = await client.fetch("*[_type == 'foo']")
+      expectTypeOf(resp).toMatchTypeOf<FooResult>()
+    })
+
+    test('ad-hoc query type', async () => {
+      const resp = await client.fetch("*[_type == 'bar']")
+      expectTypeOf(resp).toMatchTypeOf<any>()
+    })
+
+    test('ad-hoc query with a custom type', async () => {
+      type Result = {bar: string}
+      const resp = await client.fetch<Result>("*[_type == 'bar']")
+      expectTypeOf(resp).toMatchTypeOf<Result>()
+    })
+
+    test('known query type, but overriden with ad-hoc type', async () => {
+      type Result = {bar: string}
+      const resp = await client.fetch<Result>("*[_type == 'foo']")
+      expectTypeOf(resp).toMatchTypeOf<Result>()
+    })
+  })
+
+  describe('unfiltered response', () => {
+    test('known query type', async () => {
+      const resp = await client.fetch("*[_type == 'foo']", {}, {filterResponse: false})
+      expectTypeOf(resp).toMatchTypeOf<RawQueryResponse<FooResult>>()
+    })
+
+    test('ad-hoc query type', async () => {
+      const resp = await client.fetch("*[_type == 'bar']", {}, {filterResponse: false})
+      expectTypeOf(resp).toMatchTypeOf<RawQueryResponse<any>>()
+    })
+
+    test('ad-hoc query with a custom type', async () => {
+      type Result = {bar: string}
+      const resp = await client.fetch<Result>("*[_type == 'bar']", {}, {filterResponse: false})
+      expectTypeOf(resp).toMatchTypeOf<RawQueryResponse<Result>>()
+    })
+
+    test('known query type, but overriden with ad-hoc type', async () => {
+      type Result = {bar: string}
+      const resp = await client.fetch<Result>("*[_type == 'foo']", {}, {filterResponse: false})
+      expectTypeOf(resp).toMatchTypeOf<RawQueryResponse<Result>>()
+    })
+  })
+})