0pt/tests/p0/source-query-service.test.ts

import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
import { SourceService } from "@/src/services/source-service";
import type { SourceType } from "@/src/types/source";

const queryMock = vi.fn();
const closeMock = vi.fn().mockResolvedValue(undefined);

// Hoisted by Vitest above the static imports below, so
// source-query-service.ts (which statically imports FabricOneLakeAdapter)
// sees this mock too — no real mssql/AAD network calls happen in this file.
vi.mock("@/src/adapters/fabric-onelake-adapter", () => {
  return {
    // Must be a real `function` (not an arrow fn) so `new FabricOneLakeAdapter(...)`
    // in source-query-service.ts works: a constructor that explicitly returns an
    // object substitutes that object for `this`, per JS semantics.
    FabricOneLakeAdapter: vi.fn().mockImplementation(function FakeFabricOneLakeAdapter() {
      return { query: queryMock, close: closeMock };
    }),
  };
});

import {
  SourceQueryService,
  SourceQueryTimeoutError,
  SourceQueryUnsupportedError,
  SOURCE_QUERY_HARD_MAX_ROWS,
  SOURCE_QUERY_HARD_MAX_TIMEOUT_MS,
} from "@/src/services/source-query-service";

const BASE_CONNECTION_CONFIG = {
  credentials: { tenantId: "t", clientId: "c", clientSecret: "s" },
  sqlEndpoint: "ep.datawarehouse.fabric.microsoft.com",
  database: "lh_test",
};

async function registerSource(
  sourceService: SourceService,
  opts: { source_id: string; adapter_id?: string; source_type?: SourceType },
): Promise<string> {
  await sourceService.createSource({
    source_id: opts.source_id,
    tenant_id: "t1",
    project_id: "p1",
    source_type: opts.source_type ?? "database",
    name: "Fabric Warehouse",
    adapter_id: opts.adapter_id ?? "fabric_onelake_sql",
  });
  return opts.source_id;
}

describe("SourceQueryService.querySource — bounded execution", () => {
  beforeEach(() => {
    queryMock.mockReset();
    closeMock.mockClear();
  });

  it("executes a query against a registered fabric_onelake_sql database source and returns rows", async () => {
    const sourceService = new SourceService();
    const sourceId = await registerSource(sourceService, { source_id: "src_query_ok" });
    queryMock.mockResolvedValue({ rows: [{ id: 1 }, { id: 2 }], rowCount: 2, durationMs: 5 });

    const service = new SourceQueryService({ sourceService });
    const result = await service.querySource({
      source_id: sourceId,
      query: "SELECT * FROM dbo.customers",
      tenant_id: "t1",
      project_id: "p1",
      connection_config: BASE_CONNECTION_CONFIG,
    });

    expect(result.rows).toEqual([{ id: 1 }, { id: 2 }]);
    expect(result.row_count).toBe(2);
    expect(result.truncated).toBe(false);
    expect(result.bounds.max_rows).toBeGreaterThan(0);
    expect(closeMock).toHaveBeenCalledTimes(1);
  });

  it("rejects a non-database source with SourceQueryUnsupportedError", async () => {
    const sourceService = new SourceService();
    const sourceId = await registerSource(sourceService, { source_id: "src_file_1", source_type: "file", adapter_id: undefined });

    const service = new SourceQueryService({ sourceService });
    await expect(
      service.querySource({ source_id: sourceId, query: "SELECT 1", tenant_id: "t1", project_id: "p1", connection_config: BASE_CONNECTION_CONFIG }),
    ).rejects.toBeInstanceOf(SourceQueryUnsupportedError);
  });

  it("rejects a database source with an unsupported adapter_id", async () => {
    const sourceService = new SourceService();
    const sourceId = await registerSource(sourceService, { source_id: "src_bad_adapter", adapter_id: "some_other_adapter" });

    const service = new SourceQueryService({ sourceService });
    await expect(
      service.querySource({ source_id: sourceId, query: "SELECT 1", tenant_id: "t1", project_id: "p1", connection_config: BASE_CONNECTION_CONFIG }),
    ).rejects.toBeInstanceOf(SourceQueryUnsupportedError);
  });

  it("propagates a 'not found' error for an unregistered source, without ever building an adapter", async () => {
    const sourceService = new SourceService();
    const service = new SourceQueryService({ sourceService });
    await expect(
      service.querySource({ source_id: "does_not_exist", query: "SELECT 1", tenant_id: "t1", project_id: "p1", connection_config: BASE_CONNECTION_CONFIG }),
    ).rejects.toThrow(/not found/i);
    expect(queryMock).not.toHaveBeenCalled();
  });

  it("truncates rows to max_rows and flags truncated:true", async () => {
    const sourceService = new SourceService();
    const sourceId = await registerSource(sourceService, { source_id: "src_truncate" });
    const rows = Array.from({ length: 50 }, (_, i) => ({ id: i }));
    queryMock.mockResolvedValue({ rows, rowCount: rows.length, durationMs: 3 });

    const service = new SourceQueryService({ sourceService });
    const result = await service.querySource({
      source_id: sourceId,
      query: "SELECT * FROM dbo.big_table",
      tenant_id: "t1",
      project_id: "p1",
      max_rows: 10,
      connection_config: BASE_CONNECTION_CONFIG,
    });

    expect(result.rows).toHaveLength(10);
    expect(result.row_count).toBe(10);
    expect(result.truncated).toBe(true);
    expect(result.bounds.max_rows).toBe(10);
  });

  it("does not flag truncated when row count is exactly max_rows", async () => {
    const sourceService = new SourceService();
    const sourceId = await registerSource(sourceService, { source_id: "src_exact" });
    const rows = Array.from({ length: 10 }, (_, i) => ({ id: i }));
    queryMock.mockResolvedValue({ rows, rowCount: rows.length, durationMs: 1 });

    const service = new SourceQueryService({ sourceService });
    const result = await service.querySource({
      source_id: sourceId,
      query: "SELECT * FROM dbo.t",
      tenant_id: "t1",
      project_id: "p1",
      max_rows: 10,
      connection_config: BASE_CONNECTION_CONFIG,
    });

    expect(result.truncated).toBe(false);
    expect(result.rows).toHaveLength(10);
  });

  it("clamps a caller-requested max_rows above the hard ceiling instead of allowing an unbounded read", async () => {
    const sourceService = new SourceService();
    const sourceId = await registerSource(sourceService, { source_id: "src_clamp_rows" });
    queryMock.mockResolvedValue({ rows: [{ id: 1 }], rowCount: 1, durationMs: 1 });

    const service = new SourceQueryService({ sourceService });
    const result = await service.querySource({
      source_id: sourceId,
      query: "SELECT 1",
      tenant_id: "t1",
      project_id: "p1",
      max_rows: SOURCE_QUERY_HARD_MAX_ROWS * 10,
      connection_config: BASE_CONNECTION_CONFIG,
    });

    expect(result.bounds.max_rows).toBe(SOURCE_QUERY_HARD_MAX_ROWS);
  });

  it("clamps a caller-requested timeout_ms above the hard ceiling", async () => {
    const sourceService = new SourceService();
    const sourceId = await registerSource(sourceService, { source_id: "src_clamp_timeout" });
    queryMock.mockResolvedValue({ rows: [], rowCount: 0, durationMs: 1 });

    const service = new SourceQueryService({ sourceService });
    const result = await service.querySource({
      source_id: sourceId,
      query: "SELECT 1",
      tenant_id: "t1",
      project_id: "p1",
      timeout_ms: SOURCE_QUERY_HARD_MAX_TIMEOUT_MS * 100,
      connection_config: BASE_CONNECTION_CONFIG,
    });

    expect(result.bounds.timeout_ms).toBe(SOURCE_QUERY_HARD_MAX_TIMEOUT_MS);
  });

  it("times out a query that runs longer than the bounded timeout, with SourceQueryTimeoutError", async () => {
    const sourceService = new SourceService();
    const sourceId = await registerSource(sourceService, { source_id: "src_timeout" });
    queryMock.mockImplementation(
      () => new Promise((resolve) => setTimeout(() => resolve({ rows: [], rowCount: 0, durationMs: 1 }), 300)),
    );

    const service = new SourceQueryService({ sourceService });
    await expect(
      service.querySource({
        source_id: sourceId,
        query: "SELECT * FROM dbo.slow_view",
        tenant_id: "t1",
        project_id: "p1",
        timeout_ms: 20,
        connection_config: BASE_CONNECTION_CONFIG,
      }),
    ).rejects.toBeInstanceOf(SourceQueryTimeoutError);
  });

  it("still releases the adapter (calls close) after a timeout", async () => {
    const sourceService = new SourceService();
    const sourceId = await registerSource(sourceService, { source_id: "src_timeout_cleanup" });
    queryMock.mockImplementation(
      () => new Promise((resolve) => setTimeout(() => resolve({ rows: [], rowCount: 0, durationMs: 1 }), 300)),
    );

    const service = new SourceQueryService({ sourceService });
    await expect(
      service.querySource({
        source_id: sourceId,
        query: "SELECT 1",
        tenant_id: "t1",
        project_id: "p1",
        timeout_ms: 20,
        connection_config: BASE_CONNECTION_CONFIG,
      }),
    ).rejects.toThrow();

    expect(closeMock).toHaveBeenCalledTimes(1);
  });

  it("propagates adapter query errors as-is", async () => {
    const sourceService = new SourceService();
    const sourceId = await registerSource(sourceService, { source_id: "src_adapter_error" });
    queryMock.mockRejectedValue(new Error("Invalid object name 'dbo.does_not_exist'"));

    const service = new SourceQueryService({ sourceService });
    await expect(
      service.querySource({
        source_id: sourceId,
        query: "SELECT * FROM dbo.does_not_exist",
        tenant_id: "t1",
        project_id: "p1",
        connection_config: BASE_CONNECTION_CONFIG,
      }),
    ).rejects.toThrow(/Invalid object name/);
  });
});

describe("SourceQueryService — AUD0 0pt.source_queried event emission", () => {
  const originalFetch = global.fetch;

  afterEach(() => {
    global.fetch = originalFetch;
  });

  it("emits a 0pt.source_queried AUD0 event with ok:true on success", async () => {
    queryMock.mockReset();
    closeMock.mockClear();
    const sourceService = new SourceService();
    const sourceId = await registerSource(sourceService, { source_id: "src_aud0_ok" });
    queryMock.mockResolvedValue({ rows: [{ id: 1 }], rowCount: 1, durationMs: 2 });

    const calls: Array<{ url: string; body: Record<string, unknown> }> = [];
    global.fetch = vi.fn(async (url: unknown, init: unknown) => {
      calls.push({ url: String(url), body: JSON.parse(String((init as RequestInit).body)) });
      return new Response("{}", { status: 200 });
    }) as unknown as typeof fetch;

    const service = new SourceQueryService({ sourceService, aud0BaseUrl: "http://aud0.test" });
    await service.querySource({
      source_id: sourceId,
      query: "SELECT * FROM dbo.customers",
      tenant_id: "t1",
      project_id: "p1",
      actor_id: "user_42",
      correlation_id: "corr_abc",
      connection_config: BASE_CONNECTION_CONFIG,
    });

    expect(calls).toHaveLength(1);
    expect(calls[0].url).toBe("http://aud0.test/v1/aud0/events");
    expect(calls[0].body.event_type).toBe("0pt.source_queried");
    expect(calls[0].body.capability_id).toBe("0pt.query_source");
    expect(calls[0].body.app_id).toBe("0pt");
    expect(calls[0].body.resource_type).toBe("source");
    expect(calls[0].body.resource_id).toBe(sourceId);
    expect(calls[0].body.actor_id).toBe("user_42");
    expect(calls[0].body.correlation_id).toBe("corr_abc");
    const payload = calls[0].body.payload as Record<string, unknown>;
    expect(payload.ok).toBe(true);
    expect(payload.row_count).toBe(1);
    expect(payload.adapter_id).toBe("fabric_onelake_sql");
  });

  it("emits a 0pt.source_queried AUD0 event with ok:false and the error message on failure", async () => {
    queryMock.mockReset();
    closeMock.mockClear();
    const sourceService = new SourceService();
    const sourceId = await registerSource(sourceService, { source_id: "src_aud0_fail" });
    queryMock.mockRejectedValue(new Error("boom"));

    const calls: Array<{ url: string; body: Record<string, unknown> }> = [];
    global.fetch = vi.fn(async (url: unknown, init: unknown) => {
      calls.push({ url: String(url), body: JSON.parse(String((init as RequestInit).body)) });
      return new Response("{}", { status: 200 });
    }) as unknown as typeof fetch;

    const service = new SourceQueryService({ sourceService, aud0BaseUrl: "http://aud0.test" });
    await expect(
      service.querySource({ source_id: sourceId, query: "SELECT 1", tenant_id: "t1", project_id: "p1", connection_config: BASE_CONNECTION_CONFIG }),
    ).rejects.toThrow("boom");

    expect(calls).toHaveLength(1);
    const payload = calls[0].body.payload as Record<string, unknown>;
    expect(payload.ok).toBe(false);
    expect(payload.error).toBe("boom");
  });
});