ui/src/bigtrace/query/bigtrace_async_data_source.ts

// Copyright (C) 2026 The Android Open Source Project
//
// 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 {
  DataSource,
  DataSourceModel,
  DataSourceRows,
} from '../../components/widgets/datagrid/data_source';
import type {Filter} from '../../components/widgets/datagrid/model';
import type {Row, SqlValue} from '../../trace_processor/query_result';
import {AsyncMemo, type AsyncMemoResult} from '../../base/async_memo';
import {
  type BigtraceQueryClient,
  BigtraceHttpError,
  QueryCancelledError,
} from './bigtrace_query_client';
import type {BigtraceColumnSchema} from './column_types';
import {encodeFilters} from './filter_encoding';
import {FetchScheduler} from './fetch_scheduler';
import m from 'mithril';

type ModelWithColumns = DataSourceModel & {
  columns?: Array<{field: string; alias?: string}>;
};

// DataSource adapter paging `:fetch_results` into the DataGrid widget.
export class BigtraceAsyncDataSource implements DataSource {
  private loadedRows: Row[] = [];
  private error: string | null = null;
  // HTTP status of the last error, when it was an HTTP error (undefined
  // otherwise). Lets the results view tell a not-ready-yet 400 from a real
  // failure without parsing the message string.
  private errorStatus: number | undefined;
  private hasInitialFetchCompleted = false;
  // Window in `loadedRows`, for range-change detection.
  private loadedOffset = 0;
  private loadedLimit = 0;
  // AIP-132 §Ordering. Empty = default order.
  private currentOrderBy = '';
  // Aliases pre-resolved to field names. `currentFilterKey` is the JSON
  // form for cheap equality checks.
  private currentFilter: ReadonlyArray<Filter> = [];
  private currentFilterKey = '';
  // `useRows` falls back to `getTotalRows()` when undefined.
  private _filteredTotalRows: number | undefined;
  // Field-mask shipped as `:fetch_results` `columns`. Tracks the visible
  // results-grid columns so a column toggle refetches a narrower page (and
  // pulls in a metadata column when the user just enabled it).
  private currentColumns: readonly string[] = [];
  private currentColumnsKey = '';
  // availableColumnNames from the last fetch — the results-page column picker
  // reads this to know what's selectable.
  private _availableColumnNames: ReadonlyArray<string> | undefined;
  private schema?: ReadonlyArray<BigtraceColumnSchema>;
  // Owns debouncing and cancellation; see FetchScheduler.
  private readonly scheduler: FetchScheduler;
  private readonly distinctValues = new AsyncMemo<readonly SqlValue[]>();

  get filteredTotalRows(): number | undefined {
    return this._filteredTotalRows;
  }

  get availableColumnNames(): ReadonlyArray<string> | undefined {
    if (this._availableColumnNames !== undefined) {
      return this._availableColumnNames;
    }
    const initial = this.getInitialSchema?.();
    if (initial !== undefined && initial.length > 0) {
      return initial.map((s) => s.name);
    }
    return undefined;
  }

  // `signal`: owner aborts on close. `getTotalRows`: scrollbar sizing.
  constructor(
    private readonly queryUuid: string,
    private readonly queryClient: BigtraceQueryClient,
    private readonly getTotalRows: () => number,
    signal?: AbortSignal,
    private readonly getInitialSchema?: () =>
      ReadonlyArray<BigtraceColumnSchema> | undefined,
  ) {
    this.scheduler = new FetchScheduler(signal, () => m.redraw());
  }

  useRows(_model: DataSourceModel): DataSourceRows {
    const model = _model as ModelWithColumns;
    const wantedOrderBy = this.formatOrderBy(model);
    const wantedFilter = this.formatFilter(model);
    const wantedFilterKey = encodeFilters(wantedFilter);
    const wantedOffset = model.pagination?.offset ?? 0;
    const wantedLimit = model.pagination?.limit ?? 0;
    // Columns the grid is currently displaying; shipped as the `:fetch_results`
    // `columns` field-mask.
    const wantedColumns = (model.columns ?? []).map((c) => c.field);
    const wantedColumnsKey = JSON.stringify(wantedColumns);

    // Fetch on sort/filter/range/columns/initial change.
    const sortChanged = wantedOrderBy !== this.currentOrderBy;
    const filterChanged = wantedFilterKey !== this.currentFilterKey;
    const rangeChanged =
      this.hasInitialFetchCompleted &&
      (wantedOffset !== this.loadedOffset ||
        (wantedLimit > 0 && wantedLimit !== this.loadedLimit));
    const columnsChanged =
      this.hasInitialFetchCompleted &&
      wantedColumnsKey !== this.currentColumnsKey;
    const needsInitial = !this.hasInitialFetchCompleted && wantedLimit > 0;
    // No `isFetching` check: a tick arriving mid-fetch is the newest thing the
    // user has asked for, so it replaces the pending window rather than being
    // dropped. The debounce coalesces the burst and the fetch supersedes
    // whatever is still in flight.
    if (
      sortChanged ||
      filterChanged ||
      rangeChanged ||
      columnsChanged ||
      needsInitial
    ) {
      this.currentOrderBy = wantedOrderBy;
      if (filterChanged) {
        this.currentFilter = wantedFilter;
        this.currentFilterKey = wantedFilterKey;
        // Briefly oversized scrollbar > briefly collapsed while refetching.
        this._filteredTotalRows = undefined;
      }
      this.currentColumns = wantedColumns;
      this.currentColumnsKey = wantedColumnsKey;
      // First render may have limit=0; fall back so the schema comes back.
      const fetchLimit = wantedLimit > 0 ? wantedLimit : 100;
      // The first window is what the user is waiting on with nothing on
      // screen, so it goes now; later ones can afford to settle.
      this.scheduler.schedule(
        this.requestKey(wantedOffset, fetchLimit),
        needsInitial,
        (controller) =>
          this.fetchMoreRows(wantedOffset, fetchLimit, controller),
      );
    }

    const mappedRows = this.loadedRows.map((row) => {
      const mappedRow: Row = {};
      for (const key in row) {
        if (Object.prototype.hasOwnProperty.call(row, key)) {
          const col = model.columns?.find((c) => c.field === key);
          const alias =
            col !== undefined && col.alias !== undefined ? col.alias : key;
          mappedRow[alias] = row[key]!;
        }
      }
      return mappedRow;
    });

    return {
      rows: mappedRows,
      // Filtered total; falls back to unfiltered while undefined.
      totalRows: this._filteredTotalRows ?? this.getTotalRows(),
      rowOffset: this.loadedOffset,
      // Pending covers the debounce window too, so the grid doesn't flicker
      // back to "settled" between the last scroll tick and the request.
      isPending: this.scheduler.isPending,
    };
  }

  // Resolve widget alias → backend field name for the order_by wire string.
  private formatOrderBy(model: ModelWithColumns): string {
    const sort = model.sort;
    if (!sort) return '';
    const col = model.columns?.find((c) => c.alias === sort.alias);
    const field = col?.field ?? sort.alias;
    return `${field} ${sort.direction.toLowerCase()}`;
  }

  // Same alias→field remap as formatOrderBy; `fetchResults` does encoding.
  private formatFilter(model: ModelWithColumns): ReadonlyArray<Filter> {
    const filters = model.filters ?? [];
    if (filters.length === 0) return [];
    return filters.map((f) => {
      const col = model.columns?.find((c) => c.alias === f.field);
      const field = col?.field ?? f.field;
      return {...f, field};
    });
  }

  // Re-fetch the currently-loaded window, superseding anything queued or in
  // flight — an explicit refresh shouldn't wait behind a scroll.
  async refresh(): Promise<void> {
    this.distinctValues.invalidate();
    const offset = this.loadedOffset;
    const limit = this.loadedLimit > 0 ? this.loadedLimit : 100;
    await this.scheduler.runNow((controller) =>
      this.fetchMoreRows(offset, limit, controller),
    );
  }

  // What a fetch of this window would actually ask the backend for. Two
  // requests with the same key are the same request.
  private requestKey(offset: number, limit: number): string {
    return [
      offset,
      limit,
      this.currentOrderBy,
      this.currentFilterKey,
      this.currentColumnsKey,
    ].join('|');
  }

  private async fetchMoreRows(
    offset: number,
    limit: number,
    controller: AbortController,
  ) {
    this.error = null;
    this.errorStatus = undefined;
    try {
      const result = await this.queryClient.fetchResults(
        this.queryUuid,
        limit,
        offset,
        controller.signal,
        this.currentOrderBy,
        this.currentFilter,
        this.currentColumns.length > 0 ? this.currentColumns : undefined,
      );
      // A reply that raced its own abort belongs to a window the user has
      // left; adopting it would put back the page they scrolled away from.
      if (!this.scheduler.isCurrent(controller)) return;
      this.loadedRows = [...result.rows];
      this.loadedOffset = offset;
      this.loadedLimit = limit;
      this.hasInitialFetchCompleted = true;
      this._filteredTotalRows = result.totalFilteredRows;
      this._availableColumnNames = result.availableColumnNames;

      if (this.schema === undefined && result.schema.length > 0) {
        this.schema = result.schema;
      }
    } catch (e) {
      // Aborts are routine here — the tab closing, or a newer window
      // superseding this one — so they are not failures to show.
      if (e instanceof QueryCancelledError) return;
      console.error('[bigtrace] fetch_results failed:', e);
      if (e instanceof BigtraceHttpError) {
        // Surface the backend's detail (not the wrapped message) and keep the
        // status for the results view's not-ready-yet check.
        this.error = e.detail;
        this.errorStatus = e.status;
      } else {
        this.error = e instanceof Error ? e.message : String(e);
      }
    }
  }

  // Force the first window after SUCCESS without waiting for a render.
  async ensureResultsLoaded(): Promise<void> {
    if (this.hasInitialFetchCompleted || this.scheduler.isPending) return;
    await this.refresh();
  }

  getError(): string | null {
    return this.error;
  }

  // HTTP status of the last fetch error, if it was an HTTP error (undefined
  // otherwise).
  getErrorStatus(): number | undefined {
    return this.errorStatus;
  }

  getColumns(): string[] {
    return this.getSchema()?.map((s) => s.name) ?? [];
  }

  getSchema(): ReadonlyArray<BigtraceColumnSchema> | undefined {
    return this.getInitialSchema?.() ?? this.schema;
  }

  useAggregateSummaries(_model: DataSourceModel): AsyncMemoResult<Row> {
    return {data: {}, isPending: false};
  }

  useDistinctValues(
    column: string | undefined,
    search = '',
  ): AsyncMemoResult<readonly SqlValue[]> {
    if (column === undefined) return {data: [], isPending: false};
    return this.distinctValues.use({
      // Refetch as the backend inserts rows.
      key: {column, search, rows: this.getTotalRows()},
      retainOn: ['rows'],
      compute: async () => {
        try {
          return await this.queryClient.fetchColumnValues(
            this.queryUuid,
            column,
            search,
          );
        } catch (e) {
          // An empty picker beats an exception out of the grid's render.
          console.error('[bigtrace] column_values failed:', e);
          return [];
        }
      },
    });
  }

  useParameterKeys(
    _prefix: string | undefined,
  ): AsyncMemoResult<readonly string[]> {
    return {data: [], isPending: false};
  }

  async exportData(_model: DataSourceModel): Promise<readonly Row[]> {
    return this.loadedRows;
  }
}