fix the data query grpc-gw

Author: LinuxSuRen

cmd/server.go

@@ -288,6 +288,7 @@ func (o *serverOption) runE(cmd *cobra.Command, args []string) (err error) {
 		}
 		server.RegisterRunnerServer(s, remoteServer)
 		server.RegisterMockServer(s, mockServerController)
+		server.RegisterDataServerServer(s, remoteServer.(server.DataServerServer))
 		serverLogger.Info("gRPC server listening at", "addr", lis.Addr())
 		s.Serve(lis)
 	}()
@@ -322,11 +323,14 @@ func (o *serverOption) runE(cmd *cobra.Command, args []string) (err error) {
 		}
 		err = errors.Join(
 			server.RegisterRunnerHandlerFromEndpoint(ctx, mux, gRPCServerAddr, []grpc.DialOption{grpc.WithTransportCredentials(creds)}),
-			server.RegisterMockHandlerFromEndpoint(ctx, mux, gRPCServerAddr, []grpc.DialOption{grpc.WithTransportCredentials(creds)}))
+			server.RegisterMockHandlerFromEndpoint(ctx, mux, gRPCServerAddr, []grpc.DialOption{grpc.WithTransportCredentials(creds)}),
+			server.RegisterDataServerHandlerFromEndpoint(ctx, mux, gRPCServerAddr, []grpc.DialOption{grpc.WithTransportCredentials(creds)}))
 	} else {
+		dialOption := []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())}
 		err = errors.Join(
-			server.RegisterRunnerHandlerFromEndpoint(ctx, mux, gRPCServerAddr, []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())}),
-			server.RegisterMockHandlerFromEndpoint(ctx, mux, gRPCServerAddr, []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())}))
+			server.RegisterRunnerHandlerFromEndpoint(ctx, mux, gRPCServerAddr, dialOption),
+			server.RegisterMockHandlerFromEndpoint(ctx, mux, gRPCServerAddr, dialOption),
+			server.RegisterDataServerHandlerFromEndpoint(ctx, mux, gRPCServerAddr, dialOption))
 	}
 
 	if err == nil {

pkg/server/remote_server.go

@@ -64,6 +64,7 @@ var (
 
 type server struct {
 	UnimplementedRunnerServer
+	UnimplementedDataServerServer
 	loader             testing.Writer
 	storeWriterFactory testing.StoreWriterFactory
 	configDir          string
@@ -1240,6 +1241,24 @@ func (s *server) PProf(ctx context.Context, in *PProfRequest) (reply *PProfData,
 	}
 	return
 }
+func (s *server) Query(ctx context.Context, query *DataQuery) (result *DataQueryResult, err error) {
+	loader := s.getLoader(ctx)
+	defer loader.Close()
+	var dataResult map[string]string
+	if dataResult, err = loader.Query(map[string]string{
+		"sql": query.Sql,
+	}); err == nil {
+		result = &DataQueryResult{
+			Data: []*Pair{
+				{
+					Key:   "data",
+					Value: fmt.Sprintf("%v", dataResult),
+				},
+			},
+		}
+	}
+	return
+}
 
 // implement the mock server
 

pkg/testing/loader.go

@@ -26,6 +26,8 @@ type Loader interface {
 
 	Verify() (readOnly bool, err error)
 	PProf(name string) []byte
+
+	Query(query map[string]string) (result map[string]string, err error)
 }
 
 type Writer interface {

pkg/testing/loader_file.go

@@ -514,3 +514,10 @@ func (l *fileLoader) PProf(string) []byte {
 func (l *fileLoader) Close() {
 	// not support
 }
+
+func (l *fileLoader) Query(query map[string]string) (result map[string]string, err error) {
+	result = map[string]string{
+		"message": "not support",
+	}
+	return
+}

pkg/testing/loader_non.go

@@ -149,3 +149,7 @@ func (l *nonLoader) PProf(string) []byte {
 func (l *nonLoader) Close() {
 	// not support
 }
+
+func (l *nonLoader) Query(query map[string]string) (result map[string]string, err error) {
+	return
+}

pkg/testing/local/local_secret_service.go

@@ -121,6 +121,11 @@ func (s *localSecretService) UpdateSecret(ctx context.Context, in *server.Secret
 	return
 }
 
+func (s *localSecretService) Query(query map[string]string) (result map[string]string, err error) {
+	result = make(map[string]string)
+	return
+}
+
 func (s *localSecretService) getDataFilePath() string {
 	return filepath.Join(s.dataDir, "secret.yaml")
 }

pkg/testing/remote/grpc_store.go

@@ -306,6 +306,16 @@ func (g *gRPCLoader) PProf(name string) []byte {
 	return data.Data
 }
 
+func (g *gRPCLoader) Query(query map[string]string) (result map[string]string, err error) {
+	var dataResult *server.DataQueryResult
+	if dataResult, err = g.client.Query(context.Background(), &server.DataQuery{
+		Sql: query["sql"],
+	}); err == nil {
+		result = pairToMap(dataResult.Data)
+	}
+	return
+}
+
 func (g *gRPCLoader) Close() {
 	if g.conn != nil {
 		g.conn.Close()

pkg/testing/remote/loader.pb.go

@@ -32,6 +32,8 @@ service Loader {
     rpc GetVersion(server.Empty) returns (server.Version) {}
     rpc Verify(server.Empty) returns (server.ExtensionStatus) {}
     rpc PProf(server.PProfRequest) returns (server.PProfData) {}
+
+    rpc Query(server.DataQuery) returns (server.DataQueryResult) {}
 }
 
 message TestSuites {