From bf2c5def8d8d094f57a2b8c81f1c79566d885ae4 Mon Sep 17 00:00:00 2001
From: aoiasd <45024769+aoiasd@users.noreply.github.com>
Date: Thu, 25 Apr 2024 16:37:27 +0800
Subject: [PATCH] enhance: [Cherry-Pick] access log support get sdk type by
 user agent (#30760) (#32554)

Support get sdk type by user agent when we can't get sdk version by
connection in access log.

---------
pr: https://github.com/milvus-io/milvus/pull/30760

Signed-off-by: aoiasd <zhicheng.yue@zilliz.com>
---
 internal/proxy/accesslog/info.go      | 21 ++++++++++++++++++++-
 internal/proxy/accesslog/info_test.go | 20 +++++++++++++++++++-
 internal/proxy/accesslog/util.go      | 20 ++++++++++++++++++++
 internal/proxy/accesslog/util_test.go | 24 ++++++++++++++++++++++++
 pkg/util/constant.go                  |  4 +++-
 5 files changed, 86 insertions(+), 3 deletions(-)

diff --git a/internal/proxy/accesslog/info.go b/internal/proxy/accesslog/info.go
index 14275f9088..010a00f80a 100644
--- a/internal/proxy/accesslog/info.go
+++ b/internal/proxy/accesslog/info.go
@@ -32,6 +32,7 @@ import (
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
 	"github.com/milvus-io/milvus/internal/proxy/connection"
+	"github.com/milvus-io/milvus/pkg/util"
 	"github.com/milvus-io/milvus/pkg/util/merr"
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
 	"github.com/milvus-io/milvus/pkg/util/requestutil"
@@ -277,7 +278,25 @@ func getSdkVersion(i *GrpcAccessInfo) string {
 		return req.GetClientInfo().GetSdkType() + "-" + req.GetClientInfo().GetSdkVersion()
 	}
 
-	return unknownString
+	return getSdkVersionByUserAgent(i.ctx)
+}
+
+func getSdkVersionByUserAgent(ctx context.Context) string {
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		return unknownString
+	}
+	UserAgent, ok := md[util.HeaderUserAgent]
+	if !ok {
+		return unknownString
+	}
+
+	SdkType, ok := getSdkTypeByUserAgent(UserAgent)
+	if !ok {
+		return unknownString
+	}
+
+	return SdkType + "-" + unknownString
 }
 
 func getClusterPrefix(i *GrpcAccessInfo) string {
diff --git a/internal/proxy/accesslog/info_test.go b/internal/proxy/accesslog/info_test.go
index 4ed09ba4f8..f16a8c6c14 100644
--- a/internal/proxy/accesslog/info_test.go
+++ b/internal/proxy/accesslog/info_test.go
@@ -126,6 +126,24 @@ func (s *GrpcAccessInfoSuite) TestSdkInfo() {
 	result := s.info.Get("$sdk_version")
 	s.Equal(unknownString, result[0])
 
+	md := metadata.MD{}
+	ctx = metadata.NewIncomingContext(ctx, md)
+	s.info.ctx = ctx
+	result = s.info.Get("$sdk_version")
+	s.Equal(unknownString, result[0])
+
+	md = metadata.MD{util.HeaderUserAgent: []string{"invalid"}}
+	ctx = metadata.NewIncomingContext(ctx, md)
+	s.info.ctx = ctx
+	result = s.info.Get("$sdk_version")
+	s.Equal(unknownString, result[0])
+
+	md = metadata.MD{util.HeaderUserAgent: []string{"grpc-go.test"}}
+	ctx = metadata.NewIncomingContext(ctx, md)
+	s.info.ctx = ctx
+	result = s.info.Get("$sdk_version")
+	s.Equal("Golang"+"-"+unknownString, result[0])
+
 	s.info.req = &milvuspb.ConnectRequest{
 		ClientInfo: clientInfo,
 	}
@@ -133,7 +151,7 @@ func (s *GrpcAccessInfoSuite) TestSdkInfo() {
 	s.Equal(clientInfo.SdkType+"-"+clientInfo.SdkVersion, result[0])
 
 	identifier := 11111
-	md := metadata.MD{util.IdentifierKey: []string{fmt.Sprint(identifier)}}
+	md = metadata.MD{util.IdentifierKey: []string{fmt.Sprint(identifier)}}
 	ctx = metadata.NewIncomingContext(ctx, md)
 	connection.GetManager().Register(ctx, int64(identifier), clientInfo)
 
diff --git a/internal/proxy/accesslog/util.go b/internal/proxy/accesslog/util.go
index c70e7a5d94..f724bc8b21 100644
--- a/internal/proxy/accesslog/util.go
+++ b/internal/proxy/accesslog/util.go
@@ -86,3 +86,23 @@ func getCurUserFromContext(ctx context.Context) (string, error) {
 	username := secrets[0]
 	return username, nil
 }
+
+func getSdkTypeByUserAgent(userAgents []string) (string, bool) {
+	if len(userAgents) == 0 {
+		return "", false
+	}
+
+	userAgent := userAgents[0]
+	switch {
+	case strings.HasPrefix(userAgent, "grpc-node-js"):
+		return "nodejs", true
+	case strings.HasPrefix(userAgent, "grpc-python"):
+		return "Python", true
+	case strings.HasPrefix(userAgent, "grpc-go"):
+		return "Golang", true
+	case strings.HasPrefix(userAgent, "grpc-java"):
+		return "Java", true
+	default:
+		return "", false
+	}
+}
diff --git a/internal/proxy/accesslog/util_test.go b/internal/proxy/accesslog/util_test.go
index c07bc29e2c..5beebf0fb2 100644
--- a/internal/proxy/accesslog/util_test.go
+++ b/internal/proxy/accesslog/util_test.go
@@ -26,3 +26,27 @@ func TestJoin(t *testing.T) {
 	assert.Equal(t, "a/b", join("a", "b"))
 	assert.Equal(t, "a/b", join("a/", "b"))
 }
+
+func TestGetSdkTypeByUserAgent(t *testing.T) {
+	_, ok := getSdkTypeByUserAgent([]string{})
+	assert.False(t, ok)
+
+	sdk, ok := getSdkTypeByUserAgent([]string{"grpc-node-js.test"})
+	assert.True(t, ok)
+	assert.Equal(t, "nodejs", sdk)
+
+	sdk, ok = getSdkTypeByUserAgent([]string{"grpc-python.test"})
+	assert.True(t, ok)
+	assert.Equal(t, "Python", sdk)
+
+	sdk, ok = getSdkTypeByUserAgent([]string{"grpc-go.test"})
+	assert.True(t, ok)
+	assert.Equal(t, "Golang", sdk)
+
+	sdk, ok = getSdkTypeByUserAgent([]string{"grpc-java.test"})
+	assert.True(t, ok)
+	assert.Equal(t, "Java", sdk)
+
+	_, ok = getSdkTypeByUserAgent([]string{"invalid_type"})
+	assert.False(t, ok)
+}
diff --git a/pkg/util/constant.go b/pkg/util/constant.go
index 50a6ac3182..c6914d3a66 100644
--- a/pkg/util/constant.go
+++ b/pkg/util/constant.go
@@ -60,7 +60,9 @@ const (
 	AnyWord       = "*"
 
 	IdentifierKey = "identifier"
-	HeaderDBName  = "dbName"
+
+	HeaderUserAgent = "user-agent"
+	HeaderDBName    = "dbName"
 
 	RoleConfigPrivileges = "privileges"
 	RoleConfigObjectType = "object_type"