Files
dbmole-mcp/test/unit/db/postgres.test.ts
T
2026-06-11 23:58:14 +05:00

262 lines
8.9 KiB
TypeScript

import { beforeEach, describe, expect, it, vi } from 'vitest'
const pgState = vi.hoisted(() => {
const state = {
pools: [] as FakePool[],
nextResult: undefined as unknown,
resultQueue: [] as unknown[]
}
class FakePool {
options: Record<string, unknown>
queries: unknown[] = []
errorHandlers: unknown[] = []
ended = false
constructor(options: Record<string, unknown>) {
this.options = options
state.pools.push(this)
}
on(event: string, fn: unknown) {
if (event === 'error') {
this.errorHandlers.push(fn)
}
return this
}
async query(args: unknown) {
this.queries.push(args)
if (state.resultQueue.length > 0) {
return state.resultQueue.shift()
}
return state.nextResult ?? { rows: [], fields: [], rowCount: 0, command: 'SELECT' }
}
async end() {
this.ended = true
}
}
return { state, FakePool }
})
vi.mock('pg', () => ({
default: { Pool: pgState.FakePool }
}))
import type { ConnectionConfig } from '../../../src/config/types.js'
import { createPostgresDriver } from '../../../src/db/postgres.js'
const config = (extra: Partial<ConnectionConfig> = {}): ConnectionConfig => ({
name: 'pg',
type: 'postgres',
host: 'real-host',
user: 'postgres',
password: 'pw',
database: 'main',
readonly: false,
...extra
})
const target = (extra: Partial<ConnectionConfig> = {}) => ({
config: config(extra),
host: '127.0.0.1',
port: 15432
})
describe('createPostgresDriver', () => {
beforeEach(() => {
pgState.state.pools.length = 0
pgState.state.nextResult = undefined
pgState.state.resultQueue.length = 0
})
it('creates one pool per database and reuses it', async () => {
const driver = createPostgresDriver(target())
await driver.query({ sql: 'select 1', rowLimit: 10 })
await driver.query({ sql: 'select 2', rowLimit: 10 })
await driver.query({ sql: 'select 3', database: 'other', rowLimit: 10 })
expect(pgState.state.pools).toHaveLength(2)
expect(pgState.state.pools[0].options).toMatchObject({
host: '127.0.0.1',
port: 15432,
database: 'main',
max: 4
})
expect(pgState.state.pools[1].options).toMatchObject({ database: 'other' })
})
it('passes readonly as a startup option', async () => {
const driver = createPostgresDriver(target({ readonly: true }))
await driver.query({ sql: 'select 1', rowLimit: 10 })
expect(pgState.state.pools[0].options.options).toBe('-c default_transaction_read_only=on')
})
it('omits startup options when not readonly', async () => {
const driver = createPostgresDriver(target())
await driver.query({ sql: 'select 1', rowLimit: 10 })
expect(pgState.state.pools[0].options.options).toBeUndefined()
})
it('maps array-mode results with column names and normalized cells', async () => {
pgState.state.nextResult = {
rows: [[1n, new Date('2026-01-01T00:00:00Z')]],
fields: [{ name: 'id' }, { name: 'created' }],
rowCount: 1,
command: 'SELECT'
}
const driver = createPostgresDriver(target())
const result = await driver.query({ sql: 'select * from t', rowLimit: 10 })
expect(result).toEqual({
columns: ['id', 'created'],
rows: [['1', '2026-01-01T00:00:00.000Z']],
rowCount: 1,
truncated: false
})
})
it('truncates rows beyond the limit', async () => {
pgState.state.nextResult = {
rows: [[1], [2], [3]],
fields: [{ name: 'n' }],
rowCount: 3,
command: 'SELECT'
}
const driver = createPostgresDriver(target())
const result = await driver.query({ sql: 'select n', rowLimit: 2 })
expect(result.rows).toEqual([[1], [2]])
expect(result.truncated).toBe(true)
expect(result.rowCount).toBe(3)
})
it('sends positional params through', async () => {
const driver = createPostgresDriver(target())
await driver.query({ sql: 'select $1', params: [42], rowLimit: 10 })
expect(pgState.state.pools[0].queries[0]).toMatchObject({
text: 'select $1',
values: [42],
rowMode: 'array'
})
})
it('rejects multi-statement sql before touching the pool', async () => {
const driver = createPostgresDriver(target())
await expect(driver.query({ sql: 'select 1; select 2', rowLimit: 10 })).rejects.toThrow(
/one SQL statement/
)
// The guard runs before any pool is created.
expect(pgState.state.pools).toHaveLength(0)
})
it('rejects session-level statements', async () => {
const driver = createPostgresDriver(target())
await expect(driver.query({ sql: 'BEGIN', rowLimit: 1 })).rejects.toThrow(/session-level/)
expect(pgState.state.pools).toHaveLength(0)
})
it('registers a pool error listener', async () => {
const driver = createPostgresDriver(target())
await driver.query({ sql: 'select 1', rowLimit: 1 })
expect(pgState.state.pools[0].errorHandlers).toHaveLength(1)
})
it('describeTable groups composite fks and multi-column indexes', async () => {
pgState.state.resultQueue = [
// columns
{
rows: [
['id', 'integer', 'NO', null],
['name', 'text', 'YES', null]
],
fields: [{ name: 'c' }],
rowCount: 2
},
// primary key
{ rows: [['id']], fields: [{ name: 'c' }], rowCount: 1 },
// indexes: one index 'idx_ab' with two key columns
{
rows: [
['idx_ab', false, 'a'],
['idx_ab', false, 'b']
],
fields: [{ name: 'c' }],
rowCount: 2
},
// foreign keys: one composite fk 'fk1' with two column pairs
{
rows: [
['fk1', 'a', 'other', 'x'],
['fk1', 'b', 'other', 'y']
],
fields: [{ name: 'c' }],
rowCount: 2
}
]
const driver = createPostgresDriver(target())
const description = await driver.describeTable({ table: 't' })
expect(description.indexes).toEqual([
{ name: 'idx_ab', unique: false, columns: ['a', 'b'] }
])
expect(description.foreignKeys).toEqual([
{
name: 'fk1',
columns: ['a', 'b'],
referencedTable: 'other',
referencedColumns: ['x', 'y']
}
])
})
it('describeTable throws when the table has no columns', async () => {
pgState.state.resultQueue = [{ rows: [], fields: [], rowCount: 0 }]
const driver = createPostgresDriver(target())
await expect(driver.describeTable({ table: 'missing' })).rejects.toThrow(/not found/)
})
it('lists databases against the maintenance db and maps sizes', async () => {
pgState.state.nextResult = {
rows: [
['app', 1024],
['postgres', 2048]
],
fields: [{ name: 'name' }, { name: 'size_bytes' }],
rowCount: 2,
command: 'SELECT'
}
const driver = createPostgresDriver(target({ database: undefined }))
const databases = await driver.listDatabases()
expect(databases).toEqual([
{ name: 'app', sizeBytes: 1024 },
{ name: 'postgres', sizeBytes: 2048 }
])
expect(pgState.state.pools[0].options.database).toBe('postgres')
})
it('maps reltuples=-1 to null in listTables', async () => {
pgState.state.nextResult = {
rows: [
['public', 'fresh', '-1'],
['public', 'analyzed', '42']
],
fields: [{ name: 'schema' }, { name: 'name' }, { name: 'row_estimate' }],
rowCount: 2,
command: 'SELECT'
}
const driver = createPostgresDriver(target())
const tables = await driver.listTables({})
expect(tables).toEqual([
{ schema: 'public', name: 'fresh', rowEstimate: null },
{ schema: 'public', name: 'analyzed', rowEstimate: 42 }
])
})
it('dispose ends every pool', async () => {
const driver = createPostgresDriver(target())
await driver.query({ sql: 'select 1', rowLimit: 1 })
await driver.query({ sql: 'select 1', database: 'other', rowLimit: 1 })
await driver.dispose()
expect(pgState.state.pools.every((pool) => pool.ended)).toBe(true)
})
})