admin/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2025-12-08 01:58:34 +08:00
Code Issues Packages Projects Releases Wiki Activity

enhance: The Restful API supports a single page web application. (#37719)

Browse Source 
issue: #36621

Signed-off-by: jaime <yun.zhang@zilliz.com>
This commit is contained in:
jaime 2024-11-16 11:04:31 +08:00 committed by GitHub
parent 3df2c92bb0
commit ead1e7f68c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 167 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

75
internal/http/server.go
View File

@ -24,6 +24,7 @@ import (
"os" "os"
"runtime" "runtime"
"strconv" "strconv"
"strings"
"time" "time"
"go.uber.org/zap" "go.uber.org/zap"
@ -101,6 +102,8 @@ func registerDefaults() {
Path: StaticPath, Path: StaticPath,
Handler: GetStaticHandler(), Handler: GetStaticHandler(),
}) })
RegisterWebUIHandler()
} }
func RegisterStopComponent(triggerComponentStop func(role string) error) { func RegisterStopComponent(triggerComponentStop func(role string) error) {
@ -141,12 +144,78 @@ func RegisterCheckComponentReady(checkActive func(role string) error) {
w.Write([]byte(`{"msg": "OK"}`)) w.Write([]byte(`{"msg": "OK"}`))
}, },
}) })
Register(&Handler{ }
Path: RouteWebUI,
Handler: http.FileServer(http.FS(staticFiles)), func RegisterWebUIHandler() {
httpFS := http.FS(staticFiles)
fileServer := http.FileServer(httpFS)
serveIndex := serveFile(RouteWebUI+"index.html", httpFS)
http.Handle(RouteWebUI, handleNotFound(fileServer, serveIndex))
}
type responseInterceptor struct {
http.ResponseWriter
is404 bool
}
func (ri *responseInterceptor) WriteHeader(status int) {
if status == http.StatusNotFound {
ri.is404 = true
return
}
ri.ResponseWriter.WriteHeader(status)
}
func (ri *responseInterceptor) Write(p []byte) (int, error) {
if ri.is404 {
return len(p), nil // Pretend the data was written for a 404
}
return ri.ResponseWriter.Write(p)
}
// handleNotFound attempts to serve a fallback handler (on404) if the main handler returns a 404 status.
func handleNotFound(handler, on404 http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
ri := &responseInterceptor{ResponseWriter: w}
handler.ServeHTTP(ri, r)
if ri.is404 {
on404.ServeHTTP(w, r)
}
}) })
} }
// serveFile serves the specified file content (like "index.html") for HTML requests.
func serveFile(filename string, fs http.FileSystem) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
if !acceptsHTML(r) {
http.NotFound(w, r)
return
}
file, err := fs.Open(filename)
if err != nil {
http.NotFound(w, r)
return
}
defer file.Close()
fi, err := file.Stat()
if err != nil {
http.NotFound(w, r)
return
}
w.Header().Set("Content-Type", "text/html; charset=utf-8")
http.ServeContent(w, r, fi.Name(), fi.ModTime(), file)
}
}
// acceptsHTML checks if the request header specifies that HTML is acceptable.
func acceptsHTML(r *http.Request) bool {
return strings.Contains(r.Header.Get("Accept"), "text/html")
}
func Register(h *Handler) { func Register(h *Handler) {
if metricsServer == nil { if metricsServer == nil {
if paramtable.Get().HTTPCfg.EnablePprof.GetAsBool() { if paramtable.Get().HTTPCfg.EnablePprof.GetAsBool() {

95
internal/http/server_test.go
View File

@ -24,11 +24,13 @@ import (
"io" "io"
"net" "net"
"net/http" "net/http"
"net/http/httptest"
"os" "os"
"strings" "strings"
"testing" "testing"
"time" "time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/suite" "github.com/stretchr/testify/suite"
"go.uber.org/zap" "go.uber.org/zap"
@ -239,3 +241,96 @@ func (m *MockIndicator) Health(ctx context.Context) commonpb.StateCode {
func (m *MockIndicator) GetName() string { func (m *MockIndicator) GetName() string {
return m.name return m.name
} }
func TestRegisterWebUIHandler(t *testing.T) {
// Initialize the HTTP server
func() {
defer func() {
if err := recover(); err != nil {
fmt.Println("May the handler has been registered!", err)
}
}()
RegisterWebUIHandler()
}()
// Create a test server
ts := httptest.NewServer(http.DefaultServeMux)
defer ts.Close()
// Test cases
tests := []struct {
url string
expectedCode int
expectedBody string
}{
{"/webui/", http.StatusOK, "<!doctype html>"},
{"/webui/index.html", http.StatusOK, "<!doctype html>"},
{"/webui/unknown", http.StatusOK, "<!doctype html>"},
}
for _, tt := range tests {
t.Run(tt.url, func(t *testing.T) {
req, err := http.NewRequest("GET", ts.URL+tt.url, nil)
assert.NoError(t, err)
req.Header.Set("Accept", "text/html")
resp, err := ts.Client().Do(req)
assert.NoError(t, err)
defer resp.Body.Close()
assert.Equal(t, tt.expectedCode, resp.StatusCode)
body := make([]byte, len(tt.expectedBody))
_, err = resp.Body.Read(body)
assert.NoError(t, err)
assert.Contains(t, strings.ToLower(string(body)), tt.expectedBody)
})
}
}
func TestHandleNotFound(t *testing.T) {
mainHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
http.NotFound(w, r)
})
fallbackHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
w.Write([]byte("Fallback"))
})
handler := handleNotFound(mainHandler, fallbackHandler)
req := httptest.NewRequest("GET", "/", nil)
w := httptest.NewRecorder()
handler.ServeHTTP(w, req)
resp := w.Result()
body := make([]byte, 8)
resp.Body.Read(body)
assert.Equal(t, http.StatusOK, resp.StatusCode)
assert.Equal(t, "Fallback", string(body))
}
func TestServeFile(t *testing.T) {
fs := http.FS(staticFiles)
handler := serveFile("unknown", fs)
// No Accept in http header
{
req := httptest.NewRequest("GET", "/", nil)
w := httptest.NewRecorder()
handler.ServeHTTP(w, req)
resp := w.Result()
assert.Equal(t, http.StatusNotFound, resp.StatusCode)
}
// unknown request file
{
req := httptest.NewRequest("GET", "/", nil)
req.Header.Set("Accept", "text/html")
w := httptest.NewRecorder()
handler.ServeHTTP(w, req)
resp := w.Result()
assert.Equal(t, http.StatusNotFound, resp.StatusCode)
}
}