github.com/Jeffail/benthos/v3@v3.65.0/public/service/example_output_batched_plugin_test.go (about)

     1  package service_test
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"fmt"
     7  
     8  	"github.com/Jeffail/benthos/v3/public/service"
     9  
    10  	// Import all standard Benthos components
    11  	_ "github.com/Jeffail/benthos/v3/public/components/all"
    12  )
    13  
    14  type batchOfJSONWriter struct{}
    15  
    16  func (b *batchOfJSONWriter) Connect(ctx context.Context) error {
    17  	return nil
    18  }
    19  
    20  func (b *batchOfJSONWriter) WriteBatch(ctx context.Context, msgs service.MessageBatch) error {
    21  	var messageObjs []interface{}
    22  	for _, msg := range msgs {
    23  		msgObj, err := msg.AsStructured()
    24  		if err != nil {
    25  			return err
    26  		}
    27  		messageObjs = append(messageObjs, msgObj)
    28  	}
    29  	outBytes, err := json.Marshal(map[string]interface{}{
    30  		"count":   len(msgs),
    31  		"objects": messageObjs,
    32  	})
    33  	if err != nil {
    34  		return err
    35  	}
    36  	fmt.Println(string(outBytes))
    37  	return nil
    38  }
    39  
    40  func (b *batchOfJSONWriter) Close(ctx context.Context) error {
    41  	return nil
    42  }
    43  
    44  // This example demonstrates how to create a batched output plugin, which allows
    45  // us to specify a batching mechanism and implement an interface that writes a
    46  // batch of messages in one call.
    47  func Example_outputBatchedPlugin() {
    48  	spec := service.NewConfigSpec().
    49  		Field(service.NewBatchPolicyField("batching"))
    50  
    51  	// Register our new output, which doesn't require a config schema.
    52  	err := service.RegisterBatchOutput(
    53  		"batched_json_stdout", spec,
    54  		func(conf *service.ParsedConfig, mgr *service.Resources) (out service.BatchOutput, policy service.BatchPolicy, maxInFlight int, err error) {
    55  			if policy, err = conf.FieldBatchPolicy("batching"); err != nil {
    56  				return
    57  			}
    58  			maxInFlight = 1
    59  			out = &batchOfJSONWriter{}
    60  			return
    61  		})
    62  	if err != nil {
    63  		panic(err)
    64  	}
    65  
    66  	// Use the stream builder API to create a Benthos stream that uses our new
    67  	// output type.
    68  	builder := service.NewStreamBuilder()
    69  
    70  	// Set the full Benthos configuration of the stream.
    71  	err = builder.SetYAML(`
    72  input:
    73    generate:
    74      count: 5
    75      interval: 1ms
    76      mapping: |
    77        root.id = count("batched output example messages")
    78        root.text = "some stuff"
    79  
    80  output:
    81    batched_json_stdout:
    82      batching:
    83        count: 5
    84  `)
    85  	if err != nil {
    86  		panic(err)
    87  	}
    88  
    89  	// Build a stream with our configured components.
    90  	stream, err := builder.Build()
    91  	if err != nil {
    92  		panic(err)
    93  	}
    94  
    95  	// And run it, blocking until it gracefully terminates once the generate
    96  	// input has generated a message and it has flushed through the stream.
    97  	if err = stream.Run(context.Background()); err != nil {
    98  		panic(err)
    99  	}
   100  
   101  	// Output: {"count":5,"objects":[{"id":1,"text":"some stuff"},{"id":2,"text":"some stuff"},{"id":3,"text":"some stuff"},{"id":4,"text":"some stuff"},{"id":5,"text":"some stuff"}]}
   102  }