github.com/wfusion/gofusion@v1.1.14/common/infra/watermill/pubsub/kafka/context.go (about)

     1  package kafka
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  )
     7  
     8  type contextKey int
     9  
    10  const (
    11  	_ contextKey = iota
    12  	partitionContextKey
    13  	partitionOffsetContextKey
    14  	timestampContextKey
    15  	keyContextKey
    16  )
    17  
    18  func setPartitionToCtx(ctx context.Context, partition int32) context.Context {
    19  	return context.WithValue(ctx, partitionContextKey, partition)
    20  }
    21  
    22  // MessagePartitionFromCtx returns Kafka partition of the consumed message
    23  func MessagePartitionFromCtx(ctx context.Context) (int32, bool) {
    24  	partition, ok := ctx.Value(partitionContextKey).(int32)
    25  	return partition, ok
    26  }
    27  
    28  func setPartitionOffsetToCtx(ctx context.Context, offset int64) context.Context {
    29  	return context.WithValue(ctx, partitionOffsetContextKey, offset)
    30  }
    31  
    32  // MessagePartitionOffsetFromCtx returns Kafka partition offset of the consumed message
    33  func MessagePartitionOffsetFromCtx(ctx context.Context) (int64, bool) {
    34  	offset, ok := ctx.Value(partitionOffsetContextKey).(int64)
    35  	return offset, ok
    36  }
    37  
    38  func setMessageTimestampToCtx(ctx context.Context, timestamp time.Time) context.Context {
    39  	return context.WithValue(ctx, timestampContextKey, timestamp)
    40  }
    41  
    42  // MessageTimestampFromCtx returns Kafka internal timestamp of the consumed message
    43  func MessageTimestampFromCtx(ctx context.Context) (time.Time, bool) {
    44  	timestamp, ok := ctx.Value(timestampContextKey).(time.Time)
    45  	return timestamp, ok
    46  }
    47  
    48  func setMessageKeyToCtx(ctx context.Context, key []byte) context.Context {
    49  	return context.WithValue(ctx, keyContextKey, key)
    50  }
    51  
    52  // MessageKeyFromCtx returns Kafka internal key of the consumed message
    53  func MessageKeyFromCtx(ctx context.Context) ([]byte, bool) {
    54  	key, ok := ctx.Value(keyContextKey).([]byte)
    55  	return key, ok
    56  }