sink.go 5.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201
  1. // Copyright 2021-2022 EMQ Technologies Co., Ltd.
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. //go:build redisdb || !core
  15. package redis
  16. import (
  17. "encoding/json"
  18. "errors"
  19. "fmt"
  20. "github.com/lf-edge/ekuiper/pkg/ast"
  21. "time"
  22. "github.com/go-redis/redis/v7"
  23. "github.com/lf-edge/ekuiper/pkg/api"
  24. "github.com/lf-edge/ekuiper/pkg/cast"
  25. )
  26. type config struct {
  27. // host:port address.
  28. Addr string `json:"addr,omitempty"`
  29. Username string `json:"username,omitempty"`
  30. // Optional password. Must match the password specified in the
  31. Password string `json:"password,omitempty"`
  32. // Database to be selected after connecting to the server.
  33. Db int `json:"db,omitempty"`
  34. // key of field
  35. Field string `json:"field,omitempty"`
  36. // key define
  37. Key string `json:"key,omitempty"`
  38. DataType string `json:"dataType,omitempty"`
  39. Expiration time.Duration `json:"expiration,omitempty"`
  40. RowkindField string `json:"rowkindField"`
  41. DataTemplate string `json:"dataTemplate"`
  42. }
  43. type RedisSink struct {
  44. c *config
  45. cli *redis.Client
  46. }
  47. func (r *RedisSink) Configure(props map[string]interface{}) error {
  48. c := &config{DataType: "string", Expiration: -1}
  49. err := cast.MapToStruct(props, c)
  50. if err != nil {
  51. return err
  52. }
  53. if c.Key == "" && c.Field == "" {
  54. return errors.New("redis sink must have key or field")
  55. }
  56. if c.DataType != "string" && c.DataType != "list" {
  57. return errors.New("redis sink only support string or list data type")
  58. }
  59. r.c = c
  60. return nil
  61. }
  62. func (r *RedisSink) Open(ctx api.StreamContext) (err error) {
  63. logger := ctx.GetLogger()
  64. logger.Debug("Opening redis sink")
  65. r.cli = redis.NewClient(&redis.Options{
  66. Addr: r.c.Addr,
  67. Username: r.c.Username,
  68. Password: r.c.Password,
  69. DB: r.c.Db, // use default DB
  70. })
  71. return nil
  72. }
  73. func (r *RedisSink) Collect(ctx api.StreamContext, data interface{}) error {
  74. logger := ctx.GetLogger()
  75. var val string
  76. if r.c.DataTemplate != "" { // The result is a string
  77. v, _, err := ctx.TransformOutput(data)
  78. if err != nil {
  79. logger.Error(err)
  80. return err
  81. }
  82. m := make(map[string]interface{})
  83. err = json.Unmarshal(v, &m)
  84. if err != nil {
  85. return fmt.Errorf("fail to decode data %s after applying dataTemplate for error %v", string(v), err)
  86. }
  87. data = m
  88. val = string(v)
  89. }
  90. switch d := data.(type) {
  91. case []map[string]interface{}:
  92. for _, el := range d {
  93. err := r.save(ctx, el, val)
  94. if err != nil {
  95. return err
  96. }
  97. }
  98. case map[string]interface{}:
  99. err := r.save(ctx, d, val)
  100. if err != nil {
  101. return err
  102. }
  103. default:
  104. return fmt.Errorf("unrecognized format of %s", data)
  105. }
  106. logger.Debug("insert success %v", data)
  107. return nil
  108. }
  109. func (r *RedisSink) Close(ctx api.StreamContext) error {
  110. ctx.GetLogger().Infof("Closing redis sink")
  111. err := r.cli.Close()
  112. return err
  113. }
  114. func (r *RedisSink) save(ctx api.StreamContext, data map[string]interface{}, val string) error {
  115. logger := ctx.GetLogger()
  116. if val == "" {
  117. jsonBytes, err := json.Marshal(data)
  118. if err != nil {
  119. return err
  120. }
  121. val = string(jsonBytes)
  122. }
  123. key := r.c.Key
  124. var err error
  125. if r.c.Field != "" {
  126. keyval, ok := data[r.c.Field]
  127. if !ok {
  128. return fmt.Errorf("field %s does not exist in data %v", r.c.Field, data)
  129. }
  130. key, err = cast.ToString(keyval, cast.CONVERT_ALL)
  131. if err != nil {
  132. return fmt.Errorf("key must be string or convertible to string, but got %v", keyval)
  133. }
  134. }
  135. rowkind := ast.RowkindUpsert
  136. if r.c.RowkindField != "" {
  137. c, ok := data[r.c.RowkindField]
  138. if ok {
  139. rowkind, ok = c.(string)
  140. if !ok {
  141. return fmt.Errorf("rowkind field %s is not a string in data %v", r.c.RowkindField, data)
  142. }
  143. if rowkind != ast.RowkindInsert && rowkind != ast.RowkindUpdate && rowkind != ast.RowkindDelete && rowkind != ast.RowkindUpsert {
  144. return fmt.Errorf("invalid rowkind %s", rowkind)
  145. }
  146. }
  147. }
  148. switch rowkind {
  149. case ast.RowkindInsert, ast.RowkindUpdate, ast.RowkindUpsert:
  150. if r.c.DataType == "list" {
  151. err = r.cli.LPush(key, val).Err()
  152. if err != nil {
  153. return fmt.Errorf("lpush %s:%s error, %v", key, val, err)
  154. }
  155. logger.Debugf("push redis list success, key:%s data: %v", key, val)
  156. } else {
  157. err = r.cli.Set(key, val, r.c.Expiration*time.Second).Err()
  158. if err != nil {
  159. return fmt.Errorf("set %s:%s error, %v", key, val, err)
  160. }
  161. logger.Debugf("set redis string success, key:%s data: %s", key, val)
  162. }
  163. case ast.RowkindDelete:
  164. if r.c.DataType == "list" {
  165. err = r.cli.LPop(key).Err()
  166. if err != nil {
  167. return fmt.Errorf("lpop %s error, %v", key, err)
  168. }
  169. logger.Debugf("pop redis list success, key:%s data: %v", key, val)
  170. } else {
  171. err = r.cli.Del(key).Err()
  172. if err != nil {
  173. logger.Error(err)
  174. return err
  175. }
  176. logger.Debugf("delete redis string success, key:%s data: %s", key, val)
  177. }
  178. default:
  179. // never happen
  180. logger.Errorf("unexpected rowkind %s", rowkind)
  181. }
  182. return nil
  183. }
  184. func GetSink() api.Sink {
  185. return &RedisSink{}
  186. }