Merge "Treat integer args as integers" into main

ui/src/frontend/widgets/sql/table/column.ts

@@ -43,6 +43,15 @@ export type SqlColumn =
       source: SourceTable;
     };
 
+// List of columns of args, corresponding to arg values, which cause a short-form of the ID to be generated.
+// (e.g. arg_set_id[foo].int instead of args[arg_set_id,key=foo].int_value).
+const ARG_COLUMN_TO_SUFFIX: {[key: string]: string} = {
+  display_value: '',
+  int_value: '.int',
+  string_value: '.str',
+  real_value: '.real',
+};
+
 // A unique identifier for the SQL column.
 export function sqlColumnId(column: SqlColumn): string {
   if (typeof column === 'string') {
@@ -54,13 +63,13 @@ export function sqlColumnId(column: SqlColumn): string {
   }
   // Special case: args lookup. For it, we can use a simpler representation (i.e. `arg_set_id[key]`).
   if (
-    column.column === 'display_value' &&
+    column.column in ARG_COLUMN_TO_SUFFIX &&
     column.source.table === 'args' &&
     arrayEquals(Object.keys(column.source.joinOn).sort(), ['arg_set_id', 'key'])
   ) {
     const key = column.source.joinOn['key'];
     const argSetId = column.source.joinOn['arg_set_id'];
-    return `${sqlColumnId(argSetId)}[${sqlColumnId(key)}]`;
+    return `${sqlColumnId(argSetId)}[${sqlColumnId(key)}]${ARG_COLUMN_TO_SUFFIX[column.column]}`;
   }
   // Otherwise, we need to list all the join constraints.
   const lookup = Object.entries(column.source.joinOn)
@@ -137,6 +146,9 @@ export abstract class TableColumn {
   // Sometimes to display an interactive cell more than a single value is needed (e.g. "time range" corresponds to (ts, dur) pair. While we want to show the duration, we would want to highlight the interval on hover, for which both timestamp and duration are needed.
   dependentColumns?(): {[key: string]: SqlColumn};
 
+  // The set of underlying sql columns that should be sorted when this column is sorted.
+  sortColumns?(): SqlColumn[];
+
   // Render a table cell. `value` corresponds to the fetched SQL value for the primary column, `dependentColumns` are the fetched values for the dependent columns.
   abstract renderCell(
     value: SqlValue,

ui/src/frontend/widgets/sql/table/render_cell_utils.ts

@@ -157,7 +157,7 @@ export function getStandardContextMenuItems(
   return result;
 }
 
-function displayValue(value: SqlValue): m.Child {
+export function displayValue(value: SqlValue): m.Child {
   if (value === null) {
     return m('i', 'NULL');
   }

ui/src/frontend/widgets/sql/table/state.ts

@@ -21,6 +21,7 @@ import {
   SqlColumn,
   sqlColumnId,
   TableColumn,
+  tableColumnId,
 } from './column';
 import {buildSqlQuery} from './query_builder';
 import {raf} from '../../../../core/raf_scheduler';
@@ -76,7 +77,10 @@ export class SqlTableState {
   // Columns currently displayed to the user. All potential columns can be found `this.table.columns`.
   private columns: TableColumn[];
   private filters: Filter[];
-  private orderBy: ColumnOrderClause[];
+  private orderBy: {
+    column: TableColumn;
+    direction: SortDirection;
+  }[];
   private offset = 0;
   private request: Request;
   private data?: Data;
@@ -149,7 +153,7 @@ export class SqlTableState {
       table: this.config.name,
       columns,
       filters: this.filters,
-      orderBy: this.orderBy,
+      orderBy: this.getOrderedBy(),
     });
   }
 
@@ -365,10 +369,10 @@ export class SqlTableState {
     return this.filters;
   }
 
-  sortBy(clause: ColumnOrderClause) {
+  sortBy(clause: {column: TableColumn; direction: SortDirection}) {
     // Remove previous sort by the same column.
     this.orderBy = this.orderBy.filter(
-      (c) => !isSqlColumnEqual(c.column, clause.column),
+      (c) => tableColumnId(c.column) != tableColumnId(clause.column),
     );
     // Add the new sort clause to the front, so we effectively stable-sort the
     // data currently displayed to the user.
@@ -383,14 +387,23 @@ export class SqlTableState {
 
   isSortedBy(column: TableColumn): SortDirection | undefined {
     if (this.orderBy.length === 0) return undefined;
-    if (!isSqlColumnEqual(this.orderBy[0].column, column.primaryColumn())) {
+    if (tableColumnId(this.orderBy[0].column) !== tableColumnId(column)) {
       return undefined;
     }
     return this.orderBy[0].direction;
   }
 
   getOrderedBy(): ColumnOrderClause[] {
-    return this.orderBy;
+    const result: ColumnOrderClause[] = [];
+    for (const orderBy of this.orderBy) {
+      const sortColumns = orderBy.column.sortColumns?.() ?? [
+        orderBy.column.primaryColumn(),
+      ];
+      for (const column of sortColumns) {
+        result.push({column, direction: orderBy.direction});
+      }
+    }
+    return result;
   }
 
   addColumn(column: TableColumn, index: number) {
@@ -404,7 +417,7 @@ export class SqlTableState {
     // 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) => !isSqlColumnEqual(c.column, column.primaryColumn()),
+      (c) => tableColumnId(c.column) !== tableColumnId(column),
     );
     // TODO(altimin): we can avoid the fetch here if the orderBy hasn't changed.
     this.reload({offset: 'keep'});

ui/src/frontend/widgets/sql/table/state_unittest.ts

@@ -76,23 +76,23 @@ test('sqlTableState: sortedColumns', () => {
 
   // Sort by name column and verify that it is sorted by.
   state.sortBy({
-    column: nameColumn.primaryColumn(),
+    column: nameColumn,
     direction: 'ASC',
   });
   expect(state.isSortedBy(idColumn)).toBe(undefined);
   expect(state.isSortedBy(nameColumn)).toBe('ASC');
 
   // Sort by the same column in the opposite direction.
   state.sortBy({
-    column: nameColumn.primaryColumn(),
+    column: nameColumn,
     direction: 'DESC',
   });
   expect(state.isSortedBy(idColumn)).toBe(undefined);
   expect(state.isSortedBy(nameColumn)).toBe('DESC');
 
   // Sort by the id column.
   state.sortBy({
-    column: idColumn.primaryColumn(),
+    column: idColumn,
     direction: 'ASC',
   });
   expect(state.isSortedBy(idColumn)).toBe('ASC');

ui/src/frontend/widgets/sql/table/table.ts

@@ -62,8 +62,7 @@ function renderCell(
   const sqlValue = row[columns[sqlColumnId(column.primaryColumn())]];
 
   const additionalValues: {[key: string]: SqlValue} = {};
-  const dependentColumns =
-    column.dependentColumns !== undefined ? column.dependentColumns() : {};
+  const dependentColumns = column.dependentColumns?.() ?? {};
   for (const [key, col] of Object.entries(dependentColumns)) {
     additionalValues[key] = row[columns[sqlColumnId(col)]];
   }
@@ -276,7 +275,7 @@ export class SqlTable implements m.ClassComponent<SqlTableConfig> {
           icon: Icons.SortedDesc,
           onclick: () => {
             this.state.sortBy({
-              column: column.primaryColumn(),
+              column: column,
               direction: 'DESC',
             });
           },
@@ -287,7 +286,7 @@ export class SqlTable implements m.ClassComponent<SqlTableConfig> {
           icon: Icons.SortedAsc,
           onclick: () => {
             this.state.sortBy({
-              column: column.primaryColumn(),
+              column: column,
               direction: 'ASC',
             });
           },

ui/src/frontend/widgets/sql/table/well_known_columns.ts

@@ -46,6 +46,7 @@ import {
   TableManager,
 } from './column';
 import {
+  displayValue,
   getStandardContextMenuItems,
   getStandardFilters,
   renderStandardCell,
@@ -1018,6 +1019,109 @@ export class ProcessColumnSet extends TableColumnSet {
   }
 }
 
+class ArgColumn extends TableColumn {
+  private displayValue: SqlColumn;
+  private stringValue: SqlColumn;
+  private intValue: SqlColumn;
+  private realValue: SqlColumn;
+
+  constructor(
+    private argSetId: SqlColumn,
+    private key: string,
+  ) {
+    super();
+
+    const argTable: SourceTable = {
+      table: 'args',
+      joinOn: {
+        arg_set_id: argSetId,
+        key: sqliteString(key),
+      },
+    };
+
+    this.displayValue = {
+      column: 'display_value',
+      source: argTable,
+    };
+    this.stringValue = {
+      column: 'string_value',
+      source: argTable,
+    };
+    this.intValue = {
+      column: 'int_value',
+      source: argTable,
+    };
+    this.realValue = {
+      column: 'real_value',
+      source: argTable,
+    };
+  }
+
+  override primaryColumn(): SqlColumn {
+    return this.displayValue;
+  }
+
+  override sortColumns(): SqlColumn[] {
+    return [this.stringValue, this.intValue, this.realValue];
+  }
+
+  override dependentColumns() {
+    return {
+      stringValue: this.stringValue,
+      intValue: this.intValue,
+      realValue: this.realValue,
+    };
+  }
+
+  getTitle() {
+    return `${sqlColumnId(this.argSetId)}[${this.key}]`;
+  }
+
+  renderCell(
+    value: SqlValue,
+    tableManager: TableManager,
+    dependentColumns: {[key: string]: SqlValue},
+  ): m.Children {
+    const strValue = dependentColumns['stringValue'];
+    const intValue = dependentColumns['intValue'];
+    const realValue = dependentColumns['realValue'];
+
+    let contextMenuItems: m.Child[] = [];
+    if (strValue !== null) {
+      contextMenuItems = getStandardContextMenuItems(
+        strValue,
+        this.stringValue,
+        tableManager,
+      );
+    } else if (intValue !== null) {
+      contextMenuItems = getStandardContextMenuItems(
+        intValue,
+        this.intValue,
+        tableManager,
+      );
+    } else if (realValue !== null) {
+      contextMenuItems = getStandardContextMenuItems(
+        realValue,
+        this.realValue,
+        tableManager,
+      );
+    } else {
+      contextMenuItems = getStandardContextMenuItems(
+        value,
+        this.displayValue,
+        tableManager,
+      );
+    }
+    return m(
+      PopupMenu2,
+      {
+        trigger: m(Anchor, displayValue(value)),
+      },
+      ...contextMenuItems,
+    );
+  }
+}
+
 export class ArgSetColumnSet extends TableColumnSet {
   constructor(
     private column: SqlColumn,
@@ -1067,8 +1171,5 @@ export function argSqlColumn(argSetId: SqlColumn, key: string): SqlColumn {
 }
 
 export function argTableColumn(argSetId: SqlColumn, key: string) {
-  return new StandardColumn(argSqlColumn(argSetId, key), {
-    title: `${sqlColumnId(argSetId)}[${key}]`,
-    alias: `arg_${key.replace(/[^a-zA-Z0-9_]/g, '__')}`,
-  });
+  return new ArgColumn(argSetId, key);
 }