ui/src/frontend/widgets/sql/table/state.ts - third_party/perfetto - Git at Google

 // Copyright (C) 2023 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 {arrayEquals} from '../../../../base/array_utils';
 import {SortDirection} from '../../../../base/comparison_utils';
 import {isString} from '../../../../base/object_utils';
 import {sqliteString} from '../../../../base/string_utils';
 import {raf} from '../../../../core/raf_scheduler';
 import {Engine} from '../../../../trace_processor/engine';
 import {NUM, Row} from '../../../../trace_processor/query_result';
 import {
   constraintsToQueryPrefix,
   constraintsToQuerySuffix,
   SQLConstraints,
 } from '../../../../trace_processor/sql_utils';

 import {
   Column,
   columnFromSqlTableColumn,
   formatSqlProjection,
   SqlProjection,
   sqlProjectionsForColumn,
 } from './column';
 import {SqlTableDescription, startsHidden} from './table_description';

 interface ColumnOrderClause {
   // We only allow the table to be sorted by the columns which are displayed to
   // the user to avoid confusion, so we use a reference to the underlying Column
   // here and compare it by reference down the line.
   column: Column;
   direction: SortDirection;
 }

 const ROW_LIMIT = 100;

 // Result of the execution of the query.
 interface Data {
   // Rows to show, including pagination.
   rows: Row[];
   error?: string;
 }

 // In the common case, filter is an expression which evaluates to a boolean.
 // However, when filtering args, it's substantially (10x) cheaper to do a
 // join with the args table, as it means that trace processor can cache the
 // query on the key instead of invoking a function for each row of the entire
 // `slice` table.
 export type Filter =
   | string
   | {
       type: 'arg_filter';
       argSetIdColumn: string;
       argName: string;
       op: string;
     };

 interface RowCount {
   // Total number of rows in view, excluding the pagination.
   // Undefined if the query returned an error.
   count: number;
   // Filters which were used to compute this row count.
   // We need to recompute the totalRowCount only when filters change and not
   // when the set of columns / order by changes.
   filters: Filter[];
 }

 export class SqlTableState {
   private readonly engine_: Engine;
   private readonly table_: SqlTableDescription;
   private readonly additionalImports: string[];

   get engine() {
     return this.engine_;
   }
   get table() {
     return this.table_;
   }

   private filters: Filter[];
   private columns: Column[];
   private orderBy: ColumnOrderClause[];
   private offset = 0;
   private data?: Data;
   private rowCount?: RowCount;

   constructor(
     engine: Engine,
     table: SqlTableDescription,
     filters?: Filter[],
     imports?: string[],
   ) {
     this.engine_ = engine;
     this.table_ = table;
     this.additionalImports = imports || [];

     this.filters = filters || [];
     this.columns = [];
     for (const column of this.table.columns) {
       if (startsHidden(column)) continue;
       this.columns.push(columnFromSqlTableColumn(column));
     }
     this.orderBy = [];

     this.reload();
   }

   // Compute the actual columns to fetch. Some columns can appear multiple times
   // (e.g. we might need "ts" to be able to show it, as well as a dependency for
   // "slice_id" to be able to jump to it, so this function will deduplicate
   // projections by alias.
   private getSQLProjections(): SqlProjection[] {
     const projections = [];
     const aliases = new Set<string>();
     for (const column of this.columns) {
       for (const p of sqlProjectionsForColumn(column)) {
         if (aliases.has(p.alias)) continue;
         aliases.add(p.alias);
         projections.push(p);
       }
     }
     return projections;
   }

   getQueryConstraints(): SQLConstraints {
     const result: SQLConstraints = {
       commonTableExpressions: {},
       joins: [],
       filters: [],
     };
     let cteId = 0;
     for (const filter of this.filters) {
       if (isString(filter)) {
         result.filters!.push(filter);
       } else {
         const cteName = `arg_sets_${cteId++}`;
         result.commonTableExpressions![cteName] = `
           SELECT DISTINCT arg_set_id
           FROM args
           WHERE key = ${sqliteString(filter.argName)}
             AND display_value ${filter.op}
         `;
         result.joins!.push(
           `JOIN ${cteName} ON ${cteName}.arg_set_id = ${this.table.name}.${filter.argSetIdColumn}`,
         );
       }
     }
     return result;
   }

   private getSQLImports() {
     const tableImports = this.table.imports || [];
     return [...tableImports, ...this.additionalImports]
       .map((i) => `INCLUDE PERFETTO MODULE ${i};`)
       .join('\n');
   }

   private getCountRowsSQLQuery(): string {
     const constraints = this.getQueryConstraints();
     return `
       ${this.getSQLImports()}

       ${constraintsToQueryPrefix(constraints)}
       SELECT
         COUNT() AS count
       FROM ${this.table.name}
       ${constraintsToQuerySuffix(constraints)}
     `;
   }

   buildSqlSelectStatement(): {
     selectStatement: string;
     columns: string[];
   } {
     const projections = this.getSQLProjections();
     const orderBy = this.orderBy.map((c) => ({
       fieldName: c.column.alias,
       direction: c.direction,
     }));
     const constraints = this.getQueryConstraints();
     constraints.orderBy = orderBy;
     const statement = `
       ${constraintsToQueryPrefix(constraints)}
       SELECT
         ${projections.map(formatSqlProjection).join(',\n')}
       FROM ${this.table.name}
       ${constraintsToQuerySuffix(constraints)}
     `;
     return {
       selectStatement: statement,
       columns: projections.map((p) => p.alias),
     };
   }

   getNonPaginatedSQLQuery(): string {
     return `
       ${this.getSQLImports()}

       ${this.buildSqlSelectStatement().selectStatement}
     `;
   }

   getPaginatedSQLQuery(): string {
     // We fetch one more row to determine if we can go forward.
     return `
       ${this.getNonPaginatedSQLQuery()}
       LIMIT ${ROW_LIMIT + 1}
       OFFSET ${this.offset}
     `;
   }

   canGoForward(): boolean {
     if (this.data === undefined) return false;
     return this.data.rows.length > ROW_LIMIT;
   }

   canGoBack(): boolean {
     if (this.data === undefined) return false;
     return this.offset > 0;
   }

   goForward() {
     if (!this.canGoForward()) return;
     this.offset += ROW_LIMIT;
     this.reload({offset: 'keep'});
   }

   goBack() {
     if (!this.canGoBack()) return;
     this.offset -= ROW_LIMIT;
     this.reload({offset: 'keep'});
   }

   getDisplayedRange(): {from: number; to: number} | undefined {
     if (this.data === undefined) return undefined;
     return {
       from: this.offset + 1,
       to: this.offset + Math.min(this.data.rows.length, ROW_LIMIT),
     };
   }

   private async loadRowCount(): Promise<RowCount | undefined> {
     const filters = Array.from(this.filters);
     const res = await this.engine.query(this.getCountRowsSQLQuery());
     if (res.error() !== undefined) return undefined;
     return {
       count: res.firstRow({count: NUM}).count,
       filters: filters,
     };
   }

   private async loadData(): Promise<Data> {
     const queryRes = await this.engine.query(this.getPaginatedSQLQuery());
     const rows: Row[] = [];
     for (const it = queryRes.iter({}); it.valid(); it.next()) {
       const row: Row = {};
       for (const column of queryRes.columns()) {
         row[column] = it.get(column);
       }
       rows.push(row);
     }

     return {
       rows,
       error: queryRes.error(),
     };
   }

   private async reload(params?: {offset: 'reset' | 'keep'}) {
     if ((params?.offset ?? 'reset') === 'reset') {
       this.offset = 0;
     }

     const newFilters = this.rowCount?.filters;
     const filtersMatch = newFilters && arrayEquals(newFilters, this.filters);
     this.data = undefined;
     if (!filtersMatch) {
       this.rowCount = undefined;
     }

     // Delay the visual update by 50ms to avoid flickering (if the query returns
     // before the data is loaded.
     setTimeout(() => raf.scheduleFullRedraw(), 50);

     if (!filtersMatch) {
       this.rowCount = await this.loadRowCount();
     }
     this.data = await this.loadData();

     raf.scheduleFullRedraw();
   }

   getTotalRowCount(): number | undefined {
     return this.rowCount?.count;
   }

   getDisplayedRows(): Row[] {
     return this.data?.rows || [];
   }

   getQueryError(): string | undefined {
     return this.data?.error;
   }

   isLoading() {
     return this.data === undefined;
   }

   // Filters are compared by reference, so the caller is required to pass an
   // object which was previously returned by getFilters.
   removeFilter(filter: Filter) {
     this.filters = this.filters.filter((f) => f !== filter);
     this.reload();
   }

   addFilter(filter: string) {
     this.filters.push(filter);
     this.reload();
   }

   getFilters(): Filter[] {
     return this.filters;
   }

   sortBy(clause: ColumnOrderClause) {
     // Remove previous sort by the same column.
     this.orderBy = this.orderBy.filter((c) => c.column !== clause.column);
     // Add the new sort clause to the front, so we effectively stable-sort the
     // data currently displayed to the user.
     this.orderBy.unshift(clause);
     this.reload();
   }

   unsort() {
     this.orderBy = [];
     this.reload();
   }

   isSortedBy(column: Column): SortDirection | undefined {
     if (this.orderBy.length === 0) return undefined;
     if (this.orderBy[0].column !== column) return undefined;
     return this.orderBy[0].direction;
   }

   addColumn(column: Column, index: number) {
     this.columns.splice(index + 1, 0, column);
     this.reload({offset: 'keep'});
   }

   hideColumnAtIndex(index: number) {
     const column = this.columns[index];
     this.columns.splice(index, 1);
     // We can only filter by the visibile columns to avoid confusing the user,
     // so we remove order by clauses that refer to the hidden column.
     this.orderBy = this.orderBy.filter((c) => c.column !== column);
     // TODO(altimin): we can avoid the fetch here if the orderBy hasn't changed.
     this.reload({offset: 'keep'});
   }

   getSelectedColumns(): Column[] {
     return this.columns;
   }
 }
	// Copyright (C) 2023 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 {arrayEquals} from '../../../../base/array_utils';
	import {SortDirection} from '../../../../base/comparison_utils';
	import {isString} from '../../../../base/object_utils';
	import {sqliteString} from '../../../../base/string_utils';
	import {raf} from '../../../../core/raf_scheduler';
	import {Engine} from '../../../../trace_processor/engine';
	import {NUM, Row} from '../../../../trace_processor/query_result';
	import {
	constraintsToQueryPrefix,
	constraintsToQuerySuffix,
	SQLConstraints,
	} from '../../../../trace_processor/sql_utils';

	import {
	Column,
	columnFromSqlTableColumn,
	formatSqlProjection,
	SqlProjection,
	sqlProjectionsForColumn,
	} from './column';
	import {SqlTableDescription, startsHidden} from './table_description';

	interface ColumnOrderClause {
	// We only allow the table to be sorted by the columns which are displayed to
	// the user to avoid confusion, so we use a reference to the underlying Column
	// here and compare it by reference down the line.
	column: Column;
	direction: SortDirection;
	}

	const ROW_LIMIT = 100;

	// Result of the execution of the query.
	interface Data {
	// Rows to show, including pagination.
	rows: Row[];
	error?: string;
	}

	// In the common case, filter is an expression which evaluates to a boolean.
	// However, when filtering args, it's substantially (10x) cheaper to do a
	// join with the args table, as it means that trace processor can cache the
	// query on the key instead of invoking a function for each row of the entire
	// `slice` table.
	export type Filter =
	\| string
	\| {
	type: 'arg_filter';
	argSetIdColumn: string;
	argName: string;
	op: string;
	};

	interface RowCount {
	// Total number of rows in view, excluding the pagination.
	// Undefined if the query returned an error.
	count: number;
	// Filters which were used to compute this row count.
	// We need to recompute the totalRowCount only when filters change and not
	// when the set of columns / order by changes.
	filters: Filter[];
	}

	export class SqlTableState {
	private readonly engine_: Engine;
	private readonly table_: SqlTableDescription;
	private readonly additionalImports: string[];

	get engine() {
	return this.engine_;
	}
	get table() {
	return this.table_;
	}

	private filters: Filter[];
	private columns: Column[];
	private orderBy: ColumnOrderClause[];
	private offset = 0;
	private data?: Data;
	private rowCount?: RowCount;

	constructor(
	engine: Engine,
	table: SqlTableDescription,
	filters?: Filter[],
	imports?: string[],
	) {
	this.engine_ = engine;
	this.table_ = table;
	this.additionalImports = imports \|\| [];

	this.filters = filters \|\| [];
	this.columns = [];
	for (const column of this.table.columns) {
	if (startsHidden(column)) continue;
	this.columns.push(columnFromSqlTableColumn(column));
	}
	this.orderBy = [];

	this.reload();
	}

	// Compute the actual columns to fetch. Some columns can appear multiple times
	// (e.g. we might need "ts" to be able to show it, as well as a dependency for
	// "slice_id" to be able to jump to it, so this function will deduplicate
	// projections by alias.
	private getSQLProjections(): SqlProjection[] {
	const projections = [];
	const aliases = new Set<string>();
	for (const column of this.columns) {
	for (const p of sqlProjectionsForColumn(column)) {
	if (aliases.has(p.alias)) continue;
	aliases.add(p.alias);
	projections.push(p);
	}
	}
	return projections;
	}

	getQueryConstraints(): SQLConstraints {
	const result: SQLConstraints = {
	commonTableExpressions: {},
	joins: [],
	filters: [],
	};
	let cteId = 0;
	for (const filter of this.filters) {
	if (isString(filter)) {
	result.filters!.push(filter);
	} else {
	const cteName = `arg_sets_${cteId++}`;
	result.commonTableExpressions![cteName] = `
	SELECT DISTINCT arg_set_id
	FROM args
	WHERE key = ${sqliteString(filter.argName)}
	AND display_value ${filter.op}
	`;
	result.joins!.push(
	`JOIN ${cteName} ON ${cteName}.arg_set_id = ${this.table.name}.${filter.argSetIdColumn}`,
	);
	}
	}
	return result;
	}

	private getSQLImports() {
	const tableImports = this.table.imports \|\| [];
	return [...tableImports, ...this.additionalImports]
	.map((i) => `INCLUDE PERFETTO MODULE ${i};`)
	.join('\n');
	}

	private getCountRowsSQLQuery(): string {
	const constraints = this.getQueryConstraints();
	return `
	${this.getSQLImports()}

	${constraintsToQueryPrefix(constraints)}
	SELECT
	COUNT() AS count
	FROM ${this.table.name}
	${constraintsToQuerySuffix(constraints)}
	`;
	}

	buildSqlSelectStatement(): {
	selectStatement: string;
	columns: string[];
	} {
	const projections = this.getSQLProjections();
	const orderBy = this.orderBy.map((c) => ({
	fieldName: c.column.alias,
	direction: c.direction,
	}));
	const constraints = this.getQueryConstraints();
	constraints.orderBy = orderBy;
	const statement = `
	${constraintsToQueryPrefix(constraints)}
	SELECT
	${projections.map(formatSqlProjection).join(',\n')}
	FROM ${this.table.name}
	${constraintsToQuerySuffix(constraints)}
	`;
	return {
	selectStatement: statement,
	columns: projections.map((p) => p.alias),
	};
	}

	getNonPaginatedSQLQuery(): string {
	return `
	${this.getSQLImports()}

	${this.buildSqlSelectStatement().selectStatement}
	`;
	}

	getPaginatedSQLQuery(): string {
	// We fetch one more row to determine if we can go forward.
	return `
	${this.getNonPaginatedSQLQuery()}
	LIMIT ${ROW_LIMIT + 1}
	OFFSET ${this.offset}
	`;
	}

	canGoForward(): boolean {
	if (this.data === undefined) return false;
	return this.data.rows.length > ROW_LIMIT;
	}

	canGoBack(): boolean {
	if (this.data === undefined) return false;
	return this.offset > 0;
	}

	goForward() {
	if (!this.canGoForward()) return;
	this.offset += ROW_LIMIT;
	this.reload({offset: 'keep'});
	}

	goBack() {
	if (!this.canGoBack()) return;
	this.offset -= ROW_LIMIT;
	this.reload({offset: 'keep'});
	}

	getDisplayedRange(): {from: number; to: number} \| undefined {
	if (this.data === undefined) return undefined;
	return {
	from: this.offset + 1,
	to: this.offset + Math.min(this.data.rows.length, ROW_LIMIT),
	};
	}

	private async loadRowCount(): Promise<RowCount \| undefined> {
	const filters = Array.from(this.filters);
	const res = await this.engine.query(this.getCountRowsSQLQuery());
	if (res.error() !== undefined) return undefined;
	return {
	count: res.firstRow({count: NUM}).count,
	filters: filters,
	};
	}

	private async loadData(): Promise<Data> {
	const queryRes = await this.engine.query(this.getPaginatedSQLQuery());
	const rows: Row[] = [];
	for (const it = queryRes.iter({}); it.valid(); it.next()) {
	const row: Row = {};
	for (const column of queryRes.columns()) {
	row[column] = it.get(column);
	}
	rows.push(row);
	}

	return {
	rows,
	error: queryRes.error(),
	};
	}

	private async reload(params?: {offset: 'reset' \| 'keep'}) {
	if ((params?.offset ?? 'reset') === 'reset') {
	this.offset = 0;
	}

	const newFilters = this.rowCount?.filters;
	const filtersMatch = newFilters && arrayEquals(newFilters, this.filters);
	this.data = undefined;
	if (!filtersMatch) {
	this.rowCount = undefined;
	}

	// Delay the visual update by 50ms to avoid flickering (if the query returns
	// before the data is loaded.
	setTimeout(() => raf.scheduleFullRedraw(), 50);

	if (!filtersMatch) {
	this.rowCount = await this.loadRowCount();
	}
	this.data = await this.loadData();

	raf.scheduleFullRedraw();
	}

	getTotalRowCount(): number \| undefined {
	return this.rowCount?.count;
	}

	getDisplayedRows(): Row[] {
	return this.data?.rows \|\| [];
	}

	getQueryError(): string \| undefined {
	return this.data?.error;
	}

	isLoading() {
	return this.data === undefined;
	}

	// Filters are compared by reference, so the caller is required to pass an
	// object which was previously returned by getFilters.
	removeFilter(filter: Filter) {
	this.filters = this.filters.filter((f) => f !== filter);
	this.reload();
	}

	addFilter(filter: string) {
	this.filters.push(filter);
	this.reload();
	}

	getFilters(): Filter[] {
	return this.filters;
	}

	sortBy(clause: ColumnOrderClause) {
	// Remove previous sort by the same column.
	this.orderBy = this.orderBy.filter((c) => c.column !== clause.column);
	// Add the new sort clause to the front, so we effectively stable-sort the
	// data currently displayed to the user.
	this.orderBy.unshift(clause);
	this.reload();
	}

	unsort() {
	this.orderBy = [];
	this.reload();
	}

	isSortedBy(column: Column): SortDirection \| undefined {
	if (this.orderBy.length === 0) return undefined;
	if (this.orderBy[0].column !== column) return undefined;
	return this.orderBy[0].direction;
	}

	addColumn(column: Column, index: number) {
	this.columns.splice(index + 1, 0, column);
	this.reload({offset: 'keep'});
	}

	hideColumnAtIndex(index: number) {
	const column = this.columns[index];
	this.columns.splice(index, 1);
	// We can only filter by the visibile columns to avoid confusing the user,
	// so we remove order by clauses that refer to the hidden column.
	this.orderBy = this.orderBy.filter((c) => c.column !== column);
	// TODO(altimin): we can avoid the fetch here if the orderBy hasn't changed.
	this.reload({offset: 'keep'});
	}

	getSelectedColumns(): Column[] {
	return this.columns;
	}
	}