ktx/packages/context/src/scan/relationship-profiling.test.ts

import { readFile } from 'node:fs/promises';
import { join } from 'node:path';
import Database from 'better-sqlite3';
import { afterEach, describe, expect, it } from 'vitest';
import type { KloEnrichedColumn, KloEnrichedSchema, KloEnrichedTable } from './enrichment-types.js';
import { snapshotToKloEnrichedSchema } from './local-enrichment.js';
import { loadKloRelationshipBenchmarkFixture, maskKloRelationshipBenchmarkSnapshot } from './relationship-benchmarks.js';
import {
  createKloRelationshipProfileCache,
  formatKloRelationshipTableRef,
  profileKloRelationshipSchema,
  quoteKloRelationshipIdentifier,
} from './relationship-profiling.js';
import type { KloQueryResult, KloReadOnlyQueryInput, KloScanContext } from './types.js';

class InMemorySqliteExecutor {
  readonly db = new Database(':memory:');
  queryCount = 0;

  executeReadOnly(input: KloReadOnlyQueryInput, _ctx: KloScanContext): Promise<KloQueryResult> {
    this.queryCount += 1;
    const rows = this.db.prepare(input.sql).all() as Record<string, unknown>[];
    const headers = Object.keys(rows[0] ?? {});
    return Promise.resolve({
      headers,
      rows: rows.map((row) => headers.map((header) => row[header])),
      totalRows: rows.length,
      rowCount: rows.length,
    });
  }

  close(): void {
    this.db.close();
  }
}

class FileSqliteExecutor {
  readonly db: Database.Database;
  queryCount = 0;

  constructor(dataPath: string) {
    this.db = new Database(dataPath, { readonly: true, fileMustExist: true });
  }

  executeReadOnly(input: KloReadOnlyQueryInput, _ctx: KloScanContext): Promise<KloQueryResult> {
    this.queryCount += 1;
    const rows = this.db.prepare(input.sql).all() as Record<string, unknown>[];
    const headers = Object.keys(rows[0] ?? {});
    return Promise.resolve({
      headers,
      rows: rows.map((row) => headers.map((header) => row[header])),
      totalRows: rows.length,
      rowCount: rows.length,
    });
  }

  close(): void {
    this.db.close();
  }
}

function column(tableId: string, name: string, overrides: Partial<KloEnrichedColumn> = {}): KloEnrichedColumn {
  const tableRef = overrides.tableRef ?? { catalog: null, db: null, name: tableId };
  return {
    id: `${tableId}.${name}`,
    tableId,
    tableRef,
    name,
    nativeType: overrides.nativeType ?? 'INTEGER',
    normalizedType: overrides.normalizedType ?? 'integer',
    dimensionType: overrides.dimensionType ?? 'number',
    nullable: overrides.nullable ?? true,
    primaryKey: overrides.primaryKey ?? false,
    parentColumnId: null,
    descriptions: {},
    embedding: null,
    sampleValues: null,
    cardinality: null,
    ...overrides,
  };
}

function table(name: string, columns: KloEnrichedColumn[]): KloEnrichedTable {
  const ref = { catalog: null, db: null, name };
  return {
    id: name,
    ref,
    enabled: true,
    descriptions: {},
    columns: columns.map((item) => ({ ...item, tableId: name, tableRef: ref })),
  };
}

function schema(tables: KloEnrichedTable[]): KloEnrichedSchema {
  return { connectionId: 'warehouse', tables, relationships: [] };
}

describe('relationship profiling', () => {
  let executor: InMemorySqliteExecutor | null = null;

  afterEach(() => {
    executor?.close();
    executor = null;
  });

  it('keeps profiling on the batched table path', async () => {
    const source = await readFile(new URL('relationship-profiling.ts', import.meta.url), 'utf-8');

    expect(source).not.toMatch(new RegExp('queryColumn' + 'Profile'));
    expect(source).not.toMatch(/for \(const column of table\.columns\)[\s\S]*executeReadOnly/);
    expect(source).toMatch(/queryTableProfile/);
    expect(source).toMatch(/UNION ALL/);
  });

  it('quotes identifiers and formats table refs for supported local SQL drivers', () => {
    expect(quoteKloRelationshipIdentifier('sqlite', 'odd"name')).toBe('"odd""name"');
    expect(quoteKloRelationshipIdentifier('mysql', 'odd`name')).toBe('`odd``name`');
    expect(quoteKloRelationshipIdentifier('sqlserver', 'odd]name')).toBe('[odd]]name]');
    expect(formatKloRelationshipTableRef('sqlite', { catalog: null, db: null, name: 'accounts' })).toBe('"accounts"');
    expect(formatKloRelationshipTableRef('postgres', { catalog: null, db: 'analytics', name: 'accounts' })).toBe(
      '"analytics"."accounts"',
    );
  });

  it('profiles row count, null rate, uniqueness, sample values, and text lengths', async () => {
    executor = new InMemorySqliteExecutor();
    executor.db.exec(`
      CREATE TABLE accounts (id INTEGER, code TEXT, parent_id INTEGER);
      INSERT INTO accounts (id, code, parent_id) VALUES
        (1, 'A-1', NULL),
        (2, 'B-2', 1),
        (3, 'C-3', 1),
        (4, 'C-3', 2);
    `);

    const result = await profileKloRelationshipSchema({
      connectionId: 'warehouse',
      driver: 'sqlite',
      schema: schema([
        table('accounts', [
          column('accounts', 'id', { primaryKey: false, nullable: false }),
          column('accounts', 'code', { nativeType: 'TEXT', normalizedType: 'text', dimensionType: 'string' }),
          column('accounts', 'parent_id'),
        ]),
      ]),
      executor,
      ctx: { runId: 'profile-test' },
      sampleValuesPerColumn: 3,
    });

    expect(result.sqlAvailable).toBe(true);
    expect(result.queryCount).toBe(1);
    expect(executor.queryCount).toBe(1);
    expect(result.tables).toHaveLength(1);
    expect(result.tables[0]).toMatchObject({ table: { name: 'accounts' }, rowCount: 4 });
    expect(result.columns['accounts.id']).toMatchObject({
      table: { name: 'accounts' },
      column: 'id',
      rowCount: 4,
      nullCount: 0,
      distinctCount: 4,
      uniquenessRatio: 1,
      nullRate: 0,
      minTextLength: 1,
      maxTextLength: 1,
    });
    expect(result.columns['accounts.code']).toMatchObject({
      distinctCount: 3,
      uniquenessRatio: 0.75,
      sampleValues: ['C-3', 'A-1', 'B-2'],
      minTextLength: 3,
      maxTextLength: 3,
    });
    expect(result.columns['accounts.parent_id']).toMatchObject({
      nullCount: 1,
      distinctCount: 2,
      uniquenessRatio: 0.5,
      nullRate: 0.25,
    });
  });

  it('profiles each enabled table with one read-only SQL query', async () => {
    executor = new InMemorySqliteExecutor();
    executor.db.exec(`
      CREATE TABLE accounts (id INTEGER, code TEXT, parent_id INTEGER);
      CREATE TABLE users (id INTEGER, account_id INTEGER);
      INSERT INTO accounts (id, code, parent_id) VALUES
        (1, 'A-1', NULL),
        (2, 'B-2', 1),
        (3, 'C-3', 1),
        (4, 'C-3', 2);
      INSERT INTO users (id, account_id) VALUES
        (10, 1),
        (11, 1),
        (12, 2);
    `);

    const result = await profileKloRelationshipSchema({
      connectionId: 'warehouse',
      driver: 'sqlite',
      schema: schema([
        table('accounts', [
          column('accounts', 'id', { nullable: false }),
          column('accounts', 'code', { nativeType: 'TEXT', normalizedType: 'text', dimensionType: 'string' }),
          column('accounts', 'parent_id'),
        ]),
        table('users', [column('users', 'id', { nullable: false }), column('users', 'account_id')]),
      ]),
      executor,
      ctx: { runId: 'profile-batched-query-count' },
      sampleValuesPerColumn: 3,
    });

    expect(result.sqlAvailable).toBe(true);
    expect(result.queryCount).toBe(2);
    expect(executor.queryCount).toBe(2);
    expect(result.tables).toEqual([
      { table: { catalog: null, db: null, name: 'accounts' }, rowCount: 4 },
      { table: { catalog: null, db: null, name: 'users' }, rowCount: 3 },
    ]);
    expect(result.columns['accounts.code']).toMatchObject({
      distinctCount: 3,
      uniquenessRatio: 0.75,
      sampleValues: ['C-3', 'A-1', 'B-2'],
    });
    expect(result.columns['users.account_id']).toMatchObject({
      rowCount: 3,
      nullCount: 0,
      distinctCount: 2,
      uniquenessRatio: 2 / 3,
    });
  });

  it('bounds column profile statistics with profileSampleRows', async () => {
    const executor = new InMemorySqliteExecutor();
    executor.db.exec(`
      CREATE TABLE accounts (id INTEGER NOT NULL, account_code TEXT NOT NULL);
      INSERT INTO accounts VALUES (1, 'a1'), (2, 'a2'), (3, 'a3'), (4, 'a4');
    `);

    const profiles = await profileKloRelationshipSchema({
      connectionId: 'warehouse',
      driver: 'sqlite',
      schema: schema([
        table('accounts', [
          column('accounts', 'id', { nullable: false }),
          column('accounts', 'account_code', {
            nativeType: 'TEXT',
            normalizedType: 'text',
            dimensionType: 'string',
            nullable: false,
          }),
        ]),
      ]),
      executor,
      ctx: { runId: 'profile-sample-rows' },
      profileSampleRows: 2,
    });

    expect(profiles.queryCount).toBe(1);
    expect(executor.queryCount).toBe(1);
    expect(profiles.tables).toEqual([{ table: { catalog: null, db: null, name: 'accounts' }, rowCount: 4 }]);
    expect(profiles.columns['accounts.id']).toMatchObject({
      rowCount: 2,
      distinctCount: 2,
      uniquenessRatio: 1,
    });
    expect(profiles.columns['accounts.account_code']?.sampleValues).toEqual(['a1', 'a2']);

    executor.close();
  });

  it('reuses a profile cache inside one scan run but re-queries with a fresh cache', async () => {
    executor = new InMemorySqliteExecutor();
    executor.db.exec(`
      CREATE TABLE accounts (id INTEGER NOT NULL, account_code TEXT NOT NULL);
      INSERT INTO accounts VALUES (1, 'a1'), (2, 'a2'), (3, 'a2');
    `);
    const relationshipSchema = schema([
      table('accounts', [
        column('accounts', 'id', { nullable: false }),
        column('accounts', 'account_code', {
          nativeType: 'TEXT',
          normalizedType: 'text',
          dimensionType: 'string',
          nullable: false,
        }),
      ]),
    ]);
    const cache = createKloRelationshipProfileCache();

    const first = await profileKloRelationshipSchema({
      connectionId: 'warehouse',
      driver: 'sqlite',
      schema: relationshipSchema,
      executor,
      ctx: { runId: 'profile-cache-run' },
      cache,
    });
    const second = await profileKloRelationshipSchema({
      connectionId: 'warehouse',
      driver: 'sqlite',
      schema: relationshipSchema,
      executor,
      ctx: { runId: 'profile-cache-run' },
      cache,
    });
    const third = await profileKloRelationshipSchema({
      connectionId: 'warehouse',
      driver: 'sqlite',
      schema: relationshipSchema,
      executor,
      ctx: { runId: 'profile-cache-fresh-run' },
      cache: createKloRelationshipProfileCache(),
    });

    expect(first.queryCount).toBe(1);
    expect(second.queryCount).toBe(0);
    expect(third.queryCount).toBe(1);
    expect(executor.queryCount).toBe(2);
    expect(second.tables).toEqual(first.tables);
    expect(second.columns).toEqual(first.columns);
  });

  it('profiles the checked-in scale stress fixture with one query per table', async () => {
    const fixtureRoot = new URL('../../test/fixtures/relationship-benchmarks/', import.meta.url);
    const fixture = await loadKloRelationshipBenchmarkFixture(join(fixtureRoot.pathname, 'scale_stress_no_declared_constraints'));
    if (!fixture.dataPath) {
      throw new Error('scale_stress_no_declared_constraints is missing data.sqlite');
    }
    const maskedSnapshot = maskKloRelationshipBenchmarkSnapshot(
      fixture.snapshot,
      'declared_pks_and_declared_fks_removed',
    );
    const scaleExecutor = new FileSqliteExecutor(fixture.dataPath);
    try {
      const result = await profileKloRelationshipSchema({
        connectionId: fixture.snapshot.connectionId,
        driver: fixture.snapshot.driver,
        schema: snapshotToKloEnrichedSchema(maskedSnapshot, new Map()),
        executor: scaleExecutor,
        ctx: { runId: 'scale-stress-profile-query-count' },
        profileSampleRows: 3,
      });

      expect(fixture.snapshot.tables).toHaveLength(400);
      expect(result.queryCount).toBe(400);
      expect(result.queryCount).toBeLessThanOrEqual(2 * fixture.snapshot.tables.length);
      expect(scaleExecutor.queryCount).toBe(400);
    } finally {
      scaleExecutor.close();
    }
  });
});