feat: support sql lookup template sql (#3398)

Signed-off-by: Song Gao <[email protected]>
lf-edge · Nov 26, 2024 · e795d19 · e795d19
1 parent 6674454
commit e795d19
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 5 deletions.
diff --git a/extensions/impl/sql/lookupSource.go b/extensions/impl/sql/lookupSource.go
@@ -109,7 +109,20 @@ func (s *SqlLookupSource) Lookup(ctx api.StreamContext, fields []string, keys []
 			return nil, err
 		}
 	}
-	query := s.gen.buildQuery(fields, keys, values)
+	var query string
+	if s.conf.TemplateSqlQueryCfg == nil {
+		query = s.gen.buildQuery(fields, keys, values)
+	} else {
+		mapValue := make(map[string]any)
+		for index, key := range keys {
+			mapValue[key] = values[index]
+		}
+		sqlQuery, err := ctx.ParseTemplate(s.conf.TemplateSqlQueryCfg.TemplateSql, mapValue)
+		if err != nil {
+			return nil, err
+		}
+		query = sqlQuery
+	}
 	ctx.GetLogger().Debugf("Query is %s", query)
 	rows, err := s.conn.GetDB().Query(query)
 	failpoint.Inject("dbErr", func() {

diff --git a/extensions/impl/sql/lookupSource_test.go b/extensions/impl/sql/lookupSource_test.go
@@ -67,6 +67,23 @@ func TestSQLLookupSource(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, []map[string]any{{"a": int64(1), "b": int64(1)}}, got)
 	ls.Close(ctx)
+
+	props = map[string]interface{}{
+		"dburl":      fmt.Sprintf("mysql://root:@%v:%v/test", address, port),
+		"datasource": "t",
+		"templateSqlQueryCfg": map[string]interface{}{
+			"templateSql": "select * from t where b = {{.bid}}",
+		},
+	}
+	ls = &SqlLookupSource{}
+	require.NoError(t, ls.Provision(ctx, props))
+	require.NoError(t, ls.Connect(ctx, func(status string, message string) {
+		// do nothing
+	}))
+	got, err = ls.Lookup(ctx, []string{"a", "b"}, []string{"bid"}, []any{1})
+	require.NoError(t, err)
+	require.Equal(t, []map[string]any{{"a": int64(1), "b": int64(1)}}, got)
+	ls.Close(ctx)
 }
 
 func TestSQLLookupSourceProvisionErr(t *testing.T) {

diff --git a/extensions/impl/sql/source.go b/extensions/impl/sql/source.go
@@ -55,10 +55,11 @@ func (s *SQLSourceConnector) Ping(ctx api.StreamContext, m map[string]any) error
 }
 
 type SQLConf struct {
-	Interval   cast.DurationConf `json:"interval"`
-	DBUrl      string            `json:"dburl"`
-	URL        string            `json:"url,omitempty"`
-	Datasource string            `json:"datasource"`
+	Interval            cast.DurationConf           `json:"interval"`
+	DBUrl               string                      `json:"dburl"`
+	URL                 string                      `json:"url,omitempty"`
+	Datasource          string                      `json:"datasource"`
+	TemplateSqlQueryCfg *sqlgen.TemplateSqlQueryCfg `json:"templateSqlQueryCfg"`
 }
 
 func init() {