Support query cancellation

[ethanresnick]

I have some long-running analytical queries. At some point in the middle of the query running, the end user can express that they're no longer interested in getting more query results, by clicking a 'cancel' button or leaving the page. I'd like, at that point, to stop running the query in the underlying database, to save on compute costs. (For reference, my underlying database is Snowflake, which has pretty good support for cancellation in its driver.)

In some of these cases, I'm reading the query results from Kysely in a streaming fashion; in other cases, I'm just waiting for a promise with the full result.

So, this request is to somehow support query cancellation in kysely, ideally for both streaming and non-streaming queries.

I don't really know what the API should look like for this in kysely, but...

• It seems like maybe the DatabaseConnection interface could be extended to provide an optional cancelQuery method
• For the non-streaming case, execute() and friends could accept an object argument like { signal: AbortSignal }. If/when the signal is aborted before the query finishes, the underlying cancelQuery would be called.
• For the streaming case, there's already a way for the consumer of the iterator to signal disinterest in further results (with an explicit or implicit call to iterator.return()), but it would be nice for kysely to cancel the currently-running query at that time (whereas I think all that happens today is that the connection is released after the current batch of data finishes loading). It might also be nice for stream() to accept a signal too, for convenience, in case the caller already has a signal that they’re using to create one lifecycle for multiple different async operations (which could contain a mix of streaming and non-streaming queries).

Thoughts?

NB: edited to suggest AbortSignal instead (brain fart) and to clarify that this isn't limited to streaming.

[alexgleason]

+1 to { signal: AbortSignal }, like fetch. It's a non-intrusive surgery. And AbortSignals are so good (simplifies everything).

const signal = AbortSignal.timeout(1000); // Aborts after 1 second

const results = await db
  .selectFrom('users')
  .selectAll()
  .execute({ signal }); // Pass the signal

If it aborts, it throws an error, and then calls some internal stopQuery function of the dialect.

[koskimas]

This would also require this which is a breaking change. All existing 3rd party dialects would be broken unless we find some backwards compatible way to do it.

Once we had the query ID in DatabaseConnection#executeQuery we could add an optional cancelQuery?(queryId: QueryId): Promise<void> method for DatabaseConnection.

Query cancellation sounds like a good idea, but breaking all existing dialects needs to be done carefully.

[alexgleason]

Broken dialects would just need to upgrade Kysely, right? There are only about 12 of them. It would take an hour to do it.

[codetheweb]

Would love to have this, happy to help out in any way. Having .execute() return a "cancellable" Promise with .cancel(), similar to Bluebird, would be great too (e.x. would work out of the box with https://www.npmjs.com/package/p-timeout).

[alexgleason]

This is needed very badly. I've come as far as I can without it. I'm running out of database connections.

[patrickReiis]

I also need support for query cancellation +1

[igalklebanov]

Hey 👋

Let's talk about cancellation, DB-level.

PostgreSQL:
The engine has select pg_cancel_backend(<pid>) and select pg_terminate_backend(<pid>) (kills the connection), where each connection has it's own pid (retrieved with select pg_backend_pid()). You have to run them from a secondary connection.
pg, used as the underlying driver for the core PostgreSQL dialect, doesn't support it. brianc/node-postgres#773
postgres, used as the underlying driver for the organizational dialect kysely-postgres-js, supports it. https://github.com/porsager/postgres#canceling-queries-in-progress

MySQL:
The engine has kill query <pid> and kill connection <pid>. https://dev.mysql.com/doc/refman/8.4/en/kill.html, where each connection has it's own pid (retrieved with select connection_id() https://dev.mysql.com/doc/refman/8.4/en/information-functions.html#function_connection-id). You have to run them from a secondary connection.
mysql2, used as the underlying driver for the core MySQL dialect, doesn't support it. sidorares/node-mysql2#664 (comment)

SQLite:
The engine has sqlite3_interrupt() https://www.sqlite.org/c3ref/interrupt.html.
better-sqlite3, used as the underlying driver for the core SQLite dialect, doesn't support it. WiseLibs/better-sqlite3#568.

MS SQL Server:
tedious, used as the underlying driver for the core MS SQL Server dialect, supports it. We already use this functionality in streaming early exits. https://github.com/kysely-org/kysely/blob/master/src/dialect/mssql/mssql-driver.ts#L265

[oliveryasuna]

+1
This would be absurdly useful!

[benstpierre]

Same here.

[alexgleason]

Let me share what I did to work around this in Postgres, by using statement_timeout:

const query = sql`SELECT * FROM users`;
const timeout = 1_000;

const result = await kysely.transaction().execute((trx) => {
  await sql`set local statement_timeout = ${sql.raw(timeout.toString())}`.execute(trx);
  return query.execute(trx);
});

What this does:

• Starts a transaction.
• Calls set local to set a timeout for the transaction.
• Executes the query and returns the results.

Using a transaction is necessary, otherwise the timeout would apply to future queries.

I created a custom queryDatabase function in my app with this code, and updated all the code in my app to use that function instead of using Kysely methods directly.

This solved all of my problems. But boy is it ugly. I would really like to be able to cancel queries with AbortSignal.

[alexgleason]

Also, AFAIK only MySQL can cancel specific queries. Postgres will actually cancel the whole connection (so it in theory it could result in race conditions...) But anyway, that means the Postgres adapter does not even need query IDs to work.

[igalklebanov]

OK, #176 got merged. Time to figure this thing out.

[alexgleason]

That's amazing @igalklebanov !! 🚀 Thank you! 😁 🙏

[oliveryasuna]

Thank you @igalklebanov !