github-starred/better-auth
2
0
Fork 0
mirror of https://github.com/better-auth/better-auth.git synced 2026-10-19 11:24:57 -05:00
Code Issues 2.5k Packages Projects Releases 112 Wiki Activity

fix(drizzle): harden migration connections

Browse Source
This commit is contained in:
Gustavo Valverde
2026-08-26 15:53:46 -04:00
parent 744faf6176
commit abc225c2b4
2 changed files with 365 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/drizzle-adapter/src/drizzle-adapter.test.ts
+164 -1
View File
@@ -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
*/
packages/drizzle-adapter/src/drizzle-adapter.ts
+201 -39
View File
@@ -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,
@@ -107,6 +106,8 @@ function getAffectedRowCount(
function readDriverRowCount(result: unknown): unknown {
if (!result || typeof result !== "object") return undefined;
const driverResult = result as Record<string, unknown>;
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;
@@ -127,8 +128,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<string, unknown>;
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<string, unknown>[] {
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<string, unknown> =>
@@ -156,23 +172,134 @@ function getMigrationRows(result: unknown): readonly Record<string, unknown>[] {
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 ||
@@ -183,15 +310,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.",
@@ -202,8 +327,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)
@@ -218,11 +349,19 @@ function isDrizzleMigrationRead(query: MigrationDatabaseQuery): boolean {
return /^pragma\b/i.test(sql) && !/^pragma\b[^;]*=/i.test(sql);
}
function isPromiseLike(value: unknown): value is PromiseLike<unknown> {
return (
value !== null &&
(typeof value === "object" || typeof value === "function") &&
"then" in value &&
typeof value.then === "function"
);
}
function createDrizzleMigrationConnection(
db: DB,
dialect: MigrationDatabaseDialect,
drizzleSchema: Record<string, unknown> | undefined,
camelCase: boolean,
inTransaction = false,
supportsTransactions = true,
): MigrationDatabaseConnection {
@@ -236,24 +375,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<string, unknown>)[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,
},
];
}),
@@ -262,11 +410,7 @@ function createDrizzleMigrationConnection(
schemaKey,
{
...table,
modelName: hasDrizzleTable
? getTableName(drizzleTable as never)
: camelCase
? table.modelName
: toSnakeCase(table.modelName),
modelName: getTableName(drizzleTable as never),
fields,
},
];
@@ -275,18 +419,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",
@@ -297,7 +462,6 @@ function createDrizzleMigrationConnection(
db,
dialect,
drizzleSchema,
camelCase,
true,
true,
),
@@ -320,7 +484,6 @@ function createDrizzleMigrationConnection(
transactionDatabase,
dialect,
drizzleSchema,
camelCase,
true,
true,
),
@@ -1380,7 +1543,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,
),