From 8a85d17f63418fdb03d690075aa983f15acb6cbf Mon Sep 17 00:00:00 2001 From: Gustavo Valverde Date: Tue, 25 Aug 2026 07:05:31 -0400 Subject: [PATCH] fix(drizzle): harden migration connections --- .../src/drizzle-adapter.test.ts | 165 +++++++++++- .../drizzle-adapter/src/drizzle-adapter.ts | 240 +++++++++++++++--- 2 files changed, 365 insertions(+), 40 deletions(-) diff --git a/packages/drizzle-adapter/src/drizzle-adapter.test.ts b/packages/drizzle-adapter/src/drizzle-adapter.test.ts index 0ea1484015..2ed44e4cb3 100644 --- a/packages/drizzle-adapter/src/drizzle-adapter.test.ts +++ b/packages/drizzle-adapter/src/drizzle-adapter.test.ts @@ -42,6 +42,67 @@ describe("drizzle-adapter", () => { expect(resolved?.account?.fields.issuer?.fieldName).toBe("identity_issuer"); }); + it("resolves a customized model through its schema key", async () => { + const account = pgTable("auth_account", { + accountId: text("account_id").notNull(), + }); + const adapter = drizzleAdapter( + { _: { fullSchema: { account } } } as never, + { provider: "pg" }, + )({}); + + const resolved = + await adapter.options?.adapterConfig.migrationConnection?.resolvePhysicalSchema?.( + { + account: { + modelName: "accounts", + fields: { + accountId: { type: "string", fieldName: "accountId" }, + }, + }, + }, + ); + + expect(resolved?.account?.modelName).toBe("auth_account"); + expect(resolved?.account?.fields.accountId?.fieldName).toBe("account_id"); + }); + + it("fails when the configured Drizzle schema is missing migration metadata", async () => { + const account = pgTable("auth_account", { + accountId: text("account_id").notNull(), + }); + const adapter = drizzleAdapter( + { _: { fullSchema: { account } } } as never, + { provider: "pg" }, + )({}); + const migrationConnection = + adapter.options?.adapterConfig.migrationConnection; + + await expect( + migrationConnection?.resolvePhysicalSchema?.({ + account: { + modelName: "account", + fields: { + accountId: { type: "string", fieldName: "accountId" }, + issuer: { type: "string", fieldName: "issuer" }, + }, + }, + }), + ).rejects.toThrow( + 'Drizzle migration schema could not resolve field "issuer" on model "account".', + ); + await expect( + migrationConnection?.resolvePhysicalSchema?.({ + jwks: { + modelName: "jwks", + fields: { id: { type: "string", fieldName: "id" } }, + }, + }), + ).rejects.toThrow( + 'Drizzle migration schema could not resolve model "jwks".', + ); + }); + it("should create drizzle adapter", () => { const db = { _: { @@ -103,9 +164,76 @@ describe("drizzle-adapter", () => { expect(run).toHaveBeenCalledTimes(1); }); + it.each([ + { + name: "mysql2 tuple", + provider: "mysql" as const, + result: [{ affectedRows: 2 }, []], + }, + { + name: "node-postgres result", + provider: "pg" as const, + result: { rowCount: 2, rows: [] }, + }, + { + name: "postgres-js result", + provider: "pg" as const, + result: Object.assign([], { count: 2 }), + }, + ])("normalizes $name for migration writes", async ({ provider, result }) => { + const adapter = drizzleAdapter( + { + _: { fullSchema: {} }, + execute: vi.fn().mockResolvedValue(result), + } as never, + { provider }, + )({ secret: "test-secret-that-is-at-least-32-chars-long!!" }); + + await expect( + adapter.options?.adapterConfig.migrationConnection?.execute({ + parameters: ["local:credential"], + sql: + provider === "pg" + ? "UPDATE account SET issuer = $1" + : "UPDATE account SET issuer = ?", + }), + ).resolves.toEqual({ numAffectedRows: 2n, rows: [] }); + }); + + it.each([ + { + provider: "sqlite" as const, + sql: "SELECT '?' AS literal, providerId FROM account WHERE providerId = ? -- ?", + }, + { + provider: "pg" as const, + sql: "SELECT '$2' AS literal, $$ $3 $$ AS body FROM account WHERE providerId = $1 -- $4", + }, + ])("ignores parameter-like text in $provider SQL literals and comments", async ({ + provider, + sql: query, + }) => { + const rows = [{ providerId: "credential" }]; + const adapter = drizzleAdapter( + { + _: { fullSchema: {} }, + all: vi.fn().mockResolvedValue(rows), + execute: vi.fn().mockResolvedValue(rows), + } as never, + { provider }, + )({ secret: "test-secret-that-is-at-least-32-chars-long!!" }); + + await expect( + adapter.options?.adapterConfig.migrationConnection?.execute({ + parameters: ["credential"], + sql: query, + }), + ).resolves.toEqual({ rows }); + }); + it("scopes SQLite migration queries to one transaction", async () => { const statements: string[] = []; - const run = vi.fn().mockImplementation(async (statement: SQL) => { + const run = vi.fn().mockImplementation((statement: SQL) => { statements.push( statement.queryChunks .flatMap((chunk) => @@ -139,12 +267,47 @@ describe("drizzle-adapter", () => { }); expect(statements).toEqual([ + "SELECT 1", "BEGIN IMMEDIATE", "UPDATE account SET issuer = 'local:credential'", "COMMIT", ]); }); + it("uses a transaction-scoped database for asynchronous SQLite migrations", async () => { + const rootRun = vi.fn().mockResolvedValue({ changes: 0 }); + const transactionRun = vi.fn().mockResolvedValue({ changes: 1 }); + const transaction = vi.fn().mockImplementation(async (callback) => + callback({ + _: { fullSchema: {} }, + all: vi.fn(), + run: transactionRun, + }), + ); + const adapter = drizzleAdapter( + { + _: { fullSchema: {} }, + all: vi.fn(), + run: rootRun, + transaction, + } as never, + { provider: "sqlite", transaction: true }, + )({ secret: "test-secret-that-is-at-least-32-chars-long!!" }); + + await adapter.options?.adapterConfig.migrationConnection?.transaction?.( + async (connection) => { + await connection.execute({ + parameters: [], + sql: "UPDATE account SET issuer = 'local:credential'", + }); + }, + ); + + expect(transaction).toHaveBeenCalledOnce(); + expect(transactionRun).toHaveBeenCalledOnce(); + expect(rootRun).toHaveBeenCalledOnce(); + }); + /** * @see https://github.com/better-auth/better-auth/pull/10575#discussion_r3812054987 */ diff --git a/packages/drizzle-adapter/src/drizzle-adapter.ts b/packages/drizzle-adapter/src/drizzle-adapter.ts index 8f763b94e9..6123f83a38 100644 --- a/packages/drizzle-adapter/src/drizzle-adapter.ts +++ b/packages/drizzle-adapter/src/drizzle-adapter.ts @@ -13,7 +13,6 @@ import type { import { createAdapterFactory } from "@better-auth/core/db/adapter"; import { logger } from "@better-auth/core/env"; import { BetterAuthError } from "@better-auth/core/error"; -import { toSnakeCase } from "@better-auth/core/utils/string"; import type { SQL } from "drizzle-orm"; import { and, @@ -103,6 +102,8 @@ function getAffectedRowCount( function readDriverRowCount(result: unknown): unknown { if (!result || typeof result !== "object") return undefined; const driverResult = result as Record; + if ("rowCount" in driverResult) return driverResult.rowCount; + if ("count" in driverResult) return driverResult.count; if ("affectedRows" in driverResult) return driverResult.affectedRows; if ("rowsAffected" in driverResult) return driverResult.rowsAffected; if ("changes" in driverResult) return driverResult.changes; @@ -123,8 +124,23 @@ function hasDriverRowCount(result: unknown): boolean { return readDriverRowCount(result) !== undefined; } +function isDriverResultHeader(result: unknown): boolean { + if (!result || typeof result !== "object") return false; + const driverResult = result as Record; + return ( + "affectedRows" in driverResult || + "rowsAffected" in driverResult || + "changes" in driverResult || + ("meta" in driverResult && + driverResult.meta !== null && + typeof driverResult.meta === "object" && + "changes" in driverResult.meta) + ); +} + function getMigrationRows(result: unknown): readonly Record[] { if (Array.isArray(result)) { + if (result.length > 0 && isDriverResultHeader(result[0])) return []; if (Array.isArray(result[0])) { return result[0].filter( (row): row is Record => @@ -152,23 +168,134 @@ function getMigrationRows(result: unknown): readonly Record[] { return []; } +interface MigrationParameterMarker { + length: number; + offset: number; + parameterIndex: number; +} + +function findMigrationParameterMarkers( + query: string, + dialect: MigrationDatabaseDialect, +): MigrationParameterMarker[] { + const markers: MigrationParameterMarker[] = []; + let offset = 0; + let positionalParameter = 0; + while (offset < query.length) { + const character = query[offset]; + const nextCharacter = query[offset + 1]; + + if (character === "-" && nextCharacter === "-") { + offset = query.indexOf("\n", offset + 2); + if (offset === -1) break; + continue; + } + if (dialect === "mysql" && character === "#") { + offset = query.indexOf("\n", offset + 1); + if (offset === -1) break; + continue; + } + if (character === "/" && nextCharacter === "*") { + let depth = 1; + offset += 2; + while (offset < query.length && depth > 0) { + if (query[offset] === "/" && query[offset + 1] === "*") { + depth += 1; + offset += 2; + continue; + } + if (query[offset] === "*" && query[offset + 1] === "/") { + depth -= 1; + offset += 2; + continue; + } + offset += 1; + } + continue; + } + if (character === "'" || character === '"' || character === "`") { + const quote = character; + offset += 1; + while (offset < query.length) { + if (dialect === "mysql" && query[offset] === "\\") { + offset += 2; + continue; + } + if (query[offset] !== quote) { + offset += 1; + continue; + } + if (query[offset + 1] === quote) { + offset += 2; + continue; + } + offset += 1; + break; + } + continue; + } + if (character === "[") { + offset += 1; + while (offset < query.length) { + if (query[offset] !== "]") { + offset += 1; + continue; + } + if (query[offset + 1] === "]") { + offset += 2; + continue; + } + offset += 1; + break; + } + continue; + } + if (dialect === "postgres" && character === "$") { + const dollarQuote = query + .slice(offset) + .match(/^\$(?:[A-Za-z_][A-Za-z0-9_]*)?\$/)?.[0]; + if (dollarQuote) { + const closingOffset = query.indexOf( + dollarQuote, + offset + dollarQuote.length, + ); + offset = + closingOffset === -1 + ? query.length + : closingOffset + dollarQuote.length; + continue; + } + const parameterMarker = query.slice(offset).match(/^\$(\d+)/); + if (parameterMarker) { + markers.push({ + length: parameterMarker[0].length, + offset, + parameterIndex: Number.parseInt(parameterMarker[1] ?? "", 10) - 1, + }); + offset += parameterMarker[0].length; + continue; + } + } + if (dialect !== "postgres" && character === "?") { + markers.push({ length: 1, offset, parameterIndex: positionalParameter }); + positionalParameter += 1; + } + offset += 1; + } + return markers; +} + function buildMigrationStatement( query: MigrationDatabaseQuery, dialect: MigrationDatabaseDialect, ): SQL { if (query.parameters.length === 0) return sql.raw(query.sql); - const marker = dialect === "postgres" ? /\$(\d+)/g : /\?/g; + const markers = findMigrationParameterMarkers(query.sql, dialect); const chunks: SQL[] = []; let sqlOffset = 0; - let positionalParameter = 0; - for (const match of query.sql.matchAll(marker)) { - const matchOffset = match.index; - if (matchOffset === undefined) continue; - chunks.push(sql.raw(query.sql.slice(sqlOffset, matchOffset))); - const parameterIndex = - dialect === "postgres" - ? Number.parseInt(match[1] ?? "", 10) - 1 - : positionalParameter++; + for (const marker of markers) { + chunks.push(sql.raw(query.sql.slice(sqlOffset, marker.offset))); + const parameterIndex = marker.parameterIndex; if ( !Number.isInteger(parameterIndex) || parameterIndex < 0 || @@ -179,15 +306,13 @@ function buildMigrationStatement( ); } chunks.push(sql`${query.parameters[parameterIndex]}`); - sqlOffset = matchOffset + match[0].length; + sqlOffset = marker.offset + marker.length; } chunks.push(sql.raw(query.sql.slice(sqlOffset))); if ( (dialect === "postgres" - ? new Set( - [...query.sql.matchAll(marker)].map((match) => Number(match[1])), - ).size - : positionalParameter) !== query.parameters.length + ? new Set(markers.map(({ parameterIndex }) => parameterIndex)).size + : markers.length) !== query.parameters.length ) { throw new BetterAuthError( "Drizzle migration query parameter count does not match its SQL markers.", @@ -198,8 +323,14 @@ function buildMigrationStatement( function getMigrationQueryResult( driverResult: unknown, + isRead: boolean, ): MigrationDatabaseQueryResult { - const affectedRows = readDriverRowCount(driverResult); + if (isRead) return { rows: getMigrationRows(driverResult) }; + const affectedRows = + Array.isArray(driverResult) && driverResult.length > 0 + ? (readDriverRowCount(driverResult[0]) ?? + readDriverRowCount(driverResult)) + : readDriverRowCount(driverResult); return { rows: getMigrationRows(driverResult), ...(typeof affectedRows === "number" && Number.isFinite(affectedRows) @@ -214,11 +345,19 @@ function isDrizzleMigrationRead(query: MigrationDatabaseQuery): boolean { return /^pragma\b/i.test(sql) && !/^pragma\b[^;]*=/i.test(sql); } +function isPromiseLike(value: unknown): value is PromiseLike { + return ( + value !== null && + (typeof value === "object" || typeof value === "function") && + "then" in value && + typeof value.then === "function" + ); +} + function createDrizzleMigrationConnection( db: DB, dialect: MigrationDatabaseDialect, drizzleSchema: Record | undefined, - camelCase: boolean, inTransaction = false, supportsTransactions = true, ): MigrationDatabaseConnection { @@ -232,24 +371,33 @@ function createDrizzleMigrationConnection( } return Object.fromEntries( Object.entries(schema).map(([schemaKey, table]) => { - const drizzleTable = drizzleSchema[table.modelName]; - const hasDrizzleTable = - drizzleTable !== undefined && typeof drizzleTable === "object"; + const drizzleTable = + drizzleSchema[schemaKey] ?? drizzleSchema[table.modelName]; + if (!drizzleTable || typeof drizzleTable !== "object") { + throw new BetterAuthError( + `Drizzle migration schema could not resolve model "${table.modelName}".`, + ); + } const fields = Object.fromEntries( Object.entries(table.fields).map(([fieldKey, field]) => { const drizzleFieldName = field.fieldName || fieldKey; - const column = hasDrizzleTable - ? (drizzleTable as Record)[drizzleFieldName] - : undefined; + const drizzleTableFields = drizzleTable as Record< + string, + unknown + >; + const column = + drizzleTableFields[fieldKey] ?? + drizzleTableFields[drizzleFieldName]; + if (!is(column, Column)) { + throw new BetterAuthError( + `Drizzle migration schema could not resolve field "${drizzleFieldName}" on model "${table.modelName}".`, + ); + } return [ fieldKey, { ...field, - fieldName: is(column, Column) - ? column.name - : camelCase - ? drizzleFieldName - : toSnakeCase(drizzleFieldName), + fieldName: column.name, }, ]; }), @@ -258,11 +406,7 @@ function createDrizzleMigrationConnection( schemaKey, { ...table, - modelName: hasDrizzleTable - ? getTableName(drizzleTable as never) - : camelCase - ? table.modelName - : toSnakeCase(table.modelName), + modelName: getTableName(drizzleTable as never), fields, }, ]; @@ -271,18 +415,39 @@ function createDrizzleMigrationConnection( }, async execute(query) { const statement = buildMigrationStatement(query, dialect); + const isRead = isDrizzleMigrationRead(query); const driverResult = dialect === "sqlite" - ? isDrizzleMigrationRead(query) + ? isRead ? await db.all(statement) : await db.run(statement) : await db.execute(statement); - return getMigrationQueryResult(driverResult); + return getMigrationQueryResult(driverResult, isRead); }, }; if (inTransaction || !supportsTransactions) return connection; connection.transaction = async (callback) => { if (dialect === "sqlite") { + const probe = db.run(sql.raw("SELECT 1")); + if (isPromiseLike(probe)) { + await probe; + if (typeof db.transaction !== "function") { + throw new BetterAuthError( + "Drizzle migration transactions for asynchronous SQLite drivers require a native transaction-scoped database.", + ); + } + return db.transaction((transactionDatabase: DB) => + callback( + createDrizzleMigrationConnection( + transactionDatabase, + dialect, + drizzleSchema, + true, + true, + ), + ), + ); + } await connection.execute({ parameters: [], sql: "BEGIN IMMEDIATE", @@ -293,7 +458,6 @@ function createDrizzleMigrationConnection( db, dialect, drizzleSchema, - camelCase, true, true, ), @@ -316,7 +480,6 @@ function createDrizzleMigrationConnection( transactionDatabase, dialect, drizzleSchema, - camelCase, true, true, ), @@ -1350,7 +1513,6 @@ export const drizzleAdapter = (db: DB, config: DrizzleAdapterConfig) => { db, config.provider === "pg" ? "postgres" : config.provider, config.schema || db._?.fullSchema, - config.camelCase === true, false, config.transaction === true, ),