Repository: lf-edge/ekuiper
Pull request #2940 (2024-06-19 09:25:05)

feat: support rewind

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/lf-edge/ekuiper/pull/2940?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lf-edge) Report
Attention: Patch coverage is `65.38462%` with `9 lines` in your changes missing coverage. Please review.
> Please [upload](https://docs.codecov.com/docs/codecov-uploader) report for BASE (`master-v2@df8fe15`). [Learn more](https://docs.codecov.io/docs/error-reference?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lf-edge#section-missing-base-commit) about missing BASE report.

| [Files](https://app.codecov.io/gh/lf-edge/ekuiper/pull/2940?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lf-edge) | Patch % | Lines |
|---|---|---|
| [internal/topo/node/source\_node.go](https://app.codecov.io/gh/lf-edge/ekuiper/pull/2940?src=pr&el=tree&filepath=internal%2Ftopo%2Fnode%2Fsource_node.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lf-edge#diff-aW50ZXJuYWwvdG9wby9ub2RlL3NvdXJjZV9ub2RlLmdv) | 60.87% | [4 Missing and 5 partials :warning: ](https://app.codecov.io/gh/lf-edge/ekuiper/pull/2940?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lf-edge) |

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff              @@
##             master-v2    #2940   +/-   ##
============================================
  Coverage             ?   69.70%           
============================================
  Files                ?      312           
  Lines                ?    29755           
  Branches             ?        0           
============================================
  Hits                 ?    20738           
  Misses               ?     7220           
  Partials             ?     1797           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/lf-edge/ekuiper/pull/2940?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lf-edge).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lf-edge).


Patch:
diff --git a/internal/io/sql/source.go b/internal/io/sql/source.go
index 6a850458..cb3d4002 100644
--- a/internal/io/sql/source.go
+++ b/internal/io/sql/source.go
@@ -152,6 +152,23 @@ func (s *SQLSourceConnector) queryData(ctx api.StreamContext, rcvTime time.Time,
 	}
 }
 
+func (s *SQLSourceConnector) GetOffset() (interface{}, error) {
+	return s.Query.GetIndexValue(), nil
+}
+
+func (s *SQLSourceConnector) Rewind(offset interface{}) error {
+	s.Query.SetIndexValue(offset)
+	return nil
+}
+
+func (s *SQLSourceConnector) ResetOffset(input map[string]interface{}) error {
+	wrap := s.Query.GetIndexValueWrap()
+	if wrap != nil {
+		wrap.UpdateByInput(input)
+	}
+	return nil
+}
+
 func scanIntoMap(mapValue map[string]interface{}, values []interface{}, columns []string) {
 	for idx, column := range columns {
 		if reflectValue := reflect.Indirect(reflect.Indirect(reflect.ValueOf(values[idx]))); reflectValue.IsValid() {
diff --git a/internal/topo/node/source_node.go b/internal/topo/node/source_node.go
index c3f3f877..d439cd19 100644
--- a/internal/topo/node/source_node.go
+++ b/internal/topo/node/source_node.go
@@ -89,6 +89,7 @@ func (m *SourceNode) ingestBytes(ctx api.StreamContext, data []byte, meta map[st
 	m.statManager.IncTotalRecordsOut()
 	m.statManager.IncTotalMessagesProcessed(1)
 	m.statManager.ProcessTimeEnd()
+	m.updateState(ctx)
 }
 
 func (m *SourceNode) ingestAnyTuple(ctx api.StreamContext, data any, meta map[string]any, ts time.Time) {
@@ -122,6 +123,7 @@ func (m *SourceNode) ingestAnyTuple(ctx api.StreamContext, data any, meta map[st
 	}
 	m.statManager.IncTotalMessagesProcessed(1)
 	m.statManager.ProcessTimeEnd()
+	m.updateState(ctx)
 }
 
 func (m *SourceNode) ingestMap(t map[string]any, meta map[string]any, ts time.Time) {
@@ -147,6 +149,38 @@ func (m *SourceNode) ingestEof(ctx api.StreamContext) {
 	m.Broadcast(xsql.EOFTuple(0))
 }
 
+const (
+	OffsetKey = "$$offset"
+)
+
+func (m *SourceNode) Rewind(ctx api.StreamContext) error {
+	s := m.s
+	if rw, ok := s.(api.Rewindable); ok {
+		if offset, err := ctx.GetState(OffsetKey); err != nil {
+			return err
+		} else if offset != nil {
+			ctx.GetLogger().Infof("Source rewind from %v", offset)
+			err = rw.Rewind(offset)
+			if err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+func (m *SourceNode) updateState(ctx api.StreamContext) error {
+	s := m.s
+	if rw, ok := s.(api.Rewindable); ok {
+		state, err := rw.GetOffset()
+		if err != nil {
+			return err
+		}
+		return ctx.PutState(OffsetKey, state)
+	}
+	return nil
+}
+
 // Run Subscribe could be a long-running function
 func (m *SourceNode) Run(ctx api.StreamContext, ctrlCh chan<- error) {
 	defer func() {
@@ -158,6 +192,9 @@ func (m *SourceNode) Run(ctx api.StreamContext, ctrlCh chan<- error) {
 		if err != nil {
 			return err
 		}
+		if err := m.Rewind(ctx); err != nil {
+			return err
+		}
 		switch ss := m.s.(type) {
 		case api.BytesSource:
 			err = ss.Subscribe(ctx, m.ingestBytes, m.ingestError)
diff --git a/pkg/sqldatabase/sqlgen/sqlgen.go b/pkg/sqldatabase/sqlgen/sqlgen.go
index d99ef273..4602a4df 100644
--- a/pkg/sqldatabase/sqlgen/sqlgen.go
+++ b/pkg/sqldatabase/sqlgen/sqlgen.go
@@ -15,6 +15,7 @@
 package sqlgen
 
 import (
+	"errors"
 	"fmt"
 
 	"github.com/lf-edge/ekuiper/v2/pkg/cast"
@@ -77,7 +78,7 @@ func (cfg *sqlConfig) Init(props map[string]interface{}) error {
 	}
 
 	if cfg.TemplateSqlQueryCfg == nil && cfg.InternalSqlQueryCfg == nil {
-		return fmt.Errorf("read properties %v fail with error: %v", props, err)
+		return errors.New("either one of the internalSqlQueryCfg and templateSqlQueryCfg should be defined")
 	}
 
 	if cfg.TemplateSqlQueryCfg != nil {


Test patch:
diff --git a/internal/io/sql/source_test.go b/internal/io/sql/source_test.go
index f723ac74..6d0c9e98 100644
--- a/internal/io/sql/source_test.go
+++ b/internal/io/sql/source_test.go
@@ -30,6 +30,7 @@ import (
 	"github.com/lf-edge/ekuiper/v2/pkg/connection"
 	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/store"
 )
 
 func init() {
@@ -166,3 +167,62 @@ func TestSQLConnectionErr(t *testing.T) {
 	require.NoError(t, sqlSource.Provision(ctx, props))
 	require.NoError(t, sqlSource.Connect(ctx))
 }
+
+func TestSQLSourceRewind(t *testing.T) {
+	connection.InitConnectionManager4Test()
+	ctx := mockContext.NewMockContext("1", "2")
+	s, err := testx.SetupEmbeddedMysqlServer(address, port)
+	require.NoError(t, err)
+	defer func() {
+		s.Close()
+	}()
+	props := map[string]interface{}{
+		"interval": "1s",
+		"dburl":    fmt.Sprintf("mysql://root:@%v:%v/test", address, port),
+		"internalSqlQueryCfg": map[string]interface{}{
+			"table": "t",
+			"limit": 1,
+			"indexFields": []map[string]interface{}{
+				{
+					"indexField":     "a",
+					"indexValue":     0,
+					"indexFieldType": "bigint",
+				},
+				{
+					"indexField":     "b",
+					"indexValue":     0,
+					"indexFieldType": "bigint",
+				},
+			},
+		},
+	}
+	sqlSource := GetSource()
+	require.NoError(t, sqlSource.Provision(ctx, props))
+	require.NoError(t, sqlSource.Connect(ctx))
+	sqlConnector, ok := sqlSource.(*SQLSourceConnector)
+	require.True(t, ok)
+	expectedData := map[string]any{
+		"a": int64(1),
+		"b": int64(1),
+	}
+	dataChan := make(chan any, 1)
+	sqlConnector.Pull(ctx, time.Now(), func(ctx api.StreamContext, data any, meta map[string]any, ts time.Time) {
+		dataChan <- data
+	}, func(ctx api.StreamContext, err error) {})
+	require.Equal(t, expectedData, <-dataChan)
+	state, err := sqlConnector.GetOffset()
+	require.NoError(t, err)
+	expectState := store.NewIndexFieldWrap([]*store.IndexField{
+		{
+			IndexFieldName:     "a",
+			IndexFieldValue:    int64(1),
+			IndexFieldDataType: "bigint",
+		},
+		{
+			IndexFieldName:     "b",
+			IndexFieldValue:    int64(1),
+			IndexFieldDataType: "bigint",
+		},
+	}...).GetStore()
+	require.Equal(t, expectState, state)
+}
diff --git a/internal/io/sql/testx/mysql_server.go b/internal/io/sql/testx/mysql_server.go
index 3ece4d70..05d7ed7e 100644
--- a/internal/io/sql/testx/mysql_server.go
+++ b/internal/io/sql/testx/mysql_server.go
@@ -59,6 +59,9 @@ func createTestDatabase() *memory.DbProvider {
 		{Name: "b", Type: types.Int64, Nullable: false, Source: tableName},
 	}), db.GetForeignKeyCollection())
 	db.AddTable(tableName, table)
-	_ = table.Insert(ctx, sql.NewRow(1, 1))
+	err := table.Insert(ctx, sql.NewRow(int64(1), int64(1)))
+	if err != nil {
+		panic(err)
+	}
 	return pro
 }
diff --git a/internal/topo/node/source_node_test.go b/internal/topo/node/source_node_test.go
index ba14831f..8d3fda45 100644
--- a/internal/topo/node/source_node_test.go
+++ b/internal/topo/node/source_node_test.go
@@ -22,6 +22,7 @@ import (
 	"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/pkg/def"
@@ -363,3 +364,77 @@ func (m *MockPullSource) Pull(ctx api.StreamContext, trigger time.Time, ingest a
 func (m *MockPullSource) SetEofIngest(eof api.EOFIngest) {
 	m.set = true
 }
+
+type MockRewindSource struct {
+	notify chan struct{}
+	state  int
+}
+
+func (m *MockRewindSource) GetOffset() (any, error) {
+	return m.state, nil
+}
+
+func (m *MockRewindSource) Rewind(offset any) error {
+	m.state = offset.(int)
+	return nil
+}
+
+func (m *MockRewindSource) ResetOffset(input map[string]any) error {
+	return nil
+}
+
+func (m *MockRewindSource) Provision(ctx api.StreamContext, configs map[string]any) error {
+	return nil
+}
+
+func (m *MockRewindSource) Close(ctx api.StreamContext) error {
+	return nil
+}
+
+func (m *MockRewindSource) Connect(ctx api.StreamContext) error {
+	return nil
+}
+
+func (m *MockRewindSource) Subscribe(ctx api.StreamContext, ingest api.TupleIngest, ingestError api.ErrorIngest) error {
+	go func() {
+		for {
+			select {
+			case <-m.notify:
+				ingest(ctx, map[string]any{
+					"key": m.state,
+				}, nil, time.Now())
+				m.state++
+			}
+		}
+	}()
+	return nil
+}
+
+func TestMockRewind(t *testing.T) {
+	notify := make(chan struct{})
+	m := &MockRewindSource{
+		notify: notify,
+	}
+	var sc api.TupleSource = m
+	ctx := mockContext.NewMockContext("rule1", "src1")
+	// set rewind value
+	ctx.PutState(OffsetKey, 10)
+	errCh := make(chan error)
+	scn, err := NewSourceNode(ctx, "mock_connector", sc, map[string]any{"datasource": "demo"}, &def.RuleOption{
+		BufferLength: 1024,
+		SendError:    true,
+	})
+	assert.NoError(t, err)
+	result := make(chan any, 10)
+	err = scn.AddOutput(result, "testResult")
+	assert.NoError(t, err)
+	scn.Open(ctx, errCh)
+	notify <- struct{}{}
+	data := <-result
+	require.Equal(t, map[string]interface{}{"key": 10}, map[string]interface{}(data.(*xsql.Tuple).Message))
+	notify <- struct{}{}
+	data = <-result
+	require.Equal(t, map[string]interface{}{"key": 11}, map[string]interface{}(data.(*xsql.Tuple).Message))
+	v, _ := ctx.GetState(OffsetKey)
+	require.Equal(t, 11, v)
+}
diff --git a/internal/topo/topotest/mocknode/mock_source.go b/internal/topo/topotest/mocknode/mock_source.go
index 4a72802a..e247bd0d 100644
--- a/internal/topo/topotest/mocknode/mock_source.go
+++ b/internal/topo/topotest/mocknode/mock_source.go
@@ -29,7 +29,7 @@ type MockSource struct {
 	data   []*xsql.Tuple
 	offset int
 	eof    api.EOFIngest
-	sync.Mutex
+	sync.RWMutex
 }
 
 const TIMELEAP = 200
@@ -74,9 +74,9 @@ func (m *MockSource) Subscribe(ctx api.StreamContext, ingest api.TupleIngest, in
 		case <-next:
 			m.Lock()
 			m.offset = i + 1
+			m.Unlock()
 			ingest(ctx, map[string]any(d.Message), map[string]any{"topic": "mock"}, timex.GetNow())
 			log.Debugf("%d: mock source %s is sending data %d:%v", timex.GetNowInMilli(), ctx.GetOpId(), i, d)
-			m.Unlock()
 		case <-ctx.Done():
 			log.Debugf("mock source open DONE")
 			return nil
@@ -88,8 +88,8 @@ func (m *MockSource) Subscribe(ctx api.StreamContext, ingest api.TupleIngest, in
 }
 
 func (m *MockSource) GetOffset() (interface{}, error) {
-	m.Lock()
-	defer m.Unlock()
+	m.RLock()
+	defer m.RUnlock()
 	return m.offset, nil
 }