go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/milo/ui/src/common/hooks/prpc_query/use_prpc_service_client.ts

// Copyright 2024 The LUCI 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 {
  useAuthState,
  useGetAccessToken,
} from '@/common/components/auth_state_provider';
import { PrpcClient } from '@/generic_libs/tools/prpc_client';
import { Rpc } from '@/proto_utils/types';

export interface PrpcServiceClientOptions<S, Params extends unknown[] = []> {
  /**
   * The host of the pRPC server.
   */
  readonly host: string;
  /**
   * If true, use HTTP instead of HTTPS. Defaults to `false`.
   */
  readonly insecure?: boolean;
  /**
   * A generic client implementation generated by ts-proto.
   * e.g. `UserClientImpl`. If the constructor requires additional params, they
   * can be provided via `usePrpcServiceClient(opts, ...params)`.
   */
  readonly ClientImpl: Constructor<S, [Rpc, ...Params]> & {
    readonly DEFAULT_SERVICE: string;
  };
}

/**
 * A pRPC method with optional additional params.
 */
type Method<Req, Res, Params extends unknown[] = []> = (
  req: Req,
  ...params: Params
) => Res;

/**
 * The request object from a pRPC method.
 */
type MethodReq<T> = T extends Method<infer Req, infer _Res> ? Req : never;

type PagedReq<Req> = Req extends { readonly pageToken: string } ? Req : never;
type PagedRes<Res> = Awaited<Res> extends { readonly nextPageToken: string }
  ? Res
  : never;

type MethodRes<T> = T extends Method<infer _Req, infer Res> ? Res : never;

type MethodParams<T> = T extends Method<infer _Req, infer _Res, infer Params>
  ? Params
  : never;

type MethodKeys<S, Req, Res, Params extends unknown[] = []> = keyof {
  [MK in keyof S as S[MK] extends Method<Req, Res, Params>
    ? MK
    : never]: unknown;
};

export type DecoratedMethod<MK, Req, Res, Params extends unknown[] = []> = {
  (req: Req, ...params: Params): Res;
  /**
   * Builds a ReactQuery option that queries the method. The query key is
   * consisted of
   * `[userIdentity, 'prpc', serviceHost, serviceName, methodName, request]`.
   */
  query(
    req: Req,
    ...params: Params
  ): {
    queryKey: [string, 'prpc', string, string, MK, Req];
    queryFn: () => Res;
  };
  /**
   * Builds a ReactQuery option that queries the paginated method. The next page
   * param is automatically extracted from the previous response. The query key
   * is consisted of
   * `[userIdentity, 'prpc-paged', serviceHost, serviceName, methodName, request]`.
   */
  queryPaged(
    req: PagedRes<Res> extends Res ? PagedReq<Req> : never,
    ...params: Params
  ): {
    queryKey: [string, 'prpc-paged', string, string, MK, Req];
    queryFn: (ctx: { pageParam?: string }) => Res;
    getNextPageParam: (lastRes: Awaited<Res>) => string | null;
  };
};

export type DecoratedClient<S> = {
  // The request type has to be `any` because the argument type must be contra-
  // variant when sub-typing a function.
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  [MK in keyof S]: S[MK] extends Method<any, any>
    ? DecoratedMethod<
        MK,
        MethodReq<S[MK]>,
        MethodRes<S[MK]>,
        MethodParams<S[MK]>
      >
    : S[MK];
};

export interface CallOpt<MK, Request> {
  readonly method: MK;
  readonly request: Request;
}

/**
 * Construct a decorated pRPC client with the users credentials. Add helper
 * functions to each pRPC method to help generating ReactQuery options.
 *
 * Example:
 * ```typescript
 * // Constructs a client.
 * const client = usePrpcClient({
 *   // The host of the pRPC server.
 *   host: 'cr-buildbucket-dev.appspot.com',
 *   // The client implementation generated by ts-proto.
 *   ClientImpl: BuildsClientImpl,
 * })
 *
 * // Use the `.query` helper function to build a ReactQuery.
 * const {data, isLoading, ...} = useQuery({
 *   // `.query` is available on all pRPC methods. It takes the same parameters
 *   // as the method itself, and returns a ReactQuery options object with
 *   // `queryKey` and `queryFn` populated.
 *   ...client.GetBuild.query({
 *     // A type checked request object. The type of the request object is
 *     // inferred from the supplied client implementation and method name.
 *     id: "1234",
 *     ...
 *   }),
 *   // Add or override query options if needed.
 *   select: (res) => {...},
 *   ...
 * })
 *
 * // Use the `.queryPaged` helper function to build an infinite ReactQuery.
 * const {data, isLoading, ...} = useInfiniteQuery({
 *   // `.queryPaged` similar to `.query` except that
 *   // 1. it is only available on pRPC methods that have 'pageToken' in the
 *   // request object and 'nextPageToken' in the response object, and
 *   // 2. in addition to `queryKey` and `queryFn`, it also populates
 *   // `getNextPageParam`, making it suitable for `useInfiniteQuery`.
 *   ...client.SearchBuilds.queryPaged({
 *     // A type checked request object. The type of the request object is
 *     // inferred from the supplied client implementation and method name.
 *     id: "1234",
 *     ...
 *   }),
 *   // Add or override query options if needed.
 *   select: (res) => {...},
 *   ...
 * })
 * ```
 */
export function usePrpcServiceClient<
  S extends object,
  Params extends unknown[],
>(opts: PrpcServiceClientOptions<S, Params>, ...params: Params) {
  const { host, insecure, ClientImpl } = opts;

  const { identity } = useAuthState();
  const getAuthToken = useGetAccessToken();
  const client = new ClientImpl(
    new PrpcClient({ host, insecure, getAuthToken }),
    ...params,
  );

  return new Proxy(client, {
    get(target, p, receiver) {
      const mk = p as MethodKeys<S, unknown, unknown>;
      const value = Reflect.get(target, mk, receiver);
      if (typeof value !== 'function') {
        return value;
      }
      const fn = (...args: unknown[]) => value.apply(target, args);
      fn.query = (req: object, ...params: unknown[]) => ({
        queryKey: [identity, 'prpc', host, ClientImpl.DEFAULT_SERVICE, mk, req],
        queryFn: () => fn(req, ...params),
      });
      fn.queryPaged = (req: object, ...params: unknown[]) => ({
        queryKey: [
          identity,
          'prpc-paged',
          host,
          ClientImpl.DEFAULT_SERVICE,
          mk,
          req,
        ],
        queryFn: ({ pageParam = '' }) =>
          fn(pageParam ? { ...req, pageToken: pageParam } : req, ...params),
        // Return `null` when the next page token is an empty string so it get
        // treated as no next page.
        getNextPageParam: ({ nextPageToken = '' }) => nextPageToken || null,
      });
      return fn;
    },
  }) as DecoratedClient<S>;
}