github.com/apache/beam/sdks/v2@v2.48.2/go/test/integration/io/xlang/kafka/kafka.go (about)

     1  // Licensed to the Apache Software Foundation (ASF) under one or more
     2  // contributor license agreements.  See the NOTICE file distributed with
     3  // this work for additional information regarding copyright ownership.
     4  // The ASF licenses this file to You under the Apache License, Version 2.0
     5  // (the "License"); you may not use this file except in compliance with
     6  // the License.  You may obtain a copy of the License at
     7  //
     8  //    http://www.apache.org/licenses/LICENSE-2.0
     9  //
    10  // Unless required by applicable law or agreed to in writing, software
    11  // distributed under the License is distributed on an "AS IS" BASIS,
    12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  // See the License for the specific language governing permissions and
    14  // limitations under the License.
    15  
    16  // Package kafka contains integration tests for cross-language Kafka IO
    17  // transforms.
    18  package kafka
    19  
    20  import (
    21  	"bytes"
    22  	"fmt"
    23  
    24  	"github.com/apache/beam/sdks/v2/go/pkg/beam"
    25  	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder"
    26  	"github.com/apache/beam/sdks/v2/go/pkg/beam/io/xlang/kafkaio"
    27  	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/passert"
    28  	"github.com/google/uuid"
    29  )
    30  
    31  func appendUuid(prefix string) string {
    32  	return fmt.Sprintf("%v_%v", prefix, uuid.New())
    33  }
    34  
    35  // writeList encodes a list of ints and sends encoded ints to Kafka.
    36  func writeInts(s beam.Scope, expansionAddr, bootstrapAddr, topic string, inputs []int) {
    37  	s = s.Scope("kafka_test.WriteListToKafka")
    38  
    39  	ins := beam.CreateList(s, inputs)
    40  	encoded := beam.ParDo(s, func(i int) ([]byte, error) {
    41  		var buf bytes.Buffer
    42  		err := coder.EncodeVarInt(int64(i), &buf)
    43  		return buf.Bytes(), err
    44  	}, ins)
    45  	keyed := beam.ParDo(s, func(b []byte) ([]byte, []byte) {
    46  		return []byte(""), b
    47  	}, encoded)
    48  	kafkaio.Write(s, expansionAddr, bootstrapAddr, topic, keyed)
    49  }
    50  
    51  // readList reads a set number of elements from Kafka and decodes them to ints.
    52  func readInts(s beam.Scope, expansionAddr, bootstrapAddr, topic string, numRecords int64) beam.PCollection {
    53  	s = s.Scope("kafka_test.ReadListFromKafka")
    54  
    55  	reads := kafkaio.Read(s, expansionAddr, bootstrapAddr, []string{topic},
    56  		kafkaio.MaxNumRecords(numRecords),
    57  		kafkaio.ConsumerConfigs(map[string]string{"auto.offset.reset": "earliest"}))
    58  	vals := beam.DropKey(s, reads)
    59  	decoded := beam.ParDo(s, func(b []byte) (int, error) {
    60  		buf := bytes.NewBuffer(b)
    61  		i, err := coder.DecodeVarInt(buf)
    62  		return int(i), err
    63  	}, vals)
    64  	return decoded
    65  }
    66  
    67  // WritePipeline creates a pipeline that writes a given slice of ints to Kafka.
    68  func WritePipeline(expansionAddr, bootstrapAddr, topic string, inputs []int) *beam.Pipeline {
    69  	p, s := beam.NewPipelineWithRoot()
    70  	writeInts(s, expansionAddr, bootstrapAddr, topic, inputs)
    71  	return p
    72  }
    73  
    74  // ReadPipeline creates a pipeline that reads ints from Kafka and asserts that
    75  // they match a given slice of ints. This reads a number of records equal to
    76  // the length of the given slice.
    77  func ReadPipeline(expansionAddr, bootstrapAddr, topic string, inputs []int) *beam.Pipeline {
    78  	p, s := beam.NewPipelineWithRoot()
    79  	result := readInts(s, expansionAddr, bootstrapAddr, topic, int64(len(inputs)))
    80  
    81  	// Validate that records read from Kafka match the given slice.
    82  	ins := beam.CreateList(s, inputs)
    83  	passert.Equals(s, result, ins)
    84  	return p
    85  }