Feature/sink redis

go.mod

@@ -17,11 +17,12 @@ require (
 	github.com/cloudquery/plugin-sdk/v4 v4.16.1
 	github.com/go-playground/validator/v10 v10.14.0
 	github.com/goccy/go-json v0.10.2
-	github.com/google/uuid v1.3.1
 	github.com/gorilla/websocket v1.5.0
 	github.com/jackc/pgx/v5 v5.4.3
+	github.com/mehanizm/airtable v0.3.1
 	github.com/prometheus/client_golang v1.11.1
 	github.com/rcrowley/go-metrics v0.0.0-20201227073835-cf1acfcdf475
+	github.com/redis/go-redis/v9 v9.4.0
 	github.com/sashabaranov/go-openai v1.17.9
 	github.com/spf13/cobra v1.6.1
 	github.com/twmb/franz-go v1.15.2
@@ -42,6 +43,7 @@ require (
 	github.com/charmbracelet/lipgloss v0.9.1 // indirect
 	github.com/coreos/go-semver v0.3.0 // indirect
 	github.com/coreos/go-systemd/v22 v22.3.2 // indirect
+	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
 	github.com/fatih/structs v1.1.0 // indirect
 	github.com/frankban/quicktest v1.14.4 // indirect
 	github.com/gabriel-vasile/mimetype v1.4.2 // indirect
@@ -52,6 +54,7 @@ require (
 	github.com/golang/protobuf v1.5.3 // indirect
 	github.com/golang/snappy v0.0.4 // indirect
 	github.com/google/flatbuffers v23.5.26+incompatible // indirect
+	github.com/google/uuid v1.3.1 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/influxdata/line-protocol/v2 v2.2.1 // indirect
 	github.com/ivancorrales/knoa v0.0.2 // indirect
@@ -68,7 +71,6 @@ require (
 	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/mattn/go-runewidth v0.0.15 // indirect
 	github.com/matttproud/golang_protobuf_extensions v1.0.1 // indirect
-	github.com/mehanizm/airtable v0.3.1 // indirect
 	github.com/mitchellh/mapstructure v1.5.0 // indirect
 	github.com/montanaflynn/stats v0.0.0-20171201202039-1bf9dbcd8cbe // indirect
 	github.com/muesli/reflow v0.3.0 // indirect

internal/sinks/redis/config.go

@@ -0,0 +1,13 @@
+package redis
+
+type Config struct {
+	RedisAddr     string `json:"redis_addr" yaml:"redis_addr"`
+	RedisPassword string `json:"redis_password" yaml:"redis_password"`
+	// CustomNamespace can't be used alongside with NamespaceByStream, only one of the
+	// Options should be provided.
+	// In case provided both - namespace by stream will be used
+	CustomNamespace   string `json:"custom_namespace" yaml:"custom_namespace"`
+	NamespaceByStream bool   `json:"namespace_by_stream" yaml:"namespace_by_stream"`
+	KeyPrefix         string `json:"key_prefix" yaml:"key_prefix"`
+	SetWithTTL        int64  `json:"set_with_ttl" yaml:"set_with_ttl"`
+}

internal/sinks/redis/plugin.go

@@ -0,0 +1,111 @@
+package redis
+
+import (
+	"context"
+	"fmt"
+	"github.com/charmbracelet/log"
+	"github.com/goccy/go-json"
+	"github.com/redis/go-redis/v9"
+	"github.com/usedatabrew/blink/internal/schema"
+	"github.com/usedatabrew/blink/internal/sinks"
+	"github.com/usedatabrew/blink/internal/stream_context"
+	"github.com/usedatabrew/message"
+	"strings"
+	"time"
+)
+
+type SinkPlugin struct {
+	redisConn    *redis.Client
+	appCtx       *stream_context.Context
+	streamSchema []schema.StreamSchema
+	config       Config
+	logger       *log.Logger
+	pksByStream  map[string]string
+}
+
+func NewRedisSinkPlugin(config Config, schema []schema.StreamSchema, appCtx *stream_context.Context) sinks.DataSink {
+	return &SinkPlugin{
+		streamSchema: schema,
+		config:       config,
+		appCtx:       appCtx,
+		logger:       appCtx.Logger.WithPrefix("[sink]: redis"),
+		pksByStream:  map[string]string{},
+	}
+}
+
+func (s *SinkPlugin) Connect(context context.Context) error {
+	rdb := redis.NewClient(&redis.Options{
+		Addr:     s.config.RedisAddr,
+		Password: s.config.RedisPassword,
+		DB:       0,
+	})
+
+	status := rdb.Ping(context)
+	if status.Err() != nil {
+		return status.Err()
+	}
+
+	s.logger.Debug("Connect check info", "result", status)
+
+	s.redisConn = rdb
+	return nil
+}
+
+func (s *SinkPlugin) SetExpectedSchema(schema []schema.StreamSchema) {
+	for _, stream := range schema {
+		var pkCol string
+		for _, c := range stream.Columns {
+			if c.PK {
+				pkCol = c.Name
+			}
+		}
+		if strings.Index(stream.StreamName, ".") != -1 {
+			splitName := strings.Split(stream.StreamName, ".")
+			s.pksByStream[splitName[1]] = pkCol
+		} else {
+			s.pksByStream[stream.StreamName] = pkCol
+		}
+	}
+}
+
+func (s *SinkPlugin) GetType() sinks.SinkDriver {
+	return sinks.RedisSinkType
+}
+
+func (s *SinkPlugin) Write(m *message.Message) error {
+	streamName := m.GetStream()
+	namespace := ""
+	if s.config.CustomNamespace != "" {
+		namespace = s.config.CustomNamespace
+	}
+
+	if s.config.NamespaceByStream {
+		namespace = streamName
+	}
+
+	messageKey := fmt.Sprintf("%s%v", s.config.KeyPrefix, m.Data.AccessProperty(s.pksByStream[streamName]))
+	if namespace != "" {
+		messageKey = namespace + ":" + messageKey
+	}
+
+	payload := m.Data.JsonQ().First()
+	ttl := time.Duration(0)
+	if s.config.SetWithTTL > 0 {
+		ttl = time.Second * time.Duration(s.config.SetWithTTL)
+	}
+
+	data, err := json.Marshal(payload)
+	if err != nil {
+		return err
+	}
+
+	s.logger.Debug("Writing message to redis", "key", messageKey, "payload", string(data))
+	status := s.redisConn.Set(s.appCtx.GetContext(), messageKey, data, ttl)
+	return status.Err()
+}
+
+func (s *SinkPlugin) Stop() {
+	if err := s.redisConn.Close(); err != nil {
+		s.logger.Fatal("Failed to close redis client", "error", err)
+	}
+}

internal/sinks/sink_drivers.go

@@ -8,4 +8,5 @@ const (
 	KafkaSinkType     SinkDriver = "kafka"
 	PostgresSinkType  SinkDriver = "postgres"
 	MongoDBSinkType   SinkDriver = "mongodb"
+	RedisSinkType     SinkDriver = "redis"
 )

internal/sinks/websockets/config.go

@@ -1,6 +1,6 @@
 package websockets
 
 type Config struct {
-	Url     string            `json:"url"`
-	Headers map[string]string `json:"headers"`
+	Url     string            `json:"url" yaml:"url"`
+	Headers map[string]string `json:"headers" yaml:"headers"`
 }

internal/sources/airtable/plugin.go

@@ -92,7 +92,7 @@ func (s *SourcePlugin) Start() {
 				}{FieldName: s.streamPks[stream], Direction: "asc"})
 			}
 
-			getRowsRequest.PageSize(2)
+			getRowsRequest.PageSize(100)
 			result, err := getRowsRequest.Do()
 			if err != nil {
 				s.messageEvents <- sources.MessageEvent{

public/stream/sink_wrapper.go

@@ -7,6 +7,7 @@ import (
 	"github.com/usedatabrew/blink/internal/sinks/kafka"
 	"github.com/usedatabrew/blink/internal/sinks/mongodb"
 	"github.com/usedatabrew/blink/internal/sinks/postgres"
+	"github.com/usedatabrew/blink/internal/sinks/redis"
 	"github.com/usedatabrew/blink/internal/sinks/stdout"
 	websocket "github.com/usedatabrew/blink/internal/sinks/websockets"
 	"github.com/usedatabrew/blink/internal/stream_context"
@@ -71,6 +72,12 @@ func (p *SinkWrapper) LoadDriver(driver sinks.SinkDriver, cfg config.Configurati
 			panic("can read driver config")
 		}
 		return websocket.NewWebSocketSinkPlugin(driverConfig, cfg.Source.StreamSchema, p.ctx)
+	case sinks.RedisSinkType:
+		driverConfig, err := ReadDriverConfig[redis.Config](cfg.Sink.Config, redis.Config{})
+		if err != nil {
+			panic("can read driver config")
+		}
+		return redis.NewRedisSinkPlugin(driverConfig, cfg.Source.StreamSchema, p.ctx)
 	case sinks.MongoDBSinkType:
 		driverConfig, err := ReadDriverConfig[mongodb.Config](cfg.Sink.Config, mongodb.Config{})
 		if err != nil {

public/stream/source_wrapper.go

@@ -101,6 +101,8 @@ func (p *SourceWrapper) LoadDriver(driver sources.SourceDriver, config config.Co
 			}
 
 			return airtable.NewAirTableSourcePlugin(driverConfig, config.Source.StreamSchema)
+	default:
+		p.ctx.Logger.WithPrefix("Source driver loader").Fatal("Failed to load driver", "driver", driver)
 	}
 
 	return nil