Merge pull request #1479 Write message in transaction

Author: rekby

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Added write to topics within transactions
+
 ## v3.80.10
 * Added `ydb.WithSessionPoolSessionUsageLimit()` option for limitation max count of session usage  
 * Refactored experimental topic iterators in `topicsugar` package

examples/topic/topicwriter/topic_writer_transaction.go

@@ -0,0 +1,76 @@
+package topicwriter
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3"
+	"github.com/ydb-platform/ydb-go-sdk/v3/query"
+	"github.com/ydb-platform/ydb-go-sdk/v3/topic/topicreader"
+	"github.com/ydb-platform/ydb-go-sdk/v3/topic/topicwriter"
+)
+
+func CopyMessagesBetweenTopicsTxWriter(
+	ctx context.Context,
+	db *ydb.Driver,
+	reader *topicreader.Reader,
+	topic string,
+) error {
+	return db.Query().DoTx(ctx, func(ctx context.Context, tx query.TxActor) error {
+		writer, err := db.Topic().StartTransactionalWriter(tx, topic)
+		if err != nil {
+			return err
+		}
+
+		batch, err := reader.PopMessagesBatchTx(ctx, tx)
+		if err != nil {
+			return err
+		}
+		for _, mess := range batch.Messages {
+			if err = writer.Write(ctx, topicwriter.Message{Data: mess}); err != nil {
+				return err
+			}
+		}
+
+		return nil
+	}, query.WithIdempotent())
+}
+
+func TableAndTopicWithinTransaction(
+	ctx context.Context,
+	db *ydb.Driver,
+	topicPath string,
+	id int64,
+) error {
+	return db.Query().DoTx(ctx, func(ctx context.Context, t query.TxActor) error {
+		row, err := t.QueryRow(ctx, "SELECT val FROM table WHERE id=$id", query.WithParameters(
+			ydb.ParamsBuilder().
+				Param("$id").Int64(id).
+				Build()))
+		if err != nil {
+			return err
+		}
+
+		var val int64
+		if err = row.Scan(&val); err != nil {
+			return err
+		}
+
+		// the writer is dedicated for the transaction, it can't be used outside the transaction
+		// it is no needs to close or flush the messages - it happened internally on transaction commit
+		writer, err := db.Topic().StartTransactionalWriter(t, topicPath)
+		if err != nil {
+			return err
+		}
+
+		err = writer.Write(ctx, topicwriter.Message{
+			Data: strings.NewReader(fmt.Sprintf("val: %v processed", val)),
+		})
+		if err != nil {
+			return err
+		}
+
+		return nil
+	})
+}

internal/grpcwrapper/rawtopic/rawtopiccommon/transaction.go

@@ -0,0 +1,19 @@
+package rawtopiccommon
+
+import "github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Topic"
+
+type TransactionIdentity struct {
+	ID      string
+	Session string
+}
+
+func (t TransactionIdentity) ToProto() *Ydb_Topic.TransactionIdentity {
+	if t.ID == "" && t.Session == "" {
+		return nil
+	}
+
+	return &Ydb_Topic.TransactionIdentity{
+		Id:      t.ID,
+		Session: t.Session,
+	}
+}

internal/grpcwrapper/rawtopic/rawtopicwriter/messages.go

@@ -145,6 +145,7 @@ type WriteRequest struct {
 
 	Messages []MessageData
 	Codec    rawtopiccommon.Codec
+	Tx       rawtopiccommon.TransactionIdentity
 }
 
 func (r *WriteRequest) toProto() (p *Ydb_Topic.StreamWriteMessage_FromClient_WriteRequest, err error) {
@@ -161,6 +162,7 @@ func (r *WriteRequest) toProto() (p *Ydb_Topic.StreamWriteMessage_FromClient_Wri
 		WriteRequest: &Ydb_Topic.StreamWriteMessage_WriteRequest{
 			Messages: messages,
 			Codec:    int32(r.Codec.ToProto()),
+			Tx:       r.Tx.ToProto(),
 		},
 	}
 
@@ -231,11 +233,13 @@ func (r *WriteResult) GetAcks() (res traceAck) {
 	}
 	for i := range r.Acks {
 		ack := &r.Acks[i]
-		if ack.MessageWriteStatus.Type == WriteStatusTypeWritten {
+		switch ack.MessageWriteStatus.Type {
+		case WriteStatusTypeWritten:
 			res.WrittenCount++
-		}
-		if ack.MessageWriteStatus.Type == WriteStatusTypeSkipped {
+		case WriteStatusTypeSkipped:
 			res.SkipCount++
+		case WriteStatusTypeWrittenInTx:
+			res.WrittenInTxCount++
 		}
 
 		if ack.SeqNo < res.SeqNoMin {
@@ -261,6 +265,7 @@ type traceAck = struct {
 	WrittenOffsetMin int64
 	WrittenOffsetMax int64
 	WrittenCount     int
+	WrittenInTxCount int
 	SkipCount        int
 }
 
@@ -299,6 +304,12 @@ func (s *MessageWriteStatus) fromProto(status interface{}) error {
 		s.SkippedReason = WriteStatusSkipReason(v.Skipped.GetReason())
 
 		return nil
+
+	case *Ydb_Topic.StreamWriteMessage_WriteResponse_WriteAck_WrittenInTx_:
+		s.Type = WriteStatusTypeWrittenInTx
+
+		return nil
+
 	default:
 		return xerrors.WithStackTrace(xerrors.Wrap(fmt.Errorf("ydb: unexpected write status type: %v", reflect.TypeOf(v))))
 	}
@@ -307,19 +318,19 @@ func (s *MessageWriteStatus) fromProto(status interface{}) error {
 type WriteStatusType int
 
 const (
-	WriteStatusTypeUnknown WriteStatusType = iota
-	WriteStatusTypeWritten
+	WriteStatusTypeWritten WriteStatusType = iota + 1
 	WriteStatusTypeSkipped
+	WriteStatusTypeWrittenInTx
 )
 
 func (t WriteStatusType) String() string {
 	switch t {
-	case WriteStatusTypeUnknown:
-		return "Unknown"
 	case WriteStatusTypeSkipped:
 		return "Skipped"
 	case WriteStatusTypeWritten:
 		return "Written"
+	case WriteStatusTypeWrittenInTx:
+		return "WrittenInTx"
 	default:
 		return strconv.Itoa(int(t))
 	}

internal/grpcwrapper/rawtopic/update_offset_in_transaction.go

@@ -9,19 +9,16 @@ import (
 
 type UpdateOffsetsInTransactionRequest struct {
 	OperationParams rawydb.OperationParams
-	Tx              UpdateOffsetsInTransactionRequest_TransactionIdentity
+	Tx              rawtopiccommon.TransactionIdentity
 	Topics          []UpdateOffsetsInTransactionRequest_TopicOffsets
 	Consumer        string
 }
 
 func (r *UpdateOffsetsInTransactionRequest) ToProto() *Ydb_Topic.UpdateOffsetsInTransactionRequest {
 	req := &Ydb_Topic.UpdateOffsetsInTransactionRequest{
 		OperationParams: r.OperationParams.ToProto(),
-		Tx: &Ydb_Topic.TransactionIdentity{
-			Id:      r.Tx.ID,
-			Session: r.Tx.Session,
-		},
-		Consumer: r.Consumer,
+		Tx:              r.Tx.ToProto(),
+		Consumer:        r.Consumer,
 	}
 
 	req.Topics = make([]*Ydb_Topic.UpdateOffsetsInTransactionRequest_TopicOffsets, len(r.Topics))
@@ -56,11 +53,6 @@ func (r *UpdateOffsetsInTransactionRequest) ToProto() *Ydb_Topic.UpdateOffsetsIn
 	return req
 }
 
-type UpdateOffsetsInTransactionRequest_TransactionIdentity struct { //nolint:revive,stylecheck
-	ID      string
-	Session string
-}
-
 type UpdateOffsetsInTransactionRequest_TopicOffsets struct { //nolint:revive,stylecheck
 	Path       string // Topic path
 	Partitions []UpdateOffsetsInTransactionRequest_PartitionOffsets

internal/grpcwrapper/rawtopic/update_offset_in_transaction_test.go

@@ -27,7 +27,7 @@ func TestUpdateOffsetsInTransactionRequestToProto(t *testing.T) {
 				HasValue: true,
 			},
 		},
-		Tx: UpdateOffsetsInTransactionRequest_TransactionIdentity{
+		Tx: rawtopiccommon.TransactionIdentity{
 			ID:      "tx-id",
 			Session: "session-id",
 		},

internal/query/transaction.go

@@ -35,7 +35,8 @@ type (
 
 		completed bool
 
-		onCompleted xsync.Set[*baseTx.OnTransactionCompletedFunc]
+		onBeforeCommit xsync.Set[*baseTx.OnTransactionBeforeCommit]
+		onCompleted    xsync.Set[*baseTx.OnTransactionCompletedFunc]
 	}
 )
 
@@ -100,6 +101,11 @@ func (tx *Transaction) QueryResultSet(
 		}),
 	}
 	if settings.TxControl().Commit {
+		err = tx.waitOnBeforeCommit(ctx)
+		if err != nil {
+			return nil, err
+		}
+
 		// notification about complete transaction must be sended for any error or for successfully read all result if
 		// it was execution with commit flag
 		resultOpts = append(resultOpts,
@@ -144,6 +150,11 @@ func (tx *Transaction) QueryRow(
 		}),
 	}
 	if settings.TxControl().Commit {
+		err := tx.waitOnBeforeCommit(ctx)
+		if err != nil {
+			return nil, err
+		}
+
 		// notification about complete transaction must be sended for any error or for successfully read all result if
 		// it was execution with commit flag
 		resultOpts = append(resultOpts,
@@ -204,6 +215,11 @@ func (tx *Transaction) Exec(ctx context.Context, q string, opts ...options.Execu
 		}),
 	}
 	if settings.TxControl().Commit {
+		err = tx.waitOnBeforeCommit(ctx)
+		if err != nil {
+			return err
+		}
+
 		// notification about complete transaction must be sended for any error or for successfully read all result if
 		// it was execution with commit flag
 		resultOpts = append(resultOpts,
@@ -268,6 +284,11 @@ func (tx *Transaction) Query(ctx context.Context, q string, opts ...options.Exec
 		}),
 	}
 	if settings.TxControl().Commit {
+		err = tx.waitOnBeforeCommit(ctx)
+		if err != nil {
+			return nil, err
+		}
+
 		// notification about complete transaction must be sended for any error or for successfully read all result if
 		// it was execution with commit flag
 		resultOpts = append(resultOpts,
@@ -310,7 +331,12 @@ func (tx *Transaction) CommitTx(ctx context.Context) (finalErr error) {
 		tx.completed = true
 	}()
 
-	err := commitTx(ctx, tx.s.client, tx.s.ID(), tx.ID())
+	err := tx.waitOnBeforeCommit(ctx)
+	if err != nil {
+		return err
+	}
+
+	err = commitTx(ctx, tx.s.client, tx.s.ID(), tx.ID())
 	if err != nil {
 		if xerrors.IsOperationError(err, Ydb.StatusIds_BAD_SESSION) {
 			tx.s.SetStatus(session.StatusClosed)
@@ -360,10 +386,24 @@ func (tx *Transaction) Rollback(ctx context.Context) (finalErr error) {
 	return nil
 }
 
+func (tx *Transaction) OnBeforeCommit(f baseTx.OnTransactionBeforeCommit) {
+	tx.onBeforeCommit.Add(&f)
+}
+
 func (tx *Transaction) OnCompleted(f baseTx.OnTransactionCompletedFunc) {
 	tx.onCompleted.Add(&f)
 }
 
+func (tx *Transaction) waitOnBeforeCommit(ctx context.Context) (resErr error) {
+	tx.onBeforeCommit.Range(func(f *baseTx.OnTransactionBeforeCommit) bool {
+		resErr = (*f)(ctx)
+
+		return resErr == nil
+	})
+
+	return resErr
+}
+
 func (tx *Transaction) notifyOnCompleted(err error) {
 	tx.completed = true