PromQL describe_source can't list metric or label names on a large TimeSeries table: both catalog paths read every series id in the table #3648

// Repro: clickstack_describe_source on a PromQL (TimeSeries) source costs memory in proportion
// to the number of series in the table, on both label backends; a plain metric-name listing doesn't.
// Copy to packages/api/src/mcp/__tests__/ and run (REPRO_SERIES defaults to 1,000,000):
//   REPRO_SERIES=2000000 make dev-int FILE=reproPromqlCatalogScale
import { SourceKind } from '@hyperdx/common-utils/dist/types';
import { Client } from '@modelcontextprotocol/sdk/client/index.js';

import * as config from '@/config';
import {
  DEFAULT_DATABASE,
  dropTimeSeriesTable,
  executeTimeSeriesSqlCommand,
  getLoggedInAgent,
  getServer,
  getTestFixtureClickHouseClient,
} from '@/fixtures';
import { clearBackendProbeCache } from '@/controllers/promql';
import Connection from '@/models/connection';
import { Source } from '@/models/source';

import { callTool, createTestClient, getFirstText } from './mcpTestUtils';

jest.setTimeout(10 * 60 * 1000);

const mockFetch = jest.mocked(global.fetch);
const { realFetch } = globalThis as unknown as { realFetch: typeof fetch };

const TABLE = 'repro_promql_catalog';
const SERIES = Number(process.env.REPRO_SERIES ?? 1_000_000);
const METRICS = 200;

describe(`repro: describe_source on a PromQL source with ${SERIES.toLocaleString('en-US')} series`, () => {
  const server = getServer();
  let client: Client;
  let team: any;
  let connectionId: any;

  beforeAll(async () => {
    await server.start();
    await dropTimeSeriesTable({ table: TABLE });
    // Schema version 8 (the 26.10 default): series time ranges live in their own inner table.
    await executeTimeSeriesSqlCommand(`CREATE TABLE ${DEFAULT_DATABASE}.${TABLE} ENGINE = TimeSeries SETTINGS version = 8`);
    await executeTimeSeriesSqlCommand(
      `INSERT INTO TABLE FUNCTION timeSeriesTags('${DEFAULT_DATABASE}', '${TABLE}') (metric_name, tags)
       SELECT 'metric_' || toString(number % ${METRICS}), map('pod', 'pod-' || toString(number), 'job', 'job-' || toString(number % 10))
       FROM numbers(${SERIES})`,
    );
    // Every series active in the last hour, so all of them are inside describe_source's 24 h window.
    await executeTimeSeriesSqlCommand(
      `INSERT INTO TABLE FUNCTION timeSeriesTimeRanges('${DEFAULT_DATABASE}', '${TABLE}')
       SELECT id, now64(3) - 3600, now64(3) FROM timeSeriesTags('${DEFAULT_DATABASE}', '${TABLE}')`,
    );
    const agent = await getLoggedInAgent(server);
    team = agent.team;
    const user = agent.user;
    const connection = await Connection.create({
      team: team._id,
      name: 'Default',
      host: config.CLICKHOUSE_HOST,
      username: config.CLICKHOUSE_USER,
      password: config.CLICKHOUSE_PASSWORD,
    });
    connectionId = connection._id;
    client = await createTestClient({ teamId: team._id.toString(), userId: user._id.toString(), promqlEnabled: true });
  });

  afterAll(async () => {
    await client?.close();
    await dropTimeSeriesTable({ table: TABLE });
    await server.clearDBs();
    await server.stop();
  });

  afterEach(() => {
    clearBackendProbeCache();
    mockFetch.mockReset();
  });

  const ch = () => getTestFixtureClickHouseClient();
  // ClickHouse's clock, so query_log filtering doesn't depend on the test host's clock.
  const sinceMs = async () =>
    Number(
      (
        await (await ch())
          .query({ query: 'SELECT toUnixTimestamp64Milli(now64(3)) AS t', format: 'JSONEachRow' })
          .then(r => r.json<{ t: string }>())
      )[0].t,
    );

  /** Memory and rows of the label/metric-name lookups ClickHouse ran since `since`. */
  const lookups = async (since: number) => {
    await (await ch()).command({ query: 'SYSTEM FLUSH LOGS' });
    const uuid = await (await ch())
      .query({ query: `SELECT toString(uuid) AS u FROM system.tables WHERE database = '${DEFAULT_DATABASE}' AND name = '${TABLE}'`, format: 'JSONEachRow' })
      .then(r => r.json<{ u: string }>());
    const rows = await (await ch())
      .query({
        query: `
          SELECT if(query ILIKE '%timeSeriesStoreTags%', 'prometheus_api_v1 handler', 'SQL') AS path,
                 if(query ILIKE '%mapKeys%' OR query ILIKE '%tupleElement(timeSeriesIdToTags(series_id), 1)%', 'label names', 'metric names') AS lookup,
                 type, query_duration_ms AS ms, read_rows, memory_usage
          FROM system.query_log
          WHERE event_time_microseconds >= fromUnixTimestamp64Milli(${since}) AND type != 'QueryStart' AND query_kind = 'Select'
            AND (query ILIKE '%timeSeriesTimeRanges(%GROUP BY id%' OR query ILIKE '%${uuid[0].u}%')
            AND query NOT ILIKE '%system.query_log%'
          ORDER BY event_time_microseconds`,
        format: 'JSONEachRow',
      })
      .then(r => r.json<Record<string, string>>());
    return rows;
  };

  const mib = (b: string | number) => `${(Number(b) / 2 ** 20).toFixed(1)} MiB`;
  const results: string[] = [];
  let plainMemory = 0;

  it('a plain metric-name listing reads one column and stays small', async () => {
    const since = await sinceMs();
    await (await ch()).query({
      query: `SELECT DISTINCT metric_name FROM timeSeriesTags(${DEFAULT_DATABASE}, ${TABLE}) ORDER BY metric_name LIMIT 201 SETTINGS log_comment = 'repro-plain'`,
      clickhouse_settings: { allow_experimental_time_series_table: 1 },
    }).then(r => r.json());
    await (await ch()).command({ query: 'SYSTEM FLUSH LOGS' });
    const [row] = await (await ch())
      .query({ query: `SELECT query_duration_ms AS ms, read_rows, memory_usage FROM system.query_log WHERE log_comment = 'repro-plain' AND type = 'QueryFinish' AND event_time_microseconds >= fromUnixTimestamp64Milli(${since})`, format: 'JSONEachRow' })
      .then(r => r.json<Record<string, string>>());
    plainMemory = Number(row.memory_usage);
    results.push(`plain SELECT DISTINCT metric_name (tags table only)    ${row.ms} ms  rows=${row.read_rows}  memory=${mib(row.memory_usage)}`);
  });

  it.each([
    ['table-function', 'SQL path (source pinned to table-function; CH < 26.9 or no handler)'],
    [undefined, 'auto: prometheus_api_v1 handler on CH 26.9+ (#3628/#3630)'],
  ])('describe_source with promqlBackend=%s', async (promqlBackend, label) => {
    mockFetch.mockImplementation(realFetch);
    const source = await Source.create({
      kind: SourceKind.Promql,
      team: team._id,
      connection: connectionId,
      name: `PromQL ${promqlBackend ?? 'auto'}`,
      from: { databaseName: DEFAULT_DATABASE, tableName: TABLE },
      timestampValueExpression: 'timestamp',
      ...(promqlBackend ? { promqlBackend } : {}),
    });
    const since = await sinceMs();
    const t0 = Date.now();
    const result = await callTool(client, 'clickstack_describe_source', { sourceId: source._id.toString() });
    const wall = Date.now() - t0;
    const out = JSON.parse(getFirstText(result));
    const queries = await lookups(since);
    results.push(
      `describe_source, ${label}: ${wall} ms wall, metricNames=${out.metricNames?.length ?? 0}, partial=${out.partial ?? false}`,
      ...queries.map(q => `   ${q.path.padEnd(26)} ${q.lookup.padEnd(13)} ${q.type.padEnd(14)} ${String(q.ms).padStart(5)} ms  rows=${q.read_rows}  memory=${mib(q.memory_usage)}  (${(Number(q.memory_usage) / Math.max(plainMemory, 1)).toFixed(0)}x plain)`),
    );
    expect(queries.length).toBeGreaterThan(0);
    for (const q of queries) expect(Number(q.memory_usage)).toBeGreaterThan(20 * plainMemory);
  });

  afterAll(() => {
    console.log(`REPRO ${SERIES.toLocaleString('en-US')} series, ${METRICS} metric names\n${results.join('\n')}`);
  });
});