fix: infer selected data in the TanStack query integration

jest.config.json

@@ -3,6 +3,7 @@
     "testEnvironment": "node",
     "modulePathIgnorePatterns": ["<rootDir>/dist"],
     "moduleNameMapper": {
+        "^@legendapp/state/sync-plugins/tanstack-query$": "<rootDir>/src/sync-plugins/tanstack-query",
         "@legendapp/state/sync-plugins/crud": "<rootDir>/src/sync-plugins/crud",
         "@legendapp/state/sync": "<rootDir>/sync",
         "@legendapp/state/config/configureLegendState": "<rootDir>/src/config/configureLegendState",
@@ -13,4 +14,4 @@
     "transform": {
         "^.+\\.tsx?$": ["ts-jest", { "tsconfig": { "jsx": "react" } }]
     }
-}
+}

src/sync-plugins/tanstack-query.ts

@@ -15,7 +15,7 @@ import {
 let nextMutationKey = 0;
 
 export interface ObservableQueryOptions<TQueryFnData, TError, TData, TQueryKey extends QueryKey>
-    extends Omit<QueryObserverOptions<TQueryFnData, TError, TData, TData, TQueryKey>, 'queryKey'> {
+    extends Omit<QueryObserverOptions<TQueryFnData, TError, TData, TQueryFnData, TQueryKey>, 'queryKey'> {
     queryKey?: TQueryKey | (() => TQueryKey);
 }
 
@@ -43,9 +43,9 @@ export function syncedQuery<
 
     const Observer = QueryObserver;
     const defaultedOptions = queryClient!.defaultQueryOptions(
-        options as QueryObserverOptions<TQueryFnData, TError, TData, TQueryKey>,
+        options as QueryObserverOptions<TQueryFnData, TError, TData, TQueryFnData, TQueryKey>,
     );
-    let observer: QueryObserver<TQueryFnData, TError, TData, TQueryKey> | undefined = undefined;
+    let observer: QueryObserver<TQueryFnData, TError, TData, TQueryFnData, TQueryKey> | undefined = undefined;
     let latestOptions = defaultedOptions;
     let queryKeyFromFn: TQueryKey;
     let resolveInitialPromise: undefined | ((value: TData) => void) = undefined;
@@ -55,7 +55,9 @@ export function syncedQuery<
     // If the queryKey is a function, observe it and extract the raw value
     const isKeyFunction = isFunction(origQueryKey);
 
-    const updateQueryOptions = (obj: DefaultedQueryObserverOptions<TQueryFnData, TError, TData, TQueryKey>) => {
+    const updateQueryOptions = (
+        obj: DefaultedQueryObserverOptions<TQueryFnData, TError, TData, TQueryFnData, TQueryKey>,
+    ) => {
         // Since legend-state mutates the query options, we need to clone it to make Query
         // see it as changed
         const options = Object.assign({}, obj);
@@ -81,7 +83,7 @@ export function syncedQuery<
     }
 
     // Create the observer
-    observer = new Observer!<TQueryFnData, TError, TData, TQueryKey>(queryClient!, latestOptions);
+    observer = new Observer!<TQueryFnData, TError, TData, TQueryFnData, TQueryKey>(queryClient!, latestOptions);
 
     let isFirstRun = true;
     const get = (async () => {

tests/sync-plugins/tanstack-query.test.ts

@@ -1,95 +1,57 @@
-jest.mock('@tanstack/query-core', () => {
-    const refetchMock = jest.fn(() => Promise.resolve({ data: 'fresh', status: 'success' }));
-
-    class QueryObserver {
-        client: any;
-        options: any;
-
-        constructor(client: any, options: any) {
-            this.client = client;
-            this.options = options;
-        }
-
-        getOptimisticResult() {
-            return { isLoading: false, data: 'initial', status: 'success' };
-        }
-
-        refetch() {
-            return refetchMock();
-        }
-
-        subscribe() {
-            return () => {};
-        }
-
-        updateResult() {}
-    }
-
-    class MutationObserver {
-        client: any;
-        options: any;
-
-        constructor(client: any, options: any) {
-            this.client = client;
-            this.options = options;
-        }
-
-        mutate(value: any) {
-            return Promise.resolve(value);
-        }
-    }
-
-    class QueryClient {
-        defaultQueryOptions(options: any) {
-            return options;
-        }
-
-        getMutationCache() {
-            return {
-                findAll: () => [],
-                remove: () => {},
-            };
-        }
-    }
-
-    const notifyManager = {
-        batchCalls:
-            (fn: (...args: any[]) => any) =>
-            (...args: any[]) =>
-                fn(...args),
-    };
-
-    return {
-        __esModule: true,
-        QueryObserver,
-        MutationObserver,
-        QueryClient,
-        notifyManager,
-        DefaultError: Error,
-    };
-});
-
-import { Synced } from '@legendapp/state/sync';
-import { symbolLinked } from '../../src/globals';
-import { syncedQuery } from '../../src/sync-plugins/tanstack-query';
+import { observable, syncState, when } from '@legendapp/state';
 import { QueryClient } from '@tanstack/query-core';
+import { syncedQuery } from '../../src/sync-plugins/tanstack-query';
+
+type Todo = {
+    id: number;
+    title: string;
+};
 
 describe('syncedQuery', () => {
-    test('get returns refetched data after the initial run', async () => {
+    test('applies selected data from fetches and cache updates', async () => {
         const queryClient = new QueryClient();
-        const linkedFactory = syncedQuery({
-            queryClient,
-            query: {
-                queryKey: ['test'],
-            },
-        }) as () => Synced<any>;
-
-        const options = linkedFactory()[symbolLinked];
-
-        const initial = await options.get!({});
-        expect(initial).toBe('initial');
-
-        const second = await options.get!({});
-        expect(second).toBe('fresh');
+        const queryKey = ['todos'] as const;
+        let todos: Todo[] = [
+            { id: 1, title: 'First' },
+            { id: 2, title: 'Second' },
+        ];
+        const todos$ = observable(
+            syncedQuery({
+                queryClient,
+                query: {
+                    queryKey,
+                    queryFn: async () => todos,
+                    select: (data) => data.length,
+                },
+            }),
+        );
+        const state$ = syncState(todos$);
+        const unsubscribe = todos$.onChange(() => {});
+
+        try {
+            expect(todos$.get()).toBeUndefined();
+            await when(() => todos$.get() === 2);
+            await when(() => !state$.isGetting.get());
+            expect(todos$.get()).toBe(2);
+
+            todos = [{ id: 1, title: 'First' }];
+            state$.sync();
+
+            expect(state$.isGetting.get()).toBe(true);
+            await when(() => !state$.isGetting.get());
+            expect(todos$.get()).toBe(1);
+
+            queryClient.setQueryData<Todo[]>(queryKey, [
+                ...todos,
+                { id: 2, title: 'Second' },
+                { id: 3, title: 'Third' },
+            ]);
+
+            await when(() => todos$.get() === 3);
+            expect(todos$.get()).toBe(3);
+        } finally {
+            unsubscribe();
+            queryClient.clear();
+        }
     });
 });

tests/types.test.ts

@@ -7,6 +7,79 @@ import {
     ObservablePrimitive,
 } from '../src/observableTypes';
 import { observable } from '../src/observable';
+import { queryOptions, QueryClient } from '@tanstack/react-query';
+import { syncedQuery } from '../src/sync-plugins/tanstack-query';
+import { useObservableSyncedQuery } from '../src/sync-plugins/tanstack-react-query';
+
+type Todo = {
+    id: number;
+    title: string;
+};
+
+describe('TanStack Query types', () => {
+    it('infers selected data through the core and React APIs', () => {
+        const queryClient = new QueryClient();
+        const query = queryOptions({
+            queryKey: ['todos'] as const,
+            queryFn: async (): Promise<Todo[]> => [{ id: 1, title: 'First' }],
+            select: (todos) => todos.length,
+        });
+
+        function core() {
+            return syncedQuery({ queryClient, query });
+        }
+
+        function react() {
+            return useObservableSyncedQuery({ queryClient, query });
+        }
+
+        expectTypeOf<ReturnType<typeof core>>().toEqualTypeOf<number>();
+        expectTypeOf<ReturnType<typeof react>['get']>().returns.toEqualTypeOf<number>();
+    });
+
+    it('infers selected data from a plain object in the React API', () => {
+        function react() {
+            return useObservableSyncedQuery({
+                queryClient: new QueryClient(),
+                query: {
+                    queryKey: ['todos'] as const,
+                    queryFn: async () => [{ id: 1, title: 'First' }],
+                    select: (todos) => {
+                        expectTypeOf(todos).toEqualTypeOf<Todo[]>();
+                        return todos.length;
+                    },
+                },
+            });
+        }
+
+        expectTypeOf<ReturnType<typeof react>['get']>().returns.toEqualTypeOf<number>();
+    });
+
+    it('preserves raw query data and a reactive tuple key in query callbacks', () => {
+        const queryClient = new QueryClient();
+
+        function selectedQuery() {
+            return syncedQuery({
+                queryClient,
+                query: {
+                    queryKey: () => ['todo', 1] as const,
+                    queryFn: async ({ queryKey }): Promise<Todo[]> => {
+                        expectTypeOf(queryKey).toEqualTypeOf<readonly ['todo', 1]>();
+                        return [{ id: 1, title: 'First' }];
+                    },
+                    placeholderData: (previousData, previousQuery) => {
+                        expectTypeOf(previousData).toEqualTypeOf<Todo[] | undefined>();
+                        expectTypeOf(previousQuery?.state.data).toEqualTypeOf<Todo[] | undefined>();
+                        return previousData;
+                    },
+                    select: (todos) => todos.length,
+                },
+            });
+        }
+
+        expectTypeOf<ReturnType<typeof selectedQuery>>().toEqualTypeOf<number>();
+    });
+});
 
 describe('observable', () => {
     it('optional object return type when no argument is passed', () => {