[BUG]: transactions in AWS RDS Data API not working

[berenddeboer]

What version of drizzle-orm are you using?

0.28.2

What version of drizzle-kit are you using?

0.19.12

Describe the Bug

I'm using the AWS RDS Data API (yes I know), but transactions are not working. I.e. rollback does not work. The reason is that everything runs outside the transaction.

I've tried to come up with a simple sample, basically a select in a transaction, and it seems the transactionId is not correctly passed:

  return await db.transaction(async (tx: Transaction) => {
    tx.select().from(..).where(...)
  })

What happens is that this calls prepareQuery in src/aws-data-api/pg/session.ts, but transacationId is always null. So that is this code:

	prepareQuery<T extends PreparedQueryConfig = PreparedQueryConfig>(
		query: Query,
		fields: SelectedFieldsOrdered | undefined,
		transactionId?: string,
                ...) {
		return new AwsDataApiPreparedQuery(
			this.client,
			query.sql,
			query.params,
			query.typings ?? [],
			this.options,
			fields,
			transactionId,
			customResultMapper,
		);

	}

I can fix this by changing the transactionId parameter in the call that creates a new AwsDataApiPreparedQuery to be this.transactionId. Then everything is fine.

If you look at other session examples, you can see that the third parameter is called name, for example in src/pg-core/session.ts, not transactionId:

	abstract prepareQuery<T extends PreparedQueryConfig = PreparedQueryConfig>(
		query: Query,
		fields: SelectedFieldsOrdered | undefined,
		name: string | undefined,
		customResultMapper?: (rows: unknown[][], mapColumnValue?: (value: unknown) => unknown) => T['execute'],
	): PreparedQuery<T>;

The code is pretty complex, but as far as I can see it transactionId is never passed in, so a clear bug. If you look at execute in src/pg-core/query-builders/select.ts you see that it always calls _prepare() without any parameters:

	execute: ReturnType<this['prepare']>['execute'] = (placeholderValues) => {
		return tracer.startActiveSpan('drizzle.operation', () => {
			return this._prepare().execute(placeholderValues);
		});
	};

But _prepare needs a name, as that should end up as the transactionId above:

private _prepare(name?: string): PreparedQuery<
...
  session.prepareQuery(dialect.sqlToQuery(this.getSQL()), fieldsList, name)
...

As you see above, that third parameter is transactionId in session.prepareQuery, but is undefined. It cannot possibly have a transactionId the way the code is written.

I suggest prepareQuery in session gets changed to get the sessionId from itself, then everything works.

Expected behavior

Rollback works

Environment & setup

No response

[NachoVazquez]

Thanks, I'm hitting this hard as well.

[pantoninho]

Me too. Did you find a workaround @NachoVazquez?

[nacho-vazquez]

Me too. Did you find a workaround @NachoVazquez?

Yeah, I forked the repo and created a release from this PR. It is working 🙃

Once this get merge I will get back to the original package.

[pantoninho]

Sharing an alternative workaround that does not require forking the codebase.

Monkey-patch AwsDataApiSession.prepareQuery by calling the following function:

import { AwsDataApiPreparedQuery, AwsDataApiSession } from 'drizzle-orm/aws-data-api/pg';

/**
 * Patches the AWS Data API driver to fix transactions
 * @see https://github.com/drizzle-team/drizzle-orm/issues/1171
 */
function patchAwsDataApiDriver() {
    AwsDataApiSession.prototype.prepareQuery = function (
        query,
        fields,
        name,
        customResultMapper
    ) {
        return new AwsDataApiPreparedQuery(
            this.client,
            query.sql,
            query.params,
            query.typings ?? [],
            this.options,
            fields,
            this.transactionId,
            customResultMapper
        );
    };
}

[NachoVazquez]

Good stuff!

[NachoVazquez]

How do you use it @pantoninho ?

Do you call it before using drizzle? Like in a folder /drizzle/index.ts where you initialize the db object?

[pantoninho]

There are a lot of ways to use it, but I suggest the following:

• create a file like drizzle/aws-data-api-patch.js where you export that function.
• import and call it wherever you want before making the first transaction. A good place to do it is before initialising drizzle.

import { RDSDataClient } from '@aws-sdk/client-rds-data';
import {
    AwsDataApiPreparedQuery,
    AwsDataApiSession,
    drizzle
} from 'drizzle-orm/aws-data-api/pg';
import { PgDatabase } from 'drizzle-orm/pg-core';
import { patchAwsDataApiDriver } from './aws-data-api-patch.js';

/**
 * @type {PgDatabase}
 */
const db = getAWSDataApiClient();

export { db };

/**
 * Returns a drizzle client for an RDS database.
 * @returns {import('drizzle-orm/aws-data-api/pg').AwsDataApiPgDatabase} the drizzle client
 */
function getAWSDataApiClient() {
    patchAwsDataApiDriver();
    const rdsClient = new RDSDataClient({});
    return drizzle(rdsClient, {
        database,
        resourceArn,
        secretArn
    });
}

[NachoVazquez]

Yeah, thank you, that's what I had in mind!