0pt/tests/analytics/query-execution-model.test.ts

/**
 * tests/analytics/query-execution-model.test.ts
 * ==============================================
 * BI-P018 (execution model), BI-P019 (cancellation logic) and BI-P021
 * (cost/budget) — the parts provable without an engine.
 *
 * The engine-dependent claims — "cancellation actually stops a running query",
 * "the budget actually stops a running query" — are NOT asserted here. They
 * live in tests/integration/query-cancellation-budget.integration.test.ts
 * against a live ClickHouse, because a fake process-control that returns
 * `isRunning: false` on the second call would make this file pass while
 * proving nothing about the server.
 */
import { describe, expect, it } from 'vitest';
import { SiaxError } from '@/src/platform/errors';
import {
  assertLegalTransition,
  assertNoRawSql,
  compiledQueryHash,
  QueryExecutionRecorder,
  TERMINAL_STATES,
  type QueryExecutionRecord,
} from '@/src/analytics/execution/query-execution-record';
import { cancelAndVerify, type EngineProcessControl } from '@/src/analytics/execution/cancellation';
import {
  costUnits,
  DEFAULT_QUERY_BUDGET,
  InMemoryWindowStore,
  QueryBudgetEnforcer,
} from '@/src/analytics/cost/query-budget';

// ── BI-P018 ────────────────────────────────────────────────────────────────

function recorder(sink: { write: (r: QueryExecutionRecord) => void }, now = () => 1_000_000) {
  return new QueryExecutionRecorder(
    {
      tenant_id: 't_alpha',
      project_id: 'p_main',
      principal_id: 'u_1',
      workload: 'dashboard',
      engine: 'clickhouse',
      correlation_id: 'c_1',
      sink,
      now,
    },
    'opt_q1',
    {
      hash: compiledQueryHash('SELECT 1', { param_tenant: 't_alpha' }),
      datasets: ['ds_sales'],
      metrics: ['revenue'],
    },
  );
}

describe('BI-P018 — every field the prompt names, and none it forbids', () => {
  it('carries all eighteen required fields once settled', async () => {
    const written: QueryExecutionRecord[] = [];
    let clock = 1_000_000;
    const r = recorder({ write: (x) => void written.push(x) }, () => clock);
    await r.submitted();
    clock += 50;
    await r.started();
    clock += 700;
    const final = await r.finish('succeeded', {
      rows: 12,
      rows_scanned: 40_960,
      bytes_scanned: 1_234_567,
      cost_units: 0.5,
    });

    for (const field of [
      'query_id',
      'tenant_id',
      'project_id',
      'principal_id',
      'workload',
      'submitted_at',
      'started_at',
      'completed_at',
      'engine',
      'compiled_query_hash',
      'datasets',
      'metrics',
      'status',
      'rows',
      'rows_scanned',
      'bytes_scanned',
      'duration_ms',
      'cache_status',
      'cost_units',
    ]) {
      expect(final[field as keyof QueryExecutionRecord], `missing ${field}`).toBeDefined();
    }
    expect(final.duration_ms).toBe(700);
    expect(final.status).toBe('succeeded');
  });

  it('writes at EVERY transition, so an execution that never finishes is still visible', async () => {
    const written: QueryExecutionRecord[] = [];
    const r = recorder({ write: (x) => void written.push(x) });
    await r.submitted();
    await r.started();
    expect(written.map((w) => w.status)).toEqual(['queued', 'running']);
    // A recorder that only wrote at the end would have zero rows here — and
    // would lose precisely the executions worth investigating.
  });

  it('stores a hash, never the statement', async () => {
    const written: QueryExecutionRecord[] = [];
    const r = recorder({ write: (x) => void written.push(x) });
    await r.submitted();
    await r.started();
    const final = await r.finish('succeeded', { rows: 1 });
    const serialized = JSON.stringify(final);
    expect(serialized).not.toMatch(/SELECT/i);
    expect(final.compiled_query_hash).toMatch(/^[0-9a-f]{64}$/);
  });

  it('REFUSES to emit a record whose error field carries statement text', () => {
    const record: QueryExecutionRecord = {
      query_id: 'q',
      execution_id: 'e',
      tenant_id: 't_alpha',
      principal_id: 'u',
      workload: 'api',
      submitted_at: new Date().toISOString(),
      engine: 'clickhouse',
      compiled_query_hash: 'a'.repeat(64),
      datasets: [],
      metrics: [],
      status: 'failed',
      cache_status: 'miss',
      // The realistic leak: ClickHouse's error text echoes the query.
      error: 'Code: 47. Unknown identifier in scope SELECT tenant_id FROM opt_analytics.sales',
    };
    expect(() => assertNoRawSql(record)).toThrowError(SiaxError);
  });

  it('hashes the same statement with different parameters differently', () => {
    const a = compiledQueryHash('SELECT 1 WHERE t = {t:String}', { param_t: 't_alpha' });
    const b = compiledQueryHash('SELECT 1 WHERE t = {t:String}', { param_t: 't_beta' });
    expect(a).not.toBe(b);
    // ...and identically regardless of key order.
    expect(compiledQueryHash('q', { b: '2', a: '1' })).toBe(compiledQueryHash('q', { a: '1', b: '2' }));
  });

  it('refuses an illegal state transition rather than overwriting a terminal row', () => {
    for (const terminal of TERMINAL_STATES) {
      expect(() => assertLegalTransition(terminal, 'running')).toThrowError(SiaxError);
    }
    expect(() => assertLegalTransition('queued', 'succeeded')).toThrowError(SiaxError);
    expect(() => assertLegalTransition('running', 'succeeded')).not.toThrow();
  });

  it('records a cache hit as a cache hit, not as a very fast query', async () => {
    const written: QueryExecutionRecord[] = [];
    const r = recorder({ write: (x) => void written.push(x) });
    await r.submitted();
    await r.servedFromCache(42);
    const final = written.at(-1) as QueryExecutionRecord;
    expect(final.cache_status).toBe('hit');
    expect(final.rows_scanned).toBe(0);
    expect(final.bytes_scanned).toBe(0);
  });

  it('requires a verified tenant and principal', () => {
    for (const bad of [{ tenant_id: '' }, { principal_id: '' }]) {
      expect(
        () =>
          new QueryExecutionRecorder(
            {
              tenant_id: 't_alpha',
              principal_id: 'u_1',
              workload: 'api',
              engine: 'clickhouse',
              sink: { write: () => undefined },
              ...bad,
            },
            'q',
            { hash: 'h', datasets: [], metrics: [] },
          ),
      ).toThrowError(SiaxError);
    }
  });
});

// ── BI-P019 (logic only; the engine claim is in the integration suite) ─────

describe('BI-P019 — cancellation reports what it VERIFIED, not what it requested', () => {
  function control(script: boolean[]): EngineProcessControl & { kills: number } {
    let i = 0;
    return {
      engine: 'test',
      kills: 0,
      async isRunning() {
        const v = script[Math.min(i, script.length - 1)];
        i += 1;
        return v;
      },
      async requestKill() {
        (this as { kills: number }).kills += 1;
        return true;
      },
    };
  }

  it('reports verified_stopped only after OBSERVING the query gone', async () => {
    // running, running, running, gone
    const c = control([true, true, true, false]);
    const out = await cancelAndVerify(c, 'q1', { timeoutMs: 5_000, pollIntervalMs: 1 });
    expect(out.verified_stopped).toBe(true);
    expect(out.kill_accepted).toBe(true);
    expect(c.kills).toBe(1);
  });

  it('reports verified_stopped=false when the query outlives the deadline — it does not claim success', async () => {
    const c = control([true]); // never stops
    let clock = 0;
    const out = await cancelAndVerify(c, 'q1', {
      timeoutMs: 100,
      pollIntervalMs: 10,
      now: () => (clock += 30),
      sleep: async () => undefined,
    });
    expect(out.verified_stopped).toBe(false);
    expect(out.already_finished).toBe(false);
  });

  it('does not claim to have stopped a query that had already finished', async () => {
    const c = control([false]);
    const out = await cancelAndVerify(c, 'q1', { timeoutMs: 100 });
    expect(out.already_finished).toBe(true);
    expect(out.verified_stopped).toBe(true);
  });

  it('keeps polling when the engine is unreachable, then reports unverified', async () => {
    let clock = 0;
    const flaky: EngineProcessControl = {
      engine: 'test',
      isRunning: async () => {
        throw new Error('engine unreachable');
      },
      requestKill: async () => true,
    };
    // isRunning throws on the FIRST call too, which cancelAndVerify treats as
    // "assume it was running" — the fail-closed reading.
    const out = await cancelAndVerify(flaky, 'q1', {
      timeoutMs: 50,
      pollIntervalMs: 10,
      now: () => (clock += 20),
      sleep: async () => undefined,
    });
    expect(out.verified_stopped).toBe(false);
  });

  it('aborts the client side as well as the server side', async () => {
    let aborted = false;
    await cancelAndVerify(control([true, false]), 'q1', {
      timeoutMs: 1_000,
      pollIntervalMs: 1,
      abort: () => {
        aborted = true;
      },
    });
    expect(aborted).toBe(true);
  });
});

// ── BI-P021 ────────────────────────────────────────────────────────────────

describe('BI-P021 — the budget says plainly what it does and does not bound', () => {
  const ledger = {
    entries: [] as unknown[],
    record(e: unknown) {
      this.entries.push(e);
    },
  };

  it('refuses to be constructed as an in-flight enforcer without a process control', () => {
    expect(
      () =>
        new QueryBudgetEnforcer({
          budget: DEFAULT_QUERY_BUDGET,
          store: new InMemoryWindowStore(),
          ledger,
          enforceInFlight: true,
        }),
    ).toThrowError(SiaxError);
  });

  it('declares admission-only mode as advisory, in words', () => {
    const e = new QueryBudgetEnforcer({
      budget: DEFAULT_QUERY_BUDGET,
      store: new InMemoryWindowStore(),
      ledger,
      enforceInFlight: false,
    });
    expect(e.describeEnforcement().binding).toBe(false);
    expect(e.describeEnforcement().in_flight).toContain('NOT ENFORCED');
  });

  it('declares the per-process window as non-binding across replicas', () => {
    const e = new QueryBudgetEnforcer({
      budget: DEFAULT_QUERY_BUDGET,
      store: new InMemoryWindowStore(),
      ledger,
      control: {
        engine: 't',
        isRunning: async () => false,
        requestKill: async () => true,
        progress: async () => undefined,
      },
      enforceInFlight: true,
      storeIsShared: false,
    });
    expect(e.describeEnforcement().across_replicas).toContain('N times this budget');
    expect(e.describeEnforcement().binding).toBe(false);
  });

  it('rejects an estimate that does not fit the remaining window', async () => {
    const store = new InMemoryWindowStore();
    const e = new QueryBudgetEnforcer({
      budget: { ...DEFAULT_QUERY_BUDGET, maxRowsScanned: 1_000 },
      store,
      ledger,
      enforceInFlight: false,
    });
    expect((await e.admit('t_alpha', { rows_to_read: 100, bytes_to_read: 1 })).admitted).toBe(true);
    await store.add('t_alpha', { rows: 999, bytes: 0, units: 0, windowMs: DEFAULT_QUERY_BUDGET.windowMs });
    const decision = await e.admit('t_alpha', { rows_to_read: 100, bytes_to_read: 1 });
    expect(decision.admitted).toBe(false);
    expect(decision.reason).toContain('rows');
  });

  it('throttles instead of rejecting when the budget says so', async () => {
    const store = new InMemoryWindowStore();
    const e = new QueryBudgetEnforcer({
      budget: { ...DEFAULT_QUERY_BUDGET, maxRowsScanned: 10, onExhausted: 'throttle', throttleDelayMs: 250 },
      store,
      ledger,
      enforceInFlight: false,
    });
    await store.add('t_alpha', { rows: 100, bytes: 0, units: 0, windowMs: DEFAULT_QUERY_BUDGET.windowMs });
    const decision = await e.admit('t_alpha', { rows_to_read: 1, bytes_to_read: 1 });
    expect(decision.admitted).toBe(true);
    expect(decision.delay_ms).toBe(250);
  });

  it('charges for the work a cancelled query already did', async () => {
    const store = new InMemoryWindowStore();
    const entries: Array<{ terminated_by_budget: boolean; cost_units: number }> = [];
    const e = new QueryBudgetEnforcer({
      budget: DEFAULT_QUERY_BUDGET,
      store,
      ledger: { record: (x) => void entries.push(x as never) },
      enforceInFlight: false,
    });
    const settled = await e.settle({
      tenant_id: 't_alpha',
      principal_id: 'u',
      query_id: 'q',
      cost: { rows_scanned: 1e9, bytes_scanned: 2 ** 30, execution_ms: 1000, engine: 'clickhouse' },
      terminated_by_budget: true,
    });
    expect(settled.cost_units).toBeGreaterThan(0);
    expect(entries[0].terminated_by_budget).toBe(true);
    // Not charging for it is how a tenant learns the cheapest query is one
    // that gets cancelled.
    expect(entries[0].cost_units).toBe(settled.cost_units);
  });

  it('weights bytes above rows, so a wide scan is not billed like a narrow one', () => {
    const narrow = costUnits({ rows_scanned: 1e9, bytes_scanned: 1e8, execution_ms: 0, engine: 'e' });
    const wide = costUnits({ rows_scanned: 1e8, bytes_scanned: 1e10, execution_ms: 0, engine: 'e' });
    expect(wide).toBeGreaterThan(narrow);
  });

  it('expires spend out of the rolling window', async () => {
    let clock = 0;
    const store = new InMemoryWindowStore(() => clock);
    await store.add('t_alpha', { rows: 100, bytes: 0, units: 1, windowMs: 1000 });
    expect((await store.get('t_alpha', 1000)).rows_scanned).toBe(100);
    clock += 1500;
    expect((await store.get('t_alpha', 1000)).rows_scanned).toBe(0);
  });

  it('in-flight enforcement returns immediately when the query is already gone', async () => {
    const e = new QueryBudgetEnforcer({
      budget: { ...DEFAULT_QUERY_BUDGET, maxRowsScanned: 1 },
      store: new InMemoryWindowStore(),
      ledger,
      control: {
        engine: 't',
        isRunning: async () => false,
        requestKill: async () => true,
        progress: async () => undefined,
      },
      enforceInFlight: true,
    });
    const out = await e.enforceInFlight('t_alpha', 'q');
    expect(out.terminated).toBe(false);
    expect(out.polls).toBe(1);
  });
});