ktx/packages/context/src/scan/relationship-composite-candidates.ts

import type { KloEnrichedColumn, KloEnrichedSchema, KloEnrichedTable, KloRelationshipType } from './enrichment-types.js';
import {
  formatKloRelationshipTableRef,
  quoteKloRelationshipIdentifier,
  type KloRelationshipProfileArtifact,
  type KloRelationshipReadOnlyExecutor,
} from './relationship-profiling.js';
import type { KloConnectionDriver, KloQueryResult, KloScanContext, KloTableRef } from './types.js';

export type KloCompositeRelationshipStatus = 'accepted' | 'review' | 'rejected';

export interface KloCompositeRelationshipTupleEndpoint {
  tableId: string;
  columnIds: string[];
  table: KloTableRef;
  columns: string[];
}

export interface KloCompositePrimaryKeyCandidate {
  id: string;
  tableId: string;
  table: KloTableRef;
  columns: string[];
  columnIds: string[];
  score: number;
  status: KloCompositeRelationshipStatus;
  evidence: {
    rowCount: number;
    distinctCount: number;
    uniquenessRatio: number;
    nullRate: number;
    reasons: string[];
  };
}

export interface KloCompositeRelationshipValidationEvidence {
  targetUniqueness: number;
  sourceCoverage: number;
  violationCount: number;
  violationRatio: number;
  childDistinct: number;
  parentDistinct: number;
  overlap: number;
  reasons: string[];
}

export interface KloCompositeRelationshipCandidate {
  id: string;
  from: KloCompositeRelationshipTupleEndpoint;
  to: KloCompositeRelationshipTupleEndpoint;
  relationshipType: KloRelationshipType;
  confidence: number;
  status: KloCompositeRelationshipStatus;
  source: 'composite_profile_match';
  validation: KloCompositeRelationshipValidationEvidence;
}

export interface DiscoverKloCompositeRelationshipsInput {
  connectionId: string;
  driver: KloConnectionDriver;
  schema: KloEnrichedSchema;
  profiles: KloRelationshipProfileArtifact;
  executor: KloRelationshipReadOnlyExecutor | null;
  ctx: KloScanContext;
  maxCompositeWidth?: number;
  maxColumnsPerTable?: number;
  minPrimaryKeyUniqueness?: number;
  minSourceCoverage?: number;
  maxViolationRatio?: number;
}

export interface DiscoverKloCompositeRelationshipsResult {
  primaryKeys: KloCompositePrimaryKeyCandidate[];
  relationships: KloCompositeRelationshipCandidate[];
  queryCount: number;
  warnings: string[];
}

const KEY_NAME_PARTS = new Set(['id', 'key', 'code', 'number', 'num', 'line', 'warehouse', 'account', 'order']);
const DEFAULT_MAX_COMPOSITE_WIDTH = 3;
const DEFAULT_MAX_COLUMNS_PER_TABLE = 8;
const DEFAULT_MIN_PRIMARY_KEY_UNIQUENESS = 0.98;
const DEFAULT_MIN_SOURCE_COVERAGE = 0.9;
const DEFAULT_MAX_VIOLATION_RATIO = 0.01;

function enabledTables(schema: KloEnrichedSchema): KloEnrichedTable[] {
  return schema.tables.filter((table) => table.enabled);
}

function tableRowCount(profiles: KloRelationshipProfileArtifact, tableName: string): number {
  return profiles.tables.find((item) => item.table.name === tableName)?.rowCount ?? 0;
}

function profileKey(tableName: string, columnName: string): string {
  return `${tableName}.${columnName}`;
}

function profileNullRate(profiles: KloRelationshipProfileArtifact, tableName: string, columnName: string): number {
  return profiles.columns[profileKey(tableName, columnName)]?.nullRate ?? 1;
}

function normalizedColumnName(name: string): string {
  return name
    .toLowerCase()
    .replace(/[^a-z0-9]+/gu, '_')
    .replace(/^_+|_+$/gu, '');
}

function columnNameScore(column: KloEnrichedColumn): number {
  const parts = normalizedColumnName(column.name).split('_').filter(Boolean);
  if (parts.some((part) => KEY_NAME_PARTS.has(part))) {
    return 1;
  }
  return 0;
}

function nameParts(name: string): string[] {
  return normalizedColumnName(name).split('_').filter(Boolean);
}

function keyLikeTableNameParts(tableName: string): Set<string> {
  return new Set(nameParts(tableName).filter((part) => KEY_NAME_PARTS.has(part)));
}

function tupleCoversTableNameKeyParts(tableName: string, columns: readonly KloEnrichedColumn[]): boolean {
  const required = keyLikeTableNameParts(tableName);
  if (required.size === 0) {
    return true;
  }
  const columnParts = new Set(columns.flatMap((column) => nameParts(column.name)));
  return Array.from(required).every((part) => columnParts.has(part));
}

function candidateKeyColumns(input: {
  table: KloEnrichedTable;
  profiles: KloRelationshipProfileArtifact;
  maxColumnsPerTable: number;
}): KloEnrichedColumn[] {
  return input.table.columns
    .map((column, index) => ({ column, index }))
    .filter(({ column }) => {
      if (column.dimensionType === 'time' || column.dimensionType === 'boolean') {
        return false;
      }
      const profile = input.profiles.columns[profileKey(input.table.ref.name, column.name)];
      return Boolean(profile) && profile!.nullRate <= 0.02 && columnNameScore(column) > 0;
    })
    .sort(
      (left, right) =>
        columnNameScore(right.column) - columnNameScore(left.column) || left.index - right.index,
    )
    .slice(0, input.maxColumnsPerTable)
    .map(({ column }) => column);
}

function hasStrongSingleColumnKey(input: {
  table: KloEnrichedTable;
  profiles: KloRelationshipProfileArtifact;
  minPrimaryKeyUniqueness: number;
}): boolean {
  return input.table.columns.some((column) => {
    if (column.dimensionType === 'time' || column.dimensionType === 'boolean' || columnNameScore(column) === 0) {
      return false;
    }
    const profile = input.profiles.columns[profileKey(input.table.ref.name, column.name)];
    return Boolean(profile) && profile!.nullRate <= 0.02 && profile!.uniquenessRatio >= input.minPrimaryKeyUniqueness;
  });
}

function combinations<T>(values: readonly T[], width: number): T[][] {
  if (width <= 0) {
    return [[]];
  }
  if (values.length < width) {
    return [];
  }
  const output: T[][] = [];
  values.forEach((value, index) => {
    for (const tail of combinations(values.slice(index + 1), width - 1)) {
      output.push([value, ...tail]);
    }
  });
  return output;
}

function tupleKey(tableName: string, columns: readonly string[]): string {
  return `${tableName}.(${columns.join(',')})`;
}

function relationshipKey(input: {
  fromTable: string;
  fromColumns: readonly string[];
  toTable: string;
  toColumns: readonly string[];
}): string {
  return `${tupleKey(input.fromTable, input.fromColumns)}->${tupleKey(input.toTable, input.toColumns)}`;
}

function tupleEndpoint(table: KloEnrichedTable, columns: readonly KloEnrichedColumn[]): KloCompositeRelationshipTupleEndpoint {
  return {
    tableId: table.id,
    columnIds: columns.map((column) => column.id),
    table: table.ref,
    columns: columns.map((column) => column.name),
  };
}

function row(result: KloQueryResult): unknown[] {
  return result.rows[0] ?? [];
}

function numberAt(result: KloQueryResult, header: string): number {
  const index = result.headers.findIndex((candidate) => candidate.toLowerCase() === header.toLowerCase());
  const value = row(result)[index];
  if (typeof value === 'number') {
    return value;
  }
  if (typeof value === 'bigint') {
    return Number(value);
  }
  if (typeof value === 'string' && value.trim() !== '') {
    return Number(value);
  }
  return 0;
}

function topSql(driver: KloConnectionDriver, limit: number): string {
  if (driver === 'sqlserver') {
    return ` TOP (${Math.max(1, Math.floor(limit))})`;
  }
  return '';
}

function limitSql(driver: KloConnectionDriver, limit: number): string {
  if (driver === 'sqlserver') {
    return '';
  }
  return ` LIMIT ${Math.max(1, Math.floor(limit))}`;
}

function aliasedTupleSelect(driver: KloConnectionDriver, columns: readonly string[]): string {
  return columns
    .map((column, index) => `${quoteKloRelationshipIdentifier(driver, column)} AS c${index}`)
    .join(', ');
}

function nonNullPredicate(driver: KloConnectionDriver, columns: readonly string[]): string {
  return columns.map((column) => `${quoteKloRelationshipIdentifier(driver, column)} IS NOT NULL`).join(' AND ');
}

function tupleEquality(columns: number): string {
  return Array.from({ length: columns }, (_, index) => `child_values.c${index} = parent_values.c${index}`).join(
    ' AND ',
  );
}

function buildTupleDistinctSql(input: {
  driver: KloConnectionDriver;
  table: KloTableRef;
  columns: readonly string[];
}): string {
  const tableSql = formatKloRelationshipTableRef(input.driver, input.table);
  return [
    'WITH tuple_values AS (',
    `SELECT DISTINCT ${aliasedTupleSelect(input.driver, input.columns)} FROM ${tableSql}`,
    `WHERE ${nonNullPredicate(input.driver, input.columns)}`,
    ')',
    'SELECT COUNT(*) AS distinct_count FROM tuple_values',
  ].join(' ');
}

function buildCompositeCoverageSql(input: {
  driver: KloConnectionDriver;
  childTable: KloTableRef;
  childColumns: readonly string[];
  parentTable: KloTableRef;
  parentColumns: readonly string[];
  maxDistinctSourceValues: number;
}): string {
  const childTableSql = formatKloRelationshipTableRef(input.driver, input.childTable);
  const parentTableSql = formatKloRelationshipTableRef(input.driver, input.parentTable);
  const top = topSql(input.driver, input.maxDistinctSourceValues);
  const limit = limitSql(input.driver, input.maxDistinctSourceValues);
  return [
    'WITH child_values AS (',
    `SELECT DISTINCT${top} ${aliasedTupleSelect(input.driver, input.childColumns)} FROM ${childTableSql}`,
    `WHERE ${nonNullPredicate(input.driver, input.childColumns)}${limit}`,
    '), parent_values AS (',
    `SELECT DISTINCT ${aliasedTupleSelect(input.driver, input.parentColumns)} FROM ${parentTableSql}`,
    `WHERE ${nonNullPredicate(input.driver, input.parentColumns)}`,
    ')',
    'SELECT',
    '(SELECT COUNT(*) FROM child_values) AS child_distinct,',
    '(SELECT COUNT(*) FROM parent_values) AS parent_distinct,',
    'SUM(CASE WHEN parent_values.c0 IS NOT NULL THEN 1 ELSE 0 END) AS overlap,',
    'SUM(CASE WHEN parent_values.c0 IS NULL THEN 1 ELSE 0 END) AS violation_count',
    'FROM child_values',
    `LEFT JOIN parent_values ON ${tupleEquality(input.childColumns.length)}`,
  ].join(' ');
}

function relationshipStatus(input: {
  targetUniqueness: number;
  sourceCoverage: number;
  violationRatio: number;
  minSourceCoverage: number;
  maxViolationRatio: number;
}): KloCompositeRelationshipStatus {
  if (
    input.targetUniqueness >= DEFAULT_MIN_PRIMARY_KEY_UNIQUENESS &&
    input.sourceCoverage >= input.minSourceCoverage &&
    input.violationRatio <= input.maxViolationRatio
  ) {
    return 'accepted';
  }
  if (input.sourceCoverage >= 0.55) {
    return 'review';
  }
  return 'rejected';
}

function hasAcceptedSubset(
  accepted: readonly KloCompositePrimaryKeyCandidate[],
  tableName: string,
  columns: readonly string[],
): boolean {
  const columnSet = new Set(columns);
  return accepted.some(
    (candidate) =>
      candidate.table.name === tableName &&
      candidate.columns.length < columns.length &&
      candidate.columns.every((column) => columnSet.has(column)),
  );
}

async function detectCompositePrimaryKeys(input: {
  connectionId: string;
  driver: KloConnectionDriver;
  table: KloEnrichedTable;
  profiles: KloRelationshipProfileArtifact;
  executor: KloRelationshipReadOnlyExecutor;
  ctx: KloScanContext;
  maxCompositeWidth: number;
  maxColumnsPerTable: number;
  minPrimaryKeyUniqueness: number;
}): Promise<{ primaryKeys: KloCompositePrimaryKeyCandidate[]; queryCount: number }> {
  const rowCount = tableRowCount(input.profiles, input.table.ref.name);
  if (rowCount === 0) {
    return { primaryKeys: [], queryCount: 0 };
  }
  if (
    hasStrongSingleColumnKey({
      table: input.table,
      profiles: input.profiles,
      minPrimaryKeyUniqueness: input.minPrimaryKeyUniqueness,
    })
  ) {
    return { primaryKeys: [], queryCount: 0 };
  }

  const columns = candidateKeyColumns({
    table: input.table,
    profiles: input.profiles,
    maxColumnsPerTable: input.maxColumnsPerTable,
  });
  const primaryKeys: KloCompositePrimaryKeyCandidate[] = [];
  let queryCount = 0;

  for (let width = 2; width <= input.maxCompositeWidth; width += 1) {
    for (const columnTuple of combinations(columns, width)) {
      const columnNames = columnTuple.map((column) => column.name);
      if (!tupleCoversTableNameKeyParts(input.table.ref.name, columnTuple)) {
        continue;
      }
      if (hasAcceptedSubset(primaryKeys, input.table.ref.name, columnNames)) {
        continue;
      }
      const result = await input.executor.executeReadOnly(
        {
          connectionId: input.connectionId,
          sql: buildTupleDistinctSql({
            driver: input.driver,
            table: input.table.ref,
            columns: columnNames,
          }),
          maxRows: 1,
        },
        input.ctx,
      );
      queryCount += 1;
      const distinctCount = numberAt(result, 'distinct_count');
      const uniquenessRatio = rowCount === 0 ? 0 : distinctCount / rowCount;
      if (uniquenessRatio < input.minPrimaryKeyUniqueness) {
        continue;
      }
      const nullRate = Math.max(
        ...columnNames.map((columnName) => profileNullRate(input.profiles, input.table.ref.name, columnName)),
      );
      primaryKeys.push({
        id: tupleKey(input.table.ref.name, columnNames),
        tableId: input.table.id,
        table: input.table.ref,
        columns: columnNames,
        columnIds: columnTuple.map((column) => column.id),
        score: Number(Math.min(0.99, 0.72 + uniquenessRatio * 0.22 + (1 - nullRate) * 0.06).toFixed(3)),
        status: 'accepted',
        evidence: {
          rowCount,
          distinctCount,
          uniquenessRatio,
          nullRate,
          reasons: ['composite_unique_tuple', 'not_null_profile'],
        },
      });
    }
  }

  return {
    primaryKeys: primaryKeys.sort((left, right) =>
      tupleKey(left.table.name, left.columns).localeCompare(tupleKey(right.table.name, right.columns)),
    ),
    queryCount,
  };
}

function columnsByName(table: KloEnrichedTable): Map<string, KloEnrichedColumn> {
  return new Map(table.columns.map((column) => [column.name, column]));
}

function compatibleTuple(sourceColumns: readonly KloEnrichedColumn[], targetColumns: readonly KloEnrichedColumn[]): boolean {
  if (sourceColumns.length !== targetColumns.length) {
    return false;
  }
  return sourceColumns.every((source, index) => {
    const target = targetColumns[index];
    return Boolean(target) && source.dimensionType === target.dimensionType;
  });
}

async function validateCompositeRelationship(input: {
  connectionId: string;
  driver: KloConnectionDriver;
  sourceTable: KloEnrichedTable;
  sourceColumns: readonly KloEnrichedColumn[];
  targetKey: KloCompositePrimaryKeyCandidate;
  targetTable: KloEnrichedTable;
  targetColumns: readonly KloEnrichedColumn[];
  executor: KloRelationshipReadOnlyExecutor;
  ctx: KloScanContext;
  minSourceCoverage: number;
  maxViolationRatio: number;
}): Promise<{ relationship: KloCompositeRelationshipCandidate; queryCount: number }> {
  const result = await input.executor.executeReadOnly(
    {
      connectionId: input.connectionId,
      sql: buildCompositeCoverageSql({
        driver: input.driver,
        childTable: input.sourceTable.ref,
        childColumns: input.sourceColumns.map((column) => column.name),
        parentTable: input.targetTable.ref,
        parentColumns: input.targetColumns.map((column) => column.name),
        maxDistinctSourceValues: 10000,
      }),
      maxRows: 1,
    },
    input.ctx,
  );
  const childDistinct = numberAt(result, 'child_distinct');
  const parentDistinct = numberAt(result, 'parent_distinct');
  const overlap = numberAt(result, 'overlap');
  const violationCount = numberAt(result, 'violation_count');
  const sourceCoverage = childDistinct === 0 ? 0 : overlap / childDistinct;
  const violationRatio = childDistinct === 0 ? 1 : violationCount / childDistinct;
  const targetUniqueness = input.targetKey.evidence.uniquenessRatio;
  const status = relationshipStatus({
    targetUniqueness,
    sourceCoverage,
    violationRatio,
    minSourceCoverage: input.minSourceCoverage,
    maxViolationRatio: input.maxViolationRatio,
  });

  const from = tupleEndpoint(input.sourceTable, input.sourceColumns);
  const to = {
    tableId: input.targetKey.tableId,
    columnIds: input.targetKey.columnIds,
    table: input.targetKey.table,
    columns: input.targetKey.columns,
  };
  const reasons =
    status === 'accepted'
      ? ['composite_validation_passed']
      : [
          'composite_validation_failed',
          sourceCoverage < input.minSourceCoverage ? 'low_source_coverage' : '',
          violationRatio > input.maxViolationRatio ? 'excessive_violations' : '',
        ].filter(Boolean);

  return {
    queryCount: 1,
    relationship: {
      id: relationshipKey({
        fromTable: from.table.name,
        fromColumns: from.columns,
        toTable: to.table.name,
        toColumns: to.columns,
      }),
      from,
      to,
      relationshipType: 'many_to_one',
      confidence: status === 'accepted' ? 0.95 : 0.62,
      status,
      source: 'composite_profile_match',
      validation: {
        targetUniqueness,
        sourceCoverage,
        violationCount,
        violationRatio,
        childDistinct,
        parentDistinct,
        overlap,
        reasons,
      },
    },
  };
}

export async function discoverKloCompositeRelationships(
  input: DiscoverKloCompositeRelationshipsInput,
): Promise<DiscoverKloCompositeRelationshipsResult> {
  if (!input.executor || !input.profiles.sqlAvailable) {
    return {
      primaryKeys: [],
      relationships: [],
      queryCount: 0,
      warnings: ['composite_relationship_validation_unavailable'],
    };
  }

  const settings = {
    maxCompositeWidth: input.maxCompositeWidth ?? DEFAULT_MAX_COMPOSITE_WIDTH,
    maxColumnsPerTable: input.maxColumnsPerTable ?? DEFAULT_MAX_COLUMNS_PER_TABLE,
    minPrimaryKeyUniqueness: input.minPrimaryKeyUniqueness ?? DEFAULT_MIN_PRIMARY_KEY_UNIQUENESS,
    minSourceCoverage: input.minSourceCoverage ?? DEFAULT_MIN_SOURCE_COVERAGE,
    maxViolationRatio: input.maxViolationRatio ?? DEFAULT_MAX_VIOLATION_RATIO,
  };
  const tables = enabledTables(input.schema);
  const tableByName = new Map(tables.map((table) => [table.ref.name, table]));
  const primaryKeys: KloCompositePrimaryKeyCandidate[] = [];
  let queryCount = 0;

  for (const table of tables) {
    const result = await detectCompositePrimaryKeys({
      connectionId: input.connectionId,
      driver: input.driver,
      table,
      profiles: input.profiles,
      executor: input.executor,
      ctx: input.ctx,
      maxCompositeWidth: settings.maxCompositeWidth,
      maxColumnsPerTable: settings.maxColumnsPerTable,
      minPrimaryKeyUniqueness: settings.minPrimaryKeyUniqueness,
    });
    primaryKeys.push(...result.primaryKeys);
    queryCount += result.queryCount;
  }

  const relationships: KloCompositeRelationshipCandidate[] = [];
  for (const targetKey of primaryKeys) {
    const targetTable = tableByName.get(targetKey.table.name);
    if (!targetTable) {
      continue;
    }
    const targetColumnByName = columnsByName(targetTable);
    const targetColumns = targetKey.columns.flatMap((columnName) => {
      const column = targetColumnByName.get(columnName);
      return column ? [column] : [];
    });
    if (targetColumns.length !== targetKey.columns.length) {
      continue;
    }

    for (const sourceTable of tables) {
      if (sourceTable.id === targetTable.id) {
        continue;
      }
      const sourceColumnByName = columnsByName(sourceTable);
      const sourceColumns = targetKey.columns.flatMap((columnName) => {
        const column = sourceColumnByName.get(columnName);
        return column ? [column] : [];
      });
      if (sourceColumns.length !== targetKey.columns.length || !compatibleTuple(sourceColumns, targetColumns)) {
        continue;
      }

      const result = await validateCompositeRelationship({
        connectionId: input.connectionId,
        driver: input.driver,
        sourceTable,
        sourceColumns,
        targetKey,
        targetTable,
        targetColumns,
        executor: input.executor,
        ctx: input.ctx,
        minSourceCoverage: settings.minSourceCoverage,
        maxViolationRatio: settings.maxViolationRatio,
      });
      queryCount += result.queryCount;
      if (result.relationship.status !== 'rejected') {
        relationships.push(result.relationship);
      }
    }
  }

  return {
    primaryKeys: primaryKeys.sort((left, right) => left.id.localeCompare(right.id)),
    relationships: relationships.sort((left, right) => left.id.localeCompare(right.id)),
    queryCount,
    warnings: [],
  };
}