ui/src/plugins/dev.perfetto.SqlModules/sql_modules_impl.ts

// Copyright (C) 2024 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 {z} from 'zod';
import {
  SqlModules,
  SqlColumn,
  SqlFunction,
  SqlArgument,
  SqlMacro,
  SqlModule,
  SqlPackage,
  SqlTable,
  SqlTableFunction,
} from './sql_modules';
import {SqlTableDefinition} from '../../components/widgets/sql/table/table_description';
import {TableColumn} from '../../components/widgets/sql/table/table_column';
import {Trace} from '../../public/trace';
import {
  parsePerfettoSqlTypeFromString,
  PerfettoSqlType,
} from '../../trace_processor/perfetto_sql_type';
import {unwrapResult} from '../../base/result';
import {createTableColumn} from '../../components/widgets/sql/table/columns';

// Runs a data availability check query (wrapped by check_to_query in Python).
// The query is `SELECT EXISTS(...) AS has_data` which always returns one row.
// Returns true if data is present, false if not or if the query fails.
async function runDataCheck(trace: Trace, sql: string): Promise<boolean> {
  try {
    const result = await trace.engine.query(sql);
    // Use iter() to avoid type checking issues with VARINT
    const iter = result.iter({});
    iter.next();
    const hasDataValue = iter.get('has_data');
    return typeof hasDataValue === 'bigint'
      ? hasDataValue !== 0n
      : Number(hasDataValue) !== 0;
  } catch (_e) {
    // If query fails, assume no data
    return false;
  }
}

export class SqlModulesImpl implements SqlModules {
  readonly packages: SqlPackage[];
  private disabledModules: Set<string> = new Set();
  private disabledTables: Set<string> = new Set();
  private tablesWithChecks: Set<string> = new Set();
  private initPromise: Promise<void> | undefined;
  private readonly startInit: () => Promise<void>;

  constructor(trace: Trace, docs: SqlModulesDocsSchema) {
    this.packages = docs.map((json) => new StdlibPackageImpl(trace, json));
    // Capture initialization logic in a closure to avoid storing trace/docs
    this.startInit = () => this.computeDisabledModules(trace, docs);
  }

  ensureInitialized(): Promise<void> {
    if (this.initPromise === undefined) {
      this.initPromise = this.startInit();
    }
    return this.initPromise;
  }

  async waitForInit(): Promise<void> {
    await this.ensureInitialized();
  }

  private async computeDisabledModules(
    trace: Trace,
    docs: SqlModulesDocsSchema,
  ): Promise<void> {
    // Build dependency graph: module -> modules that include it
    const dependents = new Map<string, Set<string>>();
    const modulesWithChecks = new Map<string, string>();

    for (const pkg of docs) {
      for (const mod of pkg.modules) {
        const moduleName = mod.module_name;

        // Store data check SQL if present
        if (mod.data_check_sql) {
          modulesWithChecks.set(moduleName, mod.data_check_sql);
        }

        // Build reverse dependency graph
        if (mod.includes) {
          for (const includedModule of mod.includes) {
            if (!dependents.has(includedModule)) {
              dependents.set(includedModule, new Set());
            }
            dependents.get(includedModule)!.add(moduleName);
          }
        }
      }
    }

    // Check data availability for modules with checks
    const missingDataModules = new Set<string>();
    // Track modules whose trace checks explicitly passed (have data)
    const hasDataModules = new Set<string>();
    for (const [moduleName, checkSql] of modulesWithChecks) {
      const hasData = await runDataCheck(trace, checkSql);
      if (hasData) {
        hasDataModules.add(moduleName);
      } else {
        missingDataModules.add(moduleName);
      }
    }

    // Compute "protected" modules: modules with passing checks AND all modules
    // that depend on them (transitively). These should not be disabled by
    // dependency propagation since they have access to data through at least
    // one dependency with a passing trace check.
    const protectedModules = new Set(hasDataModules);
    const protectedQueue = Array.from(hasDataModules);

    while (protectedQueue.length > 0) {
      const current = protectedQueue.shift()!;
      const deps = dependents.get(current);

      if (deps) {
        for (const dependent of deps) {
          if (!protectedModules.has(dependent)) {
            protectedModules.add(dependent);
            protectedQueue.push(dependent);
          }
        }
      }
    }

    // BFS to find all transitive dependents of modules with missing data
    // Skip disabling protected modules (those with passing checks or that
    // depend on modules with passing checks)
    const queue = Array.from(missingDataModules);
    const disabled = new Set(missingDataModules);

    while (queue.length > 0) {
      const current = queue.shift()!;
      const deps = dependents.get(current);

      if (deps) {
        for (const dependent of deps) {
          if (!disabled.has(dependent) && !protectedModules.has(dependent)) {
            disabled.add(dependent);
            queue.push(dependent);
          }
        }
      }
    }

    this.disabledModules = disabled;

    // Run per-table data checks for tables that have their own check SQL.
    // This provides finer granularity than module-level checks.
    const tableChecks: Array<{name: string; sql: string}> = [];
    for (const pkg of this.packages) {
      for (const mod of pkg.modules) {
        for (const table of mod.tables) {
          if (table.dataCheckSql) {
            this.tablesWithChecks.add(table.name);
            tableChecks.push({name: table.name, sql: table.dataCheckSql});
          }
        }
      }
    }

    for (const {name, sql} of tableChecks) {
      const hasData = await runDataCheck(trace, sql);
      if (!hasData) {
        this.disabledTables.add(name);
      }
    }
  }

  isModuleDisabled(moduleName: string): boolean {
    return this.disabledModules.has(moduleName);
  }

  tablePassedDataCheck(tableName: string): boolean | undefined {
    if (!this.tablesWithChecks.has(tableName)) {
      return undefined;
    }
    return !this.disabledTables.has(tableName);
  }

  getDisabledModules(): ReadonlySet<string> {
    return this.disabledModules;
  }

  getTable(tableName: string): SqlTable | undefined {
    for (const p of this.packages) {
      const t = p.getTable(tableName);
      if (t !== undefined) {
        return t;
      }
    }
    return;
  }

  listTables(): SqlTable[] {
    return this.packages.flatMap((p) => p.listTables());
  }

  listTablesNames(): string[] {
    return this.packages.flatMap((p) => p.listTablesNames());
  }

  getModuleForTable(tableName: string): SqlModule | undefined {
    for (const stdlibPackage of this.packages) {
      const maybeTable = stdlibPackage.getModuleForTable(tableName);
      if (maybeTable) {
        return maybeTable;
      }
    }
    return undefined;
  }

  listModules(): SqlModule[] {
    return this.packages.flatMap((p) => p.modules);
  }
}

export class StdlibPackageImpl implements SqlPackage {
  readonly name: string;
  readonly modules: SqlModule[];

  constructor(trace: Trace, docs: DocsPackageSchemaType) {
    this.name = docs.name;
    this.modules = [];
    for (const moduleJson of docs.modules) {
      this.modules.push(new StdlibModuleImpl(trace, moduleJson));
    }
  }

  getTable(tableName: string): SqlTable | undefined {
    for (const module of this.modules) {
      for (const t of module.tables) {
        if (t.name == tableName) {
          return t;
        }
      }
    }
    return undefined;
  }

  listTables(): SqlTable[] {
    return this.modules.flatMap((module) => module.tables);
  }

  listTablesNames(): string[] {
    return this.listTables().map((t) => t.name);
  }

  getModuleForTable(tableName: string): SqlModule | undefined {
    for (const module of this.modules) {
      for (const t of module.tables) {
        if (t.name == tableName) {
          return module;
        }
      }
    }
    return undefined;
  }

  getSqlTableDefinition(tableName: string): SqlTableDefinition | undefined {
    for (const module of this.modules) {
      for (const t of module.tables) {
        if (t.name == tableName) {
          return module.getSqlTableDefinition(tableName);
        }
      }
    }
    return undefined;
  }
}

export class StdlibModuleImpl implements SqlModule {
  readonly includeKey: string;
  readonly tags: string[];
  readonly tables: SqlTable[];
  readonly functions: SqlFunction[];
  readonly tableFunctions: SqlTableFunction[];
  readonly macros: SqlMacro[];
  readonly dataCheckSql?: string;
  readonly includes: string[];

  constructor(trace: Trace, docs: DocsModuleSchemaType) {
    this.includeKey = docs.module_name;
    this.tags = docs.tags;
    this.dataCheckSql = docs.data_check_sql ?? undefined;
    this.includes = docs.includes ?? [];

    const neededInclude = this.includeKey.startsWith('prelude')
      ? undefined
      : this.includeKey;
    this.tables = docs.data_objects.map(
      (json) => new SqlTableImpl(trace, json, neededInclude),
    );

    this.functions = docs.functions.map((json) => new StdlibFunctionImpl(json));
    this.tableFunctions = docs.table_functions.map(
      (json) => new StdlibTableFunctionImpl(json),
    );
    this.macros = docs.macros.map((json) => new StdlibMacroImpl(json));
  }

  getTable(tableName: string): SqlTable | undefined {
    for (const t of this.tables) {
      if (t.name == tableName) {
        return t;
      }
    }
    return undefined;
  }

  getSqlTableDefinition(tableName: string): SqlTableDefinition | undefined {
    const sqlTable = this.getTable(tableName);
    if (sqlTable === undefined) {
      return undefined;
    }
    return {
      imports: [this.includeKey],
      name: sqlTable.name,
      columns: sqlTable.columns.map((col) => ({
        column: col.name,
        type: col.type,
      })),
    };
  }
}

class StdlibMacroImpl implements SqlMacro {
  readonly name: string;
  readonly summaryDesc: string;
  readonly description: string;
  readonly args: SqlArgument[];
  readonly returnType: string;

  constructor(docs: DocsMacroSchemaType) {
    this.name = docs.name;
    this.summaryDesc = docs.summary_desc;
    this.description = docs.desc;
    this.returnType = docs.return_type;
    this.args = [];
    this.args = docs.args.map((json) => new StdlibFunctionArgImpl(json));
  }
}

class StdlibTableFunctionImpl implements SqlTableFunction {
  readonly name: string;
  readonly summaryDesc: string;
  readonly description: string;
  readonly args: SqlArgument[];
  readonly returnCols: SqlColumn[];

  constructor(docs: DocsTableFunctionSchemaType) {
    this.name = docs.name;
    this.summaryDesc = docs.summary_desc;
    this.description = docs.desc;
    this.args = docs.args.map((json) => new StdlibFunctionArgImpl(json));
    this.returnCols = docs.cols.map(
      (json) => new StdlibColumnImpl(json, this.name),
    );
  }
}

class StdlibFunctionImpl implements SqlFunction {
  readonly name: string;
  readonly summaryDesc: string;
  readonly description: string;
  readonly args: SqlArgument[];
  readonly returnType: string;
  readonly returnDesc: string;

  constructor(docs: DocsFunctionSchemaType) {
    this.name = docs.name;
    this.summaryDesc = docs.summary_desc;
    this.description = docs.desc;
    this.returnType = docs.return_type;
    this.returnDesc = docs.return_desc;
    this.args = docs.args.map((json) => new StdlibFunctionArgImpl(json));
  }
}

class SqlTableImpl implements SqlTable {
  name: string;
  includeKey?: string;
  description: string;
  type: string;
  importance?: 'core' | 'high' | 'mid' | 'low';
  dataCheckSql?: string;
  columns: SqlColumn[];
  idColumn: SqlColumn | undefined;

  constructor(
    readonly trace: Trace,
    docs: DocsDataObjectSchemaType,
    includeKey: string | undefined,
  ) {
    this.name = docs.name;
    this.includeKey = includeKey;
    this.description = docs.desc;
    this.type = docs.type;
    this.importance = docs.importance ?? undefined;
    this.dataCheckSql = docs.data_check_sql ?? undefined;
    this.columns = docs.cols.map(
      (json) => new StdlibColumnImpl(json, this.name),
    );
  }

  getTableColumns(): TableColumn[] {
    return this.columns.map((col) =>
      createTableColumn({
        trace: this.trace,
        column: col.name,
        type: col.type,
      }),
    );
  }
}

class StdlibColumnImpl implements SqlColumn {
  name: string;
  type: PerfettoSqlType;
  description: string;

  constructor(docs: DocsArgOrColSchemaType, tableName: string) {
    this.name = docs.name;
    this.type = unwrapResult(
      parsePerfettoSqlTypeFromString({
        type: docs.type,
        table: tableName,
        column: this.name,
      }),
    );
    this.description = docs.desc;
  }
}

class StdlibFunctionArgImpl implements SqlArgument {
  name: string;
  description: string;
  type: string;

  constructor(docs: DocsArgOrColSchemaType) {
    this.type = docs.type;
    this.description = docs.desc;
    this.name = docs.name;
  }
}

export const ARG_OR_COL_SCHEMA = z.object({
  name: z.string(),
  type: z.string(),
  desc: z.string(),
  table: z.string().nullable(),
  column: z.string().nullable(),
});
type DocsArgOrColSchemaType = z.infer<typeof ARG_OR_COL_SCHEMA>;

export const DATA_OBJECT_SCHEMA = z.object({
  name: z.string(),
  desc: z.string(),
  summary_desc: z.string(),
  type: z.string(),
  importance: z.enum(['core', 'high', 'mid', 'low']).nullish(),
  data_check_sql: z.string().nullish(),
  cols: z.array(ARG_OR_COL_SCHEMA),
});
type DocsDataObjectSchemaType = z.infer<typeof DATA_OBJECT_SCHEMA>;

export const FUNCTION_SCHEMA = z.object({
  name: z.string(),
  desc: z.string(),
  summary_desc: z.string(),
  args: z.array(ARG_OR_COL_SCHEMA),
  return_type: z.string(),
  return_desc: z.string(),
});
type DocsFunctionSchemaType = z.infer<typeof FUNCTION_SCHEMA>;

export const TABLE_FUNCTION_SCHEMA = z.object({
  name: z.string(),
  desc: z.string(),
  summary_desc: z.string(),
  args: z.array(ARG_OR_COL_SCHEMA),
  cols: z.array(ARG_OR_COL_SCHEMA),
});
type DocsTableFunctionSchemaType = z.infer<typeof TABLE_FUNCTION_SCHEMA>;

export const MACRO_SCHEMA = z.object({
  name: z.string(),
  desc: z.string(),
  summary_desc: z.string(),
  return_desc: z.string(),
  return_type: z.string(),
  args: z.array(ARG_OR_COL_SCHEMA),
});
type DocsMacroSchemaType = z.infer<typeof MACRO_SCHEMA>;

const MODULE_SCHEMA = z.object({
  module_name: z.string(),
  tags: z.array(z.string()),
  data_objects: z.array(DATA_OBJECT_SCHEMA),
  functions: z.array(FUNCTION_SCHEMA),
  table_functions: z.array(TABLE_FUNCTION_SCHEMA),
  macros: z.array(MACRO_SCHEMA),
  data_check_sql: z.string().nullish(),
  includes: z.array(z.string()).nullish(),
});
type DocsModuleSchemaType = z.infer<typeof MODULE_SCHEMA>;

const PACKAGE_SCHEMA = z.object({
  name: z.string(),
  modules: z.array(MODULE_SCHEMA),
});
type DocsPackageSchemaType = z.infer<typeof PACKAGE_SCHEMA>;

export const SQL_MODULES_DOCS_SCHEMA = z.array(PACKAGE_SCHEMA);
export type SqlModulesDocsSchema = z.infer<typeof SQL_MODULES_DOCS_SCHEMA>;