POSTED

March 3, 2024

Kysely Retry Logic by Splitting Query Building and Execution

A description of the image.
sql
typescript

Kysely is a type-safe SQL query builder for TypeScript, and most people also use it to execute their SQL queries too. One issue with letting Kysely execute your queries it that logic you would like to apply to all of your queries can't be injected into Kysely easliy, for example retries. As Kysely mentions, it is a SQL query builder and allows you to just create queries and execute them separately.

This post shows how to split query building and execution, gives an example of retry logic, and provides two utility functions that can mimic the .execute and .takeFirstOrThrow methods while maintaining type safety.

The Challenge To Solve - Retry Logic

A common way for developers to deploy their applications now is with serverless providers such as AWS Lambda, or Cloudflare Workers (among many other providers). With these environments a new database connection must be created for each function, and from my experience with a few databases that work well with serverless applications (Neon/Turso), I have seen connection issues or timeouts that will cause your queries to fail and your site to return a 500 status response:

One way to solve for these issues is to add retry logic to your SQL queries, that way if the first connection failed the second will probably work. This also solves the issue where your serverless database has to perform a cold start and your connection request times out in the first connection attempt.

Single Failure Retry Logic

The typicaly way that you create and execute a Kysely query looks like this:

Basic Kysely Example
const selectedUser = await db.selectFrom("users").where("id", "=", id).selectAll().executeTakeFirst();

The way you can implement simple retry logic in Kysely looks like this:

Kysely Retry Example
import { InferType } from "kysely";
// Build the query separate from the execution
const selectedUserQuery = db.selectFrom("users").where("id", "=", id).selectAll();
// Run the query, and if there is an issue perform a retry with a new db connection.
let selectedUser: InferType<typeof selectedUserQuery>;
try {
selectedUser = await selectedUserQuery.executeTakeFirst();
} catch (e) {
// ... any logic you want to retry based upon
const isConnectionIssue = ["ECONNRESET", "EPIPE", "ETIMEDOUT"].includes(e.code);
if (isConnectionIssue) {
// Destroy the database connection, and then run the query again. This will force a new database
// connection to be created on the next request.
await db.destroy();
selectedUser = await selectedUserQuery.executeTakeFirst();
}
}

Multiple Retries with Exponential Backoff

The single failure retry logic will work, but there are two downsides. First, if we want this logic to apply to all database queries we have to wrap every query in this try/catch logic. This will result in a lot of duplicated code and some lengthy functions. Second, this retry logic only works for one failure right after the first failure - but what if we wanted to try a few times?

The next logical step is to extract this retry logic into a function:

Extracted Retry Logic
/**
* Adds a wait before next commands are executed.
*/
async function delay(ms: number) {
return new Promise((resolve) => setTimeout(resolve, ms));
}
/**
* A function that retries a promise up to 10 times with an exponential backoff between retries.
*/
async function retry<T>(fn: () => Promise<T>) {
let [attempt, maxAttempts] = [0, 10];
let error: unknown;
do {
try {
return await fn();
} catch (e) {
// ... any logic you want to retry based upon
const isConnectionIssue = ["ECONNRESET", "EPIPE", "ETIMEDOUT"].includes(e.code);
if (isConnectionIssue) {
await db.destroy();
await delay(5 ** attempt);
attempt = attempt + 1;
} else throw e;
}
} while (attempt < maxAttempts);
// throw the last error if all retries have failed
throw error;
}
Using Extraced Retry Logic
const selectedUser = await retry(
() => db.selectFrom("users").where("id", "=", id).selectAll().executeTakeFirst()
);

Great! Now we can run queries all over our codebase with retry logic, however passing a function just to get our type safety is a bit cumbersome. It would be great to get functions that have type safety and implement the same logic as .execute, .takeFirst, and .takeFirstOrThrow. In the next section we will implement both of these functions and show some examples.

Creating execute, takeFirst, and takeFirstOrThrow Functions

When we split the Kysely logic of query building and query execution all we need to pass to the .executeQuery function is the query and we will have type safety. A similar interface should exist for the execute, takeFirst, and takeFirstOrThrow function. Examples of this can be seen on the Kysely docs for splitting query building .

Kysely .executeQuery
// Splitting the query building and query execution
const selectedUserQuery = db.selectFrom("users").where("id", "=", id).selectAll();
const selectedUser = db.executeQuery(selectedUserQuery).

When we use the .executeQuery it accepts a wider range of inputs that may or may not have .execute, .takeFirst, or .takeFirstOrThrow. If you have a look at the type of the .executeQuery it accepts either Compilable<T> or CompiledQuery<T> as the query for execution:

Type Signature .executeQuery
Kysely<IDatabase>.executeQuery<T>(query: Compilable<T> | CompiledQuery<T>, queryId?: QueryId | undefined): Promise<QueryResult<T>>

Now that we have an idea of what our functions should accept and return, we can build the logic for these. Note that when the command.query.kind is not SelectQueryNode, the returned type has possible value of Promise<InferResult<C>> or Promise<QueryResult<T>>. The Promise<QueryResult<T>> won't return a list of items and for this reason the takeFirst and takeFirstOrThrow will need to be called on the result of an execute query. See the example below on how to use this.

Our Functions
/**
* Executes a kysely command with exponential backoff on retries.
*/
export async function execute<T, C extends Compilable<T> | CompiledQuery<T>>(cmd: C): Promise<InferResult<C>> {
let command = "compile" in cmd ? cmd.compile() : (cmd as CompiledQuery<T>);
async function fn() {
switch (command.query.kind) {
case "DeleteQueryNode":
case "UpdateQueryNode":
case "InsertQueryNode": {
if (command.query.returning)
return db
.executeQuery(command)
.then(({ rows }) => rows as InferResult<C>);
else return db.executeQuery(command) as any as Promise<InferResult<C>>;
}
case "SelectQueryNode": {
return db
.executeQuery(command)
.then(({ rows }) => rows as InferResult<C>);
}
default: {
return db.executeQuery(command) as any as Promise<InferResult<C>>;
}
}
}
return retry(fn);
}
/**
* This function is a helper that can be run on the results of an 'execute' call to extract the first item from an
* array of results. If an item isn't found, then 'undefined' will be returned.
*/
export function takeFirst<T>(input: (T | undefined)[]) {
const [result] = input;
return result;
}
/**
* This function is a helper that can be run on the results of an 'execute' call to extract the first item from an
* array of results. If an item isn't found, then and error will be thrown.
*/
export function takeFirstOrThrow<T>(input: (T | undefined)[]) {
const [result] = input;
if (result === undefined) throw new Error("No result was found for the query.");
return result;
}
Our Function's Examples
// Using just the `execute` function
const first25Users = await execute(db.selectFrom("users").selectAll().limit(25));
// Using the `execute` with `takeFirst`
const selectedUserQuery = db.selectFrom("users").where("id", "=", id).selectAll();
const selectedUser = await execute(selectedUserQuery).then((res) => takeFirst(res));
if (!selectedUser) throw new Error("No user was found!");
// Using the `execute with `takeFirstOrThrow`
const selectedUserOrThrow = await execute(selectedUserQuery).then((res) => takeFirstOrThrow(res));
if (!selectedUser) throw new Error("No user was found!");

This is how you can implement retry logic with Kysely, there is probably a way to fix the takeFirst and takeFirstOrThrow so that they don't need to work off the result of an execute statement, but this has been good enough for my purposes. I will update this blog if a better method arises, I hope you found something useful!