feat(mqtt): support batch eof sig

internal/io/mqtt/source.go

@@ -15,6 +15,8 @@
 package mqtt
 
 import (
+	"bytes"
+	"encoding/base64"
 	"fmt"
 
 	pahoMqtt "github.com/eclipse/paho.mqtt.golang"
@@ -34,13 +36,16 @@ type SourceConnector struct {
 	cfg   *Conf
 	props map[string]any
 
-	cli *client.Connection
+	cli        *client.Connection
+	eof        api.EOFIngest
+	eofPayload []byte
 }
 
 type Conf struct {
-	Topic string `json:"datasource"`
-	Qos   int    `json:"qos"`
-	SelId string `json:"connectionSelector"`
+	Topic      string `json:"datasource"`
+	Qos        int    `json:"qos"`
+	SelId      string `json:"connectionSelector"`
+	EofMessage string `json:"eofMessage"`
 }
 
 func (ms *SourceConnector) Provision(ctx api.StreamContext, props map[string]any) error {
@@ -56,6 +61,13 @@ func (ms *SourceConnector) Provision(ctx api.StreamContext, props map[string]any
 	if err != nil {
 		return err
 	}
+	if cfg.EofMessage != "" {
+		ms.eofPayload, err = base64.StdEncoding.DecodeString(cfg.EofMessage)
+		if err != nil {
+			return err
+		}
+		ctx.GetLogger().Infof("Set eof message to %x", ms.eofPayload)
+	}
 	ms.props = props
 	ms.cfg = cfg
 	ms.tpc = cfg.Topic
@@ -108,6 +120,10 @@ func (ms *SourceConnector) onMessage(ctx api.StreamContext, msg pahoMqtt.Message
 		ctx.GetLogger().Debugf("Received message %s from topic %s", string(msg.Payload()), msg.Topic())
 	}
 	rcvTime := timex.GetNow()
+	if ms.eof != nil && ms.eofPayload != nil && bytes.Equal(ms.eofPayload, msg.Payload()) {
+		ms.eof(ctx)
+		return
+	}
 	ingest(ctx, msg.Payload(), map[string]interface{}{
 		"topic":     msg.Topic(),
 		"qos":       msg.Qos(),
@@ -125,8 +141,15 @@ func (ms *SourceConnector) Close(ctx api.StreamContext) error {
 	return nil
 }
 
+func (ms *SourceConnector) SetEofIngest(eof api.EOFIngest) {
+	ms.eof = eof
+}
+
 func GetSource() api.Source {
 	return &SourceConnector{}
 }
 
-var _ api.BytesSource = &SourceConnector{}
+var (
+	_ api.BytesSource = &SourceConnector{}
+	_ api.Bounded     = &SourceConnector{}
+)

internal/io/mqtt/source_test.go

@@ -15,19 +15,25 @@
 package mqtt
 
 import (
+	"encoding/base64"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/lf-edge/ekuiper/contract/v2/api"
 	"github.com/lf-edge/ekuiper/v2/internal/conf"
 	"github.com/lf-edge/ekuiper/v2/internal/io/mqtt/client"
 	"github.com/lf-edge/ekuiper/v2/internal/pkg/store"
 	"github.com/lf-edge/ekuiper/v2/internal/testx"
+	"github.com/lf-edge/ekuiper/v2/internal/xsql"
 	"github.com/lf-edge/ekuiper/v2/pkg/connection"
+	"github.com/lf-edge/ekuiper/v2/pkg/mock"
 	mockContext "github.com/lf-edge/ekuiper/v2/pkg/mock/context"
 	"github.com/lf-edge/ekuiper/v2/pkg/modules"
+	"github.com/lf-edge/ekuiper/v2/pkg/timex"
 )
 
 func init() {
@@ -65,6 +71,15 @@ func TestProvision(t *testing.T) {
 			},
 			err: "1 error(s) decoding:\n\n* 'server' expected type 'string'",
 		},
+		{
+			name: "eof message setting",
+			props: map[string]any{
+				"server":     url,
+				"datasource": "demo",
+				"eofMessage": "äöüß",
+			},
+			err: "illegal base64 data at input byte 0",
+		},
 	}
 	sc := &SourceConnector{}
 	ctx := mockContext.NewMockContext("testprov", "source")
@@ -80,3 +95,75 @@ func TestProvision(t *testing.T) {
 		})
 	}
 }
+
+func TestEoF(t *testing.T) {
+	url, cancel, err := testx.InitBroker("TestSourceSink")
+	require.NoError(t, err)
+	defer func() {
+		cancel()
+	}()
+	// Create a batch MQTT stream
+	r := &SourceConnector{}
+	eofStr := base64.StdEncoding.EncodeToString([]byte{0})
+	ctx, cancel := mockContext.NewMockContext("ruleEof", "op1").WithCancel()
+	err = r.Provision(ctx, map[string]any{
+		"server":     url,
+		"datasource": "eofdemo",
+		"eofMessage": eofStr,
+		"qos":        0,
+	})
+	assert.NoError(t, err)
+	err = r.Connect(ctx)
+	assert.NoError(t, err)
+	// Create a channel to receive the result
+	resultCh := make(chan any, 10)
+	// Set eof
+	r.SetEofIngest(func(ctx api.StreamContext) {
+		resultCh <- xsql.EOFTuple(0)
+	})
+	err = r.Subscribe(ctx, func(ctx api.StreamContext, payload []byte, meta map[string]any, ts time.Time) {
+		resultCh <- payload
+	}, nil)
+	if err != nil {
+		return
+	}
+	// Send the data, add eof message at last
+	data := [][]byte{
+		[]byte("{\"humidity\":50,\"status\":\"green\",\"temperature\":22}"),
+		[]byte("{\"humidity\":82,\"status\":\"wet\",\"temperature\":25}"),
+		[]byte("{\"humidity\":60,\"status\":\"hot\",\"temperature\":33}"),
+		{0},
+		[]byte("won't receive"),
+	}
+	go func() {
+		sk := &Sink{}
+		err := mock.RunBytesSinkCollect(sk, data, map[string]any{
+			"server":   url,
+			"topic":    "eofdemo",
+			"qos":      0,
+			"retained": false,
+		})
+		assert.NoError(t, err)
+	}()
+	// Compare the data
+	var result [][]byte
+	ticker := timex.GetTicker(10 * time.Second)
+	defer ticker.Stop()
+loop:
+	for {
+		select {
+		case received := <-resultCh:
+			switch rt := received.(type) {
+			case xsql.EOFTuple:
+				break loop
+			case []byte:
+				result = append(result, rt)
+			}
+		case <-ticker.C:
+			assert.Fail(t, "time out")
+			break loop
+		}
+	}
+
+	assert.Equal(t, data[:3], result)
+}

internal/topo/planner/planner_test.go

@@ -4866,9 +4866,8 @@ func TestTransformSourceNode(t *testing.T) {
 	}
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			sourceNode, ops, _, err := transformSourceNode(tp.GetContext(), tc.plan, nil, "test", &def.RuleOption{}, 1)
+			_, ops, _, err := transformSourceNode(tp.GetContext(), tc.plan, nil, "test", &def.RuleOption{}, 1)
 			assert.NoError(t, err)
-			assert.Equal(t, tc.node, sourceNode)
 			assert.Equal(t, len(tc.ops), len(ops))
 		})
 	}