github.com/aristanetworks/goarista@v0.0.0-20240514173732-cca2755bbd44/kafka/gnmi/encoder.go (about)

     1  // Copyright (c) 2016 Arista Networks, Inc.
     2  // Use of this source code is governed by the Apache License 2.0
     3  // that can be found in the COPYING file.
     4  
     5  package gnmi
     6  
     7  import (
     8  	"encoding/json"
     9  	"fmt"
    10  	"time"
    11  
    12  	"github.com/aristanetworks/goarista/elasticsearch"
    13  	"github.com/aristanetworks/goarista/kafka"
    14  
    15  	"github.com/IBM/sarama"
    16  	"github.com/aristanetworks/glog"
    17  	"github.com/openconfig/gnmi/proto/gnmi"
    18  	"google.golang.org/protobuf/proto"
    19  )
    20  
    21  // UnhandledMessageError is used for proto messages not matching the handled types
    22  type UnhandledMessageError struct {
    23  	message proto.Message
    24  }
    25  
    26  func (e UnhandledMessageError) Error() string {
    27  	return fmt.Sprintf("Unexpected type %T in proto message: %#v", e.message, e.message)
    28  }
    29  
    30  // UnhandledSubscribeResponseError is used for subscribe responses not matching the handled types
    31  type UnhandledSubscribeResponseError struct {
    32  	response *gnmi.SubscribeResponse
    33  }
    34  
    35  func (e UnhandledSubscribeResponseError) Error() string {
    36  	return fmt.Sprintf("Unexpected type %T in subscribe response: %#v", e.response, e.response)
    37  }
    38  
    39  type elasticsearchMessageEncoder struct {
    40  	*kafka.BaseEncoder
    41  	topic   string
    42  	dataset string
    43  	key     sarama.Encoder
    44  }
    45  
    46  // NewEncoder creates and returns a new elasticsearch MessageEncoder
    47  func NewEncoder(topic string, key sarama.Encoder, dataset string) kafka.MessageEncoder {
    48  	baseEncoder := kafka.NewBaseEncoder("elasticsearch")
    49  	return &elasticsearchMessageEncoder{
    50  		BaseEncoder: baseEncoder,
    51  		topic:       topic,
    52  		dataset:     dataset,
    53  		key:         key,
    54  	}
    55  }
    56  
    57  func (e *elasticsearchMessageEncoder) Encode(message proto.Message) ([]*sarama.ProducerMessage,
    58  	error) {
    59  	response, ok := message.(*gnmi.SubscribeResponse)
    60  	if !ok {
    61  		return nil, UnhandledMessageError{message: message}
    62  	}
    63  	update := response.GetUpdate()
    64  	if update == nil {
    65  		return nil, UnhandledSubscribeResponseError{response: response}
    66  	}
    67  	updateMaps, err := elasticsearch.NotificationToMaps(e.dataset, update)
    68  	if err != nil {
    69  		return nil, err
    70  	}
    71  	messages := make([]*sarama.ProducerMessage, len(updateMaps))
    72  	for i, updateMap := range updateMaps {
    73  		updateJSON, err := json.Marshal(updateMap)
    74  		if err != nil {
    75  			return nil, err
    76  		}
    77  		glog.V(9).Infof("kafka: %s", updateJSON)
    78  
    79  		messages[i] = &sarama.ProducerMessage{
    80  			Topic:    e.topic,
    81  			Key:      e.key,
    82  			Value:    sarama.ByteEncoder(updateJSON),
    83  			Metadata: kafka.Metadata{StartTime: time.Unix(0, update.Timestamp), NumMessages: 1},
    84  		}
    85  	}
    86  	return messages, nil
    87  }