Repository: lf-edge/ekuiper
Pull request #2941 (2024-06-20 06:27:26)

feat: support sql reconnection

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/lf-edge/ekuiper/pull/2941?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 `59.37500%` with `13 lines` in your changes missing coverage. Please review.
> Please [upload](https://docs.codecov.com/docs/codecov-uploader) report for BASE (`master-v2@ee10892`). [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/2941?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/io/sql/client/client.go](https://app.codecov.io/gh/lf-edge/ekuiper/pull/2941?src=pr&el=tree&filepath=internal%2Fio%2Fsql%2Fclient%2Fclient.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lf-edge#diff-aW50ZXJuYWwvaW8vc3FsL2NsaWVudC9jbGllbnQuZ28=) | 57.14% | [8 Missing and 1 partial :warning: ](https://app.codecov.io/gh/lf-edge/ekuiper/pull/2941?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lf-edge) |
| [internal/io/sql/source.go](https://app.codecov.io/gh/lf-edge/ekuiper/pull/2941?src=pr&el=tree&filepath=internal%2Fio%2Fsql%2Fsource.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lf-edge#diff-aW50ZXJuYWwvaW8vc3FsL3NvdXJjZS5nbw==) | 63.64% | [3 Missing and 1 partial :warning: ](https://app.codecov.io/gh/lf-edge/ekuiper/pull/2941?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    #2941   +/-   ##
============================================
  Coverage             ?   69.70%           
============================================
  Files                ?      312           
  Lines                ?    29764           
  Branches             ?        0           
============================================
  Hits                 ?    20746           
  Misses               ?     7223           
  Partials             ?     1795           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/lf-edge/ekuiper/pull/2941?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/client/client.go b/internal/io/sql/client/client.go
index 04ea94cf..b06569f1 100644
--- a/internal/io/sql/client/client.go
+++ b/internal/io/sql/client/client.go
@@ -17,6 +17,7 @@ package client
 import (
 	"database/sql"
 	"fmt"
+	"sync"
 
 	"github.com/lf-edge/ekuiper/contract/v2/api"
 	"github.com/lf-edge/ekuiper/v2/internal/conf"
@@ -25,15 +26,37 @@ import (
 )
 
 type SQLConnection struct {
-	url string
-	db  *sql.DB
+	sync.RWMutex
+	url    string
+	db     *sql.DB
+	closed bool
+}
+
+func (s *SQLConnection) Reconnect() error {
+	s.Lock()
+	defer s.Unlock()
+	if err := s.db.Ping(); err == nil {
+		return nil
+	}
+	oldDB := s.db
+	oldDB.Close()
+	db, err := openDB(s.url)
+	if err != nil {
+		return fmt.Errorf("reconnect sql err:%v, supported drivers:%v", err, driver.GetSupportedDrivers())
+	}
+	s.db = db
+	return nil
 }
 
 func (s *SQLConnection) GetDB() *sql.DB {
+	s.RLock()
+	defer s.RUnlock()
 	return s.db
 }
 
 func (s *SQLConnection) Ping(ctx api.StreamContext) error {
+	s.RLock()
+	defer s.RUnlock()
 	return s.db.Ping()
 }
 
@@ -42,8 +65,14 @@ func (s *SQLConnection) DetachSub(ctx api.StreamContext, props map[string]any) {
 }
 
 func (s *SQLConnection) Close(ctx api.StreamContext) error {
+	s.Lock()
+	defer s.Unlock()
+	if s.closed {
+		return nil
+	}
 	conf.Log.Infof("close db with url:%v", s.url)
 	s.db.Close()
+	s.closed = true
 	return nil
 }
 
diff --git a/internal/io/sql/source.go b/internal/io/sql/source.go
index cb3d4002..54565ab4 100644
--- a/internal/io/sql/source.go
+++ b/internal/io/sql/source.go
@@ -32,10 +32,11 @@ import (
 )
 
 type SQLSourceConnector struct {
-	conf  *SQLConf
-	Query sqlgen.SqlQueryGenerator
-	conn  *client.SQLConnection
-	props map[string]any
+	conf          *SQLConf
+	Query         sqlgen.SqlQueryGenerator
+	conn          *client.SQLConnection
+	props         map[string]any
+	needReconnect bool
 }
 
 type SQLConf struct {
@@ -105,6 +106,14 @@ func (s *SQLSourceConnector) Pull(ctx api.StreamContext, recvTime time.Time, ing
 
 func (s *SQLSourceConnector) queryData(ctx api.StreamContext, rcvTime time.Time, ingest api.TupleIngest, ingestError api.ErrorIngest) {
 	logger := ctx.GetLogger()
+	if s.needReconnect {
+		err := s.conn.Reconnect()
+		if err != nil {
+			logger.Errorf("reconnect db error %v", err)
+			ingestError(ctx, err)
+			return
+		}
+	}
 	query, err := s.Query.SqlQueryStatement()
 	failpoint.Inject("StatementErr", func() {
 		err = errors.New("StatementErr")
@@ -120,8 +129,13 @@ func (s *SQLSourceConnector) queryData(ctx api.StreamContext, rcvTime time.Time,
 		err = errors.New("QueryErr")
 	})
 	if err != nil {
+		logger.Errorf("query sql error %v", err)
+		s.needReconnect = true
 		ingestError(ctx, err)
 		return
+	} else if s.needReconnect {
+		logger.Infof("reconnect sql success")
+		s.needReconnect = false
 	}
 	cols, _ := rows.Columns()
 	types, err := rows.ColumnTypes()


Test patch:
diff --git a/internal/io/sql/source_test.go b/internal/io/sql/source_test.go
index 6d0c9e98..dd53a658 100644
--- a/internal/io/sql/source_test.go
+++ b/internal/io/sql/source_test.go
@@ -226,3 +226,29 @@ func TestSQLSourceRewind(t *testing.T) {
 	}...).GetStore()
 	require.Equal(t, expectState, state)
 }
+
+func TestSQLReconnect(t *testing.T) {
+	connection.InitConnectionManager4Test()
+	ctx := mockContext.NewMockContext("1", "2")
+	props := map[string]interface{}{
+		"interval": "1s",
+		"dburl":    fmt.Sprintf("mysql://root:@%v:%v/test", address, port),
+		"templateSqlQueryCfg": map[string]interface{}{
+			"templateSql": "select a,b from t",
+		},
+	}
+	sqlSource := GetSource()
+	require.NoError(t, sqlSource.Provision(ctx, props))
+	require.NoError(t, sqlSource.Connect(ctx))
+	sqlConnector, ok := sqlSource.(*SQLSourceConnector)
+	require.True(t, ok)
+	sqlConnector.queryData(ctx, time.Now(), func(ctx api.StreamContext, data any, meta map[string]any, ts time.Time) {}, func(ctx api.StreamContext, err error) {})
+	require.True(t, sqlConnector.needReconnect)
+
+	// start server then reconnect
+	s, err := testx.SetupEmbeddedMysqlServer(address, port)
+	require.NoError(t, err)
+	defer s.Close()
+	sqlConnector.queryData(ctx, time.Now(), func(ctx api.StreamContext, data any, meta map[string]any, ts time.Time) {}, func(ctx api.StreamContext, err error) {})
+	require.False(t, sqlConnector.needReconnect)
+}