github.com/ydb-platform/ydb-go-sdk/v3@v3.89.2/topic/topicsugar/iterators.go (about)

     1  //go:build go1.23
     2  
     3  package topicsugar
     4  
     5  import (
     6  	"context"
     7  	"encoding/json"
     8  	"slices"
     9  
    10  	"google.golang.org/protobuf/proto"
    11  
    12  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xiter"
    13  	"github.com/ydb-platform/ydb-go-sdk/v3/topic/topicreader"
    14  )
    15  
    16  // TopicMessageReader is interface for topicreader.Message
    17  //
    18  // Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
    19  type TopicMessageReader interface {
    20  	ReadMessage(ctx context.Context) (*topicreader.Message, error)
    21  }
    22  
    23  // TopicMessageIterator iterator wrapper over topic reader
    24  //
    25  // Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
    26  func TopicMessageIterator(ctx context.Context, r TopicMessageReader) xiter.Seq2[*topicreader.Message, error] {
    27  	return func(yield func(*topicreader.Message, error) bool) {
    28  		for {
    29  			mess, err := r.ReadMessage(ctx)
    30  			if !yield(mess, err) {
    31  				return
    32  			}
    33  
    34  			if err != nil {
    35  				return
    36  			}
    37  		}
    38  	}
    39  }
    40  
    41  // BytesIterator produce iterator over topic messages with Data as []byte, []byte is content of the message
    42  func BytesIterator(
    43  	ctx context.Context,
    44  	r TopicMessageReader,
    45  ) xiter.Seq2[*TypedTopicMessage[[]byte], error] {
    46  	var unmarshalFunc TypedUnmarshalFunc[*[]byte] = func(data []byte, dst *[]byte) error {
    47  		*dst = slices.Clone(data)
    48  
    49  		return nil
    50  	}
    51  
    52  	return IteratorFunc[[]byte](ctx, r, unmarshalFunc)
    53  }
    54  
    55  // StringIterator produce iterator over topic messages with Data is string, created from message content
    56  func StringIterator(
    57  	ctx context.Context,
    58  	r TopicMessageReader,
    59  ) xiter.Seq2[*TypedTopicMessage[string], error] {
    60  	var unmarshalFunc TypedUnmarshalFunc[*string] = func(data []byte, dst *string) error {
    61  		*dst = string(data)
    62  
    63  		return nil
    64  	}
    65  
    66  	return IteratorFunc[string](ctx, r, unmarshalFunc)
    67  }
    68  
    69  // JSONIterator produce iterator over topic messages with Data is T, created unmarshalled from message
    70  //
    71  // Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
    72  func JSONIterator[T any](
    73  	ctx context.Context,
    74  	r TopicMessageReader,
    75  ) xiter.Seq2[*TypedTopicMessage[T], error] {
    76  	var unmarshalFunc TypedUnmarshalFunc[*T] = func(data []byte, dst *T) error {
    77  		return json.Unmarshal(data, dst)
    78  	}
    79  
    80  	return IteratorFunc[T](ctx, r, unmarshalFunc)
    81  }
    82  
    83  // ProtobufIterator produce iterator over topic messages with Data is T, created unmarshalled from message
    84  //
    85  // Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
    86  func ProtobufIterator[T proto.Message](
    87  	ctx context.Context,
    88  	r TopicMessageReader,
    89  ) xiter.Seq2[*TypedTopicMessage[T], error] {
    90  	var unmarshalFunc TypedUnmarshalFunc[*T] = func(data []byte, dst *T) error {
    91  		return proto.Unmarshal(data, *dst)
    92  	}
    93  
    94  	return IteratorFunc[T](ctx, r, unmarshalFunc)
    95  }
    96  
    97  // IteratorFunc produce iterator over topic messages with Data is T,
    98  // created unmarshalled from message by custom function
    99  //
   100  // Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
   101  func IteratorFunc[T any](
   102  	ctx context.Context,
   103  	r TopicMessageReader,
   104  	f TypedUnmarshalFunc[*T],
   105  ) xiter.Seq2[*TypedTopicMessage[T], error] {
   106  	return func(yield func(*TypedTopicMessage[T], error) bool) {
   107  		for {
   108  			mess, err := r.ReadMessage(ctx)
   109  			if err != nil {
   110  				yield(nil, err)
   111  
   112  				return
   113  			}
   114  
   115  			var res TypedTopicMessage[T]
   116  
   117  			var unmarshal UnmarshalFunc = func(data []byte, _ any) error {
   118  				return f(data, &res.Data)
   119  			}
   120  
   121  			err = UnmarshalMessageWith(mess, unmarshal, nil)
   122  			if err != nil {
   123  				yield(nil, err)
   124  
   125  				return
   126  			}
   127  
   128  			res.Message = mess
   129  
   130  			if !yield(&res, err) {
   131  				return
   132  			}
   133  		}
   134  	}
   135  }
   136  
   137  type TypedTopicMessage[T any] struct {
   138  	*topicreader.Message
   139  	Data T
   140  }
   141  
   142  type TypedUnmarshalFunc[T any] func(data []byte, dst T) error