refactor: new useQueryTable with legit types, returns the query data (#2567)

* redo useQueryTable so that the types are legit and it returns the data

* fix junky getRowId and use new QueryTable on sleds and physical disks

* get wild with a list-specific helper to make the call sites clean

* encapsulate pageSize in the query config so it is defined in one place

* do the placeholderData thing for all lists

* scroll to top when page changes

* loading spinner on page changes!

* fix the pagination test, test lovely new scroll reset logic

* fix other e2es, don't scroll reset on browser forward/back

* fix bug found while converting other tables, extract useScrollReset

Author: david-crespo

app/api/client.ts

@@ -9,6 +9,8 @@ import { QueryClient } from '@tanstack/react-query'
 
 import { Api } from './__generated__/Api'
 import {
+  getApiQueryOptions,
+  getListQueryOptionsFn,
   getUseApiMutation,
   getUseApiQueries,
   getUseApiQuery,
@@ -17,13 +19,23 @@ import {
   wrapQueryClient,
 } from './hooks'
 
+export { ensurePrefetched, PAGE_SIZE, type PaginatedQuery } from './hooks'
+
 export const api = new Api({
   // unit tests run in Node, whose fetch implementation requires a full URL
   host: process.env.NODE_ENV === 'test' ? 'http://testhost' : '',
 })
 
 export type ApiMethods = typeof api.methods
 
+/** API-specific query options helper. */
+export const apiq = getApiQueryOptions(api.methods)
+/**
+ * Query options helper that only supports list endpoints. Returns
+ * a function `(limit, pageToken) => QueryOptions` for use with
+ * `useQueryTable`.
+ */
+export const getListQFn = getListQueryOptionsFn(api.methods)
 export const useApiQuery = getUseApiQuery(api.methods)
 export const useApiQueries = getUseApiQueries(api.methods)
 /**

app/api/hooks.ts

@@ -21,13 +21,15 @@ import {
   type UseQueryOptions,
   type UseQueryResult,
 } from '@tanstack/react-query'
+import * as R from 'remeda'
 import { type SetNonNullable } from 'type-fest'
 
 import { invariant } from '~/util/invariant'
 
 import type { ApiResult } from './__generated__/Api'
 import { processServerError, type ApiError } from './errors'
 import { navToLogin } from './nav-to-login'
+import { type ResultsPage } from './util'
 
 /* eslint-disable @typescript-eslint/no-explicit-any */
 export type Params<F> = F extends (p: infer P) => any ? P : never
@@ -123,6 +125,66 @@ export const getApiQueryOptions =
       ...options,
     })
 
+// Managed here instead of at the display layer so it can be built into the
+// query options and shared between loader prefetch and QueryTable
+export const PAGE_SIZE = 25
+
+/**
+ * This primarily exists so we can have an object that encapsulates everything
+ * useQueryTable needs to know about a query. In particular, it needs the page
+ * size, and you can't pull that out of the query options object unless you
+ * stick it in `meta`, and then we don't have type safety.
+ */
+export type PaginatedQuery<TData> = {
+  optionsFn: (
+    pageToken?: string
+  ) => UseQueryOptions<TData, ApiError> & { queryKey: QueryKey }
+  pageSize: number
+}
+
+/**
+ * This is the same as getApiQueryOptions except for two things:
+ *
+ *   1. We use a type constraint on the method key to ensure it can
+ *      only be used with endpoints that return a `ResultsPage`.
+ *   2. Instead of returning the options directly, it returns a paginated
+ *      query config object containing the page size and a function that
+ *      takes `limit` and `pageToken` and merges them into the query params
+ *      so that these can be passed in by `QueryTable`.
+ */
+export const getListQueryOptionsFn =
+  <A extends ApiClient>(api: A) =>
+  <
+    M extends string &
+      {
+        // this helper can only be used with endpoints that return ResultsPage
+        [K in keyof A]: Result<A[K]> extends ResultsPage<unknown> ? K : never
+      }[keyof A],
+  >(
+    method: M,
+    params: Params<A[M]>,
+    options: UseQueryOtherOptions<Result<A[M]>, ApiError> = {}
+  ): PaginatedQuery<Result<A[M]>> => {
+    // We pull limit out of the query params rather than passing it in some
+    // other way so that there is exactly one way of specifying it. If we had
+    // some other way of doing it, and then you also passed it in as a query
+    // param, it would be hard to guess which takes precedence. (pathOr plays
+    // nice when the properties don't exist.)
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    const limit = R.pathOr(params as any, ['query', 'limit'], PAGE_SIZE)
+    return {
+      optionsFn: (pageToken?: string) => {
+        const newParams = { ...params, query: { ...params.query, limit, pageToken } }
+        return getApiQueryOptions(api)(method, newParams, {
+          ...options,
+          // identity function so current page sticks around while next loads
+          placeholderData: (x) => x,
+        })
+      },
+      pageSize: limit,
+    }
+  }
+
 export const getUseApiQuery =
   <A extends ApiClient>(api: A) =>
   <M extends string & keyof A>(
@@ -140,7 +202,7 @@ export const getUsePrefetchedApiQuery =
     options: UseQueryOtherOptions<Result<A[M]>, ApiError> = {}
   ) => {
     const qOptions = getApiQueryOptions(api)(method, params, options)
-    return ensure(useQuery(qOptions), qOptions.queryKey)
+    return ensurePrefetched(useQuery(qOptions), qOptions.queryKey)
   }
 
 const prefetchError = (key?: QueryKey) =>
@@ -152,7 +214,11 @@ Ensure the following:
 • request isn't erroring-out server-side (check the Networking tab)
 • mock API endpoint is implemented in handlers.ts`
 
-export function ensure<TData, TError>(
+/**
+ * Ensure a query result came from the cache by blowing up if `data` comes
+ * back undefined.
+ */
+export function ensurePrefetched<TData, TError>(
   result: UseQueryResult<TData, TError>,
   /**
    * Optional because if we call this manually from a component like

app/api/util.ts

@@ -23,6 +23,8 @@ import type {
   VpcFirewallRuleUpdate,
 } from './__generated__/Api'
 
+export type ResultsPage<TItem> = { items: TItem[]; nextPage?: string }
+
 // API limits encoded in https://github.com/oxidecomputer/omicron/blob/main/nexus/src/app/mod.rs
 
 export const MAX_NICS_PER_INSTANCE = 8

app/layouts/helpers.tsx

@@ -20,7 +20,12 @@ export function ContentPane() {
   const ref = useRef<HTMLDivElement>(null)
   useScrollRestoration(ref)
   return (
-    <div ref={ref} className="flex flex-col overflow-auto" data-testid="scroll-container">
+    <div
+      ref={ref}
+      className="flex flex-col overflow-auto"
+      id="scroll-container"
+      data-testid="scroll-container"
+    >
       <div className="flex grow flex-col pb-8">
         <SkipLinkTarget />
         <main className="[&>*]:gutter">

app/pages/project/snapshots/SnapshotsPage.tsx

@@ -11,6 +11,8 @@ import { Outlet, useNavigate, type LoaderFunctionArgs } from 'react-router-dom'
 
 import {
   apiQueryClient,
+  getListQFn,
+  queryClient,
   useApiMutation,
   useApiQueryClient,
   useApiQueryErrorsAllowed,
@@ -25,7 +27,7 @@ import { confirmDelete } from '~/stores/confirm-delete'
 import { SkeletonCell } from '~/table/cells/EmptyCell'
 import { useColsWithActions, type MenuAction } from '~/table/columns/action-col'
 import { Columns } from '~/table/columns/common'
-import { PAGE_SIZE, useQueryTable } from '~/table/QueryTable'
+import { useQueryTable } from '~/table/QueryTable2'
 import { Badge } from '~/ui/lib/Badge'
 import { CreateLink } from '~/ui/lib/CreateButton'
 import { EmptyMessage } from '~/ui/lib/EmptyMessage'
@@ -52,12 +54,12 @@ const EmptyState = () => (
   />
 )
 
+const snapshotList = (project: string) => getListQFn('snapshotList', { query: { project } })
+
 SnapshotsPage.loader = async ({ params }: LoaderFunctionArgs) => {
   const { project } = getProjectSelector(params)
   await Promise.all([
-    apiQueryClient.prefetchQuery('snapshotList', {
-      query: { project, limit: PAGE_SIZE },
-    }),
+    queryClient.prefetchQuery(snapshotList(project).optionsFn()),
 
     // Fetch disks and preload into RQ cache so fetches by ID in DiskNameFromId
     // can be mostly instant yet gracefully fall back to fetching individually
@@ -100,7 +102,6 @@ const staticCols = [
 export function SnapshotsPage() {
   const queryClient = useApiQueryClient()
   const { project } = useProjectSelector()
-  const { Table } = useQueryTable('snapshotList', { query: { project } })
   const navigate = useNavigate()
 
   const { mutateAsync: deleteSnapshot } = useApiMutation('snapshotDelete', {
@@ -132,6 +133,11 @@ export function SnapshotsPage() {
     [deleteSnapshot, navigate, project]
   )
   const columns = useColsWithActions(staticCols, makeActions)
+  const { table } = useQueryTable({
+    query: snapshotList(project),
+    columns,
+    emptyState: <EmptyState />,
+  })
   return (
     <>
       <PageHeader>
@@ -146,7 +152,7 @@ export function SnapshotsPage() {
       <TableActions>
         <CreateLink to={pb.snapshotsNew({ project })}>New snapshot</CreateLink>
       </TableActions>
-      <Table columns={columns} emptyState={<EmptyState />} />
+      {table}
       <Outlet />
     </>
   )

app/pages/system/inventory/DisksTab.tsx

@@ -8,14 +8,15 @@
 import { createColumnHelper } from '@tanstack/react-table'
 
 import {
-  apiQueryClient,
+  getListQFn,
+  queryClient,
   type PhysicalDisk,
   type PhysicalDiskPolicy,
   type PhysicalDiskState,
 } from '@oxide/api'
 import { Servers24Icon } from '@oxide/design-system/icons/react'
 
-import { PAGE_SIZE, useQueryTable } from '~/table/QueryTable'
+import { useQueryTable } from '~/table/QueryTable2'
 import { Badge, type BadgeColor } from '~/ui/lib/Badge'
 import { EmptyMessage } from '~/ui/lib/EmptyMessage'
 
@@ -37,8 +38,10 @@ const EmptyState = () => (
   />
 )
 
+const diskList = getListQFn('physicalDiskList', {})
+
 export async function loader() {
-  await apiQueryClient.prefetchQuery('physicalDiskList', { query: { limit: PAGE_SIZE } })
+  await queryClient.prefetchQuery(diskList.optionsFn())
   return null
 }
 
@@ -68,6 +71,7 @@ const staticCols = [
 
 Component.displayName = 'DisksTab'
 export function Component() {
-  const { Table } = useQueryTable('physicalDiskList', {})
-  return <Table emptyState={<EmptyState />} columns={staticCols} />
+  const emptyState = <EmptyState />
+  const { table } = useQueryTable({ query: diskList, columns: staticCols, emptyState })
+  return table
 }

app/pages/system/inventory/SledsTab.tsx

@@ -7,11 +7,17 @@
  */
 import { createColumnHelper } from '@tanstack/react-table'
 
-import { apiQueryClient, type Sled, type SledPolicy, type SledState } from '@oxide/api'
+import {
+  getListQFn,
+  queryClient,
+  type Sled,
+  type SledPolicy,
+  type SledState,
+} from '@oxide/api'
 import { Servers24Icon } from '@oxide/design-system/icons/react'
 
 import { makeLinkCell } from '~/table/cells/LinkCell'
-import { PAGE_SIZE, useQueryTable } from '~/table/QueryTable'
+import { useQueryTable } from '~/table/QueryTable2'
 import { Badge, type BadgeColor } from '~/ui/lib/Badge'
 import { EmptyMessage } from '~/ui/lib/EmptyMessage'
 import { pb } from '~/util/path-builder'
@@ -36,10 +42,10 @@ const EmptyState = () => {
   )
 }
 
+const sledList = getListQFn('sledList', {})
+
 export async function loader() {
-  await apiQueryClient.prefetchQuery('sledList', {
-    query: { limit: PAGE_SIZE },
-  })
+  await queryClient.prefetchQuery(sledList.optionsFn())
   return null
 }
 
@@ -69,6 +75,7 @@ const staticCols = [
 
 Component.displayName = 'SledsTab'
 export function Component() {
-  const { Table } = useQueryTable('sledList', {}, { placeholderData: (x) => x })
-  return <Table emptyState={<EmptyState />} columns={staticCols} />
+  const emptyState = <EmptyState />
+  const { table } = useQueryTable({ query: sledList, columns: staticCols, emptyState })
+  return table
 }

app/table/QueryTable.tsx

@@ -12,6 +12,7 @@ import { getCoreRowModel, useReactTable, type ColumnDef } from '@tanstack/react-
 import React, { useCallback, useMemo, type ComponentType } from 'react'
 
 import {
+  PAGE_SIZE,
   useApiQuery,
   type ApiError,
   type ApiListMethods,
@@ -27,6 +28,8 @@ import { TableEmptyBox } from '~/ui/lib/Table'
 
 import { Table } from './Table'
 
+export { PAGE_SIZE }
+
 interface UseQueryTableResult<Item extends Record<string, unknown>> {
   Table: ComponentType<QueryTableProps<Item>>
 }
@@ -59,8 +62,6 @@ type QueryTableProps<Item> = {
   columns: ColumnDef<Item, any>[]
 }
 
-export const PAGE_SIZE = 25
-
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 const makeQueryTable = <Item extends Record<string, unknown>>(
   query: any,