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

// Copyright (C) 2025 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 {LONG, STR} from '../../trace_processor/query_result';
import type {Trace} from '../../public/trace';
import {getOrCreate} from '../../base/utils';
import type {SqlModules} from './sql_modules';
import {
  SQL_MODULES_DOCS_SCHEMA,
  SqlModulesImpl,
  type ARG_OR_COL_SCHEMA,
  type DATA_OBJECT_SCHEMA,
  type FUNCTION_SCHEMA,
  type TABLE_FUNCTION_SCHEMA,
  type MACRO_SCHEMA,
} from './sql_modules_impl';

// Schema for the metadata-only JSON generated by gen_stdlib_docs_json.py.
// Contains only what the TP table functions don't expose: tags, includes,
// data availability checks, and table importance.
const TABLE_META_SCHEMA = z.object({
  importance: z.enum(['core', 'high', 'mid', 'low']).nullish(),
  data_check_sql: z.string().nullish(),
});

const MODULE_META_SCHEMA = z.object({
  tags: z.array(z.string()),
  includes: z.array(z.string()),
  data_check_sql: z.string().nullish(),
  tables: z.record(z.string(), TABLE_META_SCHEMA).optional(),
});

export const STDLIB_METADATA_SCHEMA = z.record(z.string(), MODULE_META_SCHEMA);
export type StdlibMetadata = z.infer<typeof STDLIB_METADATA_SCHEMA>;

// Types derived from the schema so any field-name mismatch is a compile error.
type ArgOrCol = z.infer<typeof ARG_OR_COL_SCHEMA>;
type DataObject = z.infer<typeof DATA_OBJECT_SCHEMA>;
type FnObject = z.infer<typeof FUNCTION_SCHEMA>;
type TblFnObject = z.infer<typeof TABLE_FUNCTION_SCHEMA>;
type MacroObject = z.infer<typeof MACRO_SCHEMA>;

// Schema for a single serialized entry in the cols/args JSON blobs from TP.
const RAW_ENTRY_SCHEMA = z.object({
  name: z.string(),
  type: z.string(),
  description: z.string(),
});
type RawEntry = z.infer<typeof RAW_ENTRY_SCHEMA>;

function parseEntries(json: string): RawEntry[] {
  try {
    const parsed: unknown = JSON.parse(json);
    if (!Array.isArray(parsed)) return [];
    return parsed.flatMap((e) => {
      const r = RAW_ENTRY_SCHEMA.safeParse(e);
      return r.success ? [r.data] : [];
    });
  } catch {
    return [];
  }
}

// Extracts the referenced table and column from a long type of the form
// "TYPE(table.column)" (e.g. "JOINID(thread.utid)"). Returns [null, null]
// when the type carries no such reference. Mirrors the parsing the Python
// generator used to do over long_type.
const LONG_TYPE_RE = /^[A-Z]*\(([a-z_]*)\.([a-z_]*)\)/;

function toArgOrCol(e: RawEntry): ArgOrCol {
  const m = LONG_TYPE_RE.exec(e.type);
  return {
    name: e.name,
    type: e.type,
    desc: e.description,
    table: m ? m[1] : null,
    column: m ? m[2] : null,
  };
}

function summaryDesc(desc: string): string {
  const trimmed = desc.replace(/\n/g, ' ');
  const dot = trimmed.indexOf('. ');
  if (dot !== -1) return trimmed.slice(0, dot);
  const lastDot = trimmed.lastIndexOf('.');
  if (lastDot !== -1) return trimmed.slice(0, lastDot);
  return trimmed;
}

// Queries the TP's unified stdlib object catalog and merges it with the
// pre-generated metadata JSON to build a SqlModules object.
export async function loadSqlModulesFromTp(
  trace: Trace,
  metadata: StdlibMetadata,
): Promise<SqlModules> {
  const result = await trace.engine.query(
    `SELECT package, module, name, object_type, description, exposed,
            return_type, return_description, args, cols
       FROM __intrinsic_stdlib_objects`,
  );

  const packageMap = new Map<string, string[]>();
  const tablesByModule = new Map<string, DataObject[]>();
  const fnsByModule = new Map<string, FnObject[]>();
  const tblFnsByModule = new Map<string, TblFnObject[]>();
  const macrosByModule = new Map<string, MacroObject[]>();
  const iter = result.iter({
    package: STR,
    module: STR,
    name: STR,
    object_type: STR,
    description: STR,
    exposed: LONG,
    return_type: STR,
    return_description: STR,
    args: STR,
    cols: STR,
  });
  for (; iter.valid(); iter.next()) {
    const modKey = iter.module;
    if (iter.object_type === 'MODULE') {
      getOrCreate(packageMap, iter.package, () => []).push(modKey);
      continue;
    }
    if (iter.exposed === 0n) continue;

    const desc = iter.description;
    const args = parseEntries(iter.args).map((entry) => toArgOrCol(entry));
    if (iter.object_type === 'TABLE' || iter.object_type === 'VIEW') {
      const tableMeta = metadata[modKey]?.tables?.[iter.name];
      getOrCreate(tablesByModule, modKey, () => []).push({
        name: iter.name,
        desc,
        summary_desc: summaryDesc(desc),
        type: iter.object_type,
        importance: tableMeta?.importance ?? null,
        data_check_sql: tableMeta?.data_check_sql ?? null,
        cols: parseEntries(iter.cols).map((entry) => toArgOrCol(entry)),
      });
    } else if (iter.object_type === 'FUNCTION') {
      getOrCreate(fnsByModule, modKey, () => []).push({
        name: iter.name,
        desc,
        summary_desc: summaryDesc(desc),
        return_type: iter.return_type,
        return_desc: iter.return_description,
        args,
      });
    } else if (iter.object_type === 'TABLE_FUNCTION') {
      getOrCreate(tblFnsByModule, modKey, () => []).push({
        name: iter.name,
        desc,
        summary_desc: summaryDesc(desc),
        args,
        cols: parseEntries(iter.cols).map((entry) => toArgOrCol(entry)),
      });
    } else if (iter.object_type === 'MACRO') {
      getOrCreate(macrosByModule, modKey, () => []).push({
        name: iter.name,
        desc,
        summary_desc: summaryDesc(desc),
        return_type: iter.return_type,
        return_desc: iter.return_description,
        args,
      });
    }
  }

  // Assemble the final docs structure.
  const docs = [];
  for (const [pkgName, moduleKeys] of packageMap) {
    const pkgModules = [];
    for (const modKey of moduleKeys) {
      const meta = metadata[modKey];
      pkgModules.push({
        module_name: modKey,
        tags: meta?.tags ?? [],
        includes: meta?.includes ?? [],
        data_check_sql: meta?.data_check_sql ?? null,
        data_objects: tablesByModule.get(modKey) ?? [],
        functions: fnsByModule.get(modKey) ?? [],
        table_functions: tblFnsByModule.get(modKey) ?? [],
        macros: macrosByModule.get(modKey) ?? [],
      });
    }
    docs.push({name: pkgName, modules: pkgModules});
  }

  const parsed = SQL_MODULES_DOCS_SCHEMA.safeParse(docs);
  if (!parsed.success) {
    throw new Error(
      `Failed to parse stdlib docs from TP: ${parsed.error.message}`,
    );
  }
  return new SqlModulesImpl(trace, parsed.data);
}