cmd/ockafka: Use gNMI client instead of openconfig.

.travis.yml

@@ -4,7 +4,7 @@ go:
 - 1.x
 - master
 before_install:
-- go get -v github.com/golang/lint/golint
+- go get -v golang.org/x/lint/golint
 - go get -v -t -d ./...
 after_success:
 - make coverdata

cmd/ockafka/README.md

@@ -9,18 +9,18 @@ Subscribe to all updates on the Arista device at `10.0.1.2` and stream to a loca
 Kafka instance:
 
 ```
-ockafka -addrs 10.0.1.2
+ockafka -addr 10.0.1.2
 ```
 
-Subscribe to temperature sensors from 2 switches and stream to a remote Kafka instance:
+Subscribe to temperature sensors and stream to a remote Kafka instance:
 
 ```
-ockafka -addrs 10.0.1.2,10.0.1.3 -kafkaaddrs kafka:9092 -subscribe /Sysdb/environment/temperature/status/tempSensor
+ockafka -addr 10.0.1.2 -kafkaaddrs kafka:9092 -subscribe /Sysdb/environment/temperature/status/tempSensor
 ```
 
 Start in a container:
 ```
-docker run aristanetworks/ockafka -addrs 10.0.1.1 -kafkaaddrs kafka:9092
+docker run aristanetworks/ockafka -addr 10.0.1.2 -kafkaaddrs kafka:9092
 ```
 
 ## Kafka/Elastic integration demo

cmd/ockafka/main.go

@@ -7,23 +7,23 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"strings"
-	"sync"
 
+	"github.com/aristanetworks/goarista/gnmi"
 	"github.com/aristanetworks/goarista/kafka"
 	"github.com/aristanetworks/goarista/kafka/openconfig"
 	"github.com/aristanetworks/goarista/kafka/producer"
-	"github.com/aristanetworks/goarista/openconfig/client"
 
 	"github.com/Shopify/sarama"
 	"github.com/aristanetworks/glog"
-	"github.com/golang/protobuf/proto"
+	pb "github.com/openconfig/gnmi/proto/gnmi"
 )
 
-var keysFlag = flag.String("kafkakeys", "",
-	"Keys for kafka messages (comma-separated, default: the value of -addrs")
+var keyFlag = flag.String("kafkakey", "",
+	"Key for kafka messages (default: the value of -addr")
 
 func newProducer(addresses []string, topic, key, dataset string) (producer.Producer, error) {
 	encodedKey := sarama.StringEncoder(key)
@@ -36,33 +36,54 @@ func newProducer(addresses []string, topic, key, dataset string) (producer.Produ
 }
 
 func main() {
-	username, password, subscriptions, grpcAddrs, opts := client.ParseFlags()
+	cfg := &gnmi.Config{}
+	flag.StringVar(&cfg.Addr, "addr", "", "Address of gNMI gRPC server with optional VRF name")
+	flag.StringVar(&cfg.CAFile, "cafile", "", "Path to server TLS certificate file")
+	flag.StringVar(&cfg.CertFile, "certfile", "", "Path to client TLS certificate file")
+	flag.StringVar(&cfg.KeyFile, "keyfile", "", "Path to client TLS private key file")
+	flag.StringVar(&cfg.Password, "password", "", "Password to authenticate with")
+	flag.StringVar(&cfg.Username, "username", "", "Username to authenticate with")
+	flag.BoolVar(&cfg.TLS, "tls", false, "Enable TLS")
+	subscribePaths := flag.String("subscribe", "/", "Comma-separated list of paths to subscribe to")
+	flag.Parse()
 
-	if *keysFlag == "" {
-		*keysFlag = strings.Join(grpcAddrs, ",")
-	}
-	keys := strings.Split(*keysFlag, ",")
-	if len(grpcAddrs) != len(keys) {
-		glog.Fatal("Please provide the same number of addresses and Kafka keys")
+	var key string
+	if key = *keyFlag; key == "" {
+		key = cfg.Addr
 	}
 	addresses := strings.Split(*kafka.Addresses, ",")
-	wg := new(sync.WaitGroup)
-	for i, grpcAddr := range grpcAddrs {
-		key := keys[i]
-		p, err := newProducer(addresses, *kafka.Topic, key, grpcAddr)
-		if err != nil {
+	p, err := newProducer(addresses, *kafka.Topic, key, cfg.Addr)
+	if err != nil {
+		glog.Fatal(err)
+	} else {
+		glog.Infof("Initialized Kafka producer for %s", cfg.Addr)
+	}
+	p.Start()
+	defer p.Stop()
+
+	subscriptions := strings.Split(*subscribePaths, ",")
+	ctx := gnmi.NewContext(context.Background(), cfg)
+	client, err := gnmi.Dial(cfg)
+	if err != nil {
+		glog.Fatal(err)
+	}
+	respChan := make(chan *pb.SubscribeResponse)
+	errChan := make(chan error)
+	defer close(errChan)
+	subscribeOptions := &gnmi.SubscribeOptions{
+		Mode:       "stream",
+		StreamMode: "target_defined",
+		Paths:      gnmi.SplitPaths(subscriptions),
+	}
+	go gnmi.Subscribe(ctx, client, subscribeOptions, respChan, errChan)
+	for {
+		select {
+		case resp := <-respChan:
+			func(response *pb.SubscribeResponse) {
+				p.Write(response)
+			}(resp)
+		case err := <-errChan:
 			glog.Fatal(err)
-		} else {
-			glog.Infof("Initialized Kafka producer for %s", grpcAddr)
-		}
-		publish := func(addr string, message proto.Message) {
-			p.Write(message)
 		}
-		wg.Add(1)
-		p.Start()
-		defer p.Stop()
-		c := client.New(username, password, grpcAddr, opts)
-		go c.Subscribe(wg, subscriptions, publish)
 	}
-	wg.Wait()
 }

kafka/encoder.go

@@ -14,13 +14,13 @@ import (
 
 	"github.com/Shopify/sarama"
 	"github.com/aristanetworks/glog"
-	"github.com/golang/protobuf/proto"
+	pb "github.com/openconfig/gnmi/proto/gnmi"
 )
 
 // MessageEncoder is an encoder interface
 // which handles encoding proto.Message to sarama.ProducerMessage
 type MessageEncoder interface {
-	Encode(proto.Message) ([]*sarama.ProducerMessage, error)
+	Encode(*pb.SubscribeResponse) ([]*sarama.ProducerMessage, error)
 	HandleSuccess(*sarama.ProducerMessage)
 	HandleError(*sarama.ProducerError)
 }
@@ -60,11 +60,9 @@ func NewBaseEncoder(typ string) *BaseEncoder {
 	return e
 }
 
-// Encode encodes the proto message to a sarama.ProducerMessage
-func (e *BaseEncoder) Encode(message proto.Message) ([]*sarama.ProducerMessage,
+// Encode does nothing, but keep it in order for BaseEncoder to implement MessageEncoder interface
+func (e *BaseEncoder) Encode(response *pb.SubscribeResponse) ([]*sarama.ProducerMessage,
 	error) {
-	// doesn't do anything, but keep it in order for BaseEncoder
-	// to implement MessageEncoder interface
 	return nil, nil
 }
 

kafka/openconfig/encoder.go

@@ -9,14 +9,13 @@ import (
 	"fmt"
 	"time"
 
-	"github.com/aristanetworks/goarista/elasticsearch"
+	"github.com/aristanetworks/goarista/gnmi"
 	"github.com/aristanetworks/goarista/kafka"
-	"github.com/aristanetworks/goarista/openconfig"
 
 	"github.com/Shopify/sarama"
 	"github.com/aristanetworks/glog"
 	"github.com/golang/protobuf/proto"
-	pb "github.com/openconfig/reference/rpc/openconfig"
+	pb "github.com/openconfig/gnmi/proto/gnmi"
 )
 
 // UnhandledMessageError is used for proto messages not matching the handled types
@@ -55,19 +54,15 @@ func NewEncoder(topic string, key sarama.Encoder, dataset string) kafka.MessageE
 	}
 }
 
-func (e *elasticsearchMessageEncoder) Encode(message proto.Message) ([]*sarama.ProducerMessage,
+func (e *elasticsearchMessageEncoder) Encode(response *pb.SubscribeResponse) (
+	[]*sarama.ProducerMessage,
 	error) {
 
-	response, ok := message.(*pb.SubscribeResponse)
-	if !ok {
-		return nil, UnhandledMessageError{message: message}
-	}
 	update := response.GetUpdate()
 	if update == nil {
 		return nil, UnhandledSubscribeResponseError{response: response}
 	}
-	updateMap, err := openconfig.NotificationToMap(e.dataset, update,
-		elasticsearch.EscapeFieldName)
+	updateMap, err := gnmi.NotificationToMap(update)
 	if err != nil {
 		return nil, err
 	}

kafka/producer/producer.go

@@ -13,18 +13,18 @@ import (
 
 	"github.com/Shopify/sarama"
 	"github.com/aristanetworks/glog"
-	"github.com/golang/protobuf/proto"
+	pb "github.com/openconfig/gnmi/proto/gnmi"
 )
 
 // Producer forwards messages recvd on a channel to kafka.
 type Producer interface {
 	Start()
-	Write(proto.Message)
+	Write(*pb.SubscribeResponse)
 	Stop()
 }
 
 type producer struct {
-	notifsChan    chan proto.Message
+	notifsChan    chan *pb.SubscribeResponse
 	kafkaProducer sarama.AsyncProducer
 	encoder       kafka.MessageEncoder
 	done          chan struct{}
@@ -53,7 +53,7 @@ func New(encoder kafka.MessageEncoder,
 	}
 
 	p := &producer{
-		notifsChan:    make(chan proto.Message),
+		notifsChan:    make(chan *pb.SubscribeResponse),
 		kafkaProducer: kafkaProducer,
 		encoder:       encoder,
 		done:          make(chan struct{}),
@@ -91,9 +91,9 @@ func (p *producer) run() {
 	}
 }
 
-func (p *producer) Write(msg proto.Message) {
+func (p *producer) Write(response *pb.SubscribeResponse) {
 	select {
-	case p.notifsChan <- msg:
+	case p.notifsChan <- response:
 	case <-p.done:
 		// TODO: This should probably return an EOF error, but that
 		// would change the API
@@ -106,8 +106,8 @@ func (p *producer) Stop() {
 	p.kafkaProducer.Close()
 }
 
-func (p *producer) produceNotifications(protoMessage proto.Message) error {
-	messages, err := p.encoder.Encode(protoMessage)
+func (p *producer) produceNotifications(response *pb.SubscribeResponse) error {
+	messages, err := p.encoder.Encode(response)
 	if err != nil {
 		return err
 	}

kafka/producer/producer_test.go

@@ -9,13 +9,14 @@ import (
 	"strings"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/aristanetworks/goarista/kafka/openconfig"
 	"github.com/aristanetworks/goarista/test"
 
 	"github.com/Shopify/sarama"
-	"github.com/golang/protobuf/proto"
-	pb "github.com/openconfig/reference/rpc/openconfig"
+	pb "github.com/openconfig/gnmi/proto/gnmi"
+	"github.com/openconfig/gnmi/value"
 )
 
 type mockAsyncProducer struct {
@@ -53,16 +54,22 @@ func (p *mockAsyncProducer) Errors() <-chan *sarama.ProducerError {
 	return p.errors
 }
 
-func newPath(path string) *pb.Path {
-	if path == "" {
-		return nil
+func newNotification(path []string, timestamp *time.Time) *pb.Notification {
+	sv, _ := value.FromScalar(timestamp.String())
+	return &pb.Notification{
+		Timestamp: timestamp.UnixNano() / 1e6,
+		Update: []*pb.Update{
+			{
+				Path: &pb.Path{Element: path},
+				Val:  sv,
+			},
+		},
 	}
-	return &pb.Path{Element: strings.Split(path, "/")}
 }
 
 func TestKafkaProducer(t *testing.T) {
 	mock := newMockAsyncProducer()
-	toDB := make(chan proto.Message)
+	toDB := make(chan *pb.SubscribeResponse)
 	topic := "occlient"
 	systemID := "Foobar"
 	toDBProducer := &producer{
@@ -75,23 +82,17 @@ func TestKafkaProducer(t *testing.T) {
 
 	toDBProducer.Start()
 
+	path := []string{"foo", "bar"}
+	timestamp := time.Now()
 	response := &pb.SubscribeResponse{
 		Response: &pb.SubscribeResponse_Update{
-			Update: &pb.Notification{
-				Timestamp: 0,
-				Prefix:    newPath("/foo/bar"),
-				Update:    []*pb.Update{},
-			},
+			Update: newNotification(path, &timestamp),
 		},
 	}
 	document := map[string]interface{}{
-		"timestamp": int64(0),
-		"update": map[string]interface{}{
-			"": map[string]interface{}{
-				"foo": map[string]interface{}{
-					"bar": map[string]interface{}{},
-				},
-			},
+		"timestamp": timestamp.UnixNano() / 1e6,
+		"updates": map[string]interface{}{
+			"/" + strings.Join(path, "/"): timestamp.String(),
 		},
 	}
 
@@ -118,11 +119,11 @@ func TestKafkaProducer(t *testing.T) {
 	if err != nil {
 		t.Errorf("Error decoding into JSON: %s", err)
 	}
-	if !test.DeepEqual(document["update"], result.(map[string]interface{})["update"]) {
-		t.Errorf("Protobuf sent from Kafka Producer does not match original.\nOriginal: %v\nNew:%v",
+	if !test.DeepEqual(document["updates"], result.(map[string]interface{})["updates"]) {
+		t.Errorf(
+			"Protobuf sent from Kafka Producer does not match original.\nOriginal: %#v\nNew:%#v",
 			document, result)
 	}
-
 	toDBProducer.Stop()
 }
 
@@ -136,7 +137,7 @@ func TestProducerStartStop(t *testing.T) {
 	// this test checks that Start() followed by Stop() doesn't cause any race conditions.
 
 	mock := newMockAsyncProducer()
-	toDB := make(chan proto.Message)
+	toDB := make(chan *pb.SubscribeResponse)
 	topic := "occlient"
 	systemID := "Foobar"
 	p := &producer{
@@ -146,13 +147,11 @@ func TestProducerStartStop(t *testing.T) {
 		done:          make(chan struct{}),
 	}
 
+	path := []string{"foo", "bar"}
+	timestamp := time.Now()
 	msg := &pb.SubscribeResponse{
 		Response: &pb.SubscribeResponse_Update{
-			Update: &pb.Notification{
-				Timestamp: 0,
-				Prefix:    newPath("/foo/bar"),
-				Update:    []*pb.Update{},
-			},
+			Update: newNotification(path, &timestamp),
 		},
 	}