From bc23b1a4472097194ac515b00b6c53ec051febc9 Mon Sep 17 00:00:00 2001
From: Andrey Avtomonov <andreybavt@gmail.com>
Date: Wed, 13 May 2026 19:05:36 +0200
Subject: [PATCH] fix(setup): prompt for postgres query history

---
 packages/cli/src/setup-databases.test.ts | 61 ++++++++++++++++++++++++
 packages/cli/src/setup-databases.ts      | 12 ++---
 2 files changed, 67 insertions(+), 6 deletions(-)

diff --git a/packages/cli/src/setup-databases.test.ts b/packages/cli/src/setup-databases.test.ts
index 5e675ebf..656d3dad 100644
--- a/packages/cli/src/setup-databases.test.ts
+++ b/packages/cli/src/setup-databases.test.ts
@@ -1454,6 +1454,67 @@ describe('setup databases step', () => {
     expect(io.stdout()).toContain('pg_stat_statements ready');
   });
 
+  it('asks interactive Postgres setup whether to enable query history', async () => {
+    await writeFile(
+      join(tempDir, 'ktx.yaml'),
+      [
+        'project: warehouse',
+        'connections:',
+        '  warehouse:',
+        '    driver: postgres',
+        '    url: env:DATABASE_URL',
+        '    readonly: true',
+        '',
+      ].join('\n'),
+      'utf-8',
+    );
+    const io = makeIo();
+    const prompts = makePromptAdapter({ selectValues: ['yes'] });
+    const historicSqlProbe = vi.fn(async () => ({ ok: true, lines: [] }));
+
+    const result = await runKtxSetupDatabasesStep(
+      {
+        projectDir: tempDir,
+        inputMode: 'auto',
+        databaseConnectionIds: ['warehouse'],
+        databaseSchemas: [],
+        skipDatabases: false,
+      },
+      io.io,
+      {
+        prompts,
+        testConnection: vi.fn(async () => 0),
+        scanConnection: vi.fn(async () => 0),
+        historicSqlProbe,
+      },
+    );
+
+    expect(result.status).toBe('ready');
+    expect(prompts.select).toHaveBeenCalledWith({
+      message: 'Enable query-history ingest for this PostgreSQL connection?',
+      options: [
+        { value: 'yes', label: 'Enable query history' },
+        { value: 'no', label: 'Do not enable query history' },
+        { value: 'back', label: 'Back' },
+      ],
+    });
+    expect(historicSqlProbe).toHaveBeenCalledWith({
+      projectDir: tempDir,
+      connectionId: 'warehouse',
+      dialect: 'postgres',
+    });
+    const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8'));
+    expect(config.connections.warehouse).toMatchObject({
+      context: {
+        queryHistory: {
+          enabled: true,
+          minExecutions: 5,
+          filters: { dropTrivialProbes: true },
+        },
+      },
+    });
+  });
+
   it('writes query history config for supported existing database connections', async () => {
     await writeFile(
       join(tempDir, 'ktx.yaml'),
diff --git a/packages/cli/src/setup-databases.ts b/packages/cli/src/setup-databases.ts
index 6b743a90..1d3623bb 100644
--- a/packages/cli/src/setup-databases.ts
+++ b/packages/cli/src/setup-databases.ts
@@ -842,7 +842,7 @@ async function maybeApplyHistoricSqlConfig(input: {
   let enabled = input.args.enableQueryHistory === true;
   if (input.args.disableQueryHistory === true) {
     enabled = false;
-  } else if (input.args.inputMode !== 'disabled' && input.args.enableQueryHistory !== true && dialect !== 'postgres') {
+  } else if (input.args.inputMode !== 'disabled' && input.args.enableQueryHistory !== true) {
     const choice = await input.prompts.select({
       message: `Enable query-history ingest for this ${driverLabel(input.driver)} connection?`,
       options: [
@@ -855,10 +855,6 @@ async function maybeApplyHistoricSqlConfig(input: {
     enabled = choice === 'yes';
   }
 
-  if (dialect === 'postgres' && input.args.enableQueryHistory !== true && input.args.disableQueryHistory !== true) {
-    return input.connection;
-  }
-
   const existingRecord = queryHistoryConfigRecord(input.connection) ?? historicSqlConfigRecord(input.connection) ?? {};
   const { dialect: _dialect, ...existing } = existingRecord;
 
@@ -1474,7 +1470,11 @@ async function applyHistoricSqlConfigToExistingConnection(input: {
   args: KtxSetupDatabasesArgs;
   prompts: KtxSetupDatabasesPromptAdapter;
 }): Promise<'back' | void> {
-  if (input.args.enableQueryHistory !== true && input.args.disableQueryHistory !== true) {
+  if (
+    input.args.inputMode === 'disabled' &&
+    input.args.enableQueryHistory !== true &&
+    input.args.disableQueryHistory !== true
+  ) {
     return;
   }