From ed47081367b9966f2837ab52620562f9ab6d0295 Mon Sep 17 00:00:00 2001 From: ferhat elmas Date: Fri, 17 Jul 2026 14:32:08 +0200 Subject: [PATCH] fix: add transport-neutral database contracts to prepare for watt adapter Signed-off-by: ferhat elmas --- src/http/plugins/db.ts | 4 +- src/http/routes/tus/lifecycle.ts | 4 +- src/internal/auth/jwks/store-pg.ts | 18 +++-- src/internal/database/connection.ts | 63 ++++++++++++++++ src/internal/database/index.ts | 1 + .../database/migration-admin-store-pg.test.ts | 9 +-- .../database/migration-admin-store-pg.ts | 6 +- src/internal/database/migrations/migrate.ts | 6 +- src/internal/database/multitenant-pg.ts | 5 +- src/internal/database/pg-connection.test.ts | 16 +++-- src/internal/database/pg-connection.ts | 71 ++++++++----------- src/internal/database/tenant-store-pg.test.ts | 14 ++-- src/internal/database/tenant-store-pg.ts | 14 ++-- src/internal/queue/database.test.ts | 6 +- src/internal/queue/database.ts | 4 +- src/internal/queue/event.ts | 5 +- src/internal/sharding/pg.ts | 17 +++-- src/storage/database/adapter.ts | 13 +--- src/storage/database/pg.test.ts | 30 ++++---- src/storage/database/pg.ts | 45 ++++++------ .../iceberg/delete-iceberg-resources.test.ts | 11 +-- .../iceberg/delete-iceberg-resources.ts | 2 +- .../events/lifecycle/bucket-deleted.ts | 2 +- src/storage/events/pgboss/move-jobs.ts | 4 +- src/storage/events/upgrades/base-event.ts | 4 +- .../iceberg/catalog/reconciler.test.ts | 2 +- .../protocols/iceberg/catalog/reconciler.ts | 4 +- src/storage/protocols/iceberg/metastore.ts | 2 +- src/storage/protocols/iceberg/pg.test.ts | 2 +- src/storage/protocols/iceberg/pg.ts | 26 +++---- .../protocols/s3/credentials/store-pg.ts | 4 +- .../vector/adapter/pgvector/index.ts | 49 +++++++------ src/storage/protocols/vector/pg.ts | 34 +++++---- src/test/database-protection.test.ts | 8 +-- src/test/iceberg.test.ts | 6 +- src/test/object.test.ts | 24 ++++--- src/test/operation-helpers.test.ts | 2 +- src/test/s3-error-code.test.ts | 2 +- src/test/s3-protocol.test.ts | 6 +- src/test/storage-pg-db.test.ts | 36 +++++----- src/test/utils/storage.ts | 17 ++--- src/test/vectors.test.ts | 6 +- src/test/webhooks.test.ts | 6 +- 43 files changed, 346 insertions(+), 264 deletions(-) create mode 100644 src/internal/database/connection.ts diff --git a/src/http/plugins/db.ts b/src/http/plugins/db.ts index a65463c41..aea9a1fb4 100644 --- a/src/http/plugins/db.ts +++ b/src/http/plugins/db.ts @@ -3,7 +3,7 @@ import { getPostgresConnection, getServiceKeyUser, getTenantConfig, - PgTenantConnection, + type TenantConnection, } from '@internal/database' import { areMigrationsUpToDate, @@ -19,7 +19,7 @@ import { getConfig, MultitenantMigrationStrategy } from '../../config' declare module 'fastify' { interface FastifyRequest { - db: PgTenantConnection + db: TenantConnection latestMigration?: keyof typeof DBMigration } } diff --git a/src/http/routes/tus/lifecycle.ts b/src/http/routes/tus/lifecycle.ts index c02fce368..1c1cd93ac 100644 --- a/src/http/routes/tus/lifecycle.ts +++ b/src/http/routes/tus/lifecycle.ts @@ -1,5 +1,5 @@ import { SIGNED_URL_SCOPE_UPLOAD } from '@internal/auth' -import { PgTenantConnection } from '@internal/database' +import type { TenantConnection } from '@internal/database' import { ERRORS, isRenderableError } from '@internal/errors' import { logSchema, RequestLogContext } from '@internal/monitoring' import { UploadId } from '@storage/protocols/tus' @@ -38,7 +38,7 @@ export type MultiPartRequest = http.IncomingMessage & { upload: RequestLogContext & { tenantId: string storage: Storage - db: PgTenantConnection + db: TenantConnection owner?: string isUpsert: boolean resources?: string[] diff --git a/src/internal/auth/jwks/store-pg.ts b/src/internal/auth/jwks/store-pg.ts index e71bced9d..d211afadd 100644 --- a/src/internal/auth/jwks/store-pg.ts +++ b/src/internal/auth/jwks/store-pg.ts @@ -1,14 +1,14 @@ import { getConfig } from '../../../config' -import { PgTransaction, PgTransactionalExecutor } from '../../database/pg-connection' +import type { DatabaseTransaction, DatabaseTransactionalExecutor } from '../../database/connection' import { logger, logSchema } from '../../monitoring' import { JWKSManagerStore, JWKStoreItem, PaginatedTenantItem } from './store' const { multitenantDatabaseQueryTimeout } = getConfig() -export class JWKSManagerStorePg implements JWKSManagerStore { - constructor(private db: PgTransactionalExecutor) {} +export class JWKSManagerStorePg implements JWKSManagerStore { + constructor(private db: DatabaseTransactionalExecutor) {} - async transaction(callback: (trx: PgTransaction) => Promise): Promise { + async transaction(callback: (trx: DatabaseTransaction) => Promise): Promise { const trx = await this.db.beginTransaction() try { @@ -34,7 +34,7 @@ export class JWKSManagerStorePg implements JWKSManagerStore { encryptedJwk: string, kind: string, idempotent = false, - trx?: PgTransaction + trx?: DatabaseTransaction ): Promise { const db = trx || this.db const insertResult = await db.query<{ id: string }>( @@ -89,7 +89,7 @@ export class JWKSManagerStorePg implements JWKSManagerStore { tenantId: string, id: string, newState: boolean, - trx?: PgTransaction + trx?: DatabaseTransaction ): Promise { const db = trx || this.db const result = await db.query( @@ -109,7 +109,11 @@ export class JWKSManagerStorePg implements JWKSManagerStore { return Boolean(result.rowCount && result.rowCount > 0) } - async listActive(tenantId: string, kind?: string, trx?: PgTransaction): Promise { + async listActive( + tenantId: string, + kind?: string, + trx?: DatabaseTransaction + ): Promise { const db = trx || this.db const result = await db.query( { diff --git a/src/internal/database/connection.ts b/src/internal/database/connection.ts new file mode 100644 index 000000000..4277bd322 --- /dev/null +++ b/src/internal/database/connection.ts @@ -0,0 +1,63 @@ +import type { QueryResult, QueryResultRow } from 'pg' + +/** + * PostgreSQL-compatible connection contracts shared by the direct pg adapter + * and the Database Watt transport adapter. + * + * Concrete pools, clients, and transport details do not belong in this module. + */ +export interface DatabaseStatement { + text: string + values?: unknown[] +} + +export interface DatabaseQueryOptions { + signal?: AbortSignal +} + +export type DatabaseQueryArgument = DatabaseQueryOptions | unknown[] + +export interface TransactionOptions { + isolation?: string + retry?: number + readOnly?: boolean + timeout?: number +} + +export interface DatabaseExecutor { + query( + statement: string | DatabaseStatement, + options?: DatabaseQueryArgument + ): Promise> +} + +export interface DatabaseTransaction extends DatabaseExecutor { + isCompleted(): boolean + commit(): Promise + rollback(): Promise +} + +export interface DatabaseTransactionalExecutor extends DatabaseExecutor { + beginTransaction(options?: TransactionOptions): Promise +} + +export interface TenantConnection extends DatabaseTransactionalExecutor { + readonly role: string + dispose(): void + setAbortSignal(signal: AbortSignal): void + getAbortSignal(): AbortSignal | undefined + asSuperUser(): TenantConnection + transaction(options?: TransactionOptions): Promise + setScope(transaction: DatabaseExecutor): Promise +} + +export function isDatabaseTransaction(executor: DatabaseExecutor): executor is DatabaseTransaction { + return ( + 'commit' in executor && + typeof executor.commit === 'function' && + 'rollback' in executor && + typeof executor.rollback === 'function' && + 'isCompleted' in executor && + typeof executor.isCompleted === 'function' + ) +} diff --git a/src/internal/database/index.ts b/src/internal/database/index.ts index 16c6ce5ed..7af5b0031 100644 --- a/src/internal/database/index.ts +++ b/src/internal/database/index.ts @@ -1,4 +1,5 @@ export * from './client' +export * from './connection' export * from './migration-admin-store-pg' export * from './multitenant-pg' export * from './pg-connection' diff --git a/src/internal/database/migration-admin-store-pg.test.ts b/src/internal/database/migration-admin-store-pg.test.ts index 22b2cac95..58ab27f3a 100644 --- a/src/internal/database/migration-admin-store-pg.test.ts +++ b/src/internal/database/migration-admin-store-pg.test.ts @@ -1,21 +1,22 @@ +import type { DatabaseExecutor, DatabaseStatement } from './connection' import { MigrationAdminStorePg } from './migration-admin-store-pg' -import type { PgExecutor, PgStatement } from './pg-connection' function createMigrationAdminStore() { const query = vi.fn().mockResolvedValue({ rows: [], rowCount: 1, }) - const store = new MigrationAdminStorePg({ query } as unknown as PgExecutor, 'pgboss') + const db = { query } as unknown as DatabaseExecutor + const store = new MigrationAdminStorePg(db, 'pgboss') return { query, store } } -function getLastStatement(query: ReturnType): PgStatement { +function getLastStatement(query: ReturnType): DatabaseStatement { const [statement] = query.mock.calls.at(-1) || [] if (!statement || typeof statement === 'string') { - throw new Error('Expected a PgStatement query') + throw new Error('Expected a DatabaseStatement query') } return statement diff --git a/src/internal/database/migration-admin-store-pg.ts b/src/internal/database/migration-admin-store-pg.ts index ba144ed17..c3ce30357 100644 --- a/src/internal/database/migration-admin-store-pg.ts +++ b/src/internal/database/migration-admin-store-pg.ts @@ -1,5 +1,5 @@ import { QueryResultRow } from 'pg' -import { PgExecutor } from './pg-connection' +import type { DatabaseExecutor } from './connection' import { quoteIdentifier } from './sql' import { TenantCursorRow } from './tenant-store-pg' @@ -9,7 +9,7 @@ export class MigrationAdminStorePg { private readonly jobTable: string constructor( - private db: PgExecutor, + private db: DatabaseExecutor, pgBossSchema: string ) { this.jobTable = `${quoteIdentifier(pgBossSchema)}.job` @@ -112,7 +112,7 @@ export class MigrationAdminStorePg { } private query( - statement: Parameters[0] + statement: Parameters[0] ) { return this.db.query(statement) } diff --git a/src/internal/database/migrations/migrate.ts b/src/internal/database/migrations/migrate.ts index 8619cc72b..f17fe3474 100644 --- a/src/internal/database/migrations/migrate.ts +++ b/src/internal/database/migrations/migrate.ts @@ -8,8 +8,8 @@ import { validateMigrationHashes } from 'postgres-migrations/dist/validation' import SQL from 'sql-template-strings' import { getConfig, MultitenantMigrationStrategy } from '../../../config' import { logger, logSchema } from '../../monitoring' +import type { DatabaseExecutor, DatabaseTransaction } from '../connection' import { multitenantPgExecutor } from '../multitenant-pg' -import { PgExecutor, PgTransaction } from '../pg-connection' import { searchPath } from '../pool' import { getSslSettings } from '../ssl' import { getTenantConfig, TenantMigrationStatus } from '../tenant' @@ -163,7 +163,7 @@ export async function updateTenantMigrationsState( options?: { migration?: keyof typeof DBMigration state: TenantMigrationStatus - tnx?: PgExecutor + tnx?: DatabaseExecutor } ) { const migrationVersion = options?.migration || (await lastLocalMigrationName()) @@ -196,7 +196,7 @@ export async function areMigrationsUpToDate(tenantId: string) { } export async function obtainLockOnMultitenantDB( - fn: (tnx: PgTransaction) => Promise, + fn: (tnx: DatabaseTransaction) => Promise, options?: { sbReqId?: string } ) { const trx = await multitenantPgExecutor.beginTransaction() diff --git a/src/internal/database/multitenant-pg.ts b/src/internal/database/multitenant-pg.ts index 439a98f97..d088a9a47 100644 --- a/src/internal/database/multitenant-pg.ts +++ b/src/internal/database/multitenant-pg.ts @@ -1,7 +1,8 @@ import { logger, logSchema } from '@internal/monitoring' import { Pool, PoolConfig } from 'pg' import { getConfig } from '../../config' -import { attachPgPoolErrorHandler, PgPoolExecutor, PgTransactionalExecutor } from './pg-connection' +import type { DatabaseTransactionalExecutor } from './connection' +import { attachPgPoolErrorHandler, PgPoolExecutor } from './pg-connection' function buildMultitenantPgPoolConfig(config: ReturnType): PoolConfig { const { @@ -149,7 +150,7 @@ function getPoolConfigSignature(config: PoolConfig): string { const multitenantPgPoolOwner = new MultitenantPgPoolOwner() -export const multitenantPgExecutor: PgTransactionalExecutor = { +export const multitenantPgExecutor: DatabaseTransactionalExecutor = { async query(statement, options) { return multitenantPgPoolOwner.getExecutor().query(statement, options) }, diff --git a/src/internal/database/pg-connection.test.ts b/src/internal/database/pg-connection.test.ts index 3762e8847..f63138814 100644 --- a/src/internal/database/pg-connection.test.ts +++ b/src/internal/database/pg-connection.test.ts @@ -5,9 +5,9 @@ import { DatabaseError, Pool as PgPool, type Pool, type PoolClient } from 'pg' // the same class to keep cancellation pending without opening real sockets. import PgConnection from 'pg/lib/connection' import { vi } from 'vitest' +import type { DatabaseExecutor } from './connection' import { getPgCancelConnectionTarget, - type PgExecutor, PgPoolExecutor, PgPoolManager, PgPoolStrategy, @@ -1536,7 +1536,7 @@ describe('PgTenantConnection', () => { ) const executor = { query: vi.fn().mockResolvedValue({ rows: [] }), - } as unknown as PgExecutor + } as unknown as DatabaseExecutor const stringifySpy = vi.spyOn(JSON, 'stringify') try { @@ -2081,7 +2081,7 @@ describe('PgTenantConnection payload serialization', () => { expect(superUserToJSON).not.toHaveBeenCalled() const query = vi.fn().mockResolvedValue({ rows: [] }) - await superUser.setScope({ query } as unknown as PgExecutor) + await superUser.setScope({ query } as unknown as DatabaseExecutor) expect(userToJSON).not.toHaveBeenCalled() expect(superUserToJSON).toHaveBeenCalledTimes(1) @@ -2119,17 +2119,19 @@ describe('PgTenantConnection payload serialization', () => { expect(stringifySpy.mock.calls.length).toBe(afterSuperUser) const parentQuery = vi.fn().mockResolvedValue({ rows: [] }) - await parent.setScope({ query: parentQuery } as unknown as PgExecutor) + await parent.setScope({ query: parentQuery } as unknown as DatabaseExecutor) const afterParentScope = stringifySpy.mock.calls.length const siblingQuery = vi.fn().mockResolvedValue({ rows: [] }) - await sibling.setScope({ query: siblingQuery } as unknown as PgExecutor) + await sibling.setScope({ query: siblingQuery } as unknown as DatabaseExecutor) expect(stringifySpy.mock.calls.length).toBe(afterParentScope) const superUserQuery = vi.fn().mockResolvedValue({ rows: [] }) - await superUser.setScope({ query: superUserQuery } as unknown as PgExecutor) + await superUser.setScope({ query: superUserQuery } as unknown as DatabaseExecutor) const afterSuperUserScope = stringifySpy.mock.calls.length const secondSuperUserQuery = vi.fn().mockResolvedValue({ rows: [] }) - await secondSuperUser.setScope({ query: secondSuperUserQuery } as unknown as PgExecutor) + await secondSuperUser.setScope({ + query: secondSuperUserQuery, + } as unknown as DatabaseExecutor) expect(stringifySpy.mock.calls.length).toBe(afterSuperUserScope) expect(stringifySpy).toHaveBeenCalledWith(userPayload) expect(stringifySpy).toHaveBeenCalledWith(superPayload) diff --git a/src/internal/database/pg-connection.ts b/src/internal/database/pg-connection.ts index 42635fb8f..8edb3eeba 100644 --- a/src/internal/database/pg-connection.ts +++ b/src/internal/database/pg-connection.ts @@ -1,9 +1,17 @@ import { ERRORS } from '@internal/errors' import { logger, logSchema } from '@internal/monitoring' -import { TransactionOptions } from '@storage/database' import pg, { DatabaseError, Pool, PoolClient, QueryResult, QueryResultRow } from 'pg' import PgConnection from 'pg/lib/connection' import { getConfig } from '../../config' +import type { + DatabaseExecutor, + DatabaseQueryArgument, + DatabaseStatement, + DatabaseTransaction, + DatabaseTransactionalExecutor, + TenantConnection, + TransactionOptions, +} from './connection' import { PoolManager, PoolRebalanceOptions, @@ -33,17 +41,6 @@ const { pg.types.setTypeParser(20, 'text', parseInt) -export interface PgStatement { - text: string - values?: unknown[] -} - -export interface PgQueryOptions { - signal?: AbortSignal -} - -type PgQueryArgument = PgQueryOptions | unknown[] - interface PgTransactionOptions { statementTimeoutMs?: number } @@ -80,17 +77,6 @@ const scopeConfigSqlWithStatementTimeout = ` set_config('statement_timeout', $10, true); ` -export interface PgExecutor { - query( - statement: string | PgStatement, - options?: PgQueryArgument - ): Promise> -} - -export interface PgTransactionalExecutor extends PgExecutor { - beginTransaction(options?: PgBeginTransactionOptions): Promise -} - interface PgPoolErrorContext { message: string tenantId?: string @@ -406,12 +392,12 @@ class PgClientErrorTracker { } } -export class PgPoolExecutor implements PgTransactionalExecutor { +export class PgPoolExecutor implements DatabaseTransactionalExecutor { constructor(private readonly pool: Pool) {} async query( - statement: string | PgStatement, - options?: PgQueryArgument + statement: string | DatabaseStatement, + options?: DatabaseQueryArgument ): Promise> { assertValidSignal(getQuerySignal(options)) @@ -482,7 +468,7 @@ export class PgPoolExecutor implements PgTransactionalExecutor { } } -export class PgTransaction implements PgExecutor { +export class PgTransaction implements DatabaseTransaction { private completed = false private statementTimeoutMs?: number @@ -499,15 +485,15 @@ export class PgTransaction implements PgExecutor { } async query( - statement: string | PgStatement, - options?: PgQueryArgument + statement: string | DatabaseStatement, + options?: DatabaseQueryArgument ): Promise> { return this.runQuery(statement, options, true) } async runSetupQuery( - statement: string | PgStatement, - options?: PgQueryArgument + statement: string | DatabaseStatement, + options?: DatabaseQueryArgument ): Promise> { // Setup statements must not consume a deferred timeout before scope can fold it in. return this.runQuery(statement, options, false) @@ -522,8 +508,8 @@ export class PgTransaction implements PgExecutor { } private async runQuery( - statement: string | PgStatement, - options: PgQueryArgument | undefined, + statement: string | DatabaseStatement, + options: DatabaseQueryArgument | undefined, applyPendingStatementTimeout: boolean ): Promise> { if (this.completed) { @@ -632,7 +618,7 @@ function serializeJwtPayload(payload: object): string { return serialized } -export class PgTenantConnection { +export class PgTenantConnection implements TenantConnection { static poolManager = new PgPoolManager() public readonly role: string private abortSignal?: AbortSignal @@ -671,8 +657,8 @@ export class PgTenantConnection { } async query( - statement: string | PgStatement, - options?: PgQueryArgument + statement: string | DatabaseStatement, + options?: DatabaseQueryArgument ): Promise> { this.assertNotDisposed() return this.pool.acquire().query(statement, options) @@ -794,7 +780,7 @@ export class PgTenantConnection { } } - async setScope(tnx: PgExecutor) { + async setScope(tnx: DatabaseExecutor) { const statementTimeoutMs = tnx instanceof PgTransaction ? tnx.takePendingStatementTimeoutMs() : undefined @@ -867,8 +853,8 @@ export function createAbortError(): Error & { code: string } { async function runPgQuery( client: PoolClient, - statement: string | PgStatement, - options?: PgQueryArgument + statement: string | DatabaseStatement, + options?: DatabaseQueryArgument ): Promise> { const signal = Array.isArray(options) ? undefined : options?.signal assertValidSignal(signal) @@ -920,7 +906,10 @@ async function runPgQuery( } } -function normalizeStatement(statement: string | PgStatement, values?: unknown[]): PgStatement { +function normalizeStatement( + statement: string | DatabaseStatement, + values?: unknown[] +): DatabaseStatement { if (typeof statement === 'string') { return { text: statement, values } } @@ -942,7 +931,7 @@ function assertValidSignal(signal?: AbortSignal): void { } } -function getQuerySignal(options?: PgQueryArgument): AbortSignal | undefined { +function getQuerySignal(options?: DatabaseQueryArgument): AbortSignal | undefined { return Array.isArray(options) ? undefined : options?.signal } diff --git a/src/internal/database/tenant-store-pg.test.ts b/src/internal/database/tenant-store-pg.test.ts index e267ce936..75d542313 100644 --- a/src/internal/database/tenant-store-pg.test.ts +++ b/src/internal/database/tenant-store-pg.test.ts @@ -1,5 +1,5 @@ import { spyOnAbortSignalAny, spyOnAbortSignalTimeout } from '../../test/utils/abort-signal' -import type { PgExecutor, PgStatement } from './pg-connection' +import type { DatabaseExecutor, DatabaseStatement } from './connection' import { TenantConfigStorePg } from './tenant-store-pg' function createTenantStore() { @@ -7,16 +7,16 @@ function createTenantStore() { rows: [], rowCount: 1, }) - const store = new TenantConfigStorePg({ query } as unknown as PgExecutor) + const store = new TenantConfigStorePg({ query } as unknown as DatabaseExecutor) return { query, store } } -function getLastStatement(query: ReturnType): PgStatement { +function getLastStatement(query: ReturnType): DatabaseStatement { const [statement] = query.mock.calls.at(-1) || [] if (!statement || typeof statement === 'string') { - throw new Error('Expected a PgStatement query') + throw new Error('Expected a DatabaseStatement query') } return statement @@ -113,7 +113,7 @@ describe('TenantConfigStorePg', () => { await ( store as unknown as { query( - statement: PgStatement, + statement: DatabaseStatement, options: { signal: AbortSignal; timeoutMs: number } ): Promise } @@ -133,7 +133,9 @@ describe('TenantConfigStorePg', () => { rows: [], rowCount: 1, }) - const store = new ConfiguredTenantConfigStorePg({ query } as unknown as PgExecutor) + const store = new ConfiguredTenantConfigStorePg({ + query, + } as unknown as DatabaseExecutor) const { timeoutSignal, timeoutSpy } = spyOnAbortSignalTimeout() try { diff --git a/src/internal/database/tenant-store-pg.ts b/src/internal/database/tenant-store-pg.ts index aea5d3758..c970f888b 100644 --- a/src/internal/database/tenant-store-pg.ts +++ b/src/internal/database/tenant-store-pg.ts @@ -1,6 +1,6 @@ import { QueryResultRow } from 'pg' import { getConfig, JwksConfigKey } from '../../config' -import { PgExecutor } from './pg-connection' +import type { DatabaseExecutor } from './connection' import { quoteIdentifier } from './sql' const { multitenantDatabaseQueryTimeout } = getConfig() @@ -75,7 +75,7 @@ export type TenantConfigRowInput = Partial & { cursor_id: number } interface TenantQueryOptions { - db?: PgExecutor + db?: DatabaseExecutor signal?: AbortSignal /** * Positive values add an internal timeout. Zero or negative values disable @@ -85,7 +85,7 @@ interface TenantQueryOptions { } export class TenantConfigStorePg { - constructor(private db: PgExecutor) {} + constructor(private db: DatabaseExecutor) {} async list(): Promise { const result = await this.query({ @@ -109,7 +109,7 @@ export class TenantConfigStorePg { return result.rows[0] } - async insert(tenantInfo: TenantConfigRowInput, db: PgExecutor = this.db): Promise { + async insert(tenantInfo: TenantConfigRowInput, db: DatabaseExecutor = this.db): Promise { const entries = getTenantEntries(tenantInfo) const columns = entries.map(([column]) => quoteIdentifier(column)) const values = entries.map(([, value]) => value) @@ -127,7 +127,7 @@ export class TenantConfigStorePg { ) } - async upsert(tenantInfo: TenantConfigRowInput, db: PgExecutor = this.db): Promise { + async upsert(tenantInfo: TenantConfigRowInput, db: DatabaseExecutor = this.db): Promise { const entries = getTenantEntries(tenantInfo) const columns = entries.map(([column]) => quoteIdentifier(column)) const values = entries.map(([, value]) => value) @@ -156,7 +156,7 @@ export class TenantConfigStorePg { async update( tenantId: string, tenantInfo: TenantConfigRowInput, - db: PgExecutor = this.db + db: DatabaseExecutor = this.db ): Promise { const entries = getTenantEntries(tenantInfo).filter(([column]) => column !== 'id') if (entries.length === 0) { @@ -289,7 +289,7 @@ export class TenantConfigStorePg { } private query( - statement: Parameters[0], + statement: Parameters[0], options: TenantQueryOptions = {} ) { const db = options.db ?? this.db diff --git a/src/internal/queue/database.test.ts b/src/internal/queue/database.test.ts index 18cd469b1..e1bb49457 100644 --- a/src/internal/queue/database.test.ts +++ b/src/internal/queue/database.test.ts @@ -1,5 +1,5 @@ import { createRequire } from 'node:module' -import type { PgExecutor, PgStatement } from '@internal/database' +import type { DatabaseExecutor, DatabaseStatement } from '@internal/database' import { PgQueueDB } from './database' const loadCjs = createRequire(__filename) @@ -158,7 +158,7 @@ describe('PgQueueDB', () => { const query = vi.fn().mockResolvedValue({ rows: [{ ok: true }], }) - const db = new PgQueueDB({ query } as unknown as PgExecutor) + const db = new PgQueueDB({ query } as unknown as DatabaseExecutor) await expect(db.executeSql('SELECT $1, $2, $3', ['queue-name', undefined, 3])).resolves.toEqual( { @@ -169,7 +169,7 @@ describe('PgQueueDB', () => { expect(query).toHaveBeenCalledWith({ text: 'SELECT $1, $2, $3', values: ['queue-name', null, 3], - } satisfies PgStatement) + } satisfies DatabaseStatement) }) it('covers pg-boss v10 generated SQL with only pg-compatible placeholders', () => { diff --git a/src/internal/queue/database.ts b/src/internal/queue/database.ts index b71789bff..e81695da3 100644 --- a/src/internal/queue/database.ts +++ b/src/internal/queue/database.ts @@ -1,8 +1,8 @@ import EventEmitter from 'node:events' +import type { DatabaseExecutor } from '@internal/database/connection' import { ERRORS } from '@internal/errors' import pg from 'pg' import { Db } from 'pg-boss' -import { PgExecutor } from '../database/pg-connection' export { quoteIdentifier } from '../database/sql' @@ -91,7 +91,7 @@ export class PgQueueDB extends EventEmitter implements Db { error: 'error', } - constructor(protected readonly db: PgExecutor) { + constructor(protected readonly db: DatabaseExecutor) { super() } diff --git a/src/internal/queue/event.ts b/src/internal/queue/event.ts index e04ddf980..d9b8d5570 100644 --- a/src/internal/queue/event.ts +++ b/src/internal/queue/event.ts @@ -1,5 +1,4 @@ -import { getTenantConfig } from '@internal/database' -import { PgExecutor } from '@internal/database/pg-connection' +import { type DatabaseExecutor, getTenantConfig } from '@internal/database' import { ERRORS } from '@internal/errors' import { logger, logSchema } from '@internal/monitoring' import { queueJobScheduled, queueJobSchedulingTime } from '@internal/monitoring/metrics' @@ -22,7 +21,7 @@ export interface BasePayload { } const { pgQueueEnable, region, isMultitenant } = getConfig() -type TransactionalQueueDb = PgExecutor +type TransactionalQueueDb = DatabaseExecutor function withPayloadVersion( payload: TPayload, diff --git a/src/internal/sharding/pg.ts b/src/internal/sharding/pg.ts index 08b2e5bb0..b8c1039ff 100644 --- a/src/internal/sharding/pg.ts +++ b/src/internal/sharding/pg.ts @@ -1,7 +1,12 @@ import { hashStringToInt } from '@internal/hashing' import { logger, logSchema } from '@internal/monitoring' import { DatabaseError, QueryResultRow } from 'pg' -import { PgExecutor, PgTransaction, PgTransactionalExecutor } from '../database/pg-connection' +import { + type DatabaseExecutor, + type DatabaseTransaction, + type DatabaseTransactionalExecutor, + isDatabaseTransaction, +} from '../database/connection' import { ReservationRow, ResourceKind, @@ -12,15 +17,15 @@ import { UniqueViolationError, } from './store' -export class PgShardStoreFactory implements ShardStoreFactory { - constructor(private db: PgTransactionalExecutor | PgTransaction) {} +export class PgShardStoreFactory implements ShardStoreFactory { + constructor(private db: DatabaseTransactionalExecutor | DatabaseTransaction) {} - withExistingTransaction(tnx: PgTransaction): ShardStoreFactory { + withExistingTransaction(tnx: DatabaseTransaction): ShardStoreFactory { return new PgShardStoreFactory(tnx) } async withTransaction(fn: (store: ShardStore) => Promise): Promise { - if (this.db instanceof PgTransaction) { + if (isDatabaseTransaction(this.db)) { return fn(new PgShardStore(this.db)) } @@ -49,7 +54,7 @@ export class PgShardStoreFactory implements ShardStoreFactory { } class PgShardStore implements ShardStore { - constructor(private db: PgExecutor) {} + constructor(private db: DatabaseExecutor) {} private query(text: string, values?: unknown[]) { return this.db.query({ text, values }) diff --git a/src/storage/database/adapter.ts b/src/storage/database/adapter.ts index 5c2b235b0..693a088fb 100644 --- a/src/storage/database/adapter.ts +++ b/src/storage/database/adapter.ts @@ -1,4 +1,4 @@ -import { PgTenantConnection } from '@internal/database' +import type { TenantConnection, TransactionOptions } from '@internal/database' import { DBMigration } from '@internal/database/migrations' import { ObjectMetadata } from '../backend' import { Bucket, IcebergCatalog, Obj, S3MultipartUpload, S3PartUpload } from '../schemas' @@ -28,13 +28,6 @@ export interface FindObjectFilters { dontErrorOnEmpty?: boolean } -export interface TransactionOptions { - isolation?: string - retry?: number - readOnly?: boolean - timeout?: number -} - export interface DatabaseOptions { tenantId: string reqId?: string @@ -43,7 +36,7 @@ export interface DatabaseOptions { host: string tnx?: TNX parentTnx?: TNX - parentConnection?: PgTenantConnection + parentConnection?: TenantConnection } export interface ListBucketOptions { @@ -65,7 +58,7 @@ export interface Database { reqId?: string sbReqId?: string role?: string - connection: PgTenantConnection + connection: TenantConnection tenant(): { ref: string; host: string } diff --git a/src/storage/database/pg.test.ts b/src/storage/database/pg.test.ts index 944a39374..d885c3eb0 100644 --- a/src/storage/database/pg.test.ts +++ b/src/storage/database/pg.test.ts @@ -1,8 +1,8 @@ import { - type PgExecutor, + type DatabaseExecutor, + type DatabaseTransaction, PgPoolExecutor, PgTenantConnection, - type PgTransaction, } from '@internal/database' import { normalizeRawError } from '@internal/errors' import { dbQueryPerformance } from '@internal/monitoring/metrics' @@ -27,7 +27,10 @@ class TestStoragePgDB extends StoragePgDB { } runUnscopedErrorMappingProbe(): Promise { - return this.runUnscopedQuery('CreateS3KeysTempTable', async () => 'ok') + return this.runUnscopedQuery('CreateS3KeysTempTable', async (db) => { + await db.query('SELECT 1') + return 'ok' + }) } } @@ -78,7 +81,7 @@ function createNestedTestPermissionFixture() { const storage = new StoragePgDB(connection, { tenantId: 'nested-test-permission-tenant', host: 'localhost', - tnx: transaction as unknown as PgTransaction, + tnx: transaction as unknown as DatabaseTransaction, }) return { connection, storage, transaction } @@ -298,7 +301,7 @@ describe('StoragePgDB healthcheck', () => { }) function createHealthcheckFixture( - executor: PgExecutor, + executor: DatabaseExecutor, options: { requestSignal?: AbortSignal StorageClass?: typeof StoragePgDB @@ -313,6 +316,7 @@ describe('StoragePgDB healthcheck', () => { } const connection = { getAbortSignal: vi.fn().mockReturnValue(requestSignal), + query: vi.fn((statement, queryOptions) => executor.query(statement, queryOptions)), pool: { acquire: vi.fn().mockReturnValue(executor), }, @@ -383,7 +387,7 @@ describe('StoragePgDB healthcheck', () => { await rejection expect(release).toHaveBeenCalledWith(expect.objectContaining(expectedAbortError)) - expect(fixture.connection.pool.acquire).toHaveBeenCalledTimes(1) + expect(fixture.connection.query).toHaveBeenCalledTimes(1) expect(fixture.connection.transaction).not.toHaveBeenCalled() expect(fixture.connection.setScope).not.toHaveBeenCalled() }) @@ -426,12 +430,12 @@ describe('StoragePgDB healthcheck', () => { test('uses the scoped readiness probe by default', async () => { const executor = { query: vi.fn().mockResolvedValue({ rows: [] }), - } as unknown as PgExecutor + } as unknown as DatabaseExecutor const fixture = createHealthcheckFixture(executor) await expect(fixture.storage.healthcheck()).resolves.toBeUndefined() - expect(fixture.connection.pool.acquire).not.toHaveBeenCalled() + expect(fixture.connection.query).not.toHaveBeenCalled() expect(fixture.connection.transaction).toHaveBeenCalledTimes(1) expect(fixture.connection.setScope).toHaveBeenCalledWith(fixture.transaction) expect(fixture.transaction.query).toHaveBeenCalledWith(probeSql, { signal: undefined }) @@ -528,16 +532,14 @@ describe('StoragePgDB error mapping', () => { }) }) - test('preserves query name for pg errors thrown while acquiring an unscoped executor', async () => { + test('preserves query name for pg errors thrown by an unscoped executor', async () => { const error = createPgError('08006', 'connection failure') error.severity = 'FATAL' const connection = { getAbortSignal: vi.fn().mockReturnValue(undefined), - pool: { - acquire: vi.fn(() => { - throw error - }), - }, + query: vi.fn(() => { + throw error + }), } as unknown as PgTenantConnection const storage = new TestStoragePgDB(connection, { tenantId: 'tenant-with-unscoped-acquire-error', diff --git a/src/storage/database/pg.ts b/src/storage/database/pg.ts index e42bef05b..a490bcab5 100644 --- a/src/storage/database/pg.ts +++ b/src/storage/database/pg.ts @@ -1,11 +1,12 @@ import { randomUUID } from 'node:crypto' import { - PgExecutor, - PgStatement, - PgTenantConnection, - PgTransaction, + type DatabaseExecutor, + type DatabaseStatement, + type DatabaseTransaction, quoteIdentifier, quoteQualifiedIdentifier, + type TenantConnection, + type TransactionOptions, } from '@internal/database' import { DBMigration, tenantHasMigrations } from '@internal/database/migrations' import { ERRORS, ErrorCode, isStorageError, StorageBackendError } from '@internal/errors' @@ -24,7 +25,6 @@ import { ListBucketOptions, ScannerS3Key, SearchObjectOption, - TransactionOptions, } from './adapter' import { DBError, mapPgTransactionAbortedError, PgErrorContext } from './errors' @@ -48,9 +48,9 @@ interface PgDatabaseOptions { latestMigration?: keyof typeof DBMigration databaseEngine?: DatabaseEngine host: string - tnx?: PgTransaction - parentTnx?: PgTransaction - parentConnection?: PgTenantConnection + tnx?: DatabaseTransaction + parentTnx?: DatabaseTransaction + parentConnection?: TenantConnection } interface UnscopedQueryOptions { @@ -63,8 +63,8 @@ const HEALTHCHECK_QUERY_OPTIONS: UnscopedQueryOptions = Object.freeze({ }) async function executeQuery( - db: PgExecutor, - statement: string | PgStatement, + db: DatabaseExecutor, + statement: string | DatabaseStatement, signal?: AbortSignal ) { try { @@ -74,7 +74,7 @@ async function executeQuery( } } -function healthcheckProbe(db: PgExecutor, signal?: AbortSignal) { +function healthcheckProbe(db: DatabaseExecutor, signal?: AbortSignal) { return executeQuery(db, HEALTHCHECK_SQL, signal) } @@ -102,7 +102,7 @@ export class StoragePgDB implements Database { public readonly latestMigration?: keyof typeof DBMigration constructor( - public readonly connection: PgTenantConnection, + public readonly connection: TenantConnection, private readonly options: PgDatabaseOptions ) { this.tenantHost = options.host @@ -1302,7 +1302,7 @@ export class StoragePgDB implements Database { } private async waitObjectLockWithTopLevelLockTimeout( - db: PgExecutor, + db: DatabaseExecutor, hash: number, lockTimeout: number, signal?: AbortSignal @@ -1567,7 +1567,10 @@ export class StoragePgDB implements Database { }) } - private async dropStaleS3KeysScratchTables(db: PgExecutor, signal?: AbortSignal): Promise { + private async dropStaleS3KeysScratchTables( + db: DatabaseExecutor, + signal?: AbortSignal + ): Promise { const staleBefore = Date.now() - S3_KEYS_SCRATCH_TABLE_MAX_AGE_MS const result = await this.query<{ table_name: string }>( db, @@ -1784,7 +1787,7 @@ export class StoragePgDB implements Database { protected async runQuery( queryName: string, - fn: (db: PgExecutor, signal?: AbortSignal) => Promise + fn: (db: DatabaseExecutor, signal?: AbortSignal) => Promise ): Promise { const startTime = performance.now() const abortSignal = this.connection.getAbortSignal() @@ -1915,7 +1918,7 @@ export class StoragePgDB implements Database { protected async runUnscopedQuery( queryName: string, - fn: (db: PgExecutor, signal?: AbortSignal) => Promise, + fn: (db: DatabaseExecutor, signal?: AbortSignal) => Promise, options?: UnscopedQueryOptions ): Promise { const startTime = performance.now() @@ -1942,7 +1945,7 @@ export class StoragePgDB implements Database { } try { - return await fn(this.connection.pool.acquire(), controller?.signal ?? requestAbortSignal) + return await fn(this.connection, controller?.signal ?? requestAbortSignal) } catch (e) { throw mapPgErrorWithQueryName(e, queryName) } finally { @@ -1957,8 +1960,8 @@ export class StoragePgDB implements Database { } protected query( - db: PgExecutor, - statement: string | PgStatement, + db: DatabaseExecutor, + statement: string | DatabaseStatement, signal?: AbortSignal ) { return executeQuery(db, statement, signal) @@ -2187,7 +2190,7 @@ function nextSavepointName(): string { return quoteIdentifier(`storage_pg_query_${randomUUID().replace(/-/g, '_')}`) } -async function createSavepoint(tnx: PgTransaction, savepoint: string): Promise { +async function createSavepoint(tnx: DatabaseTransaction, savepoint: string): Promise { const query = `SAVEPOINT ${savepoint}` try { @@ -2197,7 +2200,7 @@ async function createSavepoint(tnx: PgTransaction, savepoint: string): Promise { +async function rollbackSavepoint(tnx: DatabaseTransaction, savepoint: string): Promise { await tnx.query(`ROLLBACK TO SAVEPOINT ${savepoint}`) await tnx.query(`RELEASE SAVEPOINT ${savepoint}`) } diff --git a/src/storage/events/iceberg/delete-iceberg-resources.test.ts b/src/storage/events/iceberg/delete-iceberg-resources.test.ts index 6718462b5..0616ddf4c 100644 --- a/src/storage/events/iceberg/delete-iceberg-resources.test.ts +++ b/src/storage/events/iceberg/delete-iceberg-resources.test.ts @@ -96,9 +96,7 @@ const sharder = { freeByResource: vi.fn() } const db = { connection: { - pool: { - acquire: vi.fn(), - }, + query: vi.fn(), }, deleteAnalyticsBucket: vi.fn(), destroyConnection: vi.fn(), @@ -107,7 +105,7 @@ const db = { function expectIcebergCleanup({ multitenant }: { multitenant: boolean }) { expect(mockCreateStorage).toHaveBeenCalledWith(jobData) expect(MockPgMetastore).toHaveBeenCalledWith( - multitenant ? mockMultitenantPgExecutor : 'mock-db-connection', + multitenant ? mockMultitenantPgExecutor : db.connection, { multiTenant: multitenant, schema: multitenant ? 'public' : 'storage', @@ -201,8 +199,6 @@ describe('DeleteIcebergResources.handle', () => { metastore.transaction.mockImplementation(async (fn) => fn(store)) store.getTnx.mockReturnValue('mock-transaction') shardCatalog.withTnx.mockReturnValue(sharder) - db.connection.pool.acquire.mockReturnValue('mock-db-connection') - store.findCatalogById.mockResolvedValue({ id: 'catalog-123', deleted_at: new Date() }) store.listNamespaces.mockResolvedValue([{ id: 'ns-1', name: 'namespace-1' }]) store.listTables.mockResolvedValue([ @@ -255,7 +251,7 @@ describe('DeleteIcebergResources.handle', () => { expect(metastore.transaction).not.toHaveBeenCalled() expect(store.lockResource).not.toHaveBeenCalled() expect(store.dropCatalog).not.toHaveBeenCalled() - expect(db.connection.pool.acquire).not.toHaveBeenCalled() + expect(db.connection.query).not.toHaveBeenCalled() expect(db.deleteAnalyticsBucket).not.toHaveBeenCalled() expect(db.destroyConnection).not.toHaveBeenCalled() }) @@ -265,7 +261,6 @@ describe('DeleteIcebergResources.handle', () => { await expect(DeleteIcebergResources.handle(makeJob() as never)).resolves.toBeUndefined() - expect(db.connection.pool.acquire).toHaveBeenCalled() expectIcebergCleanup({ multitenant: false }) expect(db.deleteAnalyticsBucket).not.toHaveBeenCalled() expect(db.destroyConnection).toHaveBeenCalled() diff --git a/src/storage/events/iceberg/delete-iceberg-resources.ts b/src/storage/events/iceberg/delete-iceberg-resources.ts index c16fb0898..7c7b02c96 100644 --- a/src/storage/events/iceberg/delete-iceberg-resources.ts +++ b/src/storage/events/iceberg/delete-iceberg-resources.ts @@ -62,7 +62,7 @@ export class DeleteIcebergResources extends BaseEvent { const eventStorage = storage const metastore = new PgMetastore( - isMultitenant ? multitenantPgExecutor : eventStorage.db.connection.pool.acquire(), + isMultitenant ? multitenantPgExecutor : eventStorage.db.connection, { multiTenant: isMultitenant, schema: isMultitenant ? 'public' : 'storage', diff --git a/src/storage/events/pgboss/move-jobs.ts b/src/storage/events/pgboss/move-jobs.ts index 60e6884c6..4d9383ba5 100644 --- a/src/storage/events/pgboss/move-jobs.ts +++ b/src/storage/events/pgboss/move-jobs.ts @@ -1,4 +1,4 @@ -import { multitenantPgExecutor, PgTransaction } from '@internal/database' +import { type DatabaseTransaction, multitenantPgExecutor } from '@internal/database' import { logger, logSchema } from '@internal/monitoring' import { BasePayload, PG_BOSS_SCHEMA, Queue, SYSTEM_TENANT_REF } from '@internal/queue' import { Job, Queue as PgBossQueue, SendOptions, WorkOptions } from 'pg-boss' @@ -137,7 +137,7 @@ export class MoveJobs extends BaseEvent { } } -async function withPgTransaction(fn: (tnx: PgTransaction) => Promise): Promise { +async function withPgTransaction(fn: (tnx: DatabaseTransaction) => Promise): Promise { const tnx = await multitenantPgExecutor.beginTransaction() try { diff --git a/src/storage/events/upgrades/base-event.ts b/src/storage/events/upgrades/base-event.ts index 7d36b81af..89f614626 100644 --- a/src/storage/events/upgrades/base-event.ts +++ b/src/storage/events/upgrades/base-event.ts @@ -1,4 +1,4 @@ -import { multitenantPgExecutor, PgTransaction } from '@internal/database' +import { type DatabaseTransaction, multitenantPgExecutor } from '@internal/database' import { hashStringToInt } from '@internal/hashing' import { logger, logSchema } from '@internal/monitoring' import type { BasePayload } from '@internal/queue' @@ -9,7 +9,7 @@ import { getConfig } from '../../../config' const { isMultitenant } = getConfig() export type UpgradeBaseEventPayload = BasePayload -export type UpgradeTransaction = PgTransaction +export type UpgradeTransaction = DatabaseTransaction export abstract class UpgradeBaseEvent extends BaseEvent { static getQueueOptions(): PgBossQueue { diff --git a/src/storage/protocols/iceberg/catalog/reconciler.test.ts b/src/storage/protocols/iceberg/catalog/reconciler.test.ts index a1d982507..dccdba3ec 100644 --- a/src/storage/protocols/iceberg/catalog/reconciler.test.ts +++ b/src/storage/protocols/iceberg/catalog/reconciler.test.ts @@ -19,7 +19,7 @@ function getLastStatement(query: ReturnType): string { const [statement] = query.mock.calls.at(-1) || [] if (!statement || typeof statement === 'string') { - throw new Error('Expected a PgStatement query') + throw new Error('Expected a DatabaseStatement query') } return String((statement as { text: string }).text) diff --git a/src/storage/protocols/iceberg/catalog/reconciler.ts b/src/storage/protocols/iceberg/catalog/reconciler.ts index 36582765e..34e48cd16 100644 --- a/src/storage/protocols/iceberg/catalog/reconciler.ts +++ b/src/storage/protocols/iceberg/catalog/reconciler.ts @@ -1,4 +1,4 @@ -import { multitenantPgExecutor, PgTransaction } from '@internal/database' +import { type DatabaseTransaction, multitenantPgExecutor } from '@internal/database' import { logger, logSchema } from '@internal/monitoring' import { PgShardStoreFactory, ShardCatalog, ShardRow } from '@internal/sharding' import { @@ -10,7 +10,7 @@ import { IcebergCatalog } from '@storage/schemas' type NamespaceWithShardInfo = TableIndex & { shard_id?: string; shard_key?: string } type CatalogRow = Pick -type ReconcilerTransaction = PgTransaction +type ReconcilerTransaction = DatabaseTransaction /** * Highly experimental reconciler for iceberg catalogs diff --git a/src/storage/protocols/iceberg/metastore.ts b/src/storage/protocols/iceberg/metastore.ts index 60bb39088..2b6bc9eaa 100644 --- a/src/storage/protocols/iceberg/metastore.ts +++ b/src/storage/protocols/iceberg/metastore.ts @@ -1,4 +1,4 @@ -import { TransactionOptions } from '@storage/database' +import type { TransactionOptions } from '@internal/database' import { IcebergCatalog } from '@storage/schemas' export interface CreateNamespaceParams { diff --git a/src/storage/protocols/iceberg/pg.test.ts b/src/storage/protocols/iceberg/pg.test.ts index f00f3128e..262d05a38 100644 --- a/src/storage/protocols/iceberg/pg.test.ts +++ b/src/storage/protocols/iceberg/pg.test.ts @@ -23,7 +23,7 @@ function getLastStatement(query: ReturnType): string { return String(statement.text) } - throw new Error('Expected a PgStatement query with text property') + throw new Error('Expected a DatabaseStatement query with text property') } describe('PgMetastore.countCatalogs', () => { diff --git a/src/storage/protocols/iceberg/pg.ts b/src/storage/protocols/iceberg/pg.ts index 3d9e16c9a..d370b0ac7 100644 --- a/src/storage/protocols/iceberg/pg.ts +++ b/src/storage/protocols/iceberg/pg.ts @@ -1,14 +1,16 @@ import { randomUUID } from 'node:crypto' import { - PgStatement, - PgTransaction, - PgTransactionalExecutor, + type DatabaseStatement, + type DatabaseTransaction, + type DatabaseTransactionalExecutor, + isDatabaseTransaction, quoteIdentifier, + type TransactionOptions, } from '@internal/database' import { ERRORS, StorageBackendError } from '@internal/errors' import { hashStringToInt } from '@internal/hashing' import { logger, logSchema } from '@internal/monitoring' -import { DBError, mapPgTransactionAbortedError, TransactionOptions } from '@storage/database' +import { DBError, mapPgTransactionAbortedError } from '@storage/database' import { IcebergCatalog } from '@storage/schemas' import { DatabaseError, QueryResult, QueryResultRow } from 'pg' import { @@ -22,9 +24,9 @@ import { TableIndex, } from './metastore' -export class PgMetastore implements Metastore { +export class PgMetastore implements Metastore { constructor( - private readonly db: PgTransactionalExecutor | PgTransaction, + private readonly db: DatabaseTransactionalExecutor | DatabaseTransaction, private readonly ops: { schema: string; multiTenant?: boolean } ) {} @@ -33,8 +35,8 @@ export class PgMetastore implements Metastore { await this.query('SELECT pg_advisory_xact_lock($1::bigint)', [String(lockId)]) } - getTnx(): PgTransaction { - if (this.db instanceof PgTransaction) { + getTnx(): DatabaseTransaction { + if (isDatabaseTransaction(this.db)) { return this.db } @@ -263,7 +265,7 @@ export class PgMetastore implements Metastore { callback: (trx: PgMetastore) => Promise, opts?: TransactionOptions ): Promise { - if (this.db instanceof PgTransaction) { + if (isDatabaseTransaction(this.db)) { const savepoint = nextSavepointName() let savepointEstablished = false @@ -631,7 +633,7 @@ export class PgMetastore implements Metastore { } private queryRaw( - statement: PgStatement + statement: DatabaseStatement ): Promise> { return this.db.query(statement) } @@ -653,7 +655,7 @@ function nextSavepointName(): string { return quoteIdentifier(`iceberg_pg_transaction_${randomUUID().replace(/-/g, '_')}`) } -async function createSavepoint(tnx: PgTransaction, savepoint: string): Promise { +async function createSavepoint(tnx: DatabaseTransaction, savepoint: string): Promise { const query = `SAVEPOINT ${savepoint}` try { @@ -663,7 +665,7 @@ async function createSavepoint(tnx: PgTransaction, savepoint: string): Promise { +async function rollbackSavepoint(tnx: DatabaseTransaction, savepoint: string): Promise { await tnx.query(`ROLLBACK TO SAVEPOINT ${savepoint}`) await tnx.query(`RELEASE SAVEPOINT ${savepoint}`) } diff --git a/src/storage/protocols/s3/credentials/store-pg.ts b/src/storage/protocols/s3/credentials/store-pg.ts index a5e86dd63..7b4ce3f3b 100644 --- a/src/storage/protocols/s3/credentials/store-pg.ts +++ b/src/storage/protocols/s3/credentials/store-pg.ts @@ -1,5 +1,5 @@ +import type { DatabaseExecutor } from '@internal/database' import { getConfig } from '../../../../config' -import { PgExecutor } from '../../../../internal/database/pg-connection' import { S3Credentials, S3CredentialsManagerStore, @@ -10,7 +10,7 @@ import { const { multitenantDatabaseQueryTimeout } = getConfig() export class S3CredentialsManagerStorePg implements S3CredentialsManagerStore { - constructor(private db: PgExecutor) {} + constructor(private db: DatabaseExecutor) {} async insert(tenantId: string, credential: S3CredentialWithDescription): Promise { const credentials = await this.db.query<{ id: string }>( diff --git a/src/storage/protocols/vector/adapter/pgvector/index.ts b/src/storage/protocols/vector/adapter/pgvector/index.ts index 17a348c82..9605b2edb 100644 --- a/src/storage/protocols/vector/adapter/pgvector/index.ts +++ b/src/storage/protocols/vector/adapter/pgvector/index.ts @@ -17,10 +17,11 @@ import { QueryVectorsOutput, } from '@aws-sdk/client-s3vectors' import { - PgExecutor, + type DatabaseExecutor, + type DatabaseTransaction, + type DatabaseTransactionalExecutor, + isDatabaseTransaction, PgTenantConnection, - PgTransaction, - PgTransactionalExecutor, quoteIdentifier, } from '@internal/database' import { ERRORS } from '@internal/errors' @@ -152,7 +153,7 @@ function tableCapabilityCache(db: object): Map { const cache = tableCapabilityCache(cacheKey) let capability = cache.get(table) @@ -216,25 +217,29 @@ async function resolveTableCapability( } export type PgExecutorResolver = - | PgTransactionalExecutor - | PgTransaction + | DatabaseTransactionalExecutor + | DatabaseTransaction | { - resolve: () => PgTransactionalExecutor | PgTransaction - root?: () => PgTransactionalExecutor | PgTransaction + resolve: () => DatabaseTransactionalExecutor | DatabaseTransaction + root?: () => DatabaseTransactionalExecutor | DatabaseTransaction } function isPgExecutorProvider(r: PgExecutorResolver): r is { - resolve: () => PgTransactionalExecutor | PgTransaction - root?: () => PgTransactionalExecutor | PgTransaction + resolve: () => DatabaseTransactionalExecutor | DatabaseTransaction + root?: () => DatabaseTransactionalExecutor | DatabaseTransaction } { - return typeof (r as { resolve?: unknown }).resolve === 'function' + return 'resolve' in r && typeof r.resolve === 'function' } -function resolvePgExecutor(r: PgExecutorResolver): PgTransactionalExecutor | PgTransaction { +function resolvePgExecutor( + r: PgExecutorResolver +): DatabaseTransactionalExecutor | DatabaseTransaction { return isPgExecutorProvider(r) ? r.resolve() : r } -function resolveRootPgExecutor(r: PgExecutorResolver): PgTransactionalExecutor | PgTransaction { +function resolveRootPgExecutor( + r: PgExecutorResolver +): DatabaseTransactionalExecutor | DatabaseTransaction { return isPgExecutorProvider(r) ? (r.root?.() ?? r.resolve()) : r } @@ -243,10 +248,10 @@ function hasRootPgResolver(r: PgExecutorResolver): boolean { } async function withPgTransaction( - db: PgTransactionalExecutor | PgTransaction, - fn: (trx: PgTransaction) => Promise + db: DatabaseTransactionalExecutor | DatabaseTransaction, + fn: (trx: DatabaseTransaction) => Promise ): Promise { - if (db instanceof PgTransaction) { + if (isDatabaseTransaction(db)) { const savepoint = nextSavepointName() await db.query(`SAVEPOINT ${savepoint}`) try { @@ -377,11 +382,11 @@ export class PgVectorStore implements VectorStore { this.transactionalIndexOperations = hasRootPgResolver(executor) } - private db(): PgTransactionalExecutor | PgTransaction { + private db(): DatabaseTransactionalExecutor | DatabaseTransaction { return resolvePgExecutor(this.executor) } - private rootDb(): PgTransactionalExecutor | PgTransaction { + private rootDb(): DatabaseTransactionalExecutor | DatabaseTransaction { return resolveRootPgExecutor(this.executor) } @@ -536,7 +541,7 @@ export class PgVectorStore implements VectorStore { } private async putVectorsManually( - db: PgTransactionalExecutor | PgTransaction, + db: DatabaseTransactionalExecutor | DatabaseTransaction, table: string, serializedRows: string ): Promise { @@ -604,7 +609,7 @@ export class PgVectorStore implements VectorStore { } private async queryVectorsRaw( - db: PgTransactionalExecutor | PgTransaction, + db: DatabaseTransactionalExecutor | DatabaseTransaction, table: string, sql: string, params: unknown[], diff --git a/src/storage/protocols/vector/pg.ts b/src/storage/protocols/vector/pg.ts index 1427daa73..85c3e826e 100644 --- a/src/storage/protocols/vector/pg.ts +++ b/src/storage/protocols/vector/pg.ts @@ -2,7 +2,12 @@ import { AsyncLocalStorage } from 'node:async_hooks' import { randomUUID } from 'node:crypto' import { ListVectorBucketsInput } from '@aws-sdk/client-s3vectors' import { wait } from '@internal/concurrency' -import { PgTransaction, PgTransactionalExecutor, quoteIdentifier } from '@internal/database' +import { + type DatabaseTransaction, + type DatabaseTransactionalExecutor, + isDatabaseTransaction, + quoteIdentifier, +} from '@internal/database' import { ERRORS } from '@internal/errors' import { hashStringToInt } from '@internal/hashing' import { logger, logSchema } from '@internal/monitoring' @@ -20,13 +25,13 @@ import { } from './metadata' const vectorTransactionStorage = new AsyncLocalStorage<{ - root: PgTransactionalExecutor - transaction: PgTransaction + root: DatabaseTransactionalExecutor | DatabaseTransaction + transaction: DatabaseTransaction }>() -export function createVectorTransactionPgResolver(db: PgTransactionalExecutor): { - resolve: () => PgTransactionalExecutor | PgTransaction - root: () => PgTransactionalExecutor +export function createVectorTransactionPgResolver(db: DatabaseTransactionalExecutor): { + resolve: () => DatabaseTransactionalExecutor | DatabaseTransaction + root: () => DatabaseTransactionalExecutor | DatabaseTransaction } { return { resolve: () => { @@ -39,16 +44,17 @@ export function createVectorTransactionPgResolver(db: PgTransactionalExecutor): export class PgVectorMetadataDB implements VectorMetadataDB { constructor( - protected readonly db: PgTransactionalExecutor | PgTransaction, - private readonly rootDb: PgTransactionalExecutor = db as PgTransactionalExecutor + protected readonly db: DatabaseTransactionalExecutor | DatabaseTransaction, + private readonly rootDb: DatabaseTransactionalExecutor | DatabaseTransaction = db ) {} async withTransaction(fn: (db: VectorMetadataDB) => Promise | T): Promise { const maxRetries = 3 for (let attempt = 1; ; attempt++) { - const trx = this.db instanceof PgTransaction ? this.db : await this.db.beginTransaction() - const savepoint = this.db instanceof PgTransaction ? nextSavepointName() : undefined + const transactionExists = isDatabaseTransaction(this.db) + const trx = transactionExists ? this.db : await this.db.beginTransaction() + const savepoint = transactionExists ? nextSavepointName() : undefined let savepointEstablished = false try { @@ -366,7 +372,7 @@ function nextSavepointName(): string { return quoteIdentifier(`vector_metadata_transaction_${randomUUID().replace(/-/g, '_')}`) } -async function createSavepoint(trx: PgTransaction, savepoint: string): Promise { +async function createSavepoint(trx: DatabaseTransaction, savepoint: string): Promise { const query = `SAVEPOINT ${savepoint}` try { @@ -376,7 +382,7 @@ async function createSavepoint(trx: PgTransaction, savepoint: string): Promise { - await trx.query(`ROLLBACK TO SAVEPOINT ${savepoint}`) - await trx.query(`RELEASE SAVEPOINT ${savepoint}`) +async function rollbackSavepoint(tnx: DatabaseTransaction, savepoint: string): Promise { + await tnx.query(`ROLLBACK TO SAVEPOINT ${savepoint}`) + await tnx.query(`RELEASE SAVEPOINT ${savepoint}`) } diff --git a/src/test/database-protection.test.ts b/src/test/database-protection.test.ts index 28d210131..fa7a1ecd3 100644 --- a/src/test/database-protection.test.ts +++ b/src/test/database-protection.test.ts @@ -18,7 +18,7 @@ describe('Database Protection Triggers', () => { describe('Direct DELETE protection (migration 0050)', () => { it('should prevent direct DELETE on storage.buckets without storage.allow_delete_query', async () => { - const db = tHelper.database.connection.pool.acquire() + const db = tHelper.database.connection const testBucket = `temp-bucket-${Date.now()}` // Create a test bucket @@ -48,7 +48,7 @@ describe('Database Protection Triggers', () => { }) it('should prevent direct DELETE on storage.objects without storage.allow_delete_query', async () => { - const db = tHelper.database.connection.pool.acquire() + const db = tHelper.database.connection const testObjectName = `test-object-${Date.now()}.txt` // Create a test object @@ -87,7 +87,7 @@ describe('Database Protection Triggers', () => { }) it('should allow DELETE on storage.buckets when storage.allow_delete_query is set', async () => { - const db = tHelper.database.connection.pool.acquire() + const db = tHelper.database.connection const testBucket = `temp-bucket-allow-${Date.now()}` await withDeleteEnabled(db, async (db) => { @@ -107,7 +107,7 @@ describe('Database Protection Triggers', () => { }) it('should allow DELETE on storage.objects when storage.allow_delete_query is set', async () => { - const db = tHelper.database.connection.pool.acquire() + const db = tHelper.database.connection const testObjectName = `test-object-allow-${Date.now()}.txt` await withDeleteEnabled(db, async (db) => { diff --git a/src/test/iceberg.test.ts b/src/test/iceberg.test.ts index 5468fefb0..6e3f1048e 100644 --- a/src/test/iceberg.test.ts +++ b/src/test/iceberg.test.ts @@ -36,7 +36,7 @@ describe('Iceberg Catalog', () => { const { default: makeApp } = await import('../app') app = makeApp() - icebergMetastore = new PgMetastore(t.database.connection.pool.acquire(), { + icebergMetastore = new PgMetastore(t.database.connection, { multiTenant: false, schema: 'storage', }) @@ -48,7 +48,7 @@ describe('Iceberg Catalog', () => { afterAll(async () => { await app.close() - await t.database.connection.pool.destroy() + t.database.connection.dispose() }) it('can create an analytic bucket', async () => { @@ -108,7 +108,7 @@ describe('Iceberg Catalog', () => { expect(response.statusCode).toBe(200) - const deletedCatalog = await t.database.connection.pool.acquire().query<{ + const deletedCatalog = await t.database.connection.query<{ deleted_at: Date | string | null }>({ text: ` diff --git a/src/test/object.test.ts b/src/test/object.test.ts index 2a93ed081..ed032a999 100644 --- a/src/test/object.test.ts +++ b/src/test/object.test.ts @@ -11,7 +11,11 @@ import { signJWT, verifyJWT, } from '@internal/auth' -import { getPostgresConnection, getServiceKeyUser, PgTransaction } from '@internal/database' +import { + type DatabaseTransaction, + getPostgresConnection, + getServiceKeyUser, +} from '@internal/database' import { ErrorCode, StorageBackendError } from '@internal/errors' import { MAX_OBJECTS_PER_REQUEST } from '@storage/limits' import { randomUUID } from 'crypto' @@ -36,7 +40,7 @@ type SignedUrlResult = { signedURL: string | null } -let tnx: PgTransaction | undefined +let tnx: DatabaseTransaction | undefined async function getSuperuserPostgrestClient() { const superUser = await getServiceKeyUser(tenantId) @@ -52,7 +56,7 @@ async function getSuperuserPostgrestClient() { } async function findObject( - db: PgTransaction, + db: DatabaseTransaction, bucketId: string, name: string ): Promise { @@ -71,7 +75,7 @@ async function findObject( } async function insertObjects( - db: PgTransaction, + db: DatabaseTransaction, objects: | Array & { bucket_id: string; name: string }> | (Partial & { bucket_id: string; name: string }) @@ -90,7 +94,11 @@ async function insertObjects( } } -async function deleteObjectsByName(db: PgTransaction, bucketId: string, names: string | string[]) { +async function deleteObjectsByName( + db: DatabaseTransaction, + bucketId: string, + names: string | string[] +) { await db.query({ text: ` DELETE FROM objects @@ -101,7 +109,7 @@ async function deleteObjectsByName(db: PgTransaction, bucketId: string, names: s }) } -async function insertObjectNames(db: PgTransaction, bucketId: string, names: string[]) { +async function insertObjectNames(db: DatabaseTransaction, bucketId: string, names: string[]) { const owner = '317eadce-631a-4429-a0bb-f19a7a517b4a' const versions = names.map((_, index) => `test-version-${randomUUID()}-${index}`) @@ -116,7 +124,7 @@ async function insertObjectNames(db: PgTransaction, bucketId: string, names: str } async function insertBucket( - db: PgTransaction, + db: DatabaseTransaction, bucket: { id: string name: string @@ -1778,7 +1786,7 @@ describe('testing copy object', () => { }) await seedTx.commit() tnx = undefined - let verificationTx: PgTransaction | undefined + let verificationTx: DatabaseTransaction | undefined try { const response = await appInstance.inject({ diff --git a/src/test/operation-helpers.test.ts b/src/test/operation-helpers.test.ts index 646e19deb..d5407ab49 100644 --- a/src/test/operation-helpers.test.ts +++ b/src/test/operation-helpers.test.ts @@ -8,7 +8,7 @@ describe('Storage operation helpers', () => { bindings: unknown[] = [], currentOperation?: string ): Promise { - const db = tHelper.database.connection.pool.acquire() + const db = tHelper.database.connection const tnx = await db.beginTransaction() try { diff --git a/src/test/s3-error-code.test.ts b/src/test/s3-error-code.test.ts index 74d156b71..357a4f39f 100644 --- a/src/test/s3-error-code.test.ts +++ b/src/test/s3-error-code.test.ts @@ -40,7 +40,7 @@ describe('S3 protocol error code', () => { beforeAll(async () => { fileBackendPath = await mkdtemp(join(tmpdir(), 'storage-file-backend-')) testApp = await createFileBackedApp(fileBackendPath) - icebergMetastore = new PgMetastore(t.database.connection.pool.acquire(), { + icebergMetastore = new PgMetastore(t.database.connection, { multiTenant: false, schema: 'storage', }) diff --git a/src/test/s3-protocol.test.ts b/src/test/s3-protocol.test.ts index 9663900b2..afad1485b 100644 --- a/src/test/s3-protocol.test.ts +++ b/src/test/s3-protocol.test.ts @@ -30,7 +30,7 @@ import { Upload } from '@aws-sdk/lib-storage' import { createPresignedPost } from '@aws-sdk/s3-presigned-post' import { getSignedUrl } from '@aws-sdk/s3-request-presigner' import { wait } from '@internal/concurrency' -import { getPostgresConnection, getServiceKeyUser, PgTenantConnection } from '@internal/database' +import { getPostgresConnection, getServiceKeyUser, type TenantConnection } from '@internal/database' import { DBMigration } from '@internal/database/migrations' import { ERRORS } from '@internal/errors' import { StoragePgDB } from '@storage/database' @@ -2372,7 +2372,7 @@ describe('S3 Protocol', () => { superUser: adminUser, host: 'localhost', }) - const db = connection.pool.acquire() + const db = connection const anonKey = await anonKeyAsync const anonClient = new S3Client({ endpoint: `${baseUrl}/s3`, @@ -3401,7 +3401,7 @@ describe('S3 Protocol', () => { describe('Migration compatibility', () => { describe('integration', () => { const { tenantId } = getConfig() - let connection: PgTenantConnection + let connection: TenantConnection let bucketId: string beforeAll(async () => { diff --git a/src/test/storage-pg-db.test.ts b/src/test/storage-pg-db.test.ts index 2a10473c5..e99ae8812 100644 --- a/src/test/storage-pg-db.test.ts +++ b/src/test/storage-pg-db.test.ts @@ -1,8 +1,8 @@ import { + type DatabaseExecutor, + type DatabaseStatement, getServiceKeyUser, - type PgExecutor, PgPoolStrategy, - type PgStatement, PgTenantConnection, PgTransaction, } from '@internal/database' @@ -404,7 +404,7 @@ describe('StoragePgDB bucket metadata', () => { 'ReadSuperUserStatementTimeout', async (pg) => { await expect(readCurrentRoleFromExecutor(pg)).resolves.toBe(superUser.payload.role) - await expect(readCurrentStatementTimeoutMsFromExecutor(pg)).resolves.toBe(4321) + await expect(readStatementTimeoutMs(pg)).resolves.toBe(4321) } ) ).resolves.toBeUndefined() @@ -430,9 +430,9 @@ describe('StoragePgDB bucket metadata', () => { 'FailSuperUserStatementTimeout', async (pg) => { await expect(readCurrentRoleFromExecutor(pg)).resolves.toBe(superUser.payload.role) - await expect(readCurrentStatementTimeoutMsFromExecutor(pg)).resolves.toBe(4321) + await expect(readStatementTimeoutMs(pg)).resolves.toBe(4321) await pg.query("SELECT set_config('statement_timeout', '30s', true)") - await expect(readCurrentStatementTimeoutMsFromExecutor(pg)).resolves.toBe(30000) + await expect(readStatementTimeoutMs(pg)).resolves.toBe(30000) throw new Error('failed nested super-user timeout query') } ) @@ -754,9 +754,9 @@ describe('StoragePgDB bucket metadata', () => { it('maps PostgreSQL lock_timeout from waitObjectLock to LockTimeout', async () => { const lockTimeoutError = createPgError('55P03', 'canceling statement due to lock timeout') - const queries: Array = [] + const queries: Array = [] const transaction = { - query: vi.fn(async (statement: string | PgStatement) => { + query: vi.fn(async (statement: string | DatabaseStatement) => { queries.push(statement) const text = typeof statement === 'string' ? statement : statement.text @@ -803,9 +803,9 @@ describe('StoragePgDB bucket metadata', () => { }) it('restores the prior transaction-local lock_timeout after waitObjectLock succeeds', async () => { - const queries: Array = [] + const queries: Array = [] const transaction = { - query: vi.fn(async (statement: string | PgStatement) => { + query: vi.fn(async (statement: string | DatabaseStatement) => { queries.push(statement) const text = typeof statement === 'string' ? statement : statement.text @@ -847,9 +847,9 @@ describe('StoragePgDB bucket metadata', () => { }) it('uses top-level lock_timeout statements for Multigres waitObjectLock', async () => { - const queries: Array = [] + const queries: Array = [] const transaction = { - query: vi.fn(async (statement: string | PgStatement) => { + query: vi.fn(async (statement: string | DatabaseStatement) => { queries.push(statement) const text = typeof statement === 'string' ? statement : statement.text @@ -1742,23 +1742,23 @@ describe('StoragePgDB bucket metadata', () => { function runStorageQuery( storage: StoragePgDB, queryName: string, - fn: (db: PgExecutor, signal?: AbortSignal) => Promise + fn: (db: DatabaseExecutor, signal?: AbortSignal) => Promise ): Promise { return ( storage as unknown as { runQuery( queryName: string, - fn: (db: PgExecutor, signal?: AbortSignal) => Promise + fn: (db: DatabaseExecutor, signal?: AbortSignal) => Promise ): Promise } ).runQuery(queryName, fn) } - function statementText(statement: string | PgStatement): string { + function statementText(statement: string | DatabaseStatement): string { return typeof statement === 'string' ? statement : statement.text } - function statementValues(statement: string | PgStatement): unknown[] | undefined { + function statementValues(statement: string | DatabaseStatement): unknown[] | undefined { return typeof statement === 'string' ? undefined : statement.values } @@ -1772,7 +1772,7 @@ describe('StoragePgDB bucket metadata', () => { return runStorageQuery(storage, 'ReadCurrentRole', (pg) => readCurrentRoleFromExecutor(pg)) } - async function readCurrentRoleFromExecutor(pg: PgExecutor): Promise { + async function readCurrentRoleFromExecutor(pg: DatabaseExecutor): Promise { const result = await pg.query<{ role: string }>(` SELECT current_setting('role', true) AS role `) @@ -1782,11 +1782,11 @@ describe('StoragePgDB bucket metadata', () => { function readCurrentStatementTimeoutMs(storage: StoragePgDB): Promise { return runStorageQuery(storage, 'ReadCurrentStatementTimeout', (pg) => - readCurrentStatementTimeoutMsFromExecutor(pg) + readStatementTimeoutMs(pg) ) } - async function readCurrentStatementTimeoutMsFromExecutor(pg: PgExecutor): Promise { + async function readStatementTimeoutMs(pg: DatabaseExecutor): Promise { const result = await pg.query<{ statement_timeout: string }>(` SELECT current_setting('statement_timeout') AS statement_timeout `) diff --git a/src/test/utils/storage.ts b/src/test/utils/storage.ts index e8463eda4..16eddd077 100644 --- a/src/test/utils/storage.ts +++ b/src/test/utils/storage.ts @@ -1,10 +1,11 @@ import { CreateBucketCommand, HeadBucketCommand, S3Client } from '@aws-sdk/client-s3' import { + type DatabaseTransaction, + type DatabaseTransactionalExecutor, getPostgresConnection, getServiceKeyUser, - PgPoolExecutor, - PgTenantConnection, - PgTransaction, + isDatabaseTransaction, + type TenantConnection, } from '@internal/database' import { runMigrationsOnTenant } from '@internal/database/migrations' import { isS3Error } from '@internal/errors' @@ -23,10 +24,10 @@ const { databaseURL, tenantId, storageBackendType, storageS3Bucket } = getConfig * This is needed because raw queries bypass the normal connection scope setting */ export async function withDeleteEnabled( - db: PgPoolExecutor | PgTransaction, - fn: (db: PgTransaction) => Promise + db: DatabaseTransactionalExecutor | DatabaseTransaction, + fn: (db: DatabaseTransaction) => Promise ): Promise { - const existingTransaction = db instanceof PgTransaction + const existingTransaction = isDatabaseTransaction(db) const tnx = existingTransaction ? db : await db.beginTransaction() try { await tnx.query(`SELECT set_config('storage.allow_delete_query', 'true', true)`) @@ -44,7 +45,7 @@ export async function withDeleteEnabled( } export function useStorage(options: { ensureMigrations?: boolean } = {}) { - let connection: PgTenantConnection + let connection: TenantConnection let storage: Storage let adapter: StorageBackendAdapter let database: Database @@ -75,7 +76,7 @@ export function useStorage(options: { ensureMigrations?: boolean } = {}) { tenantId, host: 'localhost', } - database = new StoragePgDB(connection, databaseOptions) as unknown as Database + database = new StoragePgDB(connection, databaseOptions) location = new TenantLocation(storageS3Bucket) adapter = createStorageBackend(storageBackendType) storage = new Storage(adapter, database, location) diff --git a/src/test/vectors.test.ts b/src/test/vectors.test.ts index 97c982ee4..aaaf0e783 100644 --- a/src/test/vectors.test.ts +++ b/src/test/vectors.test.ts @@ -37,7 +37,7 @@ function parseJsonBody(body: string): Body { } async function findVectorIndex(bucketId: string, name: string) { - const result = await storageTest.database.connection.pool.acquire().query<{ + const result = await storageTest.database.connection.query<{ data_type: string dimension: number distance_metric: string @@ -57,7 +57,7 @@ async function findVectorIndex(bucketId: string, name: string) { } async function findVectorBucket(bucketId: string) { - const result = await storageTest.database.connection.pool.acquire().query<{ + const result = await storageTest.database.connection.query<{ id: string created_at: Date }>({ @@ -114,7 +114,7 @@ describe('Vectors API', () => { shardKey: 'test-bucket', capacity: 1000, }) - const mockVectorDB = new PgVectorMetadataDB(storageTest.database.connection.pool.acquire()) + const mockVectorDB = new PgVectorMetadataDB(storageTest.database.connection) s3Vector = new VectorStoreManager(mockVectorStore, mockVectorDB, shard, { tenantId: 'test-tenant', maxBucketCount: Infinity, diff --git a/src/test/webhooks.test.ts b/src/test/webhooks.test.ts index 1efa971f2..54e6ccab1 100644 --- a/src/test/webhooks.test.ts +++ b/src/test/webhooks.test.ts @@ -1,4 +1,4 @@ -import { PgTenantConnection } from '@internal/database' +import type { TenantConnection } from '@internal/database' import { getConfig, mergeConfig } from '../config' vi.hoisted(() => { @@ -26,7 +26,7 @@ import { mockQueue, useMockObject } from './common' describe('Webhooks', () => { useMockObject() - let pg: PgTenantConnection + let pg: TenantConnection beforeAll(async () => { const superUser = await getServiceKeyUser(tenantId) pg = await getPostgresConnection({ @@ -679,7 +679,7 @@ describe('Webhooks', () => { }) }) -async function createObject(pg: PgTenantConnection, bucketId: string, createdAt?: Date) { +async function createObject(pg: TenantConnection, bucketId: string, createdAt?: Date) { const objectName = randomUUID() const tnx = await pg.transaction() const createdAtIso = createdAt?.toISOString()