feat(query, topic): add TX lifecycle hooks and refactor tx-topic APIs

  - Introduce TX.onCommit, TX.onRollback, and TX.onClose lifecycle hooks
    (replaces registerPrecommitHook). Hooks receive AbortSignal where relevant.
  - Execute commit hooks before commit, rollback hooks on errors, and close hooks in finally.
  - Refactor Topic TX APIs to be tx-first:
      - createTxReader(tx, options) and createTxWriter(tx, options); options no longer includes tx.
  - Topic TX Reader registers a commit hook to update read offsets in the transaction.
  - Update _update_offsets_in_transaction signature to (tx, driver, consumer, updates).
  - Adjust e2e tests to the new API and parameter order.

  BREAKING CHANGE:

  - TopicClient.createTxReader and createTxWriter now accept (tx, options) instead of options with tx.
  - TX.registerPrecommitHook is removed; use tx.onCommit, tx.onRollback, and tx.onClose instead.

Signed-off-by: Vladislav Polyakov <[email protected]>

Author: polRk

e2e/topic/read-write-tx.test.ts

@@ -69,13 +69,13 @@ test('writes and reads in tx', async () => {
 
 	// Begin a transaction
 	await yql.begin({ idempotent: true }, async (tx) => {
-		await using readerTx = createTopicTxReader(driver, {
+		let readerTx = createTopicTxReader(driver, {
 			tx,
 			topic: testTopicName,
 			consumer: testConsumerName,
 		})
 
-		await using writerTx = createTopicTxWriter(driver, tx, {
+		await using writerTx = createTopicTxWriter(tx, driver, {
 			topic: testTopicName,
 			producer: testProducerName,
 		})
@@ -132,7 +132,7 @@ test('rollbacks reads', async () => {
 
 	await yql
 		.begin({ idempotent: true }, async (tx) => {
-			await using readerTx = createTopicTxReader(driver, {
+			let readerTx = createTopicTxReader(driver, {
 				tx,
 				topic: testTopicName,
 				consumer: testConsumerName,
@@ -186,7 +186,7 @@ test('rollbacks writes', async () => {
 
 	await yql
 		.begin({ idempotent: true }, async (tx) => {
-			await using writerTx = createTopicTxWriter(driver, tx, {
+			await using writerTx = createTopicTxWriter(tx, driver, {
 				topic: testTopicName,
 				producer: testProducerName,
 			})

packages/query/src/index.ts

@@ -24,7 +24,9 @@ export type TX = SQL & {
 	nodeId: bigint
 	sessionId: string
 	transactionId: string
-	registerPrecommitHook: RegisterPrecommitHook
+	onRollback: (fn: (error: unknown, signal?: AbortSignal) => Promise<void> | void) => void
+	onCommit: (fn: (signal?: AbortSignal) => Promise<void> | void) => void
+	onClose: (fn: (committed: boolean, signal?: AbortSignal) => Promise<void> | void) => void
 }
 
 interface SessionContextCallback<T> {
@@ -194,25 +196,35 @@ export function query(driver: Driver): QueryClient {
 
 			store.transactionId = beginTransactionResult.txMeta!.id
 
+			let сommitHooks: Array<(signal?: AbortSignal) => Promise<void> | void> = []
+			let rollbackHooks: Array<(error: unknown, signal?: AbortSignal) => Promise<void> | void> = []
+			let closeHooks: Array<(committed: boolean, signal?: AbortSignal) => Promise<void> | void> = []
+
+			let commited = false
 			try {
-				let precommitHooks: Array<() => Promise<void> | void> = []
 				let tx = Object.assign(yqlQuery, {
 					nodeId: store.nodeId,
 					sessionId: store.sessionId,
 					transactionId: store.transactionId,
-					registerPrecommitHook: (fn: () => Promise<void> | void) => {
-						precommitHooks.push(fn);
+					onRollback: (fn: () => Promise<void> | void) => {
+						rollbackHooks.push(fn);
+					},
+					onCommit: (fn: () => Promise<void> | void) => {
+						сommitHooks.push(fn);
+					},
+					onClose: (fn: () => Promise<void> | void) => {
+						closeHooks.push(fn);
 					},
 				}) as TX
 
 				dbg.log('executing transaction body')
 				let result = await ctx.run(store, () => caller!(tx, signal))
 
-				dbg.log('executing %d precommit hooks', precommitHooks.length)
-				await Promise.all(precommitHooks.map(async (hook, i) => {
-					dbg.log('executing precommit hook #%d', i + 1)
-					await hook()
-					dbg.log('precommit hook #%d completed', i + 1)
+				dbg.log('executing %d commit hooks', сommitHooks.length)
+				await Promise.all(сommitHooks.map(async (hook, i) => {
+					dbg.log('executing commit hook #%d', i + 1)
+					await hook(signal)
+					dbg.log('commit hook #%d completed', i + 1)
 				}))
 
 				dbg.log('committing transaction')
@@ -222,10 +234,19 @@ export function query(driver: Driver): QueryClient {
 					throw new CommitError("Transaction commit failed.", new YDBError(commitResult.status, commitResult.issues))
 				}
 
+				commited = true
 				dbg.log('transaction committed successfully')
 				return result
 			} catch (error) {
 				dbg.log('transaction error: %O', error)
+
+				dbg.log('executing %d rollback hooks', сommitHooks.length)
+				await Promise.all(rollbackHooks.map(async (hook, i) => {
+					dbg.log('executing rollback hook #%d', i + 1)
+					await hook(error, signal)
+					dbg.log('rollback hook #%d completed', i + 1)
+				}))
+
 				void client.rollbackTransaction({ sessionId: store.sessionId, txId: store.transactionId })
 
 				if (!isRetryableError(error, options.idempotent)) {
@@ -237,6 +258,13 @@ export function query(driver: Driver): QueryClient {
 			} finally {
 				dbg.log('deleting session %s', sessionResponse.sessionId)
 				void client.deleteSession({ sessionId: sessionResponse.sessionId })
+
+				dbg.log('executing %d close hooks', сommitHooks.length)
+				await Promise.all(closeHooks.map(async (hook, i) => {
+					dbg.log('executing close hook #%d', i + 1)
+					await hook(commited, signal)
+					dbg.log('close hook #%d completed', i + 1)
+				}))
 			}
 		})
 	}

packages/topic/src/index.ts

@@ -6,7 +6,7 @@ import type { TX } from "./tx.js";
 
 export interface TopicClient {
 	createReader(options: TopicReaderOptions): TopicReader;
-	createTxReader(options: TopicTxReaderOptions): TopicTxReader;
+	createTxReader(tx: TX, options: TopicTxReaderOptions): TopicTxReader;
 	createWriter(options: TopicWriterOptions): TopicWriter;
 	createTxWriter(tx: TX, options: TopicWriterOptions): TopicWriter;
 }
@@ -16,13 +16,13 @@ export function topic(driver: Driver): TopicClient {
 		createReader(options) {
 			return createTopicReader(driver, options);
 		},
-		createTxReader(options: TopicTxReaderOptions) {
-			return createTopicTxReader(driver, options);
+		createTxReader(tx: TX, options: Omit<TopicTxReaderOptions, 'tx'>) {
+			return createTopicTxReader(tx, driver, options);
 		},
 		createWriter(options: TopicWriterOptions) {
 			return createTopicWriter(driver, options);
 		},
-		createTxWriter(tx: TX, options: TopicWriterOptions) {
+		createTxWriter(tx: TX, options: Omit<TopicWriterOptions, 'tx'>) {
 			return createTopicTxWriter(tx, driver, options);
 		},
 	} as TopicClient

packages/topic/src/reader/_update_offsets_in_transaction.ts

@@ -16,8 +16,8 @@ import type { TopicPartitionSession } from "../partition-session.js"
 let dbg = loggers.topic.extend('reader')
 
 export let _update_offsets_in_transaction = async function updateOffsetsInTransaction(
-	driver: Driver,
 	tx: TX,
+	driver: Driver,
 	consumer: string,
 	updates: Array<{
 		partitionSession: TopicPartitionSession

packages/topic/src/reader/index.ts

@@ -14,6 +14,7 @@ import { _read } from "./_read.js"
 import { _commit } from "./_commit.js"
 import { _update_offsets_in_transaction } from "./_update_offsets_in_transaction.js"
 import { _create_disposal_functions, _initialize_codecs, _start_background_token_refresher } from "./_shared.js"
+import type { TX } from "../tx.js"
 
 let dbg = loggers.topic.extend('reader')
 
@@ -164,12 +165,12 @@ export const createTopicReader = function createTopicReader(driver: Driver, opti
 // Re-export types for compatibility
 export type { TopicReaderOptions, TopicReader, TopicTxReaderOptions, TopicTxReader } from "./types.js"
 
-export const createTopicTxReader = function createTopicTxReader(driver: Driver, options: TopicTxReaderOptions): TopicTxReader {
+export const createTopicTxReader = function createTopicTxReader(tx: TX, driver: Driver, options: Omit<TopicTxReaderOptions, 'tx'>): TopicTxReader {
 	options.updateTokenIntervalMs ??= 60_000 // Default is 60 seconds.
 
 	let state: TopicTxReaderState = {
 		driver,
-		options,
+		options: Object.assign(options, { tx }),
 		topicsReadSettings: _parse_topics_read_settings(options.topic),
 
 		// Control
@@ -194,7 +195,7 @@ export const createTopicTxReader = function createTopicTxReader(driver: Driver,
 	_initialize_codecs(state.codecs, options.codecMap)
 
 	// Register precommit hook to send updateOffsetsInTransaction
-	options.tx.onCommit(async () => {
+	tx.onCommit(async () => {
 		// Send updateOffsetsInTransaction for all read offsets
 		let updates = []
 
@@ -212,8 +213,8 @@ export const createTopicTxReader = function createTopicTxReader(driver: Driver,
 
 		if (updates.length > 0) {
 			await _update_offsets_in_transaction(
+				tx,
 				state.driver,
-				options.tx,
 				state.options.consumer,
 				updates
 			)
@@ -265,8 +266,8 @@ export const createTopicTxReader = function createTopicTxReader(driver: Driver,
 		if (updates.length > 0) {
 			try {
 				await _update_offsets_in_transaction(
+					tx,
 					state.driver,
-					options.tx,
 					state.options.consumer,
 					updates
 				)