✨ Add Search Query Interceptor

Signed-off-by: Rintaro Okamura <[email protected]>

examples/config/config.lua

@@ -55,7 +55,13 @@ server = {
     end
   end,
 
+  -- server-side Search Query Interceptor
+  search_query_interceptor = function (request)
+    -- print(string.format("Searching top %d neighbors", request.Config.Num))
+  end,
+
+  -- server-side Insert Data Interceptor
   insert_data_interceptor = function (request)
-    print(string.format("Inserting ID: %s", request.Vector.Id))
+    -- print(string.format("Inserting ID: %s", request.Vector.Id))
   end,
 }

pkg/alvd/cli/server/cli.go

@@ -32,7 +32,8 @@ type Opts struct {
 	CreateIndexThreshold uint
 
 	SearchResultInterceptor *lua.LFunction
-	InsertDataInterceptor *lua.LFunction
+	SearchQueryInterceptor  *lua.LFunction
+	InsertDataInterceptor   *lua.LFunction
 }
 
 var Flags = []cli.Flag{
@@ -155,6 +156,7 @@ func ToConfig(opts *Opts, agentOpts *agent.Opts) (*config.Config, error) {
 		config.WithCheckIndexInterval(opts.CheckIndexInterval),
 		config.WithCreateIndexThreshold(opts.CreateIndexThreshold),
 		config.WithSearchResultInterceptor(opts.SearchResultInterceptor),
+		config.WithSearchQueryInterceptor(opts.SearchQueryInterceptor),
 		config.WithInsertDataInterceptor(opts.InsertDataInterceptor),
 	)
 	if err != nil {

pkg/alvd/extension/lua/lua.go

@@ -103,6 +103,39 @@ func NewSearchResultInterceptorFn(sri *LFunction) SearchResultInterceptor {
 	}
 }
 
+type SearchQueryInterceptor = func(*payload.Search_Request) (
+	*payload.Search_Request,
+	error,
+)
+
+func NewSearchQueryInterceptorFn(sqi *LFunction) SearchQueryInterceptor {
+	return func(origin *payload.Search_Request) (
+		req *payload.Search_Request,
+		err error,
+	) {
+		state := lua.NewState()
+		defer state.Close()
+
+		libs.Preload(state)
+
+		req = origin
+
+		err = state.CallByParam(
+			lua.P{
+				Fn:      sqi,
+				NRet:    0,
+				Protect: true,
+			},
+			luar.New(state, req),
+		)
+		if err != nil {
+			return origin, err
+		}
+
+		return req, nil
+	}
+}
+
 type InsertDataInterceptor = func(*payload.Insert_Request) (
 	*payload.Insert_Request,
 	error,

pkg/alvd/server/config/config.go

@@ -21,6 +21,7 @@ type Config struct {
 	CreateIndexThreshold int
 
 	SearchResultInterceptor *lua.LFunction
+	SearchQueryInterceptor *lua.LFunction
 	InsertDataInterceptor   *lua.LFunction
 }
 

pkg/alvd/server/config/option.go

@@ -90,6 +90,14 @@ func WithSearchResultInterceptor(sri *lua.LFunction) OptionFunc {
 	}
 }
 
+func WithSearchQueryInterceptor(sqi *lua.LFunction) OptionFunc {
+	return func(c *Config) error {
+		c.SearchQueryInterceptor = sqi
+
+		return nil
+	}
+}
+
 func WithInsertDataInterceptor(idi *lua.LFunction) OptionFunc {
 	return func(c *Config) error {
 		c.InsertDataInterceptor = idi

pkg/alvd/server/daemon/daemon.go

@@ -65,6 +65,12 @@ func New(cfg *config.Config) (Daemon, error) {
 		)
 	}
 
+	if cfg.SearchQueryInterceptor != nil {
+		h.RegisterSearchQueryInterceptor(
+			lua.NewSearchQueryInterceptorFn(cfg.SearchQueryInterceptor),
+		)
+	}
+
 	if cfg.InsertDataInterceptor != nil {
 		h.RegisterInsertDataInterceptor(
 			lua.NewInsertDataInterceptorFn(cfg.InsertDataInterceptor),

pkg/alvd/server/service/gateway/handler/handler.go

@@ -28,12 +28,14 @@ type server struct {
 	numReplica int
 
 	searchResultInterceptor lua.SearchResultInterceptor
+	searchQueryInterceptor  lua.SearchQueryInterceptor
 	insertDataInterceptor   lua.InsertDataInterceptor
 }
 
 type Server interface {
 	vald.Server
 	RegisterSearchResultInterceptor(sri lua.SearchResultInterceptor)
+	RegisterSearchQueryInterceptor(sri lua.SearchQueryInterceptor)
 	RegisterInsertDataInterceptor(idi lua.InsertDataInterceptor)
 }
 
@@ -48,6 +50,10 @@ func (s *server) RegisterSearchResultInterceptor(sri lua.SearchResultInterceptor
 	s.searchResultInterceptor = sri
 }
 
+func (s *server) RegisterSearchQueryInterceptor(sqi lua.SearchQueryInterceptor) {
+	s.searchQueryInterceptor = sqi
+}
+
 func (s *server) RegisterInsertDataInterceptor(idi lua.InsertDataInterceptor) {
 	s.insertDataInterceptor = idi
 }
@@ -88,6 +94,13 @@ func (s *server) Search(
 	ctx context.Context,
 	req *payload.Search_Request,
 ) (res *payload.Search_Response, err error) {
+	if s.searchQueryInterceptor != nil {
+		req, err = s.searchQueryInterceptor(req)
+		if err != nil {
+			log.Warnf("an error occurred while using search query interceptor: %s", err)
+		}
+	}
+
 	cfg := req.GetConfig()
 	num := int(cfg.GetNum())