0pt/src/services/source-query-service.ts

/**
 * src/services/source-query-service.ts
 * =======================================
 * 0pt.query_source — bounded, timeout-guarded synchronous query execution
 * against a registered Source (source_type = "database",
 * adapter_id = "fabric_onelake_sql" today; extensible via `ADAPTER_REGISTRY`).
 *
 * This is deliberately NOT a general-purpose query engine: it enforces a
 * max-rows and max-time ceiling on every call and is meant for small,
 * interactive/preview-style reads. Anything heavier (large exports, long
 * scans, scheduled pulls) must go through N0D as a job — see
 * docs/P0_QUERY_SOURCE_FLOW.md — the same "synchronous vs. N0D job" split
 * `quality-check-service.ts` uses for quality checks.
 *
 * Reuses the existing SourceService (src/services/source-service.ts) purely
 * through its public `getSource()` — this file does not modify that class or
 * the Source/CreateSourceInput types, which already carry the `adapter_id`
 * field this capability depends on.
 */

import type { Source } from "../types/source";
import { SourceService } from "./source-service";
import { PiiGateService } from "./pii-gate-service";
import { createAud0Client, type Aud0Client } from "../clients/aud0-client";
import { FabricOneLakeAdapter, type FabricOneLakeConnectionConfig, type AdapterQueryResult } from "../adapters/fabric-onelake-adapter";

/** Default bound applied when the caller does not specify one. */
export const SOURCE_QUERY_DEFAULT_MAX_ROWS = 1000;
/** Absolute ceiling — a caller cannot raise max_rows above this on the synchronous path. */
export const SOURCE_QUERY_HARD_MAX_ROWS = 10_000;
/** Default timeout applied when the caller does not specify one. */
export const SOURCE_QUERY_DEFAULT_TIMEOUT_MS = 10_000;
/** Absolute ceiling — a caller cannot raise timeout_ms above this on the synchronous path. */
export const SOURCE_QUERY_HARD_MAX_TIMEOUT_MS = 30_000;

/** Adapter ids supported by the synchronous query path today. */
export type SupportedQueryAdapterId = "fabric_onelake_sql";

export class SourceQueryTimeoutError extends Error {}
export class SourceQueryUnsupportedError extends Error {}

export interface QuerySourceInput {
  source_id: string;
  query: string;
  params?: Record<string, unknown>;
  tenant_id: string;
  project_id: string;
  actor_id?: string;
  correlation_id?: string;
  /** Clamped to [1, SOURCE_QUERY_HARD_MAX_ROWS]. Defaults to SOURCE_QUERY_DEFAULT_MAX_ROWS. */
  max_rows?: number;
  /** Clamped to [1, SOURCE_QUERY_HARD_MAX_TIMEOUT_MS]. Defaults to SOURCE_QUERY_DEFAULT_TIMEOUT_MS. */
  timeout_ms?: number;
  /**
   * Connection config for the source's adapter. Never persisted by this
   * service and never defaulted from an env var — tenant-agnostic by
   * construction, same as the adapter itself. Today only the
   * "fabric_onelake_sql" shape is supported.
   */
  connection_config: FabricOneLakeConnectionConfig;
}

export interface QuerySourceResult {
  source_id: string;
  rows: Record<string, unknown>[];
  row_count: number;
  truncated: boolean;
  duration_ms: number;
  bounds: { max_rows: number; timeout_ms: number };
}

function clamp(value: number, min: number, max: number): number {
  if (!Number.isFinite(value)) return min;
  return Math.min(Math.max(value, min), max);
}

export class SourceQueryService {
  private sourceService: SourceService;
  private aud0: Aud0Client;
  /** GDPR guardrail — on by default, can be shared via opts so callers use one registry. */
  private piiGate: PiiGateService;

  constructor(opts?: { sourceService?: SourceService; aud0BaseUrl?: string; piiGate?: PiiGateService }) {
    // Accepts an existing SourceService so callers (e.g. the API route) share
    // one in-memory registry instead of each constructing an empty one.
    this.sourceService = opts?.sourceService ?? new SourceService();
    this.aud0 = createAud0Client(opts?.aud0BaseUrl);
    this.piiGate = opts?.piiGate ?? new PiiGateService({ aud0BaseUrl: opts?.aud0BaseUrl });
  }

  async querySource(input: QuerySourceInput): Promise<QuerySourceResult> {
    const startedAt = Date.now();

    // GDPR fail-closed: personal-data sources (self-flagged or known directory
    // domains such as ratsit.se) may not be processed without a PiiPolicy.
    // Runs BEFORE the capability/type gate so the lawful-basis error is the
    // one the caller sees for personal data.
    const preflight = await this.sourceService.getSource(input.source_id);
    await this.piiGate.assertCanStore(preflight, [], input.tenant_id);

    const source = await this.getQueryableSource(input.source_id);

    const maxRows = clamp(input.max_rows ?? SOURCE_QUERY_DEFAULT_MAX_ROWS, 1, SOURCE_QUERY_HARD_MAX_ROWS);
    const timeoutMs = clamp(input.timeout_ms ?? SOURCE_QUERY_DEFAULT_TIMEOUT_MS, 1, SOURCE_QUERY_HARD_MAX_TIMEOUT_MS);
    const correlationId = input.correlation_id ?? `corr_0pt_query_source_${input.source_id}_${startedAt}`;

    const adapter = this.buildAdapter(source, input.connection_config);
    try {
      const { rows, durationMs } = await this.runWithTimeout(adapter.query(input.query, input.params), timeoutMs);

      const truncated = rows.length > maxRows;
      const boundedRows = truncated ? rows.slice(0, maxRows) : rows;

      await this.emitQueriedEvent({
        source,
        input,
        correlationId,
        rowCount: boundedRows.length,
        truncated,
        durationMs,
        ok: true,
      });

      return {
        source_id: source.source_id,
        rows: boundedRows,
        row_count: boundedRows.length,
        truncated,
        duration_ms: durationMs,
        bounds: { max_rows: maxRows, timeout_ms: timeoutMs },
      };
    } catch (err) {
      const message = err instanceof Error ? err.message : String(err);
      await this.emitQueriedEvent({
        source,
        input,
        correlationId,
        rowCount: 0,
        truncated: false,
        durationMs: Date.now() - startedAt,
        ok: false,
        error: message,
      });
      throw err instanceof Error ? err : new Error(message);
    } finally {
      // Each call opens its own short-lived adapter/pool (see buildAdapter) —
      // always release it, whether the query succeeded, failed, or timed out.
      await adapter.close().catch(() => {});
    }
  }

  private async getQueryableSource(sourceId: string): Promise<Source> {
    const source = await this.sourceService.getSource(sourceId);
    if (source.source_type !== "database") {
      throw new SourceQueryUnsupportedError(
        `Source ${sourceId} has source_type "${source.source_type}", not "database" — the synchronous 0pt.query_source path only ` +
          `supports database sources. Read non-database sources through a N0D job instead.`,
      );
    }
    if (source.adapter_id !== "fabric_onelake_sql") {
      throw new SourceQueryUnsupportedError(
        `Source ${sourceId} has adapter_id "${source.adapter_id ?? "(none)"}" — the synchronous 0pt.query_source path currently only ` +
          `supports adapter_id "fabric_onelake_sql". Register the source with that adapter_id, or run the read through N0D as a job instead.`,
      );
    }
    return source;
  }

  private buildAdapter(source: Source, connectionConfig: FabricOneLakeConnectionConfig): FabricOneLakeAdapter {
    if (!connectionConfig) {
      throw new Error(`Source ${source.source_id}: connection_config is required to query a fabric_onelake_sql source`);
    }
    return new FabricOneLakeAdapter(connectionConfig);
  }

  /**
   * Races the adapter query against a timeout. Note this bounds *waiting*
   * time on the caller's side — it does not cancel an in-flight TDS request
   * on timeout (the `mssql` driver has no cooperative cancellation hook used
   * here), so a timed-out query may still complete against the warehouse
   * after this call has already thrown. That is an accepted trade-off for a
   * synchronous, best-effort preview path; anything where that matters
   * belongs on the N0D job path instead.
   */
  private async runWithTimeout<T>(promise: Promise<AdapterQueryResult<T>>, timeoutMs: number): Promise<AdapterQueryResult<T>> {
    let timer: ReturnType<typeof setTimeout> | undefined;
    const timeout = new Promise<never>((_, reject) => {
      timer = setTimeout(() => {
        reject(
          new SourceQueryTimeoutError(
            `source query exceeded the bounded synchronous timeout of ${timeoutMs}ms — run this query through N0D as a job instead of the synchronous 0pt.query_source path`,
          ),
        );
      }, timeoutMs);
    });
    try {
      return await Promise.race([promise, timeout]);
    } finally {
      clearTimeout(timer);
    }
  }

  private async emitQueriedEvent(args: {
    source: Source;
    input: QuerySourceInput;
    correlationId: string;
    rowCount: number;
    truncated: boolean;
    durationMs: number;
    ok: boolean;
    error?: string;
  }): Promise<void> {
    const { source, input, correlationId, rowCount, truncated, durationMs, ok, error } = args;
    await this.aud0.ingestEvent({
      event_id: `evt_0pt_source_queried_${source.source_id}_${Date.now()}`,
      event_type: "0pt.source_queried",
      tenant_id: input.tenant_id,
      project_id: input.project_id,
      actor_id: input.actor_id ?? "system",
      app_id: "0pt",
      capability_id: "0pt.query_source",
      resource_type: "source",
      resource_id: source.source_id,
      risk_level: "L4",
      payload: {
        adapter_id: source.adapter_id,
        ok,
        row_count: rowCount,
        truncated,
        duration_ms: durationMs,
        query_preview: input.query.slice(0, 200),
        ...(error ? { error } : {}),
      },
      created_at: new Date().toISOString(),
      correlation_id: correlationId,
    });
  }
}