ui/src/plugins/dev.perfetto.SqlModules/sql_modules_from_tp.ts - third_party/perfetto - Git at Google

 // 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 __intrinsic_stdlib_* table functions and merges with the
 // pre-generated metadata JSON to build a SqlModules object.
 // Uses 4 parallel queries (one per entity type) instead of per-module queries.
 export async function loadSqlModulesFromTp(
   trace: Trace,
   metadata: StdlibMetadata,
 ): Promise<SqlModules> {
   const engine = trace.engine;

   const [modsResult, tablesResult, fnsResult, macrosResult] = await Promise.all(
     [
       engine.query('SELECT module, package FROM __intrinsic_stdlib_modules()'),
       engine.query(
         `SELECT module, name, type, description, exposed, cols
            FROM __intrinsic_stdlib_tables('*')`,
       ),
       engine.query(
         `SELECT module, name, description, exposed, is_table_function,
                 return_type, return_description, args, cols
            FROM __intrinsic_stdlib_functions('*')`,
       ),
       engine.query(
         `SELECT module, name, description, exposed,
                 return_type, return_description, args
            FROM __intrinsic_stdlib_macros('*')`,
       ),
     ],
   );

   // Build package → [module] map.
   const packageMap = new Map<string, string[]>();
   const modsIter = modsResult.iter({module: STR, package: STR});
   for (; modsIter.valid(); modsIter.next()) {
     getOrCreate(packageMap, modsIter.package, () => []).push(modsIter.module);
   }

   // Group tables by module.
   const tablesByModule = new Map<string, DataObject[]>();
   const tIter = tablesResult.iter({
     module: STR,
     name: STR,
     type: STR,
     description: STR,
     exposed: LONG,
     cols: STR,
   });
   for (; tIter.valid(); tIter.next()) {
     if (!tIter.exposed) continue;
     const modKey = tIter.module;
     const tableMeta = metadata[modKey]?.tables?.[tIter.name];
     const desc = tIter.description;
     getOrCreate(tablesByModule, modKey, () => []).push({
       name: tIter.name,
       desc,
       summary_desc: summaryDesc(desc),
       type: tIter.type,
       importance: tableMeta?.importance ?? null,
       data_check_sql: tableMeta?.data_check_sql ?? null,
       cols: parseEntries(tIter.cols).map((e) => toArgOrCol(e)),
     });
   }

   // Group functions and table functions by module.
   const fnsByModule = new Map<string, FnObject[]>();
   const tblFnsByModule = new Map<string, TblFnObject[]>();
   const fIter = fnsResult.iter({
     module: STR,
     name: STR,
     description: STR,
     exposed: LONG,
     is_table_function: LONG,
     return_type: STR,
     return_description: STR,
     args: STR,
     cols: STR,
   });
   for (; fIter.valid(); fIter.next()) {
     if (!fIter.exposed) continue;
     const modKey = fIter.module;
     const desc = fIter.description;
     const args = parseEntries(fIter.args).map((e) => toArgOrCol(e));
     if (fIter.is_table_function) {
       getOrCreate(tblFnsByModule, modKey, () => []).push({
         name: fIter.name,
         desc,
         summary_desc: summaryDesc(desc),
         args,
         cols: parseEntries(fIter.cols).map((e) => toArgOrCol(e)),
       });
     } else {
       getOrCreate(fnsByModule, modKey, () => []).push({
         name: fIter.name,
         desc,
         summary_desc: summaryDesc(desc),
         return_type: fIter.return_type,
         return_desc: fIter.return_description,
         args,
       });
     }
   }

   // Group macros by module.
   const macrosByModule = new Map<string, MacroObject[]>();
   const mIter = macrosResult.iter({
     module: STR,
     name: STR,
     description: STR,
     exposed: LONG,
     return_type: STR,
     return_description: STR,
     args: STR,
   });
   for (; mIter.valid(); mIter.next()) {
     if (!mIter.exposed) continue;
     const modKey = mIter.module;
     const desc = mIter.description;
     getOrCreate(macrosByModule, modKey, () => []).push({
       name: mIter.name,
       desc,
       summary_desc: summaryDesc(desc),
       return_type: mIter.return_type,
       return_desc: mIter.return_description,
       args: parseEntries(mIter.args).map((e) => toArgOrCol(e)),
     });
   }

   // 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);
 }
	// 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 __intrinsic_stdlib_* table functions and merges with the
	// pre-generated metadata JSON to build a SqlModules object.
	// Uses 4 parallel queries (one per entity type) instead of per-module queries.
	export async function loadSqlModulesFromTp(
	trace: Trace,
	metadata: StdlibMetadata,
	): Promise<SqlModules> {
	const engine = trace.engine;

	const [modsResult, tablesResult, fnsResult, macrosResult] = await Promise.all(
	[
	engine.query('SELECT module, package FROM __intrinsic_stdlib_modules()'),
	engine.query(
	`SELECT module, name, type, description, exposed, cols
	FROM __intrinsic_stdlib_tables('*')`,
	),
	engine.query(
	`SELECT module, name, description, exposed, is_table_function,
	return_type, return_description, args, cols
	FROM __intrinsic_stdlib_functions('*')`,
	),
	engine.query(
	`SELECT module, name, description, exposed,
	return_type, return_description, args
	FROM __intrinsic_stdlib_macros('*')`,
	),
	],
	);

	// Build package → [module] map.
	const packageMap = new Map<string, string[]>();
	const modsIter = modsResult.iter({module: STR, package: STR});
	for (; modsIter.valid(); modsIter.next()) {
	getOrCreate(packageMap, modsIter.package, () => []).push(modsIter.module);
	}

	// Group tables by module.
	const tablesByModule = new Map<string, DataObject[]>();
	const tIter = tablesResult.iter({
	module: STR,
	name: STR,
	type: STR,
	description: STR,
	exposed: LONG,
	cols: STR,
	});
	for (; tIter.valid(); tIter.next()) {
	if (!tIter.exposed) continue;
	const modKey = tIter.module;
	const tableMeta = metadata[modKey]?.tables?.[tIter.name];
	const desc = tIter.description;
	getOrCreate(tablesByModule, modKey, () => []).push({
	name: tIter.name,
	desc,
	summary_desc: summaryDesc(desc),
	type: tIter.type,
	importance: tableMeta?.importance ?? null,
	data_check_sql: tableMeta?.data_check_sql ?? null,
	cols: parseEntries(tIter.cols).map((e) => toArgOrCol(e)),
	});
	}

	// Group functions and table functions by module.
	const fnsByModule = new Map<string, FnObject[]>();
	const tblFnsByModule = new Map<string, TblFnObject[]>();
	const fIter = fnsResult.iter({
	module: STR,
	name: STR,
	description: STR,
	exposed: LONG,
	is_table_function: LONG,
	return_type: STR,
	return_description: STR,
	args: STR,
	cols: STR,
	});
	for (; fIter.valid(); fIter.next()) {
	if (!fIter.exposed) continue;
	const modKey = fIter.module;
	const desc = fIter.description;
	const args = parseEntries(fIter.args).map((e) => toArgOrCol(e));
	if (fIter.is_table_function) {
	getOrCreate(tblFnsByModule, modKey, () => []).push({
	name: fIter.name,
	desc,
	summary_desc: summaryDesc(desc),
	args,
	cols: parseEntries(fIter.cols).map((e) => toArgOrCol(e)),
	});
	} else {
	getOrCreate(fnsByModule, modKey, () => []).push({
	name: fIter.name,
	desc,
	summary_desc: summaryDesc(desc),
	return_type: fIter.return_type,
	return_desc: fIter.return_description,
	args,
	});
	}
	}

	// Group macros by module.
	const macrosByModule = new Map<string, MacroObject[]>();
	const mIter = macrosResult.iter({
	module: STR,
	name: STR,
	description: STR,
	exposed: LONG,
	return_type: STR,
	return_description: STR,
	args: STR,
	});
	for (; mIter.valid(); mIter.next()) {
	if (!mIter.exposed) continue;
	const modKey = mIter.module;
	const desc = mIter.description;
	getOrCreate(macrosByModule, modKey, () => []).push({
	name: mIter.name,
	desc,
	summary_desc: summaryDesc(desc),
	return_type: mIter.return_type,
	return_desc: mIter.return_description,
	args: parseEntries(mIter.args).map((e) => toArgOrCol(e)),
	});
	}

	// 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);
	}