apunkt/ktx
1
0
Fork 0
mirror of https://github.com/Kaelio/ktx.git synced 2026-07-07 11:02:11 +02:00
Code Issues Projects Releases Packages Wiki Activity

feat(connectors): generalize readiness and constraint handling (#212)

Browse source 
* feat(connectors): add postgres maxConnections

* feat(connectors): add mysql maxConnections

* feat(connectors): add sqlserver maxConnections

* feat(connectors): rename snowflake pool config

* docs: document connector maxConnections

* feat(scan): add constraint discovery warning helper

* feat(scan): carry structural warnings through reports

* feat(postgres): soft-fail denied constraint discovery

* feat(mysql): soft-fail denied constraint discovery

* feat(sqlserver): soft-fail denied constraint discovery

* feat(bigquery): soft-fail denied primary key discovery

* feat(snowflake): report denied primary key discovery

* test(scan): verify constraint discovery warnings

* feat(historic-sql): use shared readiness probes

* docs: document query history readiness probes

* test(historic-sql): verify readiness probe registry

* test(ingest): account for live database warnings artifact

* Add skip option for agent setup
This commit is contained in:
Andrey Avtomonov 2026-05-24 19:30:06 +02:00 • committed by GitHub
parent cfd1749ab9
commit 78b8a0c025
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
42 changed files with 2763 additions and 554 deletions
Download patch file Download diff file Expand all files Collapse all files

93
packages/cli/src/connectors/snowflake/connector.test.ts
View file

@ -8,7 +8,7 @@ vi.mock('snowflake-sdk', () => ({
}));
import { createSnowflakeLiveDatabaseIntrospection } from '../../connectors/snowflake/live-database-introspection.js';
import { isKtxSnowflakeConnectionConfig, KtxSnowflakeScanConnector, snowflakeConnectionConfigFromConfig, type KtxSnowflakeDriver, type KtxSnowflakeDriverFactory } from '../../connectors/snowflake/connector.js';
import { isKtxSnowflakeConnectionConfig, KtxSnowflakeScanConnector, snowflakeConnectionConfigFromConfig, type KtxSnowflakeConnectionConfig, type KtxSnowflakeDriver, type KtxSnowflakeDriverFactory } from '../../connectors/snowflake/connector.js';
import { tableRefSet } from '../../context/scan/table-ref.js';
function fakeDriverFactory(): KtxSnowflakeDriverFactory {
@ -140,8 +140,8 @@ describe('KtxSnowflakeScanConnector', () => {
});
});
it('defaults and validates Snowflake maxSessions', () => {
const baseConnection = {
it('defaults and validates Snowflake maxConnections', () => {
const baseConnection: KtxSnowflakeConnectionConfig = {
driver: 'snowflake',
authMethod: 'password',
account: 'acct',
@ -150,32 +150,59 @@ describe('KtxSnowflakeScanConnector', () => {
schema_name: 'PUBLIC',
username: 'reader',
password: 'fixture-pass', // pragma: allowlist secret
} as const;
};
expect(
snowflakeConnectionConfigFromConfig({
connectionId: 'warehouse',
connection: baseConnection,
}),
).toMatchObject({ maxSessions: 4 });
).toMatchObject({ maxConnections: 4 });
expect(
snowflakeConnectionConfigFromConfig({
connectionId: 'warehouse',
connection: { ...baseConnection, maxSessions: 8 },
connection: { ...baseConnection, maxConnections: 8 },
}),
).toMatchObject({ maxSessions: 8 });
).toMatchObject({ maxConnections: 8 });
for (const maxSessions of [0, -1, 1.5, Number.NaN]) {
expect(
snowflakeConnectionConfigFromConfig({
connectionId: 'warehouse',
connection: { ...baseConnection, maxConnections: '12' as never },
}),
).toMatchObject({ maxConnections: 12 });
for (const maxConnections of [0, -1, 1.5, Number.NaN, 'abc' as never]) {
expect(() =>
snowflakeConnectionConfigFromConfig({
connectionId: 'warehouse',
connection: { ...baseConnection, maxSessions },
connection: { ...baseConnection, maxConnections },
}),
).toThrow('connections.warehouse.maxSessions must be a positive integer');
).toThrow('connections.warehouse.maxConnections must be a positive integer');
}
});
it('rejects stale Snowflake pool config key', () => {
const baseConnection: KtxSnowflakeConnectionConfig = {
driver: 'snowflake',
authMethod: 'password',
account: 'acct',
warehouse: 'WH',
database: 'ANALYTICS',
schema_name: 'PUBLIC',
username: 'reader',
password: 'fixture-pass', // pragma: allowlist secret
};
expect(() =>
snowflakeConnectionConfigFromConfig({
connectionId: 'warehouse',
connection: { ...baseConnection, maxSessions: 8 },
}),
).toThrow(/renamed to maxConnections/);
});
it('uses one lazy Snowflake pool and drains it during cleanup', async () => {
const { pool, executedSql } = installSnowflakePoolMock();
const close = vi.fn(async () => undefined);
@ -191,7 +218,7 @@ describe('KtxSnowflakeScanConnector', () => {
username: 'reader',
password: 'fixture-pass', // pragma: allowlist secret
role: 'ANALYST',
maxSessions: 3,
maxConnections: 3,
},
sdkOptionsProvider: {
resolve: vi.fn(async () => ({ sdkOptions: { application: 'ktx-test' }, close })),
@ -332,12 +359,56 @@ describe('KtxSnowflakeScanConnector', () => {
expect(snapshot.tables.map((table) => table.name).sort()).toEqual(['ORDERS', 'ORDER_SUMMARY']);
expect(snapshot.tables.every((table) => table.columns.every((column) => column.primaryKey === false))).toBe(true);
expect(snapshot.warnings).toEqual([
{
code: 'constraint_discovery_unauthorized',
message: 'Skipped primary-key discovery in PUBLIC (insufficient grants on system catalogs)',
recoverable: true,
metadata: { schema: 'PUBLIC', kind: 'primary_key' },
},
]);
expect(warn).not.toHaveBeenCalled();
} finally {
warn.mockRestore();
}
});
it('propagates non-denial Snowflake primary-key discovery errors', async () => {
const driverFactory = fakeDriverFactory();
const driver = (driverFactory.createDriver as ReturnType<typeof vi.fn>).getMockImplementation() as
| (() => KtxSnowflakeDriver)
| undefined;
if (!driver) throw new Error('driver mock missing');
const built = driver();
const networkError = new Error('network unavailable');
(built.query as ReturnType<typeof vi.fn>).mockImplementation(async (sql: string) => {
if (sql.includes('TABLE_CONSTRAINTS')) {
throw networkError;
}
throw new Error(`Unexpected SQL: ${sql}`);
});
(driverFactory.createDriver as ReturnType<typeof vi.fn>).mockReturnValue(built);
const connector = new KtxSnowflakeScanConnector({
connectionId: 'warehouse',
connection: {
driver: 'snowflake',
authMethod: 'password',
account: 'acct',
warehouse: 'WH',
database: 'ANALYTICS',
schema_name: 'PUBLIC',
username: 'reader',
password: 'fixture-pass', // pragma: allowlist secret
},
driverFactory,
});
await expect(
connector.introspect({ connectionId: 'warehouse', driver: 'snowflake' }, { runId: 'scan-run-snowflake-network' }),
).rejects.toBe(networkError);
});
it('limits introspection to tables in tableScope', async () => {
const queries: Array<{ sql: string; params?: unknown }> = [];
const getSchemaMetadata = vi.fn(async (_schemaName?: string, scopedNames?: readonly string[] | null) =>

79
packages/cli/src/connectors/snowflake/connector.ts
View file

@ -3,8 +3,28 @@ import { readFileSync } from 'node:fs';
import { homedir } from 'node:os';
import { resolve } from 'node:path';
import { assertReadOnlySql, limitSqlForExecution } from '../../context/connections/read-only-sql.js';
import { createKtxConnectorCapabilities, type KtxColumnSampleInput, type KtxColumnSampleResult, type KtxColumnStatsInput, type KtxColumnStatsResult, type KtxQueryResult, type KtxReadOnlyQueryInput, type KtxScanConnector, type KtxScanContext, type KtxScanInput, type KtxSchemaColumn, type KtxSchemaSnapshot, type KtxSchemaTable, type KtxTableRef, type KtxTableSampleInput, type KtxTableListEntry, type KtxTableSampleResult } from '../../context/scan/types.js';
import { tryConstraintQuery } from '../../context/scan/constraint-discovery.js';
import { scopedTableNames } from '../../context/scan/table-ref.js';
import {
createKtxConnectorCapabilities,
type KtxColumnSampleInput,
type KtxColumnSampleResult,
type KtxColumnStatsInput,
type KtxColumnStatsResult,
type KtxQueryResult,
type KtxReadOnlyQueryInput,
type KtxScanConnector,
type KtxScanContext,
type KtxScanInput,
type KtxScanWarning,
type KtxSchemaColumn,
type KtxSchemaSnapshot,
type KtxSchemaTable,
type KtxTableListEntry,
type KtxTableRef,
type KtxTableSampleInput,
type KtxTableSampleResult,
} from '../../context/scan/types.js';
import snowflake from 'snowflake-sdk';
import type { Bind, Binds, Connection, ConnectionOptions } from 'snowflake-sdk';
import { KtxSnowflakeDialect } from './dialect.js';
@ -24,7 +44,7 @@ export interface KtxSnowflakeConnectionConfig {
privateKey?: string;
passphrase?: string;
role?: string;
maxSessions?: number;
maxConnections?: number;
[key: string]: unknown;
}
@ -39,7 +59,7 @@ export interface KtxSnowflakeResolvedConnectionConfig {
privateKey?: string;
passphrase?: string;
role?: string;
maxSessions: number;
maxConnections: number;
}
export interface KtxSnowflakeRawColumnMetadata {
@ -166,6 +186,13 @@ function firstNumber(value: unknown): number | null {
return Number.isFinite(numberValue) ? numberValue : null;
}
function isDeniedError(error: unknown): boolean {
if (error instanceof Error) {
return /insufficient privileges|does not exist or not authorized/i.test(error.message);
}
return false;
}
function normalizeSnowflakeValue(value: unknown, columnType?: string): unknown {
if (columnType && DATE_TYPES.some((type) => columnType.toUpperCase().includes(type))) {
if (typeof value === 'number') {
@ -218,6 +245,10 @@ export function snowflakeConnectionConfigFromConfig(input: {
if (!isKtxSnowflakeConnectionConfig(input.connection)) {
throw new Error(`Native Snowflake connector cannot run driver "${inputDriver}"`);
}
const staleMaxSessionsKey = 'max' + 'Sessions';
if (Object.prototype.hasOwnProperty.call(input.connection, staleMaxSessionsKey)) {
throw new Error(`connections.${input.connectionId}.maxSessions has been renamed to maxConnections`);
}
const env = input.env ?? process.env;
const authMethod = input.connection?.authMethod ?? 'password';
const account = stringConfigValue(input.connection, 'account', env);
@ -249,9 +280,9 @@ export function snowflakeConnectionConfigFromConfig(input: {
database,
schemas: resolvedSchemas,
username,
maxSessions: positiveIntegerConfigValue({
maxConnections: positiveIntegerConfigValue({
connection: input.connection,
key: 'maxSessions',
key: 'maxConnections',
connectionId: input.connectionId,
defaultValue: 4,
}),
@ -322,7 +353,7 @@ class SnowflakeSdkDriver implements KtxSnowflakeDriver {
const message = error instanceof Error ? error.message : String(error);
if (/timeout/i.test(message) && /pool|acquire/i.test(message)) {
throw new Error(
"Snowflake session pool exhausted after 60s - consider lowering maxSessions or increasing your account's concurrent-statement limit.",
"Snowflake session pool exhausted after 60s - consider lowering maxConnections or increasing your account's concurrent-statement limit.",
);
}
throw error;
@ -432,7 +463,7 @@ class SnowflakeSdkDriver implements KtxSnowflakeDriver {
if (!this.pool) {
this.pool = snowflake.createPool(await this.resolveConnectionOptions(), {
min: 0,
max: this.resolved.maxSessions,
max: this.resolved.maxConnections,
evictionRunIntervalMillis: 30_000,
acquireTimeoutMillis: 60_000,
});
@ -540,13 +571,23 @@ export class KtxSnowflakeScanConnector implements KtxScanConnector {
async introspect(input: KtxScanInput, _ctx: KtxScanContext): Promise<KtxSchemaSnapshot> {
this.assertConnection(input.connectionId);
const tables: KtxSchemaTable[] = [];
const snapshotWarnings: KtxScanWarning[] = [];
for (const schemaName of this.resolved.schemas) {
const scopedNames = input.tableScope
? scopedTableNames(input.tableScope, { catalog: this.resolved.database, db: schemaName })
: null;
if (scopedNames && scopedNames.length === 0) continue;
const rawTables = await this.getDriver().getSchemaMetadata(schemaName, scopedNames);
const primaryKeys = await this.primaryKeys(rawTables.map((table) => table.name), schemaName);
const primaryKeysResult = await tryConstraintQuery(
{ schema: schemaName, kind: 'primary_key', isDeniedError },
() => this.primaryKeys(rawTables.map((table) => table.name), schemaName),
);
const primaryKeys = primaryKeysResult.ok
? primaryKeysResult.value
: new Map(rawTables.map((table) => [table.name, new Set<string>()]));
if (!primaryKeysResult.ok) {
snapshotWarnings.push(primaryKeysResult.warning);
}
tables.push(...rawTables.map((table) => this.toSchemaTable(table, primaryKeys)));
}
return {
@ -563,6 +604,7 @@ export class KtxSnowflakeScanConnector implements KtxScanConnector {
total_columns: tables.reduce((sum, table) => sum + table.columns.length, 0),
},
tables,
warnings: snapshotWarnings,
};
}
@ -686,9 +728,8 @@ export class KtxSnowflakeScanConnector implements KtxScanConnector {
return grouped;
}
const tableNamePlaceholders = tableNames.map(() => '?').join(', ');
try {
const result = await this.getDriver().query(
`
const result = await this.getDriver().query(
`
SELECT tc.TABLE_NAME, kcu.COLUMN_NAME
FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS tc
JOIN INFORMATION_SCHEMA.KEY_COLUMN_USAGE kcu
@ -701,16 +742,12 @@ export class KtxSnowflakeScanConnector implements KtxScanConnector {
AND tc.TABLE_NAME IN (${tableNamePlaceholders})
ORDER BY tc.TABLE_NAME, kcu.ORDINAL_POSITION
`,
[schemaName, this.resolved.database, ...tableNames],
);
for (const row of result.rows) {
const tableName = String(row[0]);
const columnName = String(row[1]);
grouped.get(tableName)?.add(columnName);
}
} catch {
// INFORMATION_SCHEMA.KEY_COLUMN_USAGE often isn't granted to read-only roles;
// continue with empty PK map and let FK inference + profiling carry the slack.
[schemaName, this.resolved.database, ...tableNames],
);
for (const row of result.rows) {
const tableName = String(row[0]);
const columnName = String(row[1]);
grouped.get(tableName)?.add(columnName);
}
return grouped;
}