diff --git a/.gitignore b/.gitignore index a77f46a2..ead027e1 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,5 @@ node_modules /packages/*/dist /packages/*/coverage +*.tsbuildinfo +.vscode diff --git a/packages/connect-query-core/src/connect-query-client.test.ts b/packages/connect-query-core/src/connect-query-client.test.ts new file mode 100644 index 00000000..65eb6cf4 --- /dev/null +++ b/packages/connect-query-core/src/connect-query-client.test.ts @@ -0,0 +1,465 @@ +// Copyright 2021-2023 The Connect Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import { type MessageShape } from "@bufbuild/protobuf"; +import type { Query } from "@tanstack/query-core"; +import { mockEliza, mockPaginatedTransport } from "test-utils"; +import type { SayResponseSchema } from "test-utils/gen/eliza_pb.js"; +import { ElizaService } from "test-utils/gen/eliza_pb.js"; +import { ListService } from "test-utils/gen/list_pb.js"; +import { describe, expect, it } from "vitest"; +import { QueryClient } from "@tanstack/query-core"; + +import { createConnectQueryKey } from "./connect-query-key.js"; +import { + createConnectQueryClient, + type ConnectQueryClient, +} from "./connect-query-client.js"; + +const sayMethodDescriptor = ElizaService.method.say; + +const mockedElizaTransport = mockEliza(); + +const paginatedTransport = mockPaginatedTransport(); + +const queryDetails = { + schema: sayMethodDescriptor, + input: { + sentence: "Pablo", + }, + transport: mockedElizaTransport, +}; + +function enhanceQueryClient(queryClient: QueryClient) { + Object.assign(queryClient, createConnectQueryClient(queryClient)); + return queryClient as ConnectQueryClient & QueryClient; +} + +describe("prefetch APIs", () => { + it("populates a single query cache", async () => { + const queryClient = enhanceQueryClient(new QueryClient()); + const currentCacheItems = queryClient.getQueryCache().findAll(); + expect(currentCacheItems).toHaveLength(0); + + await queryClient.prefetchConnectQuery( + queryDetails.schema, + queryDetails.input, + { transport: queryDetails.transport }, + ); + const item = queryClient.getConnectQueryData({ + ...queryDetails, + cardinality: "finite", + }); + expect(item?.sentence).toBe("Hello Pablo"); + + const queryState = queryClient.getConnectQueryState({ + ...queryDetails, + cardinality: "finite", + }); + expect(queryState?.status).toBe("success"); + expect(queryState?.fetchStatus).toBe("idle"); + expect(queryState?.dataUpdateCount).toBe(1); + }); + + it("populates an infinite query cache", async () => { + const queryClient = enhanceQueryClient(new QueryClient()); + const currentCacheItems = queryClient.getQueryCache().findAll(); + expect(currentCacheItems).toHaveLength(0); + + await queryClient.prefetchConnectInfiniteQuery( + ListService.method.list, + { + preview: true, + page: 0n, + }, + { + transport: paginatedTransport, + pageParamKey: "page", + getNextPageParam: (data) => data.page, + }, + ); + + const details = { + schema: ListService.method.list, + transport: paginatedTransport, + input: { + preview: true, + }, + cardinality: "infinite" as const, + }; + + const item = queryClient.getConnectQueryData(details); + + const nextItems = queryClient.getQueryCache().findAll(); + expect(nextItems).toHaveLength(1); + expect(item?.pages[0].items).toHaveLength(3); + + const queryState = queryClient.getConnectQueryState(details); + expect(queryState?.status).toBe("success"); + expect(queryState?.fetchStatus).toBe("idle"); + expect(queryState?.dataUpdateCount).toBe(1); + }); +}); + +describe("invalidateConnectQueries", () => { + it("invalidates a specific query", async () => { + const queryClient = enhanceQueryClient(new QueryClient()); + await queryClient.prefetchConnectQuery( + queryDetails.schema, + queryDetails.input, + { transport: queryDetails.transport }, + ); + await queryClient.invalidateConnectQueries(queryDetails); + const queryState = queryClient.getConnectQueryState({ + ...queryDetails, + cardinality: "finite", + }); + expect(queryState?.isInvalidated).toBe(true); + }); + + it("invalidate all methods for a given service", async () => { + const queryClient = enhanceQueryClient(new QueryClient()); + await queryClient.prefetchConnectQuery( + queryDetails.schema, + queryDetails.input, + { transport: queryDetails.transport }, + ); + + await queryClient.invalidateConnectQueries({ + schema: ElizaService, + transport: mockedElizaTransport, + }); + const queryState = queryClient.getConnectQueryState({ + ...queryDetails, + cardinality: "finite", + }); + expect(queryState?.isInvalidated).toBe(true); + }); +}); + +describe("refetchConnectQueries", () => { + it("refetch a specific query", async () => { + const queryClient = enhanceQueryClient(new QueryClient()); + await queryClient.prefetchConnectQuery( + queryDetails.schema, + queryDetails.input, + { transport: queryDetails.transport }, + ); + await queryClient.refetchConnectQueries(queryDetails); + const queryState = queryClient.getConnectQueryState({ + ...queryDetails, + cardinality: "finite", + }); + expect(queryState?.dataUpdateCount).toBe(2); + }); + + it("refetch all methods for a given service", async () => { + const queryClient = enhanceQueryClient(new QueryClient()); + await queryClient.prefetchConnectQuery( + queryDetails.schema, + queryDetails.input, + { transport: queryDetails.transport }, + ); + + await queryClient.refetchConnectQueries({ + schema: ElizaService, + transport: mockedElizaTransport, + }); + const queryState = queryClient.getConnectQueryState({ + ...queryDetails, + cardinality: "finite", + }); + expect(queryState?.dataUpdateCount).toBe(2); + }); +}); + +describe("setConnectQueryData", () => { + it("updates locally fetched data", async () => { + const queryClient = enhanceQueryClient(new QueryClient()); + await queryClient.prefetchConnectQuery( + queryDetails.schema, + queryDetails.input, + { transport: queryDetails.transport }, + ); + + const queryState = queryClient.getConnectQueryState({ + ...queryDetails, + cardinality: "finite", + }); + expect(queryState?.dataUpdateCount).toBe(1); + expect(queryState?.data?.sentence).toBe("Hello Pablo"); + + queryClient.setConnectQueryData( + { + ...queryDetails, + cardinality: "finite", + }, + { + sentence: "Hello Stu", + }, + ); + + const newQueryState = queryClient.getConnectQueryState({ + ...queryDetails, + cardinality: "finite", + }); + + expect(newQueryState?.dataUpdateCount).toBe(2); + expect(newQueryState?.data?.sentence).toBe("Hello Stu"); + }); + + it("updates locally fetched data with a callback", async () => { + const queryClient = enhanceQueryClient(new QueryClient()); + await queryClient.prefetchConnectQuery( + queryDetails.schema, + queryDetails.input, + { transport: queryDetails.transport }, + ); + + const queryState = queryClient.getConnectQueryState({ + ...queryDetails, + cardinality: "finite", + }); + expect(queryState?.dataUpdateCount).toBe(1); + expect(queryState?.data?.sentence).toBe("Hello Pablo"); + + queryClient.setConnectQueryData( + { + ...queryDetails, + cardinality: "finite", + }, + (prev) => { + if (prev === undefined) { + return undefined; + } + expect(prev.sentence).toBe("Hello Pablo"); + return { + sentence: "Hello Stu", + }; + }, + ); + + const newQueryState = queryClient.getConnectQueryState({ + ...queryDetails, + cardinality: "finite", + }); + + expect(newQueryState?.dataUpdateCount).toBe(2); + expect(newQueryState?.data?.sentence).toBe("Hello Stu"); + expect(newQueryState?.data?.$typeName).toBe( + "connectrpc.eliza.v1.SayResponse", + ); + }); + + it("can update infinite paginated data", async () => { + const queryClient = enhanceQueryClient(new QueryClient()); + await queryClient.prefetchConnectInfiniteQuery( + ListService.method.list, + { + page: 0n, + }, + { + transport: paginatedTransport, + getNextPageParam: (l) => l.page + 1n, + pageParamKey: "page", + }, + ); + + const queryState = queryClient.getConnectQueryState({ + schema: ListService.method.list, + transport: paginatedTransport, + cardinality: "infinite", + input: { + page: 0n, + }, + }); + expect(queryState?.dataUpdateCount).toBe(1); + expect(queryState?.data?.pages).toHaveLength(1); + + queryClient.setConnectQueryData( + { + schema: ListService.method.list, + transport: paginatedTransport, + cardinality: "infinite", + }, + { + pageParams: [0n, 1n], + pages: [ + { + page: 0n, + items: ["a", "b", "c"], + }, + { + page: 1n, + items: ["x", "y", "z"], + }, + ], + }, + ); + + const newQueryState = queryClient.getConnectQueryState({ + schema: ListService.method.list, + transport: paginatedTransport, + cardinality: "infinite", + input: { + page: 0n, + }, + }); + + expect(newQueryState?.dataUpdateCount).toBe(2); + expect(newQueryState?.data?.pages).toHaveLength(2); + }); +}); + +describe("setConnectQueriesData", () => { + it("update locally fetched data across multiple queries", async () => { + const queryClient = enhanceQueryClient(new QueryClient()); + await queryClient.prefetchConnectQuery( + queryDetails.schema, + queryDetails.input, + { transport: queryDetails.transport }, + ); + await queryClient.prefetchConnectQuery( + queryDetails.schema, + { + sentence: "Stu", + }, + { transport: queryDetails.transport }, + ); + + const cachedItems = queryClient.getQueryCache().findAll({ + queryKey: createConnectQueryKey({ + ...queryDetails, + input: {}, + cardinality: "finite", + }), + }); + expect(cachedItems).toHaveLength(2); + + queryClient.setConnectQueriesData( + { + schema: sayMethodDescriptor, + cardinality: "finite", + }, + (prev) => { + if (prev === undefined) { + return undefined; + } + return { + ...prev, + sentence: prev.sentence + "!", + }; + }, + ); + + const newCachedItems = queryClient.getQueryCache().findAll() as Query< + MessageShape + >[]; + expect(newCachedItems).toHaveLength(2); + expect(newCachedItems[0].state.data?.sentence).toBe("Hello Pablo!"); + expect(newCachedItems[1].state.data?.sentence).toBe("Hello Stu!"); + }); +}); + +describe("fetchConnectInfiniteQuery", () => { + it("fetches infinite data", async () => { + const queryClient = enhanceQueryClient(new QueryClient()); + const result = await queryClient.fetchConnectInfiniteQuery( + ListService.method.list, + { + preview: true, + page: 0n, + }, + { + transport: paginatedTransport, + getNextPageParam: (data) => { + return data.page + 1n; + }, + pageParamKey: "page", + }, + ); + + expect(result).toBeDefined(); + expect(result.pages).toHaveLength(1); + expect(result.pages[0].$typeName).toBe("ListResponse"); + expect(result.pages[0].items).toHaveLength(3); + }); +}); + +describe("getConnectQueryState", () => { + it("can get state for infinite queries", async () => { + const queryClient = enhanceQueryClient(new QueryClient()); + await queryClient.fetchConnectInfiniteQuery( + ListService.method.list, + { + preview: true, + page: 0n, + }, + { + transport: paginatedTransport, + getNextPageParam: (data) => { + return data.page + 1n; + }, + pageParamKey: "page", + }, + ); + + const state = queryClient.getConnectQueryState({ + schema: ListService.method.list, + transport: paginatedTransport, + input: { + preview: true, + }, + cardinality: "infinite", + }); + + expect(state?.status).toBe("success"); + expect(state?.fetchStatus).toBe("idle"); + expect(state?.dataUpdateCount).toBe(1); + }); +}); + +describe("ensure APIs", () => { + it("ensure data exists for infinite queries", async () => { + const queryClient = enhanceQueryClient(new QueryClient()); + const data = await queryClient.ensureConnectInfiniteQueryData( + ListService.method.list, + { + preview: true, + page: 0n, + }, + { + transport: paginatedTransport, + getNextPageParam: (localData) => { + return localData.page + 1n; + }, + pageParamKey: "page", + staleTime: 1000, + }, + ); + + const state = queryClient.getConnectQueryState({ + schema: ListService.method.list, + transport: paginatedTransport, + input: { + preview: true, + }, + cardinality: "infinite", + }); + expect(state?.status).toBe("success"); + expect(state?.fetchStatus).toBe("idle"); + expect(state?.dataUpdateCount).toBe(1); + expect(data.pages[0]).toBe(state?.data?.pages[0]); + }); +}); diff --git a/packages/connect-query-core/src/connect-query-client.ts b/packages/connect-query-core/src/connect-query-client.ts new file mode 100644 index 00000000..c9f7c661 --- /dev/null +++ b/packages/connect-query-core/src/connect-query-client.ts @@ -0,0 +1,726 @@ +// Copyright 2021-2023 The Connect Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import { + create, + isMessage, + type DescMessage, + type DescMethod, + type DescMethodUnary, + type DescService, + type MessageInitShape, + type MessageShape, +} from "@bufbuild/protobuf"; +import type { ConnectError, Transport } from "@connectrpc/connect"; +import type { + FetchInfiniteQueryOptions as TanstackFetchInfiniteQueryOptions, + FetchQueryOptions as TanstackFetchQueryOptions, + InfiniteData, + InvalidateOptions, + InvalidateQueryFilters, + QueryFilters, + QueryState, + RefetchOptions, + RefetchQueryFilters, + ResetOptions, + SetDataOptions, + Updater, +} from "@tanstack/query-core"; +import { QueryClient as TSQueryClient } from "@tanstack/query-core"; + +import type { ConnectQueryKey } from "./connect-query-key.js"; +import { createConnectQueryKey } from "./connect-query-key.js"; +import type { ConnectInfiniteQueryOptions } from "./create-infinite-query-options.js"; +import { createInfiniteQueryOptions } from "./create-infinite-query-options.js"; +import { createQueryOptions } from "./create-query-options.js"; + +type FetchQueryOptions< + O extends DescMessage, + SelectOutData = MessageShape, +> = Omit< + TanstackFetchQueryOptions< + MessageShape, + ConnectError, + SelectOutData, + ConnectQueryKey + >, + "queryFn" | "queryKey" +> & { + /** The transport to be used for the fetching. */ + transport?: Transport; +}; + +type FetchInfiniteQueryOptions< + I extends DescMessage, + O extends DescMessage, + ParamKey extends keyof MessageInitShape, +> = Omit< + TanstackFetchInfiniteQueryOptions< + MessageShape, + ConnectError, + MessageShape, + ConnectQueryKey, + MessageInitShape[ParamKey] + >, + "getNextPageParam" | "initialPageParam" | "queryFn" | "queryKey" +> & + ConnectInfiniteQueryOptions & { + transport: Transport; + }; + +type KeyParamsForMethod = { + /** + * Set `serviceName` and `methodName` in the key. + */ + schema: Desc; + /** + * Set `input` in the key: + * - If a SkipToken is provided, `input` is "skipped". + * - If an init shape is provided, `input` is set to a message key. + * - If omitted or undefined, `input` is not set in the key. + */ + input?: MessageInitShape | undefined; + /** + * Set `transport` in the key. + */ + transport?: Transport; + /** + * Set `cardinality` in the key - undefined is used for filters to match both finite and infinite queries. + */ + cardinality?: "finite" | "infinite"; + /** + * If omit the field with this name from the key for infinite queries. + */ + pageParamKey?: keyof MessageInitShape; +}; + +type KeyParamsForService = { + /** + * Set `serviceName` in the key, and omit `methodName`. + */ + schema: Desc; + /** + * Set `transport` in the key. + */ + transport?: Transport; + /** + * Set `cardinality` in the key - undefined is used for filters to match both finite and infinite queries. + */ + cardinality?: "finite" | "infinite"; +}; + +/** + * A custom query client that adds some useful methods to access typesafe query data and other shortcuts. + */ +export interface ConnectQueryClient { + /** + * Invalidate and refetch all queries that match the given schema. This + * can include all queries for a service (and sub methods), or all queries for a method. + * + * @see {@link https://tanstack.com/query/latest/docs/reference/QueryClient/#queryclientinvalidatequeries} + */ + invalidateConnectQueries< + I extends DescMessage, + O extends DescMessage, + Desc extends DescService, + >( + params: + | KeyParamsForMethod> + | KeyParamsForService, + filterOptions?: Omit, + options?: InvalidateOptions, + ): Promise; + + /** + * Refetch all queries that match the given schema. This can include all queries for a service (and sub methods), + * or all queries for a method. + * + * @see {@link https://tanstack.com/query/latest/docs/reference/QueryClient/#queryclientrefetchqueries} + */ + refetchConnectQueries< + I extends DescMessage, + O extends DescMessage, + Desc extends DescService, + >( + params: + | KeyParamsForMethod> + | KeyParamsForService, + filterOptions?: Omit, + options?: RefetchOptions, + ): Promise; + + /** + * Set the data for a single query. The query must match exactly the input provided, as well + * as the transport and cardinality (whether it was a finite or infinite query). + * + * @see {@link https://tanstack.com/query/latest/docs/reference/QueryClient/#queryclientsetquerydata} + */ + setConnectQueryData( + keyDescriptor: { + schema: DescMethodUnary; + input?: MessageInitShape; + transport: Transport; + cardinality: "infinite"; + }, + updater: ConnectInfiniteUpdater, + options?: SetDataOptions, + ): O; + setConnectQueryData( + keyDescriptor: { + schema: DescMethodUnary; + input?: MessageInitShape; + transport: Transport; + cardinality: "finite"; + }, + updater: ConnectUpdater, + options?: SetDataOptions, + ): O; + + /** + * Get the data for a single query. The query must match exactly the input provided, as well + * as the transport and cardinality (whether it was a finite or infinite query). + * + * @see {@link https://tanstack.com/query/latest/docs/reference/QueryClient/#queryclientgetquerydata} + */ + getConnectQueryData< + I extends DescMessage, + O extends DescMessage, + >(keyDescriptor: { + schema: DescMethodUnary; + input?: MessageInitShape; + transport: Transport; + cardinality: "finite"; + }): MessageShape | undefined; + getConnectQueryData< + I extends DescMessage, + O extends DescMessage, + >(keyDescriptor: { + schema: DescMethodUnary; + input?: MessageInitShape; + transport: Transport; + cardinality: "infinite"; + }): InfiniteData> | undefined; + + /** + * Sets the data for any matching queries for a given method. The input is optional, and anything left + * as undefined will greedily match queries. + * + * @see {@link https://tanstack.com/query/latest/docs/reference/QueryClient/#queryclientsetqueriesdata} + */ + setConnectQueriesData( + keyDescriptor: { + schema: DescMethodUnary; + input?: MessageInitShape; + transport?: Transport; + cardinality: "finite"; + }, + updater: ConnectUpdater, + options?: SetDataOptions & { + exact?: boolean; + }, + ): [readonly unknown[], unknown][]; + setConnectQueriesData( + keyDescriptor: { + schema: DescMethodUnary; + input?: MessageInitShape; + transport?: Transport; + cardinality: "infinite"; + }, + updater: ConnectInfiniteUpdater, + options?: SetDataOptions & { + exact?: boolean; + }, + ): [readonly unknown[], unknown][]; + + /** + * Fetch a single query and return the result. + * + * @see {@link https://tanstack.com/query/latest/docs/reference/QueryClient/#queryclientfetchquery} + */ + fetchConnectQuery< + I extends DescMessage, + O extends DescMessage, + SelectOutData = MessageShape, + >( + schema: DescMethodUnary, + input: MessageInitShape | undefined, + { + transport, + ...queryOptions + }: { + transport: Transport; + } & FetchQueryOptions, + ): Promise; + + /** + * Fetch a single infinite query and return the result. + * + * @see {@link https://tanstack.com/query/latest/docs/reference/QueryClient/#queryclientfetchinfinitequery} + */ + fetchConnectInfiniteQuery< + I extends DescMessage, + O extends DescMessage, + ParamKey extends keyof MessageInitShape, + >( + schema: DescMethodUnary, + input: MessageInitShape & Required, ParamKey>>, + { + transport, + getNextPageParam, + pageParamKey, + ...queryOptions + }: FetchInfiniteQueryOptions, + ): Promise>>; + + /** + * Prefetch a single query and discard the result. + * + * @see {@link https://tanstack.com/query/latest/docs/reference/QueryClient/#queryclientprefetchquery} + */ + prefetchConnectQuery< + I extends DescMessage, + O extends DescMessage, + SelectOutData = MessageShape, + >( + schema: DescMethodUnary, + input: MessageInitShape | undefined, + { + transport, + ...queryOptions + }: { + transport: Transport; + } & FetchQueryOptions, + ): Promise; + + /** + * Prefetch a single infinite query and discard the result. + * + * @see {@link https://tanstack.com/query/latest/docs/reference/QueryClient/#queryclientprefetchinfinitequery} + */ + prefetchConnectInfiniteQuery< + I extends DescMessage, + O extends DescMessage, + ParamKey extends keyof MessageInitShape, + >( + schema: DescMethodUnary, + input: MessageInitShape & Required, ParamKey>>, + { + transport, + getNextPageParam, + pageParamKey, + ...queryOptions + }: FetchInfiniteQueryOptions, + ): Promise; + + /** + * Get the query state for a single query. The query must match exactly the input provided, as well + * as the transport and cardinality (whether it was a finite or infinite query). + * + * @see {@link https://tanstack.com/query/latest/docs/reference/QueryClient/#queryclientgetquerystate} + */ + getConnectQueryState(keyDescriptor: { + schema: Desc; + input?: MessageInitShape; + transport: Transport; + cardinality: "finite"; + }): QueryState, ConnectError> | undefined; + getConnectQueryState(keyDescriptor: { + schema: Desc; + input?: MessageInitShape; + transport: Transport; + cardinality: "infinite"; + }): + | QueryState>, ConnectError> + | undefined; + + /** + * Ensure the query data for a single query. The query must match exactly the input provided, as well + * as the transport and cardinality (whether it was a finite or infinite query). + * + * @see {@link https://tanstack.com/query/latest/docs/reference/QueryClient/#queryclientensurequerydata} + */ + ensureConnectQueryData< + I extends DescMessage, + O extends DescMessage, + SelectOutData = MessageShape, + >( + schema: DescMethodUnary, + input: MessageInitShape | undefined, + { + transport, + ...queryOptions + }: { + transport: Transport; + } & FetchQueryOptions, + ): Promise; + + /** + * Ensure the query data for a single infinite query. The query must match exactly the input provided, as well + * as the transport and cardinality (whether it was a finite or infinite query). + * + * @see {@link https://tanstack.com/query/latest/docs/reference/QueryClient/#queryclientensureinfinitequerydata} + */ + ensureConnectInfiniteQueryData< + I extends DescMessage, + O extends DescMessage, + ParamKey extends keyof MessageInitShape, + >( + schema: DescMethodUnary, + input: MessageInitShape & Required, ParamKey>>, + { + transport, + getNextPageParam, + pageParamKey, + ...queryOptions + }: FetchInfiniteQueryOptions, + ): Promise>>; + + /** + * Get all data entries that match the given schema. This + * can include all queries for a service (and sub methods), or all queries for a method. + * + * @see {@link https://tanstack.com/query/latest/docs/reference/QueryClient/#queryclientgetqueriesdata} + */ + getConnectQueriesData< + I extends DescMessage, + O extends DescMessage, + Desc extends DescService, + >( + params: + | KeyParamsForMethod> + | KeyParamsForService, + filterOptions?: Omit, + ): unknown; + + /** + * Cancels any outgoing queries that match the given schema. This + * can include all queries for a service (and sub methods), or all queries for a method. + * + * @see {@link https://tanstack.com/query/latest/docs/reference/QueryClient/#queryclientcancelqueries} + */ + cancelConnectQueries< + I extends DescMessage, + O extends DescMessage, + Desc extends DescService, + >( + params: + | KeyParamsForMethod> + | KeyParamsForService, + filterOptions?: Omit, + ): Promise; + + /** + * Removes any queries from the cache that match the given schema. This + * can include all queries for a service (and sub methods), or all queries for a method. + * + * @see {@link https://tanstack.com/query/latest/docs/reference/QueryClient/#queryclientremovequeries} + */ + removeConnectQueries< + I extends DescMessage, + O extends DescMessage, + Desc extends DescService, + >( + params: + | KeyParamsForMethod> + | KeyParamsForService, + filterOptions?: Omit, + ): void; + + /** + * Resets any queries that match the given schema. This + * can include all queries for a service (and sub methods), or all queries for a method. + * + * @see {@link https://tanstack.com/query/latest/docs/reference/QueryClient/#queryclientresetqueries} + */ + resetConnectQueries< + I extends DescMessage, + O extends DescMessage, + Desc extends DescService, + >( + params: + | KeyParamsForMethod> + | KeyParamsForService, + filterOptions?: Omit, + options?: ResetOptions, + ): Promise; +} + +export function createConnectQueryClient(queryClient: TSQueryClient) { + const connectQueryClient: ConnectQueryClient = { + async invalidateConnectQueries(params, filterOptions, options) { + return queryClient.invalidateQueries( + { + ...filterOptions, + queryKey: createConnectQueryKey({ + ...params, + cardinality: params.cardinality, + }), + }, + options, + ); + }, + + async refetchConnectQueries(params, filterOptions, options) { + return queryClient.refetchQueries( + { + ...filterOptions, + queryKey: createConnectQueryKey({ + ...params, + cardinality: params.cardinality, + }), + }, + options, + ); + }, + + setConnectQueryData( + keyDescriptor: { + schema: DescMethodUnary; + input?: MessageInitShape; + transport: Transport; + cardinality: "finite" | "infinite"; + }, + updater: ConnectUpdater | ConnectInfiniteUpdater, + options?: SetDataOptions, + ): O | undefined { + const safeUpdater = + keyDescriptor.cardinality === "infinite" + ? createProtobufSafeInfiniteUpdater( + keyDescriptor.schema, + updater as ConnectInfiniteUpdater, + ) + : createProtobufSafeUpdater( + keyDescriptor.schema, + updater as ConnectUpdater, + ); + return queryClient.setQueryData( + createConnectQueryKey({ + ...keyDescriptor, + // Since we are matching on the exact input, we match what connect-query does in createQueryOptions + input: keyDescriptor.input ?? ({} as MessageInitShape), + }), + // eslint-disable-next-line @typescript-eslint/no-unsafe-argument, @typescript-eslint/no-explicit-any -- Need this override to keep avoid issues with NoInfer type messing with type checks + safeUpdater as any, + options, + ); + }, + + getConnectQueryData(keyDescriptor) { + return queryClient.getQueryData(createConnectQueryKey(keyDescriptor)); + }, + + setConnectQueriesData( + keyDescriptor: { + schema: DescMethodUnary; + input?: MessageInitShape; + transport?: Transport; + cardinality: "infinite" | "finite"; + }, + updater: ConnectInfiniteUpdater | ConnectUpdater, + options?: SetDataOptions & { + exact?: boolean; + }, + ): [readonly unknown[], unknown][] { + const safeUpdater = + keyDescriptor.cardinality === "finite" + ? createProtobufSafeUpdater( + keyDescriptor.schema, + updater as ConnectUpdater, + ) + : createProtobufSafeInfiniteUpdater( + keyDescriptor.schema, + updater as ConnectInfiniteUpdater, + ); + return queryClient.setQueriesData( + { + queryKey: createConnectQueryKey({ + ...keyDescriptor, + cardinality: keyDescriptor.cardinality, + }), + exact: options?.exact ?? false, + }, + safeUpdater, + options, + ); + }, + + fetchConnectQuery(schema, input, { transport, ...queryOptions }) { + return queryClient.fetchQuery({ + ...createQueryOptions(schema, input, { transport }), + ...queryOptions, + }); + }, + fetchConnectInfiniteQuery( + schema, + input, + { transport, getNextPageParam, pageParamKey, ...queryOptions }, + ) { + return queryClient.fetchInfiniteQuery({ + ...createInfiniteQueryOptions(schema, input, { + transport, + pageParamKey, + getNextPageParam, + }), + ...queryOptions, + }); + }, + prefetchConnectQuery(schema, input, { transport, ...queryOptions }) { + return queryClient.prefetchQuery({ + ...createQueryOptions(schema, input, { transport }), + ...queryOptions, + }); + }, + + prefetchConnectInfiniteQuery( + schema, + input, + { transport, getNextPageParam, pageParamKey, ...queryOptions }, + ) { + return queryClient.prefetchInfiniteQuery({ + ...createInfiniteQueryOptions(schema, input, { + transport, + pageParamKey, + getNextPageParam, + }), + ...queryOptions, + }); + }, + getConnectQueryState(keyDescriptor) { + return queryClient.getQueryState(createConnectQueryKey(keyDescriptor)); + }, + + ensureConnectQueryData(schema, input, { transport, ...queryOptions }) { + return queryClient.ensureQueryData({ + ...createQueryOptions(schema, input, { transport }), + ...queryOptions, + }); + }, + ensureConnectInfiniteQueryData( + schema, + input, + { transport, getNextPageParam, pageParamKey, ...queryOptions }, + ) { + return queryClient.ensureInfiniteQueryData({ + ...createInfiniteQueryOptions(schema, input, { + transport, + pageParamKey, + getNextPageParam, + }), + ...queryOptions, + }); + }, + getConnectQueriesData(params, filterOptions) { + return queryClient.getQueriesData({ + ...filterOptions, + queryKey: createConnectQueryKey({ + ...params, + cardinality: params.cardinality, + }), + }); + }, + cancelConnectQueries(params, filterOptions) { + return queryClient.cancelQueries({ + ...filterOptions, + queryKey: createConnectQueryKey({ + ...params, + cardinality: params.cardinality, + }), + }); + }, + removeConnectQueries(params, filterOptions) { + queryClient.removeQueries({ + ...filterOptions, + queryKey: createConnectQueryKey({ + ...params, + cardinality: params.cardinality, + }), + }); + return; + }, + resetConnectQueries(params, filterOptions, options) { + return queryClient.resetQueries( + { + ...filterOptions, + queryKey: createConnectQueryKey({ + ...params, + cardinality: params.cardinality, + }), + }, + options, + ); + }, + }; + return connectQueryClient; +} + +type ConnectInfiniteUpdater = + | InfiniteData> + | undefined + | (( + prev?: InfiniteData>, + ) => InfiniteData> | undefined); + +const createProtobufSafeInfiniteUpdater = + ( + schema: Pick, "output">, + updater: ConnectInfiniteUpdater, + ) => + ( + prev: InfiniteData> | undefined, + ): InfiniteData> | undefined => { + if (typeof updater !== "function") { + if (updater === undefined) { + return undefined; + } + return { + pageParams: updater.pageParams, + pages: updater.pages.map((i) => create(schema.output, i)), + }; + } + const result = updater(prev); + if (result === undefined) { + return undefined; + } + return { + pageParams: result.pageParams, + pages: result.pages.map((i) => create(schema.output, i)), + }; + }; + +type ConnectUpdater = + | MessageInitShape + | undefined + | ((prev?: MessageShape) => MessageInitShape | undefined); + +/** + * This method makes sure that the object returned + * is of the message type. If an init shape is returned, + * we'll run it through create again. + */ +const createProtobufSafeUpdater: ( + schema: Pick, "output">, + updater: ConnectUpdater, +) => Updater, MessageShape | undefined> = + (schema, updater) => (prev) => { + if (typeof updater !== "function") { + if (updater === undefined) { + return undefined; + } + if (isMessage(updater, schema.output)) { + return updater; + } + return create(schema.output, updater); + } + return create(schema.output, updater(prev)); + }; diff --git a/packages/connect-query-core/src/index.ts b/packages/connect-query-core/src/index.ts index 7222a945..3ef338af 100644 --- a/packages/connect-query-core/src/index.ts +++ b/packages/connect-query-core/src/index.ts @@ -23,3 +23,5 @@ export { createQueryOptions } from "./create-query-options.js"; export { addStaticKeyToTransport } from "./transport-key.js"; export type { SkipToken } from "@tanstack/query-core"; export { skipToken } from "@tanstack/query-core"; +export type { ConnectQueryClient } from "./connect-query-client.js"; +export { createConnectQueryClient } from "./connect-query-client.js"; diff --git a/packages/connect-query/src/index.ts b/packages/connect-query/src/index.ts index 4f4dfe89..ba46f4d1 100644 --- a/packages/connect-query/src/index.ts +++ b/packages/connect-query/src/index.ts @@ -23,3 +23,4 @@ export type { UseMutationOptions } from "./use-mutation.js"; export { useMutation } from "./use-mutation.js"; export type { UseInfiniteQueryOptions } from "./use-infinite-query.js"; export type { UseQueryOptions } from "./use-query.js"; +export { useConnectQueryClient } from "./use-connect-query-client.js"; diff --git a/packages/connect-query/src/use-connect-query-client.ts b/packages/connect-query/src/use-connect-query-client.ts new file mode 100644 index 00000000..a22e3ecc --- /dev/null +++ b/packages/connect-query/src/use-connect-query-client.ts @@ -0,0 +1,25 @@ +// Copyright 2021-2023 The Connect Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import { useQueryClient } from "@tanstack/react-query"; +import { useMemo } from "react"; +import { createConnectQueryClient } from "@connectrpc/connect-query-core"; +import type { ConnectQueryClient } from "@connectrpc/connect-query-core"; + +export function useConnectQueryClient(): ConnectQueryClient { + const queryClient = useQueryClient(); + return useMemo(() => { + return createConnectQueryClient(queryClient); + }, [queryClient]); +}