mirror of
https://github.com/zeromicro/go-zero.git
synced 2026-05-11 08:50:00 +08:00
Compare commits
140 Commits
v1.7.4
...
tools/goct
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
aeceb3cfbe | ||
|
|
15ea07aad1 | ||
|
|
98bebbc74f | ||
|
|
eafd11d949 | ||
|
|
b251ce346e | ||
|
|
812140ba36 | ||
|
|
44735e949c | ||
|
|
bf313c3c56 | ||
|
|
94e7753262 | ||
|
|
9c478626d2 | ||
|
|
801c283478 | ||
|
|
2a54faf997 | ||
|
|
ecd98f3653 | ||
|
|
61641581eb | ||
|
|
6f2730d5ae | ||
|
|
0eff777b62 | ||
|
|
cafbf535f7 | ||
|
|
6edfce63e3 | ||
|
|
cdb0098b18 | ||
|
|
620c7f9693 | ||
|
|
dba444a382 | ||
|
|
b24fb3ebf7 | ||
|
|
967f0926eb | ||
|
|
e68c683df9 | ||
|
|
247985a065 | ||
|
|
80573af0d8 | ||
|
|
c0394b631a | ||
|
|
68d1aba377 | ||
|
|
3315e60272 | ||
|
|
327ef73700 | ||
|
|
eb11521655 | ||
|
|
4c37545e55 | ||
|
|
2f47c1fba4 | ||
|
|
16d54d0ace | ||
|
|
9925bcbf99 | ||
|
|
38a5ecb796 | ||
|
|
af78fc7c5f | ||
|
|
790302b486 | ||
|
|
6a0672b801 | ||
|
|
560c61612c | ||
|
|
6a988dc4a9 | ||
|
|
15842c3c7a | ||
|
|
f2914a74df | ||
|
|
f113d512e8 | ||
|
|
7a4818da59 | ||
|
|
48d0709ca6 | ||
|
|
f747585518 | ||
|
|
507ff96546 | ||
|
|
651eabb4c6 | ||
|
|
e6b4372056 | ||
|
|
24073969a1 | ||
|
|
ca797ed22c | ||
|
|
e347d3f8f8 | ||
|
|
396393b336 | ||
|
|
1f0531b254 | ||
|
|
77fb271a06 | ||
|
|
af7cf79963 | ||
|
|
7926d396d7 | ||
|
|
080cd3df84 | ||
|
|
c4e1a6a2d8 | ||
|
|
4e71e95e44 | ||
|
|
84db9bcd15 | ||
|
|
b28f79ac11 | ||
|
|
e134e77b2b | ||
|
|
f669d84ce8 | ||
|
|
9213b8ac27 | ||
|
|
ae09d0e56d | ||
|
|
0bc4206d08 | ||
|
|
39ce17bfd2 | ||
|
|
d415ba39e2 | ||
|
|
c71829c8de | ||
|
|
a32f6d7642 | ||
|
|
64e8c94198 | ||
|
|
7d05a4bc93 | ||
|
|
44504e8df7 | ||
|
|
114311e51b | ||
|
|
4307ce45fc | ||
|
|
37b54d1fc7 | ||
|
|
00e0db5def | ||
|
|
cbcacf31c1 | ||
|
|
238c92aaa9 | ||
|
|
520d2a2075 | ||
|
|
1023800b02 | ||
|
|
030c859171 | ||
|
|
e6d1b47a43 | ||
|
|
6138f85470 | ||
|
|
bf883101d7 | ||
|
|
33011c7ed1 | ||
|
|
17d98f69e0 | ||
|
|
b650c8c425 | ||
|
|
3d931d7030 | ||
|
|
68da9ed51a | ||
|
|
b25c45b352 | ||
|
|
f05234a967 | ||
|
|
12071d17b4 | ||
|
|
11c47d23df | ||
|
|
024f285f86 | ||
|
|
fa4674611a | ||
|
|
730c3c5246 | ||
|
|
2c9310ac3a | ||
|
|
74ba0bcd50 | ||
|
|
5f4190b6c6 | ||
|
|
e1787b4ccb | ||
|
|
4ac8b492ef | ||
|
|
cdd068575c | ||
|
|
e89e2d8a75 | ||
|
|
acd2b94bd9 | ||
|
|
6a0c8047f4 | ||
|
|
cfe03ea9e1 | ||
|
|
48d21ef8ad | ||
|
|
28a001c5f9 | ||
|
|
22a41cacc7 | ||
|
|
fcc246933c | ||
|
|
5c3679ffe7 | ||
|
|
eaa01ccb9f | ||
|
|
b8206fb46a | ||
|
|
1c3876810e | ||
|
|
1d9159ea39 | ||
|
|
2159d112c3 | ||
|
|
f57874a51f | ||
|
|
8625864d43 | ||
|
|
8f9ba3ec11 | ||
|
|
a1d9bc08f0 | ||
|
|
b9d7f1cc77 | ||
|
|
6700910f64 | ||
|
|
9c4ed394a7 | ||
|
|
fd07a9c6e4 | ||
|
|
b8c239630c | ||
|
|
672ea55736 | ||
|
|
f7097866bf | ||
|
|
796b2bd1b0 | ||
|
|
e1e5fb2071 | ||
|
|
89ecb50005 | ||
|
|
dbed1ea042 | ||
|
|
ad291daf78 | ||
|
|
13746a3706 | ||
|
|
f03b13f632 | ||
|
|
f6f64b1286 | ||
|
|
300a415f5d | ||
|
|
c5de546f8a |
18
.github/workflows/issue-translator.yml
vendored
18
.github/workflows/issue-translator.yml
vendored
@@ -1,18 +0,0 @@
|
||||
name: 'issue-translator'
|
||||
on:
|
||||
issue_comment:
|
||||
types: [created]
|
||||
issues:
|
||||
types: [opened]
|
||||
|
||||
jobs:
|
||||
build:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: usthe/issues-translate-action@v2.7
|
||||
with:
|
||||
IS_MODIFY_TITLE: true
|
||||
# not require, default false, . Decide whether to modify the issue title
|
||||
# if true, the robot account @Issues-translate-bot must have modification permissions, invite @Issues-translate-bot to your project or use your custom bot.
|
||||
CUSTOM_BOT_NOTE: Bot detected the issue body's language is not English, translate it automatically. 👯👭🏻🧑🤝🧑👫🧑🏿🤝🧑🏻👩🏾🤝👨🏿👬🏿
|
||||
# not require. Customize the translation robot prefix message.
|
||||
4
.github/workflows/release.yaml
vendored
4
.github/workflows/release.yaml
vendored
@@ -22,7 +22,7 @@ jobs:
|
||||
github_token: ${{ secrets.GITHUB_TOKEN }}
|
||||
goos: ${{ matrix.goos }}
|
||||
goarch: ${{ matrix.goarch }}
|
||||
goversion: "https://dl.google.com/go/go1.20.14.linux-amd64.tar.gz"
|
||||
goversion: "https://dl.google.com/go/go1.21.13.linux-amd64.tar.gz"
|
||||
project_path: "tools/goctl"
|
||||
binary_name: "goctl"
|
||||
extra_files: tools/goctl/readme.md tools/goctl/readme-cn.md
|
||||
extra_files: tools/goctl/readme.md tools/goctl/readme-cn.md
|
||||
|
||||
2
.github/workflows/reviewdog.yml
vendored
2
.github/workflows/reviewdog.yml
vendored
@@ -14,6 +14,6 @@ jobs:
|
||||
# Report all results.
|
||||
filter_mode: nofilter
|
||||
# Exit with 1 when it find at least one finding.
|
||||
fail_on_error: true
|
||||
fail_level: any
|
||||
# Set staticcheck flags
|
||||
staticcheck_flags: -checks=inherit,-SA1019,-SA1029,-SA5008
|
||||
|
||||
42
.github/workflows/version-check.yml
vendored
Normal file
42
.github/workflows/version-check.yml
vendored
Normal file
@@ -0,0 +1,42 @@
|
||||
name: Release Version Check
|
||||
|
||||
on:
|
||||
push:
|
||||
tags:
|
||||
- 'tools/goctl/v*'
|
||||
workflow_dispatch:
|
||||
|
||||
jobs:
|
||||
version-check:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Go
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: '1.21'
|
||||
|
||||
- name: Extract tag version
|
||||
id: get_version
|
||||
run: |
|
||||
# Extract version from tools/goctl/v* format
|
||||
VERSION="${GITHUB_REF#refs/tags/tools/goctl/v}"
|
||||
echo "VERSION=$VERSION" >> $GITHUB_ENV
|
||||
echo "Extracted version: $VERSION"
|
||||
|
||||
- name: Check version in goctl source code
|
||||
run: |
|
||||
# Change to goctl directory
|
||||
cd tools/goctl
|
||||
|
||||
# Check version in BuildVersion constant
|
||||
VERSION_IN_CODE=$(grep -r "const BuildVersion =" . | grep -o '".*"' | tr -d '"')
|
||||
echo "Version in code: $VERSION_IN_CODE"
|
||||
echo "Expected version: $VERSION"
|
||||
|
||||
if [ "$VERSION_IN_CODE" != "$VERSION" ]; then
|
||||
echo "Version mismatch: Version in code ($VERSION_IN_CODE) doesn't match tag version ($VERSION)"
|
||||
exit 1
|
||||
fi
|
||||
echo "✅ Version check passed!"
|
||||
@@ -11,9 +11,9 @@ import (
|
||||
func TestNopBreaker(t *testing.T) {
|
||||
b := NopBreaker()
|
||||
assert.Equal(t, nopBreakerName, b.Name())
|
||||
p, err := b.Allow()
|
||||
_, err := b.Allow()
|
||||
assert.Nil(t, err)
|
||||
p, err = b.AllowCtx(context.Background())
|
||||
p, err := b.AllowCtx(context.Background())
|
||||
assert.Nil(t, err)
|
||||
p.Accept()
|
||||
for i := 0; i < 1000; i++ {
|
||||
|
||||
@@ -62,7 +62,11 @@ func Load(file string, v any, opts ...Option) error {
|
||||
return loader([]byte(os.ExpandEnv(string(content))), v)
|
||||
}
|
||||
|
||||
return loader(content, v)
|
||||
if err = loader(content, v); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return validate(v)
|
||||
}
|
||||
|
||||
// LoadConfig loads config into v from file, .json, .yaml and .yml are acceptable.
|
||||
@@ -85,7 +89,12 @@ func LoadFromJsonBytes(content []byte, v any) error {
|
||||
|
||||
lowerCaseKeyMap := toLowerCaseKeyMap(m, info)
|
||||
|
||||
return mapping.UnmarshalJsonMap(lowerCaseKeyMap, v, mapping.WithCanonicalKeyFunc(toLowerCase))
|
||||
if err = mapping.UnmarshalJsonMap(lowerCaseKeyMap, v,
|
||||
mapping.WithCanonicalKeyFunc(toLowerCase)); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return validate(v)
|
||||
}
|
||||
|
||||
// LoadConfigFromJsonBytes loads config into v from content json bytes.
|
||||
@@ -192,7 +201,7 @@ func buildFieldsInfo(tp reflect.Type, fullName string) (*fieldInfo, error) {
|
||||
case reflect.Array, reflect.Slice, reflect.Map:
|
||||
return buildFieldsInfo(mapping.Deref(tp.Elem()), fullName)
|
||||
case reflect.Chan, reflect.Func:
|
||||
return nil, fmt.Errorf("unsupported type: %s", tp.Kind())
|
||||
return nil, fmt.Errorf("unsupported type: %s, fullName: %s", tp.Kind(), fullName)
|
||||
default:
|
||||
return &fieldInfo{
|
||||
children: make(map[string]*fieldInfo),
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package conf
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"os"
|
||||
"reflect"
|
||||
"testing"
|
||||
@@ -40,9 +41,8 @@ func TestConfigJson(t *testing.T) {
|
||||
for _, test := range tests {
|
||||
test := test
|
||||
t.Run(test, func(t *testing.T) {
|
||||
tmpfile, err := createTempFile(test, text)
|
||||
tmpfile, err := createTempFile(t, test, text)
|
||||
assert.Nil(t, err)
|
||||
defer os.Remove(tmpfile)
|
||||
|
||||
var val struct {
|
||||
A string `json:"a"`
|
||||
@@ -82,9 +82,8 @@ c = "${FOO}"
|
||||
d = "abcd!@#$112"
|
||||
`
|
||||
t.Setenv("FOO", "2")
|
||||
tmpfile, err := createTempFile(".toml", text)
|
||||
tmpfile, err := createTempFile(t, ".toml", text)
|
||||
assert.Nil(t, err)
|
||||
defer os.Remove(tmpfile)
|
||||
|
||||
var val struct {
|
||||
A string `json:"a"`
|
||||
@@ -105,9 +104,8 @@ b = 1
|
||||
c = "FOO"
|
||||
d = "abcd"
|
||||
`
|
||||
tmpfile, err := createTempFile(".toml", text)
|
||||
tmpfile, err := createTempFile(t, ".toml", text)
|
||||
assert.Nil(t, err)
|
||||
defer os.Remove(tmpfile)
|
||||
|
||||
var val struct {
|
||||
A string `json:"a"`
|
||||
@@ -127,9 +125,8 @@ func TestConfigWithLower(t *testing.T) {
|
||||
text := `a = "foo"
|
||||
b = 1
|
||||
`
|
||||
tmpfile, err := createTempFile(".toml", text)
|
||||
tmpfile, err := createTempFile(t, ".toml", text)
|
||||
assert.Nil(t, err)
|
||||
defer os.Remove(tmpfile)
|
||||
|
||||
var val struct {
|
||||
A string `json:"a"`
|
||||
@@ -207,9 +204,8 @@ c = "${FOO}"
|
||||
d = "abcd!@#112"
|
||||
`
|
||||
t.Setenv("FOO", "2")
|
||||
tmpfile, err := createTempFile(".toml", text)
|
||||
tmpfile, err := createTempFile(t, ".toml", text)
|
||||
assert.Nil(t, err)
|
||||
defer os.Remove(tmpfile)
|
||||
|
||||
var val struct {
|
||||
A string `json:"a"`
|
||||
@@ -241,9 +237,8 @@ func TestConfigJsonEnv(t *testing.T) {
|
||||
for _, test := range tests {
|
||||
test := test
|
||||
t.Run(test, func(t *testing.T) {
|
||||
tmpfile, err := createTempFile(test, text)
|
||||
tmpfile, err := createTempFile(t, test, text)
|
||||
assert.Nil(t, err)
|
||||
defer os.Remove(tmpfile)
|
||||
|
||||
var val struct {
|
||||
A string `json:"a"`
|
||||
@@ -1217,11 +1212,44 @@ Name = "bar"
|
||||
})
|
||||
}
|
||||
|
||||
func Test_LoadBadConfig(t *testing.T) {
|
||||
type Config struct {
|
||||
Name string `json:"name,options=foo|bar"`
|
||||
}
|
||||
|
||||
file, err := createTempFile(t, ".json", `{"name": "baz"}`)
|
||||
assert.NoError(t, err)
|
||||
|
||||
var c Config
|
||||
err = Load(file, &c)
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
func Test_getFullName(t *testing.T) {
|
||||
assert.Equal(t, "a.b", getFullName("a", "b"))
|
||||
assert.Equal(t, "a", getFullName("", "a"))
|
||||
}
|
||||
|
||||
func TestValidate(t *testing.T) {
|
||||
t.Run("normal config", func(t *testing.T) {
|
||||
var c mockConfig
|
||||
err := LoadFromJsonBytes([]byte(`{"val": "hello", "number": 8}`), &c)
|
||||
assert.NoError(t, err)
|
||||
})
|
||||
|
||||
t.Run("error no int", func(t *testing.T) {
|
||||
var c mockConfig
|
||||
err := LoadFromJsonBytes([]byte(`{"val": "hello"}`), &c)
|
||||
assert.Error(t, err)
|
||||
})
|
||||
|
||||
t.Run("error no string", func(t *testing.T) {
|
||||
var c mockConfig
|
||||
err := LoadFromJsonBytes([]byte(`{"number": 8}`), &c)
|
||||
assert.Error(t, err)
|
||||
})
|
||||
}
|
||||
|
||||
func Test_buildFieldsInfo(t *testing.T) {
|
||||
type ParentSt struct {
|
||||
Name string
|
||||
@@ -1311,13 +1339,13 @@ func Test_buildFieldsInfo(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func createTempFile(ext, text string) (string, error) {
|
||||
func createTempFile(t *testing.T, ext, text string) (string, error) {
|
||||
tmpFile, err := os.CreateTemp(os.TempDir(), hash.Md5Hex([]byte(text))+"*"+ext)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
if err := os.WriteFile(tmpFile.Name(), []byte(text), os.ModeTemporary); err != nil {
|
||||
if err = os.WriteFile(tmpFile.Name(), []byte(text), os.ModeTemporary); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
@@ -1326,5 +1354,26 @@ func createTempFile(ext, text string) (string, error) {
|
||||
return "", err
|
||||
}
|
||||
|
||||
t.Cleanup(func() {
|
||||
_ = os.Remove(filename)
|
||||
})
|
||||
|
||||
return filename, nil
|
||||
}
|
||||
|
||||
type mockConfig struct {
|
||||
Val string
|
||||
Number int
|
||||
}
|
||||
|
||||
func (m mockConfig) Validate() error {
|
||||
if len(m.Val) == 0 {
|
||||
return errors.New("val is empty")
|
||||
}
|
||||
|
||||
if m.Number == 0 {
|
||||
return errors.New("number is zero")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
12
core/conf/validate.go
Normal file
12
core/conf/validate.go
Normal file
@@ -0,0 +1,12 @@
|
||||
package conf
|
||||
|
||||
import "github.com/zeromicro/go-zero/core/validation"
|
||||
|
||||
// validate validates the value if it implements the Validator interface.
|
||||
func validate(v any) error {
|
||||
if val, ok := v.(validation.Validator); ok {
|
||||
return val.Validate()
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
81
core/conf/validate_test.go
Normal file
81
core/conf/validate_test.go
Normal file
@@ -0,0 +1,81 @@
|
||||
package conf
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
type mockType int
|
||||
|
||||
func (m mockType) Validate() error {
|
||||
if m < 10 {
|
||||
return errors.New("invalid value")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
type anotherMockType int
|
||||
|
||||
func Test_validate(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
v any
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "invalid",
|
||||
v: mockType(5),
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "valid",
|
||||
v: mockType(10),
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "not validator",
|
||||
v: anotherMockType(5),
|
||||
wantErr: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
err := validate(tt.v)
|
||||
assert.Equal(t, tt.wantErr, err != nil)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
type mockVal struct {
|
||||
}
|
||||
|
||||
func (m mockVal) Validate() error {
|
||||
return errors.New("invalid value")
|
||||
}
|
||||
|
||||
func Test_validateValPtr(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
v any
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "invalid",
|
||||
v: mockVal{},
|
||||
},
|
||||
{
|
||||
name: "invalid value",
|
||||
v: &mockVal{},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
assert.Error(t, validate(tt.v))
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -10,22 +10,24 @@ import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"go.etcd.io/etcd/api/v3/v3rpc/rpctypes"
|
||||
clientv3 "go.etcd.io/etcd/client/v3"
|
||||
|
||||
"github.com/zeromicro/go-zero/core/contextx"
|
||||
"github.com/zeromicro/go-zero/core/lang"
|
||||
"github.com/zeromicro/go-zero/core/logx"
|
||||
"github.com/zeromicro/go-zero/core/logc"
|
||||
"github.com/zeromicro/go-zero/core/mathx"
|
||||
"github.com/zeromicro/go-zero/core/syncx"
|
||||
"github.com/zeromicro/go-zero/core/threading"
|
||||
"go.etcd.io/etcd/api/v3/v3rpc/rpctypes"
|
||||
clientv3 "go.etcd.io/etcd/client/v3"
|
||||
)
|
||||
|
||||
const coolDownDeviation = 0.05
|
||||
|
||||
var (
|
||||
registry = Registry{
|
||||
clusters: make(map[string]*cluster),
|
||||
}
|
||||
connManager = syncx.NewResourceManager()
|
||||
errClosed = errors.New("etcd monitor chan has been closed")
|
||||
connManager = syncx.NewResourceManager()
|
||||
coolDownUnstable = mathx.NewUnstable(coolDownDeviation)
|
||||
errClosed = errors.New("etcd monitor chan has been closed")
|
||||
)
|
||||
|
||||
// A Registry is a registry that manages the etcd client connections.
|
||||
@@ -41,33 +43,92 @@ func GetRegistry() *Registry {
|
||||
|
||||
// GetConn returns an etcd client connection associated with given endpoints.
|
||||
func (r *Registry) GetConn(endpoints []string) (EtcdClient, error) {
|
||||
c, _ := r.getCluster(endpoints)
|
||||
c, _ := r.getOrCreateCluster(endpoints)
|
||||
return c.getClient()
|
||||
}
|
||||
|
||||
// Monitor monitors the key on given etcd endpoints, notify with the given UpdateListener.
|
||||
func (r *Registry) Monitor(endpoints []string, key string, l UpdateListener, exactMatch bool) error {
|
||||
c, exists := r.getCluster(endpoints)
|
||||
func (r *Registry) Monitor(endpoints []string, key string, exactMatch bool, l UpdateListener) error {
|
||||
wkey := watchKey{
|
||||
key: key,
|
||||
exactMatch: exactMatch,
|
||||
}
|
||||
|
||||
c, exists := r.getOrCreateCluster(endpoints)
|
||||
// if exists, the existing values should be updated to the listener.
|
||||
if exists {
|
||||
kvs := c.getCurrent(key)
|
||||
for _, kv := range kvs {
|
||||
l.OnAdd(kv)
|
||||
c.lock.Lock()
|
||||
watcher, ok := c.watchers[wkey]
|
||||
if ok {
|
||||
watcher.listeners = append(watcher.listeners, l)
|
||||
}
|
||||
c.lock.Unlock()
|
||||
|
||||
if ok {
|
||||
kvs := c.getCurrent(wkey)
|
||||
for _, kv := range kvs {
|
||||
l.OnAdd(kv)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
return c.monitor(key, l, exactMatch)
|
||||
return c.monitor(wkey, l)
|
||||
}
|
||||
|
||||
func (r *Registry) getCluster(endpoints []string) (c *cluster, exists bool) {
|
||||
func (r *Registry) Unmonitor(endpoints []string, key string, exactMatch bool, l UpdateListener) {
|
||||
c, exists := r.getCluster(endpoints)
|
||||
if !exists {
|
||||
return
|
||||
}
|
||||
|
||||
wkey := watchKey{
|
||||
key: key,
|
||||
exactMatch: exactMatch,
|
||||
}
|
||||
|
||||
c.lock.Lock()
|
||||
defer c.lock.Unlock()
|
||||
|
||||
watcher, ok := c.watchers[wkey]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
for i, listener := range watcher.listeners {
|
||||
if listener == l {
|
||||
watcher.listeners = append(watcher.listeners[:i], watcher.listeners[i+1:]...)
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if len(watcher.listeners) == 0 {
|
||||
if watcher.cancel != nil {
|
||||
watcher.cancel()
|
||||
}
|
||||
delete(c.watchers, wkey)
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Registry) getCluster(endpoints []string) (*cluster, bool) {
|
||||
clusterKey := getClusterKey(endpoints)
|
||||
|
||||
r.lock.RLock()
|
||||
c, exists = r.clusters[clusterKey]
|
||||
c, ok := r.clusters[clusterKey]
|
||||
r.lock.RUnlock()
|
||||
|
||||
return c, ok
|
||||
}
|
||||
|
||||
func (r *Registry) getOrCreateCluster(endpoints []string) (c *cluster, exists bool) {
|
||||
c, exists = r.getCluster(endpoints)
|
||||
if !exists {
|
||||
clusterKey := getClusterKey(endpoints)
|
||||
|
||||
r.lock.Lock()
|
||||
defer r.lock.Unlock()
|
||||
|
||||
// double-check locking
|
||||
c, exists = r.clusters[clusterKey]
|
||||
if !exists {
|
||||
@@ -79,30 +140,51 @@ func (r *Registry) getCluster(endpoints []string) (c *cluster, exists bool) {
|
||||
return
|
||||
}
|
||||
|
||||
type cluster struct {
|
||||
endpoints []string
|
||||
key string
|
||||
values map[string]map[string]string
|
||||
listeners map[string][]UpdateListener
|
||||
watchGroup *threading.RoutineGroup
|
||||
done chan lang.PlaceholderType
|
||||
lock sync.RWMutex
|
||||
exactMatch bool
|
||||
}
|
||||
type (
|
||||
watchKey struct {
|
||||
key string
|
||||
exactMatch bool
|
||||
}
|
||||
|
||||
watchValue struct {
|
||||
listeners []UpdateListener
|
||||
values map[string]string
|
||||
cancel context.CancelFunc
|
||||
}
|
||||
|
||||
cluster struct {
|
||||
endpoints []string
|
||||
key string
|
||||
watchers map[watchKey]*watchValue
|
||||
watchGroup *threading.RoutineGroup
|
||||
done chan lang.PlaceholderType
|
||||
lock sync.RWMutex
|
||||
}
|
||||
)
|
||||
|
||||
func newCluster(endpoints []string) *cluster {
|
||||
return &cluster{
|
||||
endpoints: endpoints,
|
||||
key: getClusterKey(endpoints),
|
||||
values: make(map[string]map[string]string),
|
||||
listeners: make(map[string][]UpdateListener),
|
||||
watchers: make(map[watchKey]*watchValue),
|
||||
watchGroup: threading.NewRoutineGroup(),
|
||||
done: make(chan lang.PlaceholderType),
|
||||
}
|
||||
}
|
||||
|
||||
func (c *cluster) context(cli EtcdClient) context.Context {
|
||||
return contextx.ValueOnlyFrom(cli.Ctx())
|
||||
func (c *cluster) addListener(key watchKey, l UpdateListener) {
|
||||
c.lock.Lock()
|
||||
defer c.lock.Unlock()
|
||||
|
||||
watcher, ok := c.watchers[key]
|
||||
if ok {
|
||||
watcher.listeners = append(watcher.listeners, l)
|
||||
return
|
||||
}
|
||||
|
||||
val := newWatchValue()
|
||||
val.listeners = []UpdateListener{l}
|
||||
c.watchers[key] = val
|
||||
}
|
||||
|
||||
func (c *cluster) getClient() (EtcdClient, error) {
|
||||
@@ -116,12 +198,17 @@ func (c *cluster) getClient() (EtcdClient, error) {
|
||||
return val.(EtcdClient), nil
|
||||
}
|
||||
|
||||
func (c *cluster) getCurrent(key string) []KV {
|
||||
func (c *cluster) getCurrent(key watchKey) []KV {
|
||||
c.lock.RLock()
|
||||
defer c.lock.RUnlock()
|
||||
|
||||
watcher, ok := c.watchers[key]
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
var kvs []KV
|
||||
for k, v := range c.values[key] {
|
||||
for k, v := range watcher.values {
|
||||
kvs = append(kvs, KV{
|
||||
Key: k,
|
||||
Val: v,
|
||||
@@ -131,43 +218,23 @@ func (c *cluster) getCurrent(key string) []KV {
|
||||
return kvs
|
||||
}
|
||||
|
||||
func (c *cluster) handleChanges(key string, kvs []KV) {
|
||||
var add []KV
|
||||
var remove []KV
|
||||
|
||||
func (c *cluster) handleChanges(key watchKey, kvs []KV) {
|
||||
c.lock.Lock()
|
||||
listeners := append([]UpdateListener(nil), c.listeners[key]...)
|
||||
vals, ok := c.values[key]
|
||||
watcher, ok := c.watchers[key]
|
||||
if !ok {
|
||||
add = kvs
|
||||
vals = make(map[string]string)
|
||||
for _, kv := range kvs {
|
||||
vals[kv.Key] = kv.Val
|
||||
}
|
||||
c.values[key] = vals
|
||||
} else {
|
||||
m := make(map[string]string)
|
||||
for _, kv := range kvs {
|
||||
m[kv.Key] = kv.Val
|
||||
}
|
||||
for k, v := range vals {
|
||||
if val, ok := m[k]; !ok || v != val {
|
||||
remove = append(remove, KV{
|
||||
Key: k,
|
||||
Val: v,
|
||||
})
|
||||
}
|
||||
}
|
||||
for k, v := range m {
|
||||
if val, ok := vals[k]; !ok || v != val {
|
||||
add = append(add, KV{
|
||||
Key: k,
|
||||
Val: v,
|
||||
})
|
||||
}
|
||||
}
|
||||
c.values[key] = m
|
||||
c.lock.Unlock()
|
||||
return
|
||||
}
|
||||
|
||||
listeners := append([]UpdateListener(nil), watcher.listeners...)
|
||||
// watcher.values cannot be nil
|
||||
vals := watcher.values
|
||||
newVals := make(map[string]string, len(kvs)+len(vals))
|
||||
for _, kv := range kvs {
|
||||
newVals[kv.Key] = kv.Val
|
||||
}
|
||||
add, remove := calculateChanges(vals, newVals)
|
||||
watcher.values = newVals
|
||||
c.lock.Unlock()
|
||||
|
||||
for _, kv := range add {
|
||||
@@ -182,20 +249,22 @@ func (c *cluster) handleChanges(key string, kvs []KV) {
|
||||
}
|
||||
}
|
||||
|
||||
func (c *cluster) handleWatchEvents(key string, events []*clientv3.Event) {
|
||||
func (c *cluster) handleWatchEvents(ctx context.Context, key watchKey, events []*clientv3.Event) {
|
||||
c.lock.RLock()
|
||||
listeners := append([]UpdateListener(nil), c.listeners[key]...)
|
||||
watcher, ok := c.watchers[key]
|
||||
if !ok {
|
||||
c.lock.RUnlock()
|
||||
return
|
||||
}
|
||||
|
||||
listeners := append([]UpdateListener(nil), watcher.listeners...)
|
||||
c.lock.RUnlock()
|
||||
|
||||
for _, ev := range events {
|
||||
switch ev.Type {
|
||||
case clientv3.EventTypePut:
|
||||
c.lock.Lock()
|
||||
if vals, ok := c.values[key]; ok {
|
||||
vals[string(ev.Kv.Key)] = string(ev.Kv.Value)
|
||||
} else {
|
||||
c.values[key] = map[string]string{string(ev.Kv.Key): string(ev.Kv.Value)}
|
||||
}
|
||||
watcher.values[string(ev.Kv.Key)] = string(ev.Kv.Value)
|
||||
c.lock.Unlock()
|
||||
for _, l := range listeners {
|
||||
l.OnAdd(KV{
|
||||
@@ -205,9 +274,7 @@ func (c *cluster) handleWatchEvents(key string, events []*clientv3.Event) {
|
||||
}
|
||||
case clientv3.EventTypeDelete:
|
||||
c.lock.Lock()
|
||||
if vals, ok := c.values[key]; ok {
|
||||
delete(vals, string(ev.Kv.Key))
|
||||
}
|
||||
delete(watcher.values, string(ev.Kv.Key))
|
||||
c.lock.Unlock()
|
||||
for _, l := range listeners {
|
||||
l.OnDelete(KV{
|
||||
@@ -216,20 +283,20 @@ func (c *cluster) handleWatchEvents(key string, events []*clientv3.Event) {
|
||||
})
|
||||
}
|
||||
default:
|
||||
logx.Errorf("Unknown event type: %v", ev.Type)
|
||||
logc.Errorf(ctx, "Unknown event type: %v", ev.Type)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *cluster) load(cli EtcdClient, key string) int64 {
|
||||
func (c *cluster) load(cli EtcdClient, key watchKey) int64 {
|
||||
var resp *clientv3.GetResponse
|
||||
for {
|
||||
var err error
|
||||
ctx, cancel := context.WithTimeout(c.context(cli), RequestTimeout)
|
||||
if c.exactMatch {
|
||||
resp, err = cli.Get(ctx, key)
|
||||
ctx, cancel := context.WithTimeout(cli.Ctx(), RequestTimeout)
|
||||
if key.exactMatch {
|
||||
resp, err = cli.Get(ctx, key.key)
|
||||
} else {
|
||||
resp, err = cli.Get(ctx, makeKeyPrefix(key), clientv3.WithPrefix())
|
||||
resp, err = cli.Get(ctx, makeKeyPrefix(key.key), clientv3.WithPrefix())
|
||||
}
|
||||
|
||||
cancel()
|
||||
@@ -237,8 +304,8 @@ func (c *cluster) load(cli EtcdClient, key string) int64 {
|
||||
break
|
||||
}
|
||||
|
||||
logx.Errorf("%s, key is %s", err.Error(), key)
|
||||
time.Sleep(coolDownInterval)
|
||||
logc.Errorf(cli.Ctx(), "%s, key: %s, exactMatch: %t", err.Error(), key.key, key.exactMatch)
|
||||
time.Sleep(coolDownUnstable.AroundDuration(coolDownInterval))
|
||||
}
|
||||
|
||||
var kvs []KV
|
||||
@@ -254,17 +321,13 @@ func (c *cluster) load(cli EtcdClient, key string) int64 {
|
||||
return resp.Header.Revision
|
||||
}
|
||||
|
||||
func (c *cluster) monitor(key string, l UpdateListener, exactMatch bool) error {
|
||||
c.lock.Lock()
|
||||
c.listeners[key] = append(c.listeners[key], l)
|
||||
c.exactMatch = exactMatch
|
||||
c.lock.Unlock()
|
||||
|
||||
func (c *cluster) monitor(key watchKey, l UpdateListener) error {
|
||||
cli, err := c.getClient()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
c.addListener(key, l)
|
||||
rev := c.load(cli, key)
|
||||
c.watchGroup.Run(func() {
|
||||
c.watch(cli, key, rev)
|
||||
@@ -286,16 +349,22 @@ func (c *cluster) newClient() (EtcdClient, error) {
|
||||
|
||||
func (c *cluster) reload(cli EtcdClient) {
|
||||
c.lock.Lock()
|
||||
// cancel the previous watches
|
||||
close(c.done)
|
||||
c.watchGroup.Wait()
|
||||
var keys []watchKey
|
||||
for wk, wval := range c.watchers {
|
||||
keys = append(keys, wk)
|
||||
if wval.cancel != nil {
|
||||
wval.cancel()
|
||||
}
|
||||
}
|
||||
|
||||
c.done = make(chan lang.PlaceholderType)
|
||||
c.watchGroup = threading.NewRoutineGroup()
|
||||
var keys []string
|
||||
for k := range c.listeners {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
c.lock.Unlock()
|
||||
|
||||
// start new watches
|
||||
for _, key := range keys {
|
||||
k := key
|
||||
c.watchGroup.Run(func() {
|
||||
@@ -305,7 +374,7 @@ func (c *cluster) reload(cli EtcdClient) {
|
||||
}
|
||||
}
|
||||
|
||||
func (c *cluster) watch(cli EtcdClient, key string, rev int64) {
|
||||
func (c *cluster) watch(cli EtcdClient, key watchKey, rev int64) {
|
||||
for {
|
||||
err := c.watchStream(cli, key, rev)
|
||||
if err == nil {
|
||||
@@ -313,30 +382,17 @@ func (c *cluster) watch(cli EtcdClient, key string, rev int64) {
|
||||
}
|
||||
|
||||
if rev != 0 && errors.Is(err, rpctypes.ErrCompacted) {
|
||||
logx.Errorf("etcd watch stream has been compacted, try to reload, rev %d", rev)
|
||||
logc.Errorf(cli.Ctx(), "etcd watch stream has been compacted, try to reload, rev %d", rev)
|
||||
rev = c.load(cli, key)
|
||||
}
|
||||
|
||||
// log the error and retry
|
||||
logx.Error(err)
|
||||
logc.Error(cli.Ctx(), err)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *cluster) watchStream(cli EtcdClient, key string, rev int64) error {
|
||||
var (
|
||||
rch clientv3.WatchChan
|
||||
ops []clientv3.OpOption
|
||||
watchKey = key
|
||||
)
|
||||
if !c.exactMatch {
|
||||
watchKey = makeKeyPrefix(key)
|
||||
ops = append(ops, clientv3.WithPrefix())
|
||||
}
|
||||
if rev != 0 {
|
||||
ops = append(ops, clientv3.WithRev(rev+1))
|
||||
}
|
||||
|
||||
rch = cli.Watch(clientv3.WithRequireLeader(c.context(cli)), watchKey, ops...)
|
||||
func (c *cluster) watchStream(cli EtcdClient, key watchKey, rev int64) error {
|
||||
ctx, rch := c.setupWatch(cli, key, rev)
|
||||
|
||||
for {
|
||||
select {
|
||||
@@ -351,13 +407,47 @@ func (c *cluster) watchStream(cli EtcdClient, key string, rev int64) error {
|
||||
return fmt.Errorf("etcd monitor chan error: %w", wresp.Err())
|
||||
}
|
||||
|
||||
c.handleWatchEvents(key, wresp.Events)
|
||||
c.handleWatchEvents(ctx, key, wresp.Events)
|
||||
case <-ctx.Done():
|
||||
return nil
|
||||
case <-c.done:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *cluster) setupWatch(cli EtcdClient, key watchKey, rev int64) (context.Context, clientv3.WatchChan) {
|
||||
var (
|
||||
rch clientv3.WatchChan
|
||||
ops []clientv3.OpOption
|
||||
wkey = key.key
|
||||
)
|
||||
|
||||
if !key.exactMatch {
|
||||
wkey = makeKeyPrefix(key.key)
|
||||
ops = append(ops, clientv3.WithPrefix())
|
||||
}
|
||||
if rev != 0 {
|
||||
ops = append(ops, clientv3.WithRev(rev+1))
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(cli.Ctx())
|
||||
if watcher, ok := c.watchers[key]; ok {
|
||||
watcher.cancel = cancel
|
||||
} else {
|
||||
val := newWatchValue()
|
||||
val.cancel = cancel
|
||||
|
||||
c.lock.Lock()
|
||||
c.watchers[key] = val
|
||||
c.lock.Unlock()
|
||||
}
|
||||
|
||||
rch = cli.Watch(clientv3.WithRequireLeader(ctx), wkey, ops...)
|
||||
|
||||
return ctx, rch
|
||||
}
|
||||
|
||||
func (c *cluster) watchConnState(cli EtcdClient) {
|
||||
watcher := newStateWatcher()
|
||||
watcher.addListener(func() {
|
||||
@@ -386,6 +476,28 @@ func DialClient(endpoints []string) (EtcdClient, error) {
|
||||
return clientv3.New(cfg)
|
||||
}
|
||||
|
||||
func calculateChanges(oldVals, newVals map[string]string) (add, remove []KV) {
|
||||
for k, v := range newVals {
|
||||
if val, ok := oldVals[k]; !ok || v != val {
|
||||
add = append(add, KV{
|
||||
Key: k,
|
||||
Val: v,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
for k, v := range oldVals {
|
||||
if val, ok := newVals[k]; !ok || v != val {
|
||||
remove = append(remove, KV{
|
||||
Key: k,
|
||||
Val: v,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return add, remove
|
||||
}
|
||||
|
||||
func getClusterKey(endpoints []string) string {
|
||||
sort.Strings(endpoints)
|
||||
return strings.Join(endpoints, endpointsSeparator)
|
||||
@@ -394,3 +506,10 @@ func getClusterKey(endpoints []string) string {
|
||||
func makeKeyPrefix(key string) string {
|
||||
return fmt.Sprintf("%s%c", key, Delimiter)
|
||||
}
|
||||
|
||||
// NewClient returns a watchValue that make sure values are not nil.
|
||||
func newWatchValue() *watchValue {
|
||||
return &watchValue{
|
||||
values: make(map[string]string),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
"github.com/zeromicro/go-zero/core/lang"
|
||||
"github.com/zeromicro/go-zero/core/logx"
|
||||
"github.com/zeromicro/go-zero/core/stringx"
|
||||
"github.com/zeromicro/go-zero/core/threading"
|
||||
"go.etcd.io/etcd/api/v3/etcdserverpb"
|
||||
"go.etcd.io/etcd/api/v3/mvccpb"
|
||||
clientv3 "go.etcd.io/etcd/client/v3"
|
||||
@@ -38,9 +39,9 @@ func setMockClient(cli EtcdClient) func() {
|
||||
|
||||
func TestGetCluster(t *testing.T) {
|
||||
AddAccount([]string{"first"}, "foo", "bar")
|
||||
c1, _ := GetRegistry().getCluster([]string{"first"})
|
||||
c2, _ := GetRegistry().getCluster([]string{"second"})
|
||||
c3, _ := GetRegistry().getCluster([]string{"first"})
|
||||
c1, _ := GetRegistry().getOrCreateCluster([]string{"first"})
|
||||
c2, _ := GetRegistry().getOrCreateCluster([]string{"second"})
|
||||
c3, _ := GetRegistry().getOrCreateCluster([]string{"first"})
|
||||
assert.Equal(t, c1, c3)
|
||||
assert.NotEqual(t, c1, c2)
|
||||
}
|
||||
@@ -50,6 +51,36 @@ func TestGetClusterKey(t *testing.T) {
|
||||
getClusterKey([]string{"remotehost:5678", "localhost:1234"}))
|
||||
}
|
||||
|
||||
func TestUnmonitor(t *testing.T) {
|
||||
t.Run("no listener", func(t *testing.T) {
|
||||
reg := &Registry{
|
||||
clusters: map[string]*cluster{},
|
||||
}
|
||||
assert.NotPanics(t, func() {
|
||||
reg.Unmonitor([]string{"any"}, "any", false, nil)
|
||||
})
|
||||
})
|
||||
|
||||
t.Run("no value", func(t *testing.T) {
|
||||
reg := &Registry{
|
||||
clusters: map[string]*cluster{
|
||||
"any": {
|
||||
watchers: map[watchKey]*watchValue{
|
||||
{
|
||||
key: "any",
|
||||
}: {
|
||||
values: map[string]string{},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
assert.NotPanics(t, func() {
|
||||
reg.Unmonitor([]string{"any"}, "another", false, nil)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func TestCluster_HandleChanges(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
l := NewMockUpdateListener(ctrl)
|
||||
@@ -78,8 +109,14 @@ func TestCluster_HandleChanges(t *testing.T) {
|
||||
Val: "4",
|
||||
})
|
||||
c := newCluster([]string{"any"})
|
||||
c.listeners["any"] = []UpdateListener{l}
|
||||
c.handleChanges("any", []KV{
|
||||
key := watchKey{
|
||||
key: "any",
|
||||
exactMatch: false,
|
||||
}
|
||||
c.watchers[key] = &watchValue{
|
||||
listeners: []UpdateListener{l},
|
||||
}
|
||||
c.handleChanges(key, []KV{
|
||||
{
|
||||
Key: "first",
|
||||
Val: "1",
|
||||
@@ -92,8 +129,8 @@ func TestCluster_HandleChanges(t *testing.T) {
|
||||
assert.EqualValues(t, map[string]string{
|
||||
"first": "1",
|
||||
"second": "2",
|
||||
}, c.values["any"])
|
||||
c.handleChanges("any", []KV{
|
||||
}, c.watchers[key].values)
|
||||
c.handleChanges(key, []KV{
|
||||
{
|
||||
Key: "third",
|
||||
Val: "3",
|
||||
@@ -106,7 +143,7 @@ func TestCluster_HandleChanges(t *testing.T) {
|
||||
assert.EqualValues(t, map[string]string{
|
||||
"third": "3",
|
||||
"fourth": "4",
|
||||
}, c.values["any"])
|
||||
}, c.watchers[key].values)
|
||||
}
|
||||
|
||||
func TestCluster_Load(t *testing.T) {
|
||||
@@ -126,9 +163,11 @@ func TestCluster_Load(t *testing.T) {
|
||||
}, nil)
|
||||
cli.EXPECT().Ctx().Return(context.Background())
|
||||
c := &cluster{
|
||||
values: make(map[string]map[string]string),
|
||||
watchers: make(map[watchKey]*watchValue),
|
||||
}
|
||||
c.load(cli, "any")
|
||||
c.load(cli, watchKey{
|
||||
key: "any",
|
||||
})
|
||||
}
|
||||
|
||||
func TestCluster_Watch(t *testing.T) {
|
||||
@@ -160,11 +199,16 @@ func TestCluster_Watch(t *testing.T) {
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
c := &cluster{
|
||||
listeners: make(map[string][]UpdateListener),
|
||||
values: make(map[string]map[string]string),
|
||||
watchers: make(map[watchKey]*watchValue),
|
||||
}
|
||||
key := watchKey{
|
||||
key: "any",
|
||||
}
|
||||
listener := NewMockUpdateListener(ctrl)
|
||||
c.listeners["any"] = []UpdateListener{listener}
|
||||
c.watchers[key] = &watchValue{
|
||||
listeners: []UpdateListener{listener},
|
||||
values: make(map[string]string),
|
||||
}
|
||||
listener.EXPECT().OnAdd(gomock.Any()).Do(func(kv KV) {
|
||||
assert.Equal(t, "hello", kv.Key)
|
||||
assert.Equal(t, "world", kv.Val)
|
||||
@@ -173,7 +217,7 @@ func TestCluster_Watch(t *testing.T) {
|
||||
listener.EXPECT().OnDelete(gomock.Any()).Do(func(_ any) {
|
||||
wg.Done()
|
||||
}).MaxTimes(1)
|
||||
go c.watch(cli, "any", 0)
|
||||
go c.watch(cli, key, 0)
|
||||
ch <- clientv3.WatchResponse{
|
||||
Events: []*clientv3.Event{
|
||||
{
|
||||
@@ -211,17 +255,111 @@ func TestClusterWatch_RespFailures(t *testing.T) {
|
||||
ch := make(chan clientv3.WatchResponse)
|
||||
cli.EXPECT().Watch(gomock.Any(), "any/", gomock.Any()).Return(ch).AnyTimes()
|
||||
cli.EXPECT().Ctx().Return(context.Background()).AnyTimes()
|
||||
c := new(cluster)
|
||||
c := &cluster{
|
||||
watchers: make(map[watchKey]*watchValue),
|
||||
}
|
||||
c.done = make(chan lang.PlaceholderType)
|
||||
go func() {
|
||||
ch <- resp
|
||||
close(c.done)
|
||||
}()
|
||||
c.watch(cli, "any", 0)
|
||||
key := watchKey{
|
||||
key: "any",
|
||||
}
|
||||
c.watch(cli, key, 0)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCluster_getCurrent(t *testing.T) {
|
||||
t.Run("no value", func(t *testing.T) {
|
||||
c := &cluster{
|
||||
watchers: map[watchKey]*watchValue{
|
||||
{
|
||||
key: "any",
|
||||
}: {
|
||||
values: map[string]string{},
|
||||
},
|
||||
},
|
||||
}
|
||||
assert.Nil(t, c.getCurrent(watchKey{
|
||||
key: "another",
|
||||
}))
|
||||
})
|
||||
}
|
||||
|
||||
func TestCluster_handleWatchEvents(t *testing.T) {
|
||||
t.Run("no value", func(t *testing.T) {
|
||||
c := &cluster{
|
||||
watchers: map[watchKey]*watchValue{
|
||||
{
|
||||
key: "any",
|
||||
}: {
|
||||
values: map[string]string{},
|
||||
},
|
||||
},
|
||||
}
|
||||
assert.NotPanics(t, func() {
|
||||
c.handleWatchEvents(context.Background(), watchKey{
|
||||
key: "another",
|
||||
}, nil)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func TestCluster_addListener(t *testing.T) {
|
||||
t.Run("has listener", func(t *testing.T) {
|
||||
c := &cluster{
|
||||
watchers: map[watchKey]*watchValue{
|
||||
{
|
||||
key: "any",
|
||||
}: {
|
||||
listeners: make([]UpdateListener, 0),
|
||||
},
|
||||
},
|
||||
}
|
||||
assert.NotPanics(t, func() {
|
||||
c.addListener(watchKey{
|
||||
key: "any",
|
||||
}, nil)
|
||||
})
|
||||
})
|
||||
|
||||
t.Run("no listener", func(t *testing.T) {
|
||||
c := &cluster{
|
||||
watchers: map[watchKey]*watchValue{
|
||||
{
|
||||
key: "any",
|
||||
}: {
|
||||
listeners: make([]UpdateListener, 0),
|
||||
},
|
||||
},
|
||||
}
|
||||
assert.NotPanics(t, func() {
|
||||
c.addListener(watchKey{
|
||||
key: "another",
|
||||
}, nil)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func TestCluster_reload(t *testing.T) {
|
||||
c := &cluster{
|
||||
watchers: map[watchKey]*watchValue{},
|
||||
watchGroup: threading.NewRoutineGroup(),
|
||||
done: make(chan lang.PlaceholderType),
|
||||
}
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
defer ctrl.Finish()
|
||||
cli := NewMockEtcdClient(ctrl)
|
||||
restore := setMockClient(cli)
|
||||
defer restore()
|
||||
assert.NotPanics(t, func() {
|
||||
c.reload(cli)
|
||||
})
|
||||
}
|
||||
|
||||
func TestClusterWatch_CloseChan(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
defer ctrl.Finish()
|
||||
@@ -231,13 +369,17 @@ func TestClusterWatch_CloseChan(t *testing.T) {
|
||||
ch := make(chan clientv3.WatchResponse)
|
||||
cli.EXPECT().Watch(gomock.Any(), "any/", gomock.Any()).Return(ch).AnyTimes()
|
||||
cli.EXPECT().Ctx().Return(context.Background()).AnyTimes()
|
||||
c := new(cluster)
|
||||
c := &cluster{
|
||||
watchers: make(map[watchKey]*watchValue),
|
||||
}
|
||||
c.done = make(chan lang.PlaceholderType)
|
||||
go func() {
|
||||
close(ch)
|
||||
close(c.done)
|
||||
}()
|
||||
c.watch(cli, "any", 0)
|
||||
c.watch(cli, watchKey{
|
||||
key: "any",
|
||||
}, 0)
|
||||
}
|
||||
|
||||
func TestValueOnlyContext(t *testing.T) {
|
||||
@@ -280,16 +422,59 @@ func TestRegistry_Monitor(t *testing.T) {
|
||||
GetRegistry().lock.Lock()
|
||||
GetRegistry().clusters = map[string]*cluster{
|
||||
getClusterKey(endpoints): {
|
||||
listeners: map[string][]UpdateListener{},
|
||||
values: map[string]map[string]string{
|
||||
"foo": {
|
||||
"bar": "baz",
|
||||
watchers: map[watchKey]*watchValue{
|
||||
watchKey{
|
||||
key: "foo",
|
||||
exactMatch: true,
|
||||
}: {
|
||||
values: map[string]string{
|
||||
"bar": "baz",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
GetRegistry().lock.Unlock()
|
||||
assert.Error(t, GetRegistry().Monitor(endpoints, "foo", new(mockListener), false))
|
||||
assert.Error(t, GetRegistry().Monitor(endpoints, "foo", false, new(mockListener)))
|
||||
}
|
||||
|
||||
func TestRegistry_Unmonitor(t *testing.T) {
|
||||
svr, err := mockserver.StartMockServers(1)
|
||||
assert.NoError(t, err)
|
||||
svr.StartAt(0)
|
||||
|
||||
_, cancel := context.WithCancel(context.Background())
|
||||
endpoints := []string{svr.Servers[0].Address}
|
||||
GetRegistry().lock.Lock()
|
||||
GetRegistry().clusters = map[string]*cluster{
|
||||
getClusterKey(endpoints): {
|
||||
watchers: map[watchKey]*watchValue{
|
||||
watchKey{
|
||||
key: "foo",
|
||||
exactMatch: true,
|
||||
}: {
|
||||
values: map[string]string{
|
||||
"bar": "baz",
|
||||
},
|
||||
cancel: cancel,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
GetRegistry().lock.Unlock()
|
||||
l := new(mockListener)
|
||||
assert.NoError(t, GetRegistry().Monitor(endpoints, "foo", true, l))
|
||||
watchVals := GetRegistry().clusters[getClusterKey(endpoints)].watchers[watchKey{
|
||||
key: "foo",
|
||||
exactMatch: true,
|
||||
}]
|
||||
assert.Equal(t, 1, len(watchVals.listeners))
|
||||
GetRegistry().Unmonitor(endpoints, "foo", true, l)
|
||||
watchVals = GetRegistry().clusters[getClusterKey(endpoints)].watchers[watchKey{
|
||||
key: "foo",
|
||||
exactMatch: true,
|
||||
}]
|
||||
assert.Nil(t, watchVals)
|
||||
}
|
||||
|
||||
type mockListener struct {
|
||||
|
||||
@@ -10,6 +10,7 @@ type (
|
||||
}
|
||||
|
||||
// UpdateListener wraps the OnAdd and OnDelete methods.
|
||||
// The implementation should be thread-safe and idempotent.
|
||||
UpdateListener interface {
|
||||
OnAdd(kv KV)
|
||||
OnDelete(kv KV)
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
|
||||
"github.com/zeromicro/go-zero/core/discov/internal"
|
||||
"github.com/zeromicro/go-zero/core/lang"
|
||||
"github.com/zeromicro/go-zero/core/logc"
|
||||
"github.com/zeromicro/go-zero/core/logx"
|
||||
"github.com/zeromicro/go-zero/core/proc"
|
||||
"github.com/zeromicro/go-zero/core/syncx"
|
||||
@@ -91,12 +92,12 @@ func (p *Publisher) doKeepAlive() error {
|
||||
default:
|
||||
cli, err := p.doRegister()
|
||||
if err != nil {
|
||||
logx.Errorf("etcd publisher doRegister: %s", err.Error())
|
||||
logc.Errorf(cli.Ctx(), "etcd publisher doRegister: %s", err.Error())
|
||||
break
|
||||
}
|
||||
|
||||
if err := p.keepAliveAsync(cli); err != nil {
|
||||
logx.Errorf("etcd publisher keepAliveAsync: %s", err.Error())
|
||||
logc.Errorf(cli.Ctx(), "etcd publisher keepAliveAsync: %s", err.Error())
|
||||
break
|
||||
}
|
||||
|
||||
@@ -130,17 +131,17 @@ func (p *Publisher) keepAliveAsync(cli internal.EtcdClient) error {
|
||||
if !ok {
|
||||
p.revoke(cli)
|
||||
if err := p.doKeepAlive(); err != nil {
|
||||
logx.Errorf("etcd publisher KeepAlive: %s", err.Error())
|
||||
logc.Errorf(cli.Ctx(), "etcd publisher KeepAlive: %s", err.Error())
|
||||
}
|
||||
return
|
||||
}
|
||||
case <-p.pauseChan:
|
||||
logx.Infof("paused etcd renew, key: %s, value: %s", p.key, p.value)
|
||||
logc.Infof(cli.Ctx(), "paused etcd renew, key: %s, value: %s", p.key, p.value)
|
||||
p.revoke(cli)
|
||||
select {
|
||||
case <-p.resumeChan:
|
||||
if err := p.doKeepAlive(); err != nil {
|
||||
logx.Errorf("etcd publisher KeepAlive: %s", err.Error())
|
||||
logc.Errorf(cli.Ctx(), "etcd publisher KeepAlive: %s", err.Error())
|
||||
}
|
||||
return
|
||||
case <-p.quit.Done():
|
||||
@@ -175,7 +176,7 @@ func (p *Publisher) register(client internal.EtcdClient) (clientv3.LeaseID, erro
|
||||
|
||||
func (p *Publisher) revoke(cli internal.EtcdClient) {
|
||||
if _, err := cli.Revoke(cli.Ctx(), p.lease); err != nil {
|
||||
logx.Errorf("etcd publisher revoke: %s", err.Error())
|
||||
logc.Errorf(cli.Ctx(), "etcd publisher revoke: %s", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -17,6 +17,7 @@ type (
|
||||
Subscriber struct {
|
||||
endpoints []string
|
||||
exclusive bool
|
||||
key string
|
||||
exactMatch bool
|
||||
items *container
|
||||
}
|
||||
@@ -29,13 +30,14 @@ type (
|
||||
func NewSubscriber(endpoints []string, key string, opts ...SubOption) (*Subscriber, error) {
|
||||
sub := &Subscriber{
|
||||
endpoints: endpoints,
|
||||
key: key,
|
||||
}
|
||||
for _, opt := range opts {
|
||||
opt(sub)
|
||||
}
|
||||
sub.items = newContainer(sub.exclusive)
|
||||
|
||||
if err := internal.GetRegistry().Monitor(endpoints, key, sub.items, sub.exactMatch); err != nil {
|
||||
if err := internal.GetRegistry().Monitor(endpoints, key, sub.exactMatch, sub.items); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
@@ -47,6 +49,11 @@ func (s *Subscriber) AddListener(listener func()) {
|
||||
s.items.addListener(listener)
|
||||
}
|
||||
|
||||
// Close closes the subscriber.
|
||||
func (s *Subscriber) Close() {
|
||||
internal.GetRegistry().Unmonitor(s.endpoints, s.key, s.exactMatch, s.items)
|
||||
}
|
||||
|
||||
// Values returns all the subscription values.
|
||||
func (s *Subscriber) Values() []string {
|
||||
return s.items.getValues()
|
||||
|
||||
@@ -225,3 +225,28 @@ func TestWithSubEtcdAccount(t *testing.T) {
|
||||
assert.Equal(t, user, account.User)
|
||||
assert.Equal(t, "bar", account.Pass)
|
||||
}
|
||||
|
||||
func TestWithExactMatch(t *testing.T) {
|
||||
sub := new(Subscriber)
|
||||
WithExactMatch()(sub)
|
||||
sub.items = newContainer(sub.exclusive)
|
||||
var count int32
|
||||
sub.AddListener(func() {
|
||||
atomic.AddInt32(&count, 1)
|
||||
})
|
||||
sub.items.notifyChange()
|
||||
assert.Empty(t, sub.Values())
|
||||
assert.Equal(t, int32(1), atomic.LoadInt32(&count))
|
||||
}
|
||||
|
||||
func TestSubscriberClose(t *testing.T) {
|
||||
l := newContainer(false)
|
||||
sub := &Subscriber{
|
||||
endpoints: []string{"localhost:12379"},
|
||||
key: "foo",
|
||||
items: l,
|
||||
}
|
||||
assert.NotPanics(t, func() {
|
||||
sub.Close()
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//go:build linux || darwin
|
||||
//go:build linux || darwin || freebsd
|
||||
|
||||
package fs
|
||||
|
||||
|
||||
@@ -37,6 +37,13 @@ func Debugf(ctx context.Context, format string, v ...interface{}) {
|
||||
getLogger(ctx).Debugf(format, v...)
|
||||
}
|
||||
|
||||
// Debugfn writes fn result into access log.
|
||||
// This is useful when the function is expensive to compute,
|
||||
// and we want to log it only when necessary.
|
||||
func Debugfn(ctx context.Context, fn func() any) {
|
||||
getLogger(ctx).Debugfn(fn)
|
||||
}
|
||||
|
||||
// Debugv writes v into access log with json content.
|
||||
func Debugv(ctx context.Context, v interface{}) {
|
||||
getLogger(ctx).Debugv(v)
|
||||
@@ -57,6 +64,13 @@ func Errorf(ctx context.Context, format string, v ...any) {
|
||||
getLogger(ctx).Errorf(fmt.Errorf(format, v...).Error())
|
||||
}
|
||||
|
||||
// Errorfn writes fn result into error log.
|
||||
// This is useful when the function is expensive to compute,
|
||||
// and we want to log it only when necessary.
|
||||
func Errorfn(ctx context.Context, fn func() any) {
|
||||
getLogger(ctx).Errorfn(fn)
|
||||
}
|
||||
|
||||
// Errorv writes v into error log with json content.
|
||||
// No call stack attached, because not elegant to pack the messages.
|
||||
func Errorv(ctx context.Context, v any) {
|
||||
@@ -83,6 +97,13 @@ func Infof(ctx context.Context, format string, v ...any) {
|
||||
getLogger(ctx).Infof(format, v...)
|
||||
}
|
||||
|
||||
// Infofn writes fn result into access log.
|
||||
// This is useful when the function is expensive to compute,
|
||||
// and we want to log it only when necessary.
|
||||
func Infofn(ctx context.Context, fn func() any) {
|
||||
getLogger(ctx).Infofn(fn)
|
||||
}
|
||||
|
||||
// Infov writes v into access log with json content.
|
||||
func Infov(ctx context.Context, v any) {
|
||||
getLogger(ctx).Infov(v)
|
||||
@@ -127,6 +148,13 @@ func Slowf(ctx context.Context, format string, v ...any) {
|
||||
getLogger(ctx).Slowf(format, v...)
|
||||
}
|
||||
|
||||
// Slowfn writes fn result into slow log.
|
||||
// This is useful when the function is expensive to compute,
|
||||
// and we want to log it only when necessary.
|
||||
func Slowfn(ctx context.Context, fn func() any) {
|
||||
getLogger(ctx).Slowfn(fn)
|
||||
}
|
||||
|
||||
// Slowv writes v into slow log with json content.
|
||||
func Slowv(ctx context.Context, v any) {
|
||||
getLogger(ctx).Slowv(v)
|
||||
|
||||
@@ -49,6 +49,15 @@ func TestErrorf(t *testing.T) {
|
||||
assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
|
||||
}
|
||||
|
||||
func TestErrorfn(t *testing.T) {
|
||||
buf := logtest.NewCollector(t)
|
||||
file, line := getFileLine()
|
||||
Errorfn(context.Background(), func() any {
|
||||
return fmt.Sprintf("foo %s", "bar")
|
||||
})
|
||||
assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
|
||||
}
|
||||
|
||||
func TestErrorv(t *testing.T) {
|
||||
buf := logtest.NewCollector(t)
|
||||
file, line := getFileLine()
|
||||
@@ -77,6 +86,15 @@ func TestInfof(t *testing.T) {
|
||||
assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
|
||||
}
|
||||
|
||||
func TestInfofn(t *testing.T) {
|
||||
buf := logtest.NewCollector(t)
|
||||
file, line := getFileLine()
|
||||
Infofn(context.Background(), func() any {
|
||||
return fmt.Sprintf("foo %s", "bar")
|
||||
})
|
||||
assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
|
||||
}
|
||||
|
||||
func TestInfov(t *testing.T) {
|
||||
buf := logtest.NewCollector(t)
|
||||
file, line := getFileLine()
|
||||
@@ -105,6 +123,15 @@ func TestDebugf(t *testing.T) {
|
||||
assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
|
||||
}
|
||||
|
||||
func TestDebugfn(t *testing.T) {
|
||||
buf := logtest.NewCollector(t)
|
||||
file, line := getFileLine()
|
||||
Debugfn(context.Background(), func() any {
|
||||
return fmt.Sprintf("foo %s", "bar")
|
||||
})
|
||||
assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
|
||||
}
|
||||
|
||||
func TestDebugv(t *testing.T) {
|
||||
buf := logtest.NewCollector(t)
|
||||
file, line := getFileLine()
|
||||
@@ -148,6 +175,15 @@ func TestSlowf(t *testing.T) {
|
||||
assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)), buf.String())
|
||||
}
|
||||
|
||||
func TestSlowfn(t *testing.T) {
|
||||
buf := logtest.NewCollector(t)
|
||||
file, line := getFileLine()
|
||||
Slowfn(context.Background(), func() any {
|
||||
return fmt.Sprintf("foo %s", "bar")
|
||||
})
|
||||
assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)), buf.String())
|
||||
}
|
||||
|
||||
func TestSlowv(t *testing.T) {
|
||||
buf := logtest.NewCollector(t)
|
||||
file, line := getFileLine()
|
||||
|
||||
@@ -11,6 +11,8 @@ type Logger interface {
|
||||
Debug(...any)
|
||||
// Debugf logs a message at debug level.
|
||||
Debugf(string, ...any)
|
||||
// Debugfn logs a message at debug level.
|
||||
Debugfn(func() any)
|
||||
// Debugv logs a message at debug level.
|
||||
Debugv(any)
|
||||
// Debugw logs a message at debug level.
|
||||
@@ -19,6 +21,8 @@ type Logger interface {
|
||||
Error(...any)
|
||||
// Errorf logs a message at error level.
|
||||
Errorf(string, ...any)
|
||||
// Errorfn logs a message at error level.
|
||||
Errorfn(func() any)
|
||||
// Errorv logs a message at error level.
|
||||
Errorv(any)
|
||||
// Errorw logs a message at error level.
|
||||
@@ -27,6 +31,8 @@ type Logger interface {
|
||||
Info(...any)
|
||||
// Infof logs a message at info level.
|
||||
Infof(string, ...any)
|
||||
// Infofn logs a message at info level.
|
||||
Infofn(func() any)
|
||||
// Infov logs a message at info level.
|
||||
Infov(any)
|
||||
// Infow logs a message at info level.
|
||||
@@ -35,6 +41,8 @@ type Logger interface {
|
||||
Slow(...any)
|
||||
// Slowf logs a message at slow level.
|
||||
Slowf(string, ...any)
|
||||
// Slowfn logs a message at slow level.
|
||||
Slowfn(func() any)
|
||||
// Slowv logs a message at slow level.
|
||||
Slowv(any)
|
||||
// Sloww logs a message at slow level.
|
||||
|
||||
@@ -100,6 +100,14 @@ func Debugf(format string, v ...any) {
|
||||
}
|
||||
}
|
||||
|
||||
// Debugfn writes function result into access log if debug level enabled.
|
||||
// This is useful when the function is expensive to call and debug level disabled.
|
||||
func Debugfn(fn func() any) {
|
||||
if shallLog(DebugLevel) {
|
||||
writeDebug(fn())
|
||||
}
|
||||
}
|
||||
|
||||
// Debugv writes v into access log with json content.
|
||||
func Debugv(v any) {
|
||||
if shallLog(DebugLevel) {
|
||||
@@ -139,6 +147,13 @@ func Errorf(format string, v ...any) {
|
||||
}
|
||||
}
|
||||
|
||||
// Errorfn writes function result into error log.
|
||||
func Errorfn(fn func() any) {
|
||||
if shallLog(ErrorLevel) {
|
||||
writeError(fn())
|
||||
}
|
||||
}
|
||||
|
||||
// ErrorStack writes v along with call stack into error log.
|
||||
func ErrorStack(v ...any) {
|
||||
if shallLog(ErrorLevel) {
|
||||
@@ -222,6 +237,14 @@ func Infof(format string, v ...any) {
|
||||
}
|
||||
}
|
||||
|
||||
// Infofn writes function result into access log.
|
||||
// This is useful when the function is expensive to call and info level disabled.
|
||||
func Infofn(fn func() any) {
|
||||
if shallLog(InfoLevel) {
|
||||
writeInfo(fn())
|
||||
}
|
||||
}
|
||||
|
||||
// Infov writes v into access log with json content.
|
||||
func Infov(v any) {
|
||||
if shallLog(InfoLevel) {
|
||||
@@ -348,6 +371,14 @@ func Slowf(format string, v ...any) {
|
||||
}
|
||||
}
|
||||
|
||||
// Slowfn writes function result into slow log.
|
||||
// This is useful when the function is expensive to call and slow level disabled.
|
||||
func Slowfn(fn func() any) {
|
||||
if shallLog(ErrorLevel) {
|
||||
writeSlow(fn())
|
||||
}
|
||||
}
|
||||
|
||||
// Slowv writes v into slow log with json content.
|
||||
func Slowv(v any) {
|
||||
if shallLog(ErrorLevel) {
|
||||
@@ -529,7 +560,7 @@ func shallLogStat() bool {
|
||||
// If we check shallLog here, the fmt.Sprint might be called even if the log level is not enabled.
|
||||
// The caller should check shallLog before calling this function.
|
||||
func writeDebug(val any, fields ...LogField) {
|
||||
getWriter().Debug(val, addCaller(fields...)...)
|
||||
getWriter().Debug(val, mergeGlobalFields(addCaller(fields...))...)
|
||||
}
|
||||
|
||||
// writeError writes v into the error log.
|
||||
@@ -537,7 +568,7 @@ func writeDebug(val any, fields ...LogField) {
|
||||
// If we check shallLog here, the fmt.Sprint might be called even if the log level is not enabled.
|
||||
// The caller should check shallLog before calling this function.
|
||||
func writeError(val any, fields ...LogField) {
|
||||
getWriter().Error(val, addCaller(fields...)...)
|
||||
getWriter().Error(val, mergeGlobalFields(addCaller(fields...))...)
|
||||
}
|
||||
|
||||
// writeInfo writes v into info log.
|
||||
@@ -545,7 +576,7 @@ func writeError(val any, fields ...LogField) {
|
||||
// If we check shallLog here, the fmt.Sprint might be called even if the log level is not enabled.
|
||||
// The caller should check shallLog before calling this function.
|
||||
func writeInfo(val any, fields ...LogField) {
|
||||
getWriter().Info(val, addCaller(fields...)...)
|
||||
getWriter().Info(val, mergeGlobalFields(addCaller(fields...))...)
|
||||
}
|
||||
|
||||
// writeSevere writes v into severe log.
|
||||
@@ -561,7 +592,7 @@ func writeSevere(msg string) {
|
||||
// If we check shallLog here, the fmt.Sprint might be called even if the log level is not enabled.
|
||||
// The caller should check shallLog before calling this function.
|
||||
func writeSlow(val any, fields ...LogField) {
|
||||
getWriter().Slow(val, addCaller(fields...)...)
|
||||
getWriter().Slow(val, mergeGlobalFields(addCaller(fields...))...)
|
||||
}
|
||||
|
||||
// writeStack writes v into stack log.
|
||||
@@ -577,5 +608,5 @@ func writeStack(msg string) {
|
||||
// If we check shallLog here, the fmt.Sprint might be called even if the log level is not enabled.
|
||||
// The caller should check shallLog before calling this function.
|
||||
func writeStat(msg string) {
|
||||
getWriter().Stat(msg, addCaller()...)
|
||||
getWriter().Stat(msg, mergeGlobalFields(addCaller())...)
|
||||
}
|
||||
|
||||
@@ -248,6 +248,32 @@ func TestStructedLogDebugf(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestStructedLogDebugfn(t *testing.T) {
|
||||
t.Run("debugfn with output", func(t *testing.T) {
|
||||
w := new(mockWriter)
|
||||
old := writer.Swap(w)
|
||||
defer writer.Store(old)
|
||||
|
||||
doTestStructedLog(t, levelDebug, w, func(v ...any) {
|
||||
Debugfn(func() any {
|
||||
return fmt.Sprint(v...)
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
t.Run("debugfn without output", func(t *testing.T) {
|
||||
w := new(mockWriter)
|
||||
old := writer.Swap(w)
|
||||
defer writer.Store(old)
|
||||
|
||||
doTestStructedLogEmpty(t, w, InfoLevel, func(v ...any) {
|
||||
Debugfn(func() any {
|
||||
return fmt.Sprint(v...)
|
||||
})
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func TestStructedLogDebugv(t *testing.T) {
|
||||
w := new(mockWriter)
|
||||
old := writer.Swap(w)
|
||||
@@ -288,6 +314,32 @@ func TestStructedLogErrorf(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestStructedLogErrorfn(t *testing.T) {
|
||||
t.Run("errorfn with output", func(t *testing.T) {
|
||||
w := new(mockWriter)
|
||||
old := writer.Swap(w)
|
||||
defer writer.Store(old)
|
||||
|
||||
doTestStructedLog(t, levelError, w, func(v ...any) {
|
||||
Errorfn(func() any {
|
||||
return fmt.Sprint(v...)
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
t.Run("errorfn without output", func(t *testing.T) {
|
||||
w := new(mockWriter)
|
||||
old := writer.Swap(w)
|
||||
defer writer.Store(old)
|
||||
|
||||
doTestStructedLogEmpty(t, w, SevereLevel, func(v ...any) {
|
||||
Errorfn(func() any {
|
||||
return fmt.Sprint(v...)
|
||||
})
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func TestStructedLogErrorv(t *testing.T) {
|
||||
w := new(mockWriter)
|
||||
old := writer.Swap(w)
|
||||
@@ -328,6 +380,32 @@ func TestStructedLogInfof(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestStructedInfofn(t *testing.T) {
|
||||
t.Run("infofn with output", func(t *testing.T) {
|
||||
w := new(mockWriter)
|
||||
old := writer.Swap(w)
|
||||
defer writer.Store(old)
|
||||
|
||||
doTestStructedLog(t, levelInfo, w, func(v ...any) {
|
||||
Infofn(func() any {
|
||||
return fmt.Sprint(v...)
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
t.Run("infofn without output", func(t *testing.T) {
|
||||
w := new(mockWriter)
|
||||
old := writer.Swap(w)
|
||||
defer writer.Store(old)
|
||||
|
||||
doTestStructedLogEmpty(t, w, ErrorLevel, func(v ...any) {
|
||||
Infofn(func() any {
|
||||
return fmt.Sprint(v...)
|
||||
})
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func TestStructedLogInfov(t *testing.T) {
|
||||
w := new(mockWriter)
|
||||
old := writer.Swap(w)
|
||||
@@ -451,6 +529,17 @@ func TestStructedLogInfoConsoleText(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestInfofnWithErrorLevel(t *testing.T) {
|
||||
called := false
|
||||
SetLevel(ErrorLevel)
|
||||
defer SetLevel(DebugLevel)
|
||||
Infofn(func() any {
|
||||
called = true
|
||||
return "info log"
|
||||
})
|
||||
assert.False(t, called)
|
||||
}
|
||||
|
||||
func TestStructedLogSlow(t *testing.T) {
|
||||
w := new(mockWriter)
|
||||
old := writer.Swap(w)
|
||||
@@ -471,6 +560,32 @@ func TestStructedLogSlowf(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestStructedLogSlowfn(t *testing.T) {
|
||||
t.Run("slowfn with output", func(t *testing.T) {
|
||||
w := new(mockWriter)
|
||||
old := writer.Swap(w)
|
||||
defer writer.Store(old)
|
||||
|
||||
doTestStructedLog(t, levelSlow, w, func(v ...any) {
|
||||
Slowfn(func() any {
|
||||
return fmt.Sprint(v...)
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
t.Run("slowfn without output", func(t *testing.T) {
|
||||
w := new(mockWriter)
|
||||
old := writer.Swap(w)
|
||||
defer writer.Store(old)
|
||||
|
||||
doTestStructedLogEmpty(t, w, SevereLevel, func(v ...any) {
|
||||
Slowfn(func() any {
|
||||
return fmt.Sprint(v...)
|
||||
})
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func TestStructedLogSlowv(t *testing.T) {
|
||||
w := new(mockWriter)
|
||||
old := writer.Swap(w)
|
||||
@@ -847,6 +962,16 @@ func doTestStructedLogConsole(t *testing.T, w *mockWriter, write func(...any)) {
|
||||
assert.True(t, strings.Contains(w.String(), message))
|
||||
}
|
||||
|
||||
func doTestStructedLogEmpty(t *testing.T, w *mockWriter, level uint32, write func(...any)) {
|
||||
olevel := atomic.LoadUint32(&logLevel)
|
||||
SetLevel(level)
|
||||
defer SetLevel(olevel)
|
||||
|
||||
const message = "hello there"
|
||||
write(message)
|
||||
assert.Empty(t, w.String())
|
||||
}
|
||||
|
||||
func testSetLevelTwiceWithMode(t *testing.T, mode string, w *mockWriter) {
|
||||
writer.Store(nil)
|
||||
SetUp(LogConf{
|
||||
|
||||
@@ -52,6 +52,12 @@ func (l *richLogger) Debugf(format string, v ...any) {
|
||||
}
|
||||
}
|
||||
|
||||
func (l *richLogger) Debugfn(fn func() any) {
|
||||
if shallLog(DebugLevel) {
|
||||
l.debug(fn())
|
||||
}
|
||||
}
|
||||
|
||||
func (l *richLogger) Debugv(v any) {
|
||||
if shallLog(DebugLevel) {
|
||||
l.debug(v)
|
||||
@@ -76,6 +82,12 @@ func (l *richLogger) Errorf(format string, v ...any) {
|
||||
}
|
||||
}
|
||||
|
||||
func (l *richLogger) Errorfn(fn func() any) {
|
||||
if shallLog(ErrorLevel) {
|
||||
l.err(fn())
|
||||
}
|
||||
}
|
||||
|
||||
func (l *richLogger) Errorv(v any) {
|
||||
if shallLog(ErrorLevel) {
|
||||
l.err(v)
|
||||
@@ -100,6 +112,12 @@ func (l *richLogger) Infof(format string, v ...any) {
|
||||
}
|
||||
}
|
||||
|
||||
func (l *richLogger) Infofn(fn func() any) {
|
||||
if shallLog(InfoLevel) {
|
||||
l.info(fn())
|
||||
}
|
||||
}
|
||||
|
||||
func (l *richLogger) Infov(v any) {
|
||||
if shallLog(InfoLevel) {
|
||||
l.info(v)
|
||||
@@ -124,6 +142,12 @@ func (l *richLogger) Slowf(format string, v ...any) {
|
||||
}
|
||||
}
|
||||
|
||||
func (l *richLogger) Slowfn(fn func() any) {
|
||||
if shallLog(ErrorLevel) {
|
||||
l.slow(fn())
|
||||
}
|
||||
}
|
||||
|
||||
func (l *richLogger) Slowv(v any) {
|
||||
if shallLog(ErrorLevel) {
|
||||
l.slow(v)
|
||||
@@ -182,7 +206,9 @@ func (l *richLogger) WithFields(fields ...LogField) Logger {
|
||||
|
||||
func (l *richLogger) buildFields(fields ...LogField) []LogField {
|
||||
fields = append(l.fields, fields...)
|
||||
// caller field should always appear together with global fields
|
||||
fields = append(fields, Field(callerKey, getCaller(callerDepth+l.callerSkip)))
|
||||
fields = mergeGlobalFields(fields)
|
||||
|
||||
if l.ctx == nil {
|
||||
return fields
|
||||
|
||||
@@ -63,6 +63,11 @@ func TestTraceDebug(t *testing.T) {
|
||||
l.WithDuration(time.Second).Debugf(testlog)
|
||||
validate(t, w.String(), true, true)
|
||||
w.Reset()
|
||||
l.WithDuration(time.Second).Debugfn(func() any {
|
||||
return testlog
|
||||
})
|
||||
validate(t, w.String(), true, true)
|
||||
w.Reset()
|
||||
l.WithDuration(time.Second).Debugv(testlog)
|
||||
validate(t, w.String(), true, true)
|
||||
w.Reset()
|
||||
@@ -103,6 +108,11 @@ func TestTraceError(t *testing.T) {
|
||||
l.WithDuration(time.Second).Errorf(testlog)
|
||||
validate(t, w.String(), true, true)
|
||||
w.Reset()
|
||||
l.WithDuration(time.Second).Errorfn(func() any {
|
||||
return testlog
|
||||
})
|
||||
validate(t, w.String(), true, true)
|
||||
w.Reset()
|
||||
l.WithDuration(time.Second).Errorv(testlog)
|
||||
validate(t, w.String(), true, true)
|
||||
w.Reset()
|
||||
@@ -140,6 +150,11 @@ func TestTraceInfo(t *testing.T) {
|
||||
l.WithDuration(time.Second).Infof(testlog)
|
||||
validate(t, w.String(), true, true)
|
||||
w.Reset()
|
||||
l.WithDuration(time.Second).Infofn(func() any {
|
||||
return testlog
|
||||
})
|
||||
validate(t, w.String(), true, true)
|
||||
w.Reset()
|
||||
l.WithDuration(time.Second).Infov(testlog)
|
||||
validate(t, w.String(), true, true)
|
||||
w.Reset()
|
||||
@@ -213,6 +228,11 @@ func TestTraceSlow(t *testing.T) {
|
||||
l.WithDuration(time.Second).Slowf(testlog)
|
||||
validate(t, w.String(), true, true)
|
||||
w.Reset()
|
||||
l.WithDuration(time.Second).Slowfn(func() any {
|
||||
return testlog
|
||||
})
|
||||
validate(t, w.String(), true, true)
|
||||
w.Reset()
|
||||
l.WithDuration(time.Second).Slowv(testlog)
|
||||
validate(t, w.String(), true, true)
|
||||
w.Reset()
|
||||
|
||||
@@ -17,15 +17,27 @@ import (
|
||||
)
|
||||
|
||||
type (
|
||||
// Writer is the interface for writing logs.
|
||||
// It's designed to let users customize their own log writer,
|
||||
// such as writing logs to a kafka, a database, or using third-party loggers.
|
||||
Writer interface {
|
||||
// Alert sends an alert message, if your writer implemented alerting functionality.
|
||||
Alert(v any)
|
||||
// Close closes the writer.
|
||||
Close() error
|
||||
// Debug logs a message at debug level.
|
||||
Debug(v any, fields ...LogField)
|
||||
// Error logs a message at error level.
|
||||
Error(v any, fields ...LogField)
|
||||
// Info logs a message at info level.
|
||||
Info(v any, fields ...LogField)
|
||||
// Severe logs a message at severe level.
|
||||
Severe(v any)
|
||||
// Slow logs a message at slow level.
|
||||
Slow(v any, fields ...LogField)
|
||||
// Stack logs a message at error level.
|
||||
Stack(v any)
|
||||
// Stat logs a message at stat level.
|
||||
Stat(v any, fields ...LogField)
|
||||
}
|
||||
|
||||
@@ -324,20 +336,6 @@ func buildPlainFields(fields logEntry) []string {
|
||||
return items
|
||||
}
|
||||
|
||||
func combineGlobalFields(fields []LogField) []LogField {
|
||||
globals := globalFields.Load()
|
||||
if globals == nil {
|
||||
return fields
|
||||
}
|
||||
|
||||
gf := globals.([]LogField)
|
||||
ret := make([]LogField, 0, len(gf)+len(fields))
|
||||
ret = append(ret, gf...)
|
||||
ret = append(ret, fields...)
|
||||
|
||||
return ret
|
||||
}
|
||||
|
||||
func marshalJson(t interface{}) ([]byte, error) {
|
||||
var buf bytes.Buffer
|
||||
encoder := json.NewEncoder(&buf)
|
||||
@@ -352,6 +350,20 @@ func marshalJson(t interface{}) ([]byte, error) {
|
||||
return buf.Bytes(), err
|
||||
}
|
||||
|
||||
func mergeGlobalFields(fields []LogField) []LogField {
|
||||
globals := globalFields.Load()
|
||||
if globals == nil {
|
||||
return fields
|
||||
}
|
||||
|
||||
gf := globals.([]LogField)
|
||||
ret := make([]LogField, 0, len(gf)+len(fields))
|
||||
ret = append(ret, gf...)
|
||||
ret = append(ret, fields...)
|
||||
|
||||
return ret
|
||||
}
|
||||
|
||||
func output(writer io.Writer, level string, val any, fields ...LogField) {
|
||||
// only truncate string content, don't know how to truncate the values of other types.
|
||||
if v, ok := val.(string); ok {
|
||||
@@ -362,7 +374,6 @@ func output(writer io.Writer, level string, val any, fields ...LogField) {
|
||||
}
|
||||
}
|
||||
|
||||
fields = combineGlobalFields(fields)
|
||||
// +3 for timestamp, level and content
|
||||
entry := make(logEntry, len(fields)+3)
|
||||
for _, field := range fields {
|
||||
|
||||
@@ -13,6 +13,15 @@ const (
|
||||
|
||||
// Marshal marshals the given val and returns the map that contains the fields.
|
||||
// optional=another is not implemented, and it's hard to implement and not commonly used.
|
||||
// support anonymous field, e.g.:
|
||||
//
|
||||
// type Foo struct {
|
||||
// Token string `header:"token"`
|
||||
// }
|
||||
// type FooB struct {
|
||||
// Foo
|
||||
// Bar string `json:"bar"`
|
||||
// }
|
||||
func Marshal(val any) (map[string]map[string]any, error) {
|
||||
ret := make(map[string]map[string]any)
|
||||
tp := reflect.TypeOf(val)
|
||||
@@ -44,6 +53,16 @@ func getTag(field reflect.StructField) (string, bool) {
|
||||
return strings.TrimSpace(tag), false
|
||||
}
|
||||
|
||||
func insertValue(collector map[string]map[string]any, tag string, key string, val any) {
|
||||
if m, ok := collector[tag]; ok {
|
||||
m[key] = val
|
||||
} else {
|
||||
collector[tag] = map[string]any{
|
||||
key: val,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func processMember(field reflect.StructField, value reflect.Value,
|
||||
collector map[string]map[string]any) error {
|
||||
var key string
|
||||
@@ -69,15 +88,20 @@ func processMember(field reflect.StructField, value reflect.Value,
|
||||
val = fmt.Sprint(val)
|
||||
}
|
||||
|
||||
m, ok := collector[tag]
|
||||
if ok {
|
||||
m[key] = val
|
||||
} else {
|
||||
m = map[string]any{
|
||||
key: val,
|
||||
if field.Anonymous {
|
||||
anonCollector, err := Marshal(val)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for anonTag, anonMap := range anonCollector {
|
||||
for anonKey, anonVal := range anonMap {
|
||||
insertValue(collector, anonTag, anonKey, anonVal)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
insertValue(collector, tag, key, val)
|
||||
}
|
||||
collector[tag] = m
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -27,6 +27,124 @@ func TestMarshal(t *testing.T) {
|
||||
assert.True(t, m[emptyTag]["Anonymous"].(bool))
|
||||
}
|
||||
|
||||
func TestMarshal_Anonymous(t *testing.T) {
|
||||
t.Run("anonymous", func(t *testing.T) {
|
||||
type BaseHeader struct {
|
||||
Token string `header:"token"`
|
||||
}
|
||||
v := struct {
|
||||
Name string `json:"name"`
|
||||
Address string `json:"address,options=[beijing,shanghai]"`
|
||||
Age int `json:"age"`
|
||||
BaseHeader
|
||||
}{
|
||||
Name: "kevin",
|
||||
Address: "shanghai",
|
||||
Age: 20,
|
||||
BaseHeader: BaseHeader{
|
||||
Token: "token_xxx",
|
||||
},
|
||||
}
|
||||
m, err := Marshal(v)
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, "kevin", m["json"]["name"])
|
||||
assert.Equal(t, "shanghai", m["json"]["address"])
|
||||
assert.Equal(t, 20, m["json"]["age"].(int))
|
||||
assert.Equal(t, "token_xxx", m["header"]["token"])
|
||||
|
||||
v1 := struct {
|
||||
Name string `json:"name"`
|
||||
Address string `json:"address,options=[beijing,shanghai]"`
|
||||
Age int `json:"age"`
|
||||
BaseHeader
|
||||
}{
|
||||
Name: "kevin",
|
||||
Address: "shanghai",
|
||||
Age: 20,
|
||||
}
|
||||
m1, err1 := Marshal(v1)
|
||||
assert.Nil(t, err1)
|
||||
assert.Equal(t, "kevin", m1["json"]["name"])
|
||||
assert.Equal(t, "shanghai", m1["json"]["address"])
|
||||
assert.Equal(t, 20, m1["json"]["age"].(int))
|
||||
|
||||
type AnotherHeader struct {
|
||||
Version string `header:"version"`
|
||||
}
|
||||
v2 := struct {
|
||||
Name string `json:"name"`
|
||||
Address string `json:"address,options=[beijing,shanghai]"`
|
||||
Age int `json:"age"`
|
||||
BaseHeader
|
||||
AnotherHeader
|
||||
}{
|
||||
Name: "kevin",
|
||||
Address: "shanghai",
|
||||
Age: 20,
|
||||
BaseHeader: BaseHeader{
|
||||
Token: "token_xxx",
|
||||
},
|
||||
AnotherHeader: AnotherHeader{
|
||||
Version: "v1.0",
|
||||
},
|
||||
}
|
||||
m2, err2 := Marshal(v2)
|
||||
assert.Nil(t, err2)
|
||||
assert.Equal(t, "kevin", m2["json"]["name"])
|
||||
assert.Equal(t, "shanghai", m2["json"]["address"])
|
||||
assert.Equal(t, 20, m2["json"]["age"].(int))
|
||||
assert.Equal(t, "token_xxx", m2["header"]["token"])
|
||||
assert.Equal(t, "v1.0", m2["header"]["version"])
|
||||
|
||||
type PointerHeader struct {
|
||||
Ref *string `header:"ref"`
|
||||
}
|
||||
ref := "reference"
|
||||
v3 := struct {
|
||||
Name string `json:"name"`
|
||||
Address string `json:"address,options=[beijing,shanghai]"`
|
||||
Age int `json:"age"`
|
||||
PointerHeader
|
||||
}{
|
||||
Name: "kevin",
|
||||
Address: "shanghai",
|
||||
Age: 20,
|
||||
PointerHeader: PointerHeader{
|
||||
Ref: &ref,
|
||||
},
|
||||
}
|
||||
m3, err3 := Marshal(v3)
|
||||
assert.Nil(t, err3)
|
||||
assert.Equal(t, "kevin", m3["json"]["name"])
|
||||
assert.Equal(t, "shanghai", m3["json"]["address"])
|
||||
assert.Equal(t, 20, m3["json"]["age"].(int))
|
||||
assert.Equal(t, "reference", *m3["header"]["ref"].(*string))
|
||||
})
|
||||
|
||||
t.Run("bad anonymous", func(t *testing.T) {
|
||||
type BaseHeader struct {
|
||||
Token string `json:"token,options=[a,b]"`
|
||||
}
|
||||
|
||||
v := struct {
|
||||
Name string `json:"name"`
|
||||
Address string `json:"address,options=[beijing,shanghai]"`
|
||||
Age int `json:"age"`
|
||||
BaseHeader
|
||||
}{
|
||||
Name: "kevin",
|
||||
Address: "shanghai",
|
||||
Age: 20,
|
||||
BaseHeader: BaseHeader{
|
||||
Token: "c",
|
||||
},
|
||||
}
|
||||
|
||||
_, err := Marshal(v)
|
||||
assert.NotNil(t, err)
|
||||
})
|
||||
}
|
||||
|
||||
func TestMarshal_Ptr(t *testing.T) {
|
||||
v := &struct {
|
||||
Name string `path:"name"`
|
||||
|
||||
@@ -2,6 +2,7 @@ package mapping
|
||||
|
||||
import (
|
||||
"encoding"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
@@ -80,40 +81,11 @@ func (u *Unmarshaler) Unmarshal(i, v any) error {
|
||||
return u.unmarshal(i, v, "")
|
||||
}
|
||||
|
||||
func (u *Unmarshaler) unmarshal(i, v any, fullName string) error {
|
||||
valueType := reflect.TypeOf(v)
|
||||
if valueType.Kind() != reflect.Ptr {
|
||||
return errValueNotSettable
|
||||
}
|
||||
|
||||
elemType := Deref(valueType)
|
||||
switch iv := i.(type) {
|
||||
case map[string]any:
|
||||
if elemType.Kind() != reflect.Struct {
|
||||
return errTypeMismatch
|
||||
}
|
||||
|
||||
return u.unmarshalValuer(mapValuer(iv), v, fullName)
|
||||
case []any:
|
||||
if elemType.Kind() != reflect.Slice {
|
||||
return errTypeMismatch
|
||||
}
|
||||
|
||||
return u.fillSlice(elemType, reflect.ValueOf(v).Elem(), iv, fullName)
|
||||
default:
|
||||
return errUnsupportedType
|
||||
}
|
||||
}
|
||||
|
||||
// UnmarshalValuer unmarshals m into v.
|
||||
func (u *Unmarshaler) UnmarshalValuer(m Valuer, v any) error {
|
||||
return u.unmarshalValuer(simpleValuer{current: m}, v, "")
|
||||
}
|
||||
|
||||
func (u *Unmarshaler) unmarshalValuer(m Valuer, v any, fullName string) error {
|
||||
return u.unmarshalWithFullName(simpleValuer{current: m}, v, fullName)
|
||||
}
|
||||
|
||||
func (u *Unmarshaler) fillMap(fieldType reflect.Type, value reflect.Value,
|
||||
mapValue any, fullName string) error {
|
||||
if !value.CanSet() {
|
||||
@@ -173,13 +145,14 @@ func (u *Unmarshaler) fillSlice(fieldType reflect.Type, value reflect.Value,
|
||||
baseType := fieldType.Elem()
|
||||
dereffedBaseType := Deref(baseType)
|
||||
dereffedBaseKind := dereffedBaseType.Kind()
|
||||
conv := reflect.MakeSlice(reflect.SliceOf(baseType), refValue.Len(), refValue.Cap())
|
||||
if refValue.Len() == 0 {
|
||||
value.Set(conv)
|
||||
value.Set(reflect.MakeSlice(reflect.SliceOf(baseType), 0, 0))
|
||||
return nil
|
||||
}
|
||||
|
||||
var valid bool
|
||||
conv := reflect.MakeSlice(reflect.SliceOf(baseType), refValue.Len(), refValue.Cap())
|
||||
|
||||
for i := 0; i < refValue.Len(); i++ {
|
||||
ithValue := refValue.Index(i).Interface()
|
||||
if ithValue == nil {
|
||||
@@ -191,17 +164,9 @@ func (u *Unmarshaler) fillSlice(fieldType reflect.Type, value reflect.Value,
|
||||
|
||||
switch dereffedBaseKind {
|
||||
case reflect.Struct:
|
||||
target := reflect.New(dereffedBaseType)
|
||||
val, ok := ithValue.(map[string]any)
|
||||
if !ok {
|
||||
return errTypeMismatch
|
||||
}
|
||||
|
||||
if err := u.unmarshal(val, target.Interface(), sliceFullName); err != nil {
|
||||
if err := u.fillStructElement(baseType, conv.Index(i), ithValue, sliceFullName); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
SetValue(fieldType.Elem(), conv.Index(i), target.Elem())
|
||||
case reflect.Slice:
|
||||
if err := u.fillSlice(dereffedBaseType, conv.Index(i), ithValue, sliceFullName); err != nil {
|
||||
return err
|
||||
@@ -236,7 +201,7 @@ func (u *Unmarshaler) fillSliceFromString(fieldType reflect.Type, value reflect.
|
||||
return errUnsupportedType
|
||||
}
|
||||
|
||||
baseFieldType := Deref(fieldType.Elem())
|
||||
baseFieldType := fieldType.Elem()
|
||||
baseFieldKind := baseFieldType.Kind()
|
||||
conv := reflect.MakeSlice(reflect.SliceOf(baseFieldType), len(slice), cap(slice))
|
||||
|
||||
@@ -257,29 +222,39 @@ func (u *Unmarshaler) fillSliceValue(slice reflect.Value, index int,
|
||||
}
|
||||
|
||||
ithVal := slice.Index(index)
|
||||
ithValType := ithVal.Type()
|
||||
|
||||
switch v := value.(type) {
|
||||
case fmt.Stringer:
|
||||
return setValueFromString(baseKind, ithVal, v.String())
|
||||
case string:
|
||||
return setValueFromString(baseKind, ithVal, v)
|
||||
case map[string]any:
|
||||
return u.fillMap(ithVal.Type(), ithVal, value, fullName)
|
||||
// deref to handle both pointer and non-pointer types.
|
||||
switch Deref(ithValType).Kind() {
|
||||
case reflect.Struct:
|
||||
return u.fillStructElement(ithValType, ithVal, v, fullName)
|
||||
case reflect.Map:
|
||||
return u.fillMap(ithValType, ithVal, value, fullName)
|
||||
default:
|
||||
return errTypeMismatch
|
||||
}
|
||||
default:
|
||||
// don't need to consider the difference between int, int8, int16, int32, int64,
|
||||
// uint, uint8, uint16, uint32, uint64, because they're handled as json.Number.
|
||||
if ithVal.Kind() == reflect.Ptr {
|
||||
baseType := Deref(ithVal.Type())
|
||||
baseType := Deref(ithValType)
|
||||
if !reflect.TypeOf(value).AssignableTo(baseType) {
|
||||
return errTypeMismatch
|
||||
}
|
||||
|
||||
target := reflect.New(baseType).Elem()
|
||||
target.Set(reflect.ValueOf(value))
|
||||
SetValue(ithVal.Type(), ithVal, target)
|
||||
SetValue(ithValType, ithVal, target)
|
||||
return nil
|
||||
}
|
||||
|
||||
if !reflect.TypeOf(value).AssignableTo(ithVal.Type()) {
|
||||
if !reflect.TypeOf(value).AssignableTo(ithValType) {
|
||||
return errTypeMismatch
|
||||
}
|
||||
|
||||
@@ -310,6 +285,23 @@ func (u *Unmarshaler) fillSliceWithDefault(derefedType reflect.Type, value refle
|
||||
return u.fillSlice(derefedType, value, slice, fullName)
|
||||
}
|
||||
|
||||
func (u *Unmarshaler) fillStructElement(baseType reflect.Type, target reflect.Value,
|
||||
value any, fullName string) error {
|
||||
val, ok := value.(map[string]any)
|
||||
if !ok {
|
||||
return errTypeMismatch
|
||||
}
|
||||
|
||||
// use Deref(baseType) to get the base type in case the type is a pointer type.
|
||||
ptr := reflect.New(Deref(baseType))
|
||||
if err := u.unmarshal(val, ptr.Interface(), fullName); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
SetValue(baseType, target, ptr.Elem())
|
||||
return nil
|
||||
}
|
||||
|
||||
func (u *Unmarshaler) fillUnmarshalerStruct(fieldType reflect.Type,
|
||||
value reflect.Value, targetValue string) error {
|
||||
if !value.CanSet() {
|
||||
@@ -612,6 +604,22 @@ func (u *Unmarshaler) processFieldNotFromString(fieldType reflect.Type, value re
|
||||
case valueKind == reflect.String && typeKind == reflect.Map:
|
||||
return u.fillMapFromString(value, mapValue)
|
||||
case valueKind == reflect.String && typeKind == reflect.Slice:
|
||||
// try to find out if it's a byte slice,
|
||||
// more details https://pkg.go.dev/encoding/json#Marshal
|
||||
// array and slice values encode as JSON arrays,
|
||||
// except that []byte encodes as a base64-encoded string,
|
||||
// and a nil slice encoded as the null JSON value.
|
||||
// https://stackoverflow.com/questions/34089750/marshal-byte-to-json-giving-a-strange-string
|
||||
if fieldType.Elem().Kind() == reflect.Uint8 {
|
||||
// check whether string type, because the kind of some other types can be string
|
||||
if strVal, ok := mapValue.(string); ok {
|
||||
if decodedBytes, err := base64.StdEncoding.DecodeString(strVal); err == nil {
|
||||
value.Set(reflect.ValueOf(decodedBytes))
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return u.fillSliceFromString(fieldType, value, mapValue, fullName)
|
||||
case valueKind == reflect.String && derefedFieldType == durationType:
|
||||
return fillDurationValue(fieldType, value, mapValue.(string))
|
||||
@@ -952,6 +960,35 @@ func (u *Unmarshaler) processNamedFieldWithoutValue(fieldType reflect.Type, valu
|
||||
return nil
|
||||
}
|
||||
|
||||
func (u *Unmarshaler) unmarshal(i, v any, fullName string) error {
|
||||
valueType := reflect.TypeOf(v)
|
||||
if valueType.Kind() != reflect.Ptr {
|
||||
return errValueNotSettable
|
||||
}
|
||||
|
||||
elemType := Deref(valueType)
|
||||
switch iv := i.(type) {
|
||||
case map[string]any:
|
||||
if elemType.Kind() != reflect.Struct {
|
||||
return errTypeMismatch
|
||||
}
|
||||
|
||||
return u.unmarshalValuer(mapValuer(iv), v, fullName)
|
||||
case []any:
|
||||
if elemType.Kind() != reflect.Slice {
|
||||
return errTypeMismatch
|
||||
}
|
||||
|
||||
return u.fillSlice(elemType, reflect.ValueOf(v).Elem(), iv, fullName)
|
||||
default:
|
||||
return errUnsupportedType
|
||||
}
|
||||
}
|
||||
|
||||
func (u *Unmarshaler) unmarshalValuer(m Valuer, v any, fullName string) error {
|
||||
return u.unmarshalWithFullName(simpleValuer{current: m}, v, fullName)
|
||||
}
|
||||
|
||||
func (u *Unmarshaler) unmarshalWithFullName(m valuerWithParent, v any, fullName string) error {
|
||||
rv := reflect.ValueOf(v)
|
||||
if err := ValidatePtr(rv); err != nil {
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/zeromicro/go-zero/core/jsonx"
|
||||
"github.com/zeromicro/go-zero/core/stringx"
|
||||
)
|
||||
|
||||
@@ -351,7 +352,7 @@ func TestUnmarshalIntSliceOfPtr(t *testing.T) {
|
||||
assert.Error(t, UnmarshalKey(m, &in))
|
||||
})
|
||||
|
||||
t.Run("int slice with nil", func(t *testing.T) {
|
||||
t.Run("int slice with nil element", func(t *testing.T) {
|
||||
type inner struct {
|
||||
Ints []int `key:"ints"`
|
||||
}
|
||||
@@ -365,6 +366,21 @@ func TestUnmarshalIntSliceOfPtr(t *testing.T) {
|
||||
assert.Empty(t, in.Ints)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("int slice with nil", func(t *testing.T) {
|
||||
type inner struct {
|
||||
Ints []int `key:"ints"`
|
||||
}
|
||||
|
||||
m := map[string]any{
|
||||
"ints": []any(nil),
|
||||
}
|
||||
|
||||
var in inner
|
||||
if assert.NoError(t, UnmarshalKey(m, &in)) {
|
||||
assert.Empty(t, in.Ints)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestUnmarshalIntWithDefault(t *testing.T) {
|
||||
@@ -1374,20 +1390,80 @@ func TestUnmarshalWithFloatPtr(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestUnmarshalIntSlice(t *testing.T) {
|
||||
var v struct {
|
||||
Ages []int `key:"ages"`
|
||||
Slice []int `key:"slice"`
|
||||
}
|
||||
m := map[string]any{
|
||||
"ages": []int{1, 2},
|
||||
"slice": []any{},
|
||||
}
|
||||
t.Run("int slice from int", func(t *testing.T) {
|
||||
var v struct {
|
||||
Ages []int `key:"ages"`
|
||||
Slice []int `key:"slice"`
|
||||
}
|
||||
m := map[string]any{
|
||||
"ages": []int{1, 2},
|
||||
"slice": []any{},
|
||||
}
|
||||
|
||||
ast := assert.New(t)
|
||||
if ast.NoError(UnmarshalKey(m, &v)) {
|
||||
ast.ElementsMatch([]int{1, 2}, v.Ages)
|
||||
ast.Equal([]int{}, v.Slice)
|
||||
}
|
||||
ast := assert.New(t)
|
||||
if ast.NoError(UnmarshalKey(m, &v)) {
|
||||
ast.ElementsMatch([]int{1, 2}, v.Ages)
|
||||
ast.Equal([]int{}, v.Slice)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("int slice from one int", func(t *testing.T) {
|
||||
var v struct {
|
||||
Ages []int `key:"ages"`
|
||||
}
|
||||
m := map[string]any{
|
||||
"ages": []int{2},
|
||||
}
|
||||
|
||||
ast := assert.New(t)
|
||||
unmarshaler := NewUnmarshaler(defaultKeyName, WithFromArray())
|
||||
if ast.NoError(unmarshaler.Unmarshal(m, &v)) {
|
||||
ast.ElementsMatch([]int{2}, v.Ages)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("int slice from one int string", func(t *testing.T) {
|
||||
var v struct {
|
||||
Ages []int `key:"ages"`
|
||||
}
|
||||
m := map[string]any{
|
||||
"ages": []string{"2"},
|
||||
}
|
||||
|
||||
ast := assert.New(t)
|
||||
unmarshaler := NewUnmarshaler(defaultKeyName, WithFromArray())
|
||||
if ast.NoError(unmarshaler.Unmarshal(m, &v)) {
|
||||
ast.ElementsMatch([]int{2}, v.Ages)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("int slice from one json.Number", func(t *testing.T) {
|
||||
var v struct {
|
||||
Ages []int `key:"ages"`
|
||||
}
|
||||
m := map[string]any{
|
||||
"ages": []json.Number{"2"},
|
||||
}
|
||||
|
||||
ast := assert.New(t)
|
||||
unmarshaler := NewUnmarshaler(defaultKeyName, WithFromArray())
|
||||
if ast.NoError(unmarshaler.Unmarshal(m, &v)) {
|
||||
ast.ElementsMatch([]int{2}, v.Ages)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("int slice from one int strings", func(t *testing.T) {
|
||||
var v struct {
|
||||
Ages []int `key:"ages"`
|
||||
}
|
||||
m := map[string]any{
|
||||
"ages": []string{"1,2"},
|
||||
}
|
||||
|
||||
ast := assert.New(t)
|
||||
unmarshaler := NewUnmarshaler(defaultKeyName, WithFromArray())
|
||||
ast.Error(unmarshaler.Unmarshal(m, &v))
|
||||
})
|
||||
}
|
||||
|
||||
func TestUnmarshalString(t *testing.T) {
|
||||
@@ -1442,6 +1518,51 @@ func TestUnmarshalStringSliceFromString(t *testing.T) {
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("slice from empty string", func(t *testing.T) {
|
||||
var v struct {
|
||||
Names []string `key:"names"`
|
||||
}
|
||||
m := map[string]any{
|
||||
"names": []string{""},
|
||||
}
|
||||
|
||||
ast := assert.New(t)
|
||||
unmarshaler := NewUnmarshaler(defaultKeyName, WithFromArray())
|
||||
if ast.NoError(unmarshaler.Unmarshal(m, &v)) {
|
||||
ast.ElementsMatch([]string{""}, v.Names)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("slice from empty and valid string", func(t *testing.T) {
|
||||
var v struct {
|
||||
Names []string `key:"names"`
|
||||
}
|
||||
m := map[string]any{
|
||||
"names": []string{","},
|
||||
}
|
||||
|
||||
ast := assert.New(t)
|
||||
unmarshaler := NewUnmarshaler(defaultKeyName, WithFromArray())
|
||||
if ast.NoError(unmarshaler.Unmarshal(m, &v)) {
|
||||
ast.ElementsMatch([]string{","}, v.Names)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("slice from valid strings with comma", func(t *testing.T) {
|
||||
var v struct {
|
||||
Names []string `key:"names"`
|
||||
}
|
||||
m := map[string]any{
|
||||
"names": []string{"aa,bb"},
|
||||
}
|
||||
|
||||
ast := assert.New(t)
|
||||
unmarshaler := NewUnmarshaler(defaultKeyName, WithFromArray())
|
||||
if ast.NoError(unmarshaler.Unmarshal(m, &v)) {
|
||||
ast.ElementsMatch([]string{"aa,bb"}, v.Names)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("slice from string with slice error", func(t *testing.T) {
|
||||
var v struct {
|
||||
Names []int `key:"names"`
|
||||
@@ -4761,14 +4882,28 @@ func TestUnmarshal_EnvWithOptionsWrongValueString(t *testing.T) {
|
||||
|
||||
func TestUnmarshalJsonReaderMultiArray(t *testing.T) {
|
||||
t.Run("reader multi array", func(t *testing.T) {
|
||||
var res struct {
|
||||
type testRes struct {
|
||||
A string `json:"a"`
|
||||
B [][]string `json:"b"`
|
||||
C []byte `json:"c"`
|
||||
}
|
||||
payload := `{"a": "133", "b": [["add", "cccd"], ["eeee"]]}`
|
||||
|
||||
var res testRes
|
||||
marshal := testRes{
|
||||
A: "133",
|
||||
B: [][]string{
|
||||
{"add", "cccd"},
|
||||
{"eeee"},
|
||||
},
|
||||
C: []byte("11122344wsss"),
|
||||
}
|
||||
bytes, err := jsonx.Marshal(marshal)
|
||||
assert.NoError(t, err)
|
||||
payload := string(bytes)
|
||||
reader := strings.NewReader(payload)
|
||||
if assert.NoError(t, UnmarshalJsonReader(reader, &res)) {
|
||||
assert.Equal(t, 2, len(res.B))
|
||||
assert.Equal(t, string(marshal.C), string(res.C))
|
||||
}
|
||||
})
|
||||
|
||||
@@ -5862,6 +5997,38 @@ func TestUnmarshal_Unmarshaler(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestParseJsonStringValue(t *testing.T) {
|
||||
t.Run("string", func(t *testing.T) {
|
||||
type GoodsInfo struct {
|
||||
Sku int64 `json:"sku,optional"`
|
||||
}
|
||||
|
||||
type GetReq struct {
|
||||
GoodsList []*GoodsInfo `json:"goods_list"`
|
||||
}
|
||||
|
||||
input := map[string]any{"goods_list": "[{\"sku\":11},{\"sku\":22}]"}
|
||||
var v GetReq
|
||||
assert.NotPanics(t, func() {
|
||||
assert.NoError(t, UnmarshalJsonMap(input, &v))
|
||||
assert.Equal(t, 2, len(v.GoodsList))
|
||||
assert.ElementsMatch(t, []int64{11, 22}, []int64{v.GoodsList[0].Sku, v.GoodsList[1].Sku})
|
||||
})
|
||||
})
|
||||
|
||||
t.Run("string with invalid type", func(t *testing.T) {
|
||||
type GetReq struct {
|
||||
GoodsList []*int `json:"goods_list"`
|
||||
}
|
||||
|
||||
input := map[string]any{"goods_list": "[{\"sku\":11},{\"sku\":22}]"}
|
||||
var v GetReq
|
||||
assert.NotPanics(t, func() {
|
||||
assert.Error(t, UnmarshalJsonMap(input, &v))
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func BenchmarkDefaultValue(b *testing.B) {
|
||||
for i := 0; i < b.N; i++ {
|
||||
var a struct {
|
||||
|
||||
@@ -142,89 +142,6 @@ func MapReduceChan[T, U, V any](source <-chan T, mapper MapperFunc[T, U], reduce
|
||||
return mapReduceWithPanicChan(source, panicChan, mapper, reducer, opts...)
|
||||
}
|
||||
|
||||
// mapReduceWithPanicChan maps all elements from source, and reduce the output elements with given reducer.
|
||||
func mapReduceWithPanicChan[T, U, V any](source <-chan T, panicChan *onceChan, mapper MapperFunc[T, U],
|
||||
reducer ReducerFunc[U, V], opts ...Option) (val V, err error) {
|
||||
options := buildOptions(opts...)
|
||||
// output is used to write the final result
|
||||
output := make(chan V)
|
||||
defer func() {
|
||||
// reducer can only write once, if more, panic
|
||||
for range output {
|
||||
panic("more than one element written in reducer")
|
||||
}
|
||||
}()
|
||||
|
||||
// collector is used to collect data from mapper, and consume in reducer
|
||||
collector := make(chan U, options.workers)
|
||||
// if done is closed, all mappers and reducer should stop processing
|
||||
done := make(chan struct{})
|
||||
writer := newGuardedWriter(options.ctx, output, done)
|
||||
var closeOnce sync.Once
|
||||
// use atomic type to avoid data race
|
||||
var retErr errorx.AtomicError
|
||||
finish := func() {
|
||||
closeOnce.Do(func() {
|
||||
close(done)
|
||||
close(output)
|
||||
})
|
||||
}
|
||||
cancel := once(func(err error) {
|
||||
if err != nil {
|
||||
retErr.Set(err)
|
||||
} else {
|
||||
retErr.Set(ErrCancelWithNil)
|
||||
}
|
||||
|
||||
drain(source)
|
||||
finish()
|
||||
})
|
||||
|
||||
go func() {
|
||||
defer func() {
|
||||
drain(collector)
|
||||
if r := recover(); r != nil {
|
||||
panicChan.write(r)
|
||||
}
|
||||
finish()
|
||||
}()
|
||||
|
||||
reducer(collector, writer, cancel)
|
||||
}()
|
||||
|
||||
go executeMappers(mapperContext[T, U]{
|
||||
ctx: options.ctx,
|
||||
mapper: func(item T, w Writer[U]) {
|
||||
mapper(item, w, cancel)
|
||||
},
|
||||
source: source,
|
||||
panicChan: panicChan,
|
||||
collector: collector,
|
||||
doneChan: done,
|
||||
workers: options.workers,
|
||||
})
|
||||
|
||||
select {
|
||||
case <-options.ctx.Done():
|
||||
cancel(context.DeadlineExceeded)
|
||||
err = context.DeadlineExceeded
|
||||
case v := <-panicChan.channel:
|
||||
// drain output here, otherwise for loop panic in defer
|
||||
drain(output)
|
||||
panic(v)
|
||||
case v, ok := <-output:
|
||||
if e := retErr.Load(); e != nil {
|
||||
err = e
|
||||
} else if ok {
|
||||
val = v
|
||||
} else {
|
||||
err = ErrReduceNoOutput
|
||||
}
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
// MapReduceVoid maps all elements generated from given generate,
|
||||
// and reduce the output elements with given reducer.
|
||||
func MapReduceVoid[T, U any](generate GenerateFunc[T], mapper MapperFunc[T, U],
|
||||
@@ -330,6 +247,89 @@ func executeMappers[T, U any](mCtx mapperContext[T, U]) {
|
||||
}
|
||||
}
|
||||
|
||||
// mapReduceWithPanicChan maps all elements from source, and reduce the output elements with given reducer.
|
||||
func mapReduceWithPanicChan[T, U, V any](source <-chan T, panicChan *onceChan, mapper MapperFunc[T, U],
|
||||
reducer ReducerFunc[U, V], opts ...Option) (val V, err error) {
|
||||
options := buildOptions(opts...)
|
||||
// output is used to write the final result
|
||||
output := make(chan V)
|
||||
defer func() {
|
||||
// reducer can only write once, if more, panic
|
||||
for range output {
|
||||
panic("more than one element written in reducer")
|
||||
}
|
||||
}()
|
||||
|
||||
// collector is used to collect data from mapper, and consume in reducer
|
||||
collector := make(chan U, options.workers)
|
||||
// if done is closed, all mappers and reducer should stop processing
|
||||
done := make(chan struct{})
|
||||
writer := newGuardedWriter(options.ctx, output, done)
|
||||
var closeOnce sync.Once
|
||||
// use atomic type to avoid data race
|
||||
var retErr errorx.AtomicError
|
||||
finish := func() {
|
||||
closeOnce.Do(func() {
|
||||
close(done)
|
||||
close(output)
|
||||
})
|
||||
}
|
||||
cancel := once(func(err error) {
|
||||
if err != nil {
|
||||
retErr.Set(err)
|
||||
} else {
|
||||
retErr.Set(ErrCancelWithNil)
|
||||
}
|
||||
|
||||
drain(source)
|
||||
finish()
|
||||
})
|
||||
|
||||
go func() {
|
||||
defer func() {
|
||||
drain(collector)
|
||||
if r := recover(); r != nil {
|
||||
panicChan.write(r)
|
||||
}
|
||||
finish()
|
||||
}()
|
||||
|
||||
reducer(collector, writer, cancel)
|
||||
}()
|
||||
|
||||
go executeMappers(mapperContext[T, U]{
|
||||
ctx: options.ctx,
|
||||
mapper: func(item T, w Writer[U]) {
|
||||
mapper(item, w, cancel)
|
||||
},
|
||||
source: source,
|
||||
panicChan: panicChan,
|
||||
collector: collector,
|
||||
doneChan: done,
|
||||
workers: options.workers,
|
||||
})
|
||||
|
||||
select {
|
||||
case <-options.ctx.Done():
|
||||
cancel(context.DeadlineExceeded)
|
||||
err = context.DeadlineExceeded
|
||||
case v := <-panicChan.channel:
|
||||
// drain output here, otherwise for loop panic in defer
|
||||
drain(output)
|
||||
panic(v)
|
||||
case v, ok := <-output:
|
||||
if e := retErr.Load(); e != nil {
|
||||
err = e
|
||||
} else if ok {
|
||||
val = v
|
||||
} else {
|
||||
err = ErrReduceNoOutput
|
||||
}
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func newOptions() *mapReduceOptions {
|
||||
return &mapReduceOptions{
|
||||
ctx: context.Background(),
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//go:build linux || darwin
|
||||
//go:build linux || darwin || freebsd
|
||||
|
||||
package proc
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//go:build linux || darwin
|
||||
//go:build linux || darwin || freebsd
|
||||
|
||||
package proc
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//go:build linux || darwin
|
||||
//go:build linux || darwin || freebsd
|
||||
|
||||
package proc
|
||||
|
||||
|
||||
@@ -4,6 +4,9 @@ package proc
|
||||
|
||||
import "time"
|
||||
|
||||
// ShutdownConf is empty on windows.
|
||||
type ShutdownConf struct{}
|
||||
|
||||
// AddShutdownListener returns fn itself on windows, lets callers call fn on their own.
|
||||
func AddShutdownListener(fn func()) func() {
|
||||
return fn
|
||||
@@ -18,6 +21,10 @@ func AddWrapUpListener(fn func()) func() {
|
||||
func SetTimeToForceQuit(duration time.Duration) {
|
||||
}
|
||||
|
||||
// Setup does nothing on windows.
|
||||
func Setup(conf ShutdownConf) {
|
||||
}
|
||||
|
||||
// Shutdown does nothing on windows.
|
||||
func Shutdown() {
|
||||
}
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//go:build linux || darwin
|
||||
//go:build linux || darwin || freebsd
|
||||
|
||||
package proc
|
||||
|
||||
@@ -14,17 +14,29 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
wrapUpTime = time.Second
|
||||
// why we use 5500 milliseconds is because most of our queue are blocking mode with 5 seconds
|
||||
waitTime = 5500 * time.Millisecond
|
||||
// defaultWrapUpTime is the default time to wait before calling wrap up listeners.
|
||||
defaultWrapUpTime = time.Second
|
||||
// defaultWaitTime is the default time to wait before force quitting.
|
||||
// why we use 5500 milliseconds is because most of our queues are blocking mode with 5 seconds
|
||||
defaultWaitTime = 5500 * time.Millisecond
|
||||
)
|
||||
|
||||
var (
|
||||
wrapUpListeners = new(listenerManager)
|
||||
shutdownListeners = new(listenerManager)
|
||||
delayTimeBeforeForceQuit = waitTime
|
||||
wrapUpListeners = new(listenerManager)
|
||||
shutdownListeners = new(listenerManager)
|
||||
wrapUpTime = defaultWrapUpTime
|
||||
waitTime = defaultWaitTime
|
||||
shutdownLock sync.Mutex
|
||||
)
|
||||
|
||||
// ShutdownConf defines the shutdown configuration for the process.
|
||||
type ShutdownConf struct {
|
||||
// WrapUpTime is the time to wait before calling shutdown listeners.
|
||||
WrapUpTime time.Duration `json:",default=1s"`
|
||||
// WaitTime is the time to wait before force quitting.
|
||||
WaitTime time.Duration `json:",default=5.5s"`
|
||||
}
|
||||
|
||||
// AddShutdownListener adds fn as a shutdown listener.
|
||||
// The returned func can be used to wait for fn getting called.
|
||||
func AddShutdownListener(fn func()) (waitForCalled func()) {
|
||||
@@ -39,7 +51,21 @@ func AddWrapUpListener(fn func()) (waitForCalled func()) {
|
||||
|
||||
// SetTimeToForceQuit sets the waiting time before force quitting.
|
||||
func SetTimeToForceQuit(duration time.Duration) {
|
||||
delayTimeBeforeForceQuit = duration
|
||||
shutdownLock.Lock()
|
||||
defer shutdownLock.Unlock()
|
||||
waitTime = duration
|
||||
}
|
||||
|
||||
func Setup(conf ShutdownConf) {
|
||||
shutdownLock.Lock()
|
||||
defer shutdownLock.Unlock()
|
||||
|
||||
if conf.WrapUpTime > 0 {
|
||||
wrapUpTime = conf.WrapUpTime
|
||||
}
|
||||
if conf.WaitTime > 0 {
|
||||
waitTime = conf.WaitTime
|
||||
}
|
||||
}
|
||||
|
||||
// Shutdown calls the registered shutdown listeners, only for test purpose.
|
||||
@@ -61,8 +87,12 @@ func gracefulStop(signals chan os.Signal, sig syscall.Signal) {
|
||||
time.Sleep(wrapUpTime)
|
||||
go shutdownListeners.notifyListeners()
|
||||
|
||||
time.Sleep(delayTimeBeforeForceQuit - wrapUpTime)
|
||||
logx.Infof("Still alive after %v, going to force kill the process...", delayTimeBeforeForceQuit)
|
||||
shutdownLock.Lock()
|
||||
remainingTime := waitTime - wrapUpTime
|
||||
shutdownLock.Unlock()
|
||||
|
||||
time.Sleep(remainingTime)
|
||||
logx.Infof("Still alive after %v, going to force kill the process...", waitTime)
|
||||
_ = syscall.Kill(syscall.Getpid(), sig)
|
||||
}
|
||||
|
||||
@@ -82,6 +112,9 @@ func (lm *listenerManager) addListener(fn func()) (waitForCalled func()) {
|
||||
})
|
||||
lm.lock.Unlock()
|
||||
|
||||
// we can return lm.waitGroup.Wait directly,
|
||||
// but we want to make the returned func more readable.
|
||||
// creating an extra closure would be negligible in practice.
|
||||
return func() {
|
||||
lm.waitGroup.Wait()
|
||||
}
|
||||
|
||||
@@ -1,8 +1,9 @@
|
||||
//go:build linux || darwin
|
||||
//go:build linux || darwin || freebsd
|
||||
|
||||
package proc
|
||||
|
||||
import (
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -10,8 +11,12 @@ import (
|
||||
)
|
||||
|
||||
func TestShutdown(t *testing.T) {
|
||||
t.Cleanup(restoreSettings)
|
||||
|
||||
SetTimeToForceQuit(time.Hour)
|
||||
assert.Equal(t, time.Hour, delayTimeBeforeForceQuit)
|
||||
shutdownLock.Lock()
|
||||
assert.Equal(t, time.Hour, waitTime)
|
||||
shutdownLock.Unlock()
|
||||
|
||||
var val int
|
||||
called := AddWrapUpListener(func() {
|
||||
@@ -29,7 +34,53 @@ func TestShutdown(t *testing.T) {
|
||||
assert.Equal(t, 3, val)
|
||||
}
|
||||
|
||||
func TestShutdownWithMultipleServices(t *testing.T) {
|
||||
t.Cleanup(restoreSettings)
|
||||
|
||||
SetTimeToForceQuit(time.Hour)
|
||||
shutdownLock.Lock()
|
||||
assert.Equal(t, time.Hour, waitTime)
|
||||
shutdownLock.Unlock()
|
||||
|
||||
var val int32
|
||||
called1 := AddShutdownListener(func() {
|
||||
atomic.AddInt32(&val, 1)
|
||||
})
|
||||
called2 := AddShutdownListener(func() {
|
||||
atomic.AddInt32(&val, 2)
|
||||
})
|
||||
Shutdown()
|
||||
called1()
|
||||
called2()
|
||||
|
||||
assert.Equal(t, int32(3), atomic.LoadInt32(&val))
|
||||
}
|
||||
|
||||
func TestWrapUpWithMultipleServices(t *testing.T) {
|
||||
t.Cleanup(restoreSettings)
|
||||
|
||||
SetTimeToForceQuit(time.Hour)
|
||||
shutdownLock.Lock()
|
||||
assert.Equal(t, time.Hour, waitTime)
|
||||
shutdownLock.Unlock()
|
||||
|
||||
var val int32
|
||||
called1 := AddWrapUpListener(func() {
|
||||
atomic.AddInt32(&val, 1)
|
||||
})
|
||||
called2 := AddWrapUpListener(func() {
|
||||
atomic.AddInt32(&val, 2)
|
||||
})
|
||||
WrapUp()
|
||||
called1()
|
||||
called2()
|
||||
|
||||
assert.Equal(t, int32(3), atomic.LoadInt32(&val))
|
||||
}
|
||||
|
||||
func TestNotifyMoreThanOnce(t *testing.T) {
|
||||
t.Cleanup(restoreSettings)
|
||||
|
||||
ch := make(chan struct{}, 1)
|
||||
|
||||
go func() {
|
||||
@@ -58,3 +109,38 @@ func TestNotifyMoreThanOnce(t *testing.T) {
|
||||
t.Fatal("timeout, check error logs")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetup(t *testing.T) {
|
||||
t.Run("valid time", func(t *testing.T) {
|
||||
defer restoreSettings()
|
||||
|
||||
Setup(ShutdownConf{
|
||||
WrapUpTime: time.Second * 2,
|
||||
WaitTime: time.Second * 30,
|
||||
})
|
||||
|
||||
shutdownLock.Lock()
|
||||
assert.Equal(t, time.Second*2, wrapUpTime)
|
||||
assert.Equal(t, time.Second*30, waitTime)
|
||||
shutdownLock.Unlock()
|
||||
})
|
||||
|
||||
t.Run("valid time", func(t *testing.T) {
|
||||
defer restoreSettings()
|
||||
|
||||
Setup(ShutdownConf{})
|
||||
|
||||
shutdownLock.Lock()
|
||||
assert.Equal(t, defaultWrapUpTime, wrapUpTime)
|
||||
assert.Equal(t, defaultWaitTime, waitTime)
|
||||
shutdownLock.Unlock()
|
||||
})
|
||||
}
|
||||
|
||||
func restoreSettings() {
|
||||
shutdownLock.Lock()
|
||||
defer shutdownLock.Unlock()
|
||||
|
||||
wrapUpTime = defaultWrapUpTime
|
||||
waitTime = defaultWaitTime
|
||||
}
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//go:build linux || darwin
|
||||
//go:build linux || darwin || freebsd
|
||||
|
||||
package proc
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//go:build linux || darwin
|
||||
//go:build linux || darwin || freebsd
|
||||
|
||||
package proc
|
||||
|
||||
|
||||
@@ -37,6 +37,7 @@ type (
|
||||
Prometheus prometheus.Config `json:",optional"`
|
||||
Telemetry trace.Config `json:",optional"`
|
||||
DevServer DevServerConfig `json:",optional"`
|
||||
Shutdown proc.ShutdownConf `json:",optional"`
|
||||
}
|
||||
)
|
||||
|
||||
@@ -61,6 +62,7 @@ func (sc ServiceConf) SetUp() error {
|
||||
sc.Telemetry.Name = sc.Name
|
||||
}
|
||||
trace.StartAgent(sc.Telemetry)
|
||||
proc.Setup(sc.Shutdown)
|
||||
proc.AddShutdownListener(func() {
|
||||
trace.StopAgent()
|
||||
})
|
||||
|
||||
@@ -76,9 +76,14 @@ func (sg *ServiceGroup) doStart() {
|
||||
}
|
||||
|
||||
func (sg *ServiceGroup) doStop() {
|
||||
group := threading.NewRoutineGroup()
|
||||
for _, service := range sg.services {
|
||||
service.Stop()
|
||||
// new variable to avoid closure problems, can be removed after go 1.22
|
||||
// see https://golang.org/doc/faq#closures_and_goroutines
|
||||
service := service
|
||||
group.Run(service.Stop)
|
||||
}
|
||||
group.Wait()
|
||||
}
|
||||
|
||||
// WithStart wraps a start func as a Service.
|
||||
|
||||
@@ -19,6 +19,7 @@ type (
|
||||
RedisConf struct {
|
||||
Host string
|
||||
Type string `json:",default=node,options=node|cluster"`
|
||||
User string `json:",optional"`
|
||||
Pass string `json:",optional"`
|
||||
Tls bool `json:",optional"`
|
||||
NonBlock bool `json:",default=true"`
|
||||
@@ -40,6 +41,9 @@ func (rc RedisConf) NewRedis() *Redis {
|
||||
if rc.Type == ClusterType {
|
||||
opts = append(opts, Cluster())
|
||||
}
|
||||
if len(rc.User) > 0 {
|
||||
opts = append(opts, WithUser(rc.User))
|
||||
}
|
||||
if len(rc.Pass) > 0 {
|
||||
opts = append(opts, WithPass(rc.Pass))
|
||||
}
|
||||
|
||||
@@ -55,6 +55,7 @@ type (
|
||||
Redis struct {
|
||||
Addr string
|
||||
Type string
|
||||
User string
|
||||
Pass string
|
||||
tls bool
|
||||
brk breaker.Breaker
|
||||
@@ -126,6 +127,9 @@ func NewRedis(conf RedisConf, opts ...Option) (*Redis, error) {
|
||||
if conf.Type == ClusterType {
|
||||
opts = append([]Option{Cluster()}, opts...)
|
||||
}
|
||||
if len(conf.User) > 0 {
|
||||
opts = append([]Option{WithUser(conf.User)}, opts...)
|
||||
}
|
||||
if len(conf.Pass) > 0 {
|
||||
opts = append([]Option{WithPass(conf.Pass)}, opts...)
|
||||
}
|
||||
@@ -605,6 +609,28 @@ func (s *Redis) GetBitCtx(ctx context.Context, key string, offset int64) (int, e
|
||||
return int(v), nil
|
||||
}
|
||||
|
||||
// GetDel is the implementation of redis getdel command.
|
||||
// Available since: redis version 6.2.0
|
||||
func (s *Redis) GetDel(key string) (string, error) {
|
||||
return s.GetDelCtx(context.Background(), key)
|
||||
}
|
||||
|
||||
// GetDelCtx is the implementation of redis getdel command.
|
||||
// Available since: redis version 6.2.0
|
||||
func (s *Redis) GetDelCtx(ctx context.Context, key string) (string, error) {
|
||||
conn, err := getRedis(s)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
val, err := conn.GetDel(ctx, key).Result()
|
||||
if errors.Is(err, red.Nil) {
|
||||
return "", nil
|
||||
}
|
||||
|
||||
return val, err
|
||||
}
|
||||
|
||||
// GetSet is the implementation of redis getset command.
|
||||
func (s *Redis) GetSet(key, value string) (string, error) {
|
||||
return s.GetSetCtx(context.Background(), key, value)
|
||||
@@ -2405,6 +2431,13 @@ func SetSlowThreshold(threshold time.Duration) {
|
||||
slowThreshold.Set(threshold)
|
||||
}
|
||||
|
||||
// WithHook customizes the given Redis with given durationHook.
|
||||
func WithHook(hook Hook) Option {
|
||||
return func(r *Redis) {
|
||||
r.hooks = append(r.hooks, hook)
|
||||
}
|
||||
}
|
||||
|
||||
// WithPass customizes the given Redis with given password.
|
||||
func WithPass(pass string) Option {
|
||||
return func(r *Redis) {
|
||||
@@ -2419,11 +2452,10 @@ func WithTLS() Option {
|
||||
}
|
||||
}
|
||||
|
||||
// WithHook customizes the given Redis with given durationHook, only for private use now,
|
||||
// maybe expose later.
|
||||
func WithHook(hook Hook) Option {
|
||||
// WithUser customizes the given Redis with given username.
|
||||
func WithUser(user string) Option {
|
||||
return func(r *Redis) {
|
||||
r.hooks = append(r.hooks, hook)
|
||||
r.User = user
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1071,6 +1071,34 @@ func TestRedis_Set(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestRedis_GetDel(t *testing.T) {
|
||||
t.Run("get_del", func(t *testing.T) {
|
||||
runOnRedis(t, func(client *Redis) {
|
||||
val, err := newRedis(client.Addr).GetDel("hello")
|
||||
assert.Equal(t, "", val)
|
||||
assert.Nil(t, err)
|
||||
err = client.Set("hello", "world")
|
||||
assert.Nil(t, err)
|
||||
val, err = client.Get("hello")
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, "world", val)
|
||||
val, err = client.GetDel("hello")
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, "world", val)
|
||||
val, err = client.Get("hello")
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, "", val)
|
||||
})
|
||||
})
|
||||
|
||||
t.Run("get_del_with_error", func(t *testing.T) {
|
||||
runOnRedisWithError(t, func(client *Redis) {
|
||||
_, err := newRedis(client.Addr, badType()).GetDel("hello")
|
||||
assert.Error(t, err)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func TestRedis_GetSet(t *testing.T) {
|
||||
t.Run("set_get", func(t *testing.T) {
|
||||
runOnRedis(t, func(client *Redis) {
|
||||
@@ -1996,9 +2024,9 @@ func TestSetSlowThreshold(t *testing.T) {
|
||||
assert.Equal(t, time.Second, slowThreshold.Load())
|
||||
}
|
||||
|
||||
func TestRedis_WithPass(t *testing.T) {
|
||||
func TestRedis_WithUserPass(t *testing.T) {
|
||||
runOnRedis(t, func(client *Redis) {
|
||||
err := newRedis(client.Addr, WithPass("any")).Ping()
|
||||
err := newRedis(client.Addr, WithUser("any"), WithPass("any")).Ping()
|
||||
assert.NotNil(t, err)
|
||||
})
|
||||
}
|
||||
@@ -2119,9 +2147,9 @@ func TestRedisUnlink(t *testing.T) {
|
||||
func TestRedisTxPipeline(t *testing.T) {
|
||||
runOnRedis(t, func(client *Redis) {
|
||||
ctx := context.Background()
|
||||
pipe, err := newRedis(client.Addr, badType()).TxPipeline()
|
||||
_, err := newRedis(client.Addr, badType()).TxPipeline()
|
||||
assert.NotNil(t, err)
|
||||
pipe, err = client.TxPipeline()
|
||||
pipe, err := client.TxPipeline()
|
||||
assert.Nil(t, err)
|
||||
key := "key"
|
||||
hashKey := "field"
|
||||
|
||||
@@ -21,6 +21,7 @@ func CreateBlockingNode(r *Redis) (ClosableNode, error) {
|
||||
case NodeType:
|
||||
client := red.NewClient(&red.Options{
|
||||
Addr: r.Addr,
|
||||
Username: r.User,
|
||||
Password: r.Pass,
|
||||
DB: defaultDatabase,
|
||||
MaxRetries: maxRetries,
|
||||
@@ -32,6 +33,7 @@ func CreateBlockingNode(r *Redis) (ClosableNode, error) {
|
||||
case ClusterType:
|
||||
client := red.NewClusterClient(&red.ClusterOptions{
|
||||
Addrs: splitClusterAddrs(r.Addr),
|
||||
Username: r.User,
|
||||
Password: r.Pass,
|
||||
MaxRetries: maxRetries,
|
||||
PoolSize: 1,
|
||||
|
||||
@@ -31,6 +31,7 @@ func getClient(r *Redis) (*red.Client, error) {
|
||||
}
|
||||
store := red.NewClient(&red.Options{
|
||||
Addr: r.Addr,
|
||||
Username: r.User,
|
||||
Password: r.Pass,
|
||||
DB: defaultDatabase,
|
||||
MaxRetries: maxRetries,
|
||||
|
||||
@@ -28,6 +28,7 @@ func getCluster(r *Redis) (*red.ClusterClient, error) {
|
||||
}
|
||||
store := red.NewClusterClient(&red.ClusterOptions{
|
||||
Addrs: splitClusterAddrs(r.Addr),
|
||||
Username: r.User,
|
||||
Password: r.Pass,
|
||||
MaxRetries: maxRetries,
|
||||
MinIdleConns: idleConns,
|
||||
|
||||
@@ -41,6 +41,7 @@ var (
|
||||
|
||||
type (
|
||||
statGetter struct {
|
||||
host string
|
||||
dbName string
|
||||
hash string
|
||||
poolStats func() sql.DBStats
|
||||
|
||||
@@ -267,6 +267,20 @@ func TestUnmarshalRowStruct(t *testing.T) {
|
||||
}, "select name, age from users where user=?", "anyone"), ErrNotMatchDestination)
|
||||
})
|
||||
|
||||
dbtest.RunTest(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
|
||||
value := new(struct {
|
||||
Name string
|
||||
age int
|
||||
})
|
||||
|
||||
rs := sqlmock.NewRows([]string{"name", "age"}).FromCSVString("liao,5")
|
||||
mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
|
||||
|
||||
assert.ErrorIs(t, query(context.Background(), db, func(rows *sql.Rows) error {
|
||||
return unmarshalRow(value, rows, true)
|
||||
}, "select name, age from users where user=?", "anyone"), ErrNotMatchDestination)
|
||||
})
|
||||
|
||||
dbtest.RunTest(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
|
||||
rs := sqlmock.NewRows([]string{"value"}).FromCSVString("8")
|
||||
mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
|
||||
@@ -310,6 +324,20 @@ func TestUnmarshalRowStructWithTags(t *testing.T) {
|
||||
}, "select name, age from users where user=?", "anyone"), ErrNotReadableValue)
|
||||
})
|
||||
|
||||
dbtest.RunTest(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
|
||||
value := new(struct {
|
||||
age int `db:"age"`
|
||||
Name string `db:"name"`
|
||||
})
|
||||
|
||||
rs := sqlmock.NewRows([]string{"name", "age"}).FromCSVString("liao,5")
|
||||
mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
|
||||
|
||||
assert.ErrorIs(t, query(context.Background(), db, func(rows *sql.Rows) error {
|
||||
return unmarshalRow(value, rows, true)
|
||||
}, "select name, age from users where user=?", "anyone"), ErrNotReadableValue)
|
||||
})
|
||||
|
||||
dbtest.RunTest(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
|
||||
var value struct {
|
||||
Age *int `db:"age"`
|
||||
@@ -1307,25 +1335,25 @@ func TestAnonymousStructPr(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestAnonymousStructPrError(t *testing.T) {
|
||||
type Score struct {
|
||||
Discipline string `db:"discipline"`
|
||||
score uint `db:"score"`
|
||||
}
|
||||
type ClassType struct {
|
||||
Grade sql.NullString `db:"grade"`
|
||||
ClassName *string `db:"class_name"`
|
||||
}
|
||||
type Class struct {
|
||||
*ClassType
|
||||
Score
|
||||
}
|
||||
var value []*struct {
|
||||
Age int64 `db:"age"`
|
||||
Class
|
||||
Name string `db:"name"`
|
||||
}
|
||||
|
||||
dbtest.RunTest(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
|
||||
type Score struct {
|
||||
Discipline string `db:"discipline"`
|
||||
score uint `db:"score"`
|
||||
}
|
||||
type ClassType struct {
|
||||
Grade sql.NullString `db:"grade"`
|
||||
ClassName *string `db:"class_name"`
|
||||
}
|
||||
type Class struct {
|
||||
*ClassType
|
||||
Score
|
||||
}
|
||||
|
||||
var value []*struct {
|
||||
Age int64 `db:"age"`
|
||||
Class
|
||||
Name string `db:"name"`
|
||||
}
|
||||
rs := sqlmock.NewRows([]string{
|
||||
"name",
|
||||
"age",
|
||||
@@ -1338,10 +1366,50 @@ func TestAnonymousStructPrError(t *testing.T) {
|
||||
AddRow("second", 3, "grade one", "chinese", "class three grade two", 99)
|
||||
mock.ExpectQuery("select (.+) from users where user=?").
|
||||
WithArgs("anyone").WillReturnRows(rs)
|
||||
assert.Error(t, query(context.Background(), db, func(rows *sql.Rows) error {
|
||||
assert.ErrorIs(t, query(context.Background(), db, func(rows *sql.Rows) error {
|
||||
return unmarshalRows(&value, rows, true)
|
||||
}, "select name, age, grade, discipline, class_name, score from users where user=?",
|
||||
"anyone"))
|
||||
"anyone"), ErrNotReadableValue)
|
||||
if len(value) > 0 {
|
||||
assert.Equal(t, value[0].score, 0)
|
||||
}
|
||||
})
|
||||
|
||||
dbtest.RunTest(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
|
||||
type Score struct {
|
||||
Discipline string
|
||||
score uint
|
||||
}
|
||||
type ClassType struct {
|
||||
Grade sql.NullString
|
||||
ClassName *string
|
||||
}
|
||||
type Class struct {
|
||||
*ClassType
|
||||
Score
|
||||
}
|
||||
|
||||
var value []*struct {
|
||||
Age int64
|
||||
Class
|
||||
Name string
|
||||
}
|
||||
rs := sqlmock.NewRows([]string{
|
||||
"name",
|
||||
"age",
|
||||
"grade",
|
||||
"discipline",
|
||||
"class_name",
|
||||
"score",
|
||||
}).
|
||||
AddRow("first", 2, nil, "math", "experimental class", 100).
|
||||
AddRow("second", 3, "grade one", "chinese", "class three grade two", 99)
|
||||
mock.ExpectQuery("select (.+) from users where user=?").
|
||||
WithArgs("anyone").WillReturnRows(rs)
|
||||
assert.ErrorIs(t, query(context.Background(), db, func(rows *sql.Rows) error {
|
||||
return unmarshalRows(&value, rows, true)
|
||||
}, "select name, age, grade, discipline, class_name, score from users where user=?",
|
||||
"anyone"), ErrNotMatchDestination)
|
||||
if len(value) > 0 {
|
||||
assert.Equal(t, value[0].score, 0)
|
||||
}
|
||||
|
||||
@@ -1,10 +1,14 @@
|
||||
package sqlx
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"database/sql"
|
||||
"encoding/hex"
|
||||
"io"
|
||||
"time"
|
||||
|
||||
"github.com/go-sql-driver/mysql"
|
||||
"github.com/zeromicro/go-zero/core/logx"
|
||||
"github.com/zeromicro/go-zero/core/syncx"
|
||||
)
|
||||
|
||||
@@ -23,6 +27,23 @@ func getCachedSqlConn(driverName, server string) (*sql.DB, error) {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if driverName != mysqlDriverName {
|
||||
if cfg, e := mysql.ParseDSN(server); e != nil {
|
||||
// if cannot parse, don't collect the metrics
|
||||
logx.Error(e)
|
||||
} else {
|
||||
checksum := sha256.Sum256([]byte(server))
|
||||
connCollector.registerClient(&statGetter{
|
||||
host: cfg.Addr,
|
||||
dbName: cfg.DBName,
|
||||
hash: hex.EncodeToString(checksum[:]),
|
||||
poolStats: func() sql.DBStats {
|
||||
return conn.Stats()
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return conn, nil
|
||||
})
|
||||
if err != nil {
|
||||
|
||||
@@ -76,7 +76,7 @@ func FuzzNodeFind(f *testing.F) {
|
||||
fmt.Fprintf(&buf, "text:\n\t%s\n", str)
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
t.Errorf(buf.String())
|
||||
t.Error(buf.String())
|
||||
}
|
||||
}()
|
||||
assert.ElementsMatchf(t, scopes, n.find([]rune(str)), buf.String())
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"runtime"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
const factor = 10
|
||||
@@ -100,6 +101,6 @@ func (r *StableRunner[I, O]) Wait() {
|
||||
close(r.done)
|
||||
r.runner.Wait()
|
||||
for atomic.LoadUint64(&r.consumedIndex) < atomic.LoadUint64(&r.writtenIndex) {
|
||||
runtime.Gosched()
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -12,14 +12,22 @@ type (
|
||||
Upstreams []Upstream
|
||||
}
|
||||
|
||||
// HttpClientConf is the configuration for an HTTP client.
|
||||
HttpClientConf struct {
|
||||
Target string
|
||||
Prefix string `json:",optional"`
|
||||
Timeout int64 `json:",default=3000"`
|
||||
}
|
||||
|
||||
// RouteMapping is a mapping between a gateway route and an upstream rpc method.
|
||||
RouteMapping struct {
|
||||
// Method is the HTTP method, like GET, POST, PUT, DELETE.
|
||||
Method string
|
||||
// Path is the HTTP path.
|
||||
Path string
|
||||
// RpcPath is the gRPC rpc method, with format of package.service/method
|
||||
RpcPath string
|
||||
// RpcPath is the gRPC rpc method, with format of package.service/method, optional.
|
||||
// If the mapping is for HTTP, it's not necessary.
|
||||
RpcPath string `json:",optional"`
|
||||
}
|
||||
|
||||
// Upstream is the configuration for an upstream.
|
||||
@@ -27,12 +35,14 @@ type (
|
||||
// Name is the name of the upstream.
|
||||
Name string `json:",optional"`
|
||||
// Grpc is the target of the upstream.
|
||||
Grpc zrpc.RpcClientConf
|
||||
Grpc *zrpc.RpcClientConf `json:",optional"`
|
||||
// Http is the target of the upstream.
|
||||
Http *HttpClientConf `json:",optional=!grpc"`
|
||||
// ProtoSets is the file list of proto set, like [hello.pb].
|
||||
// if your proto file import another proto file, you need to write multi-file slice,
|
||||
// like [hello.pb, common.pb].
|
||||
ProtoSets []string `json:",optional"`
|
||||
// Mappings is the mapping between gateway routes and Upstream rpc methods.
|
||||
// Mappings is the mapping between gateway routes and Upstream methods.
|
||||
// Keep it blank if annotations are added in rpc methods.
|
||||
Mappings []RouteMapping `json:",optional"`
|
||||
}
|
||||
|
||||
@@ -3,26 +3,35 @@ package gateway
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/fullstorydev/grpcurl"
|
||||
"github.com/golang/protobuf/jsonpb"
|
||||
"github.com/jhump/protoreflect/grpcreflect"
|
||||
"github.com/zeromicro/go-zero/core/logc"
|
||||
"github.com/zeromicro/go-zero/core/logx"
|
||||
"github.com/zeromicro/go-zero/core/mr"
|
||||
"github.com/zeromicro/go-zero/core/threading"
|
||||
"github.com/zeromicro/go-zero/gateway/internal"
|
||||
"github.com/zeromicro/go-zero/rest"
|
||||
"github.com/zeromicro/go-zero/rest/httpc"
|
||||
"github.com/zeromicro/go-zero/rest/httpx"
|
||||
"github.com/zeromicro/go-zero/zrpc"
|
||||
"google.golang.org/grpc/codes"
|
||||
)
|
||||
|
||||
const defaultHttpScheme = "http"
|
||||
|
||||
type (
|
||||
// Server is a gateway server.
|
||||
Server struct {
|
||||
*rest.Server
|
||||
upstreams []Upstream
|
||||
conns []zrpc.Client
|
||||
processHeader func(http.Header) []string
|
||||
dialer func(conf zrpc.RpcClientConf) zrpc.Client
|
||||
}
|
||||
@@ -51,8 +60,24 @@ func (s *Server) Start() {
|
||||
}
|
||||
|
||||
// Stop stops the gateway server.
|
||||
// To get a graceful shutdown, it stops the HTTP server first, then closes gRPC connections.
|
||||
func (s *Server) Stop() {
|
||||
// stop the HTTP server first, then close gRPC connections.
|
||||
// in case the gRPC server is stopped first,
|
||||
// the HTTP server may still be running to accept requests.
|
||||
s.Server.Stop()
|
||||
|
||||
group := threading.NewRoutineGroup()
|
||||
for _, conn := range s.conns {
|
||||
// new variable to avoid closure problems, can be removed after go 1.22
|
||||
// see https://golang.org/doc/faq#closures_and_goroutines
|
||||
conn := conn
|
||||
group.Run(func() {
|
||||
// ignore the error when closing the connection
|
||||
_ = conn.Conn().Close()
|
||||
})
|
||||
}
|
||||
group.Wait()
|
||||
}
|
||||
|
||||
func (s *Server) build() error {
|
||||
@@ -65,51 +90,11 @@ func (s *Server) build() error {
|
||||
source <- up
|
||||
}
|
||||
}, func(up Upstream, writer mr.Writer[rest.Route], cancel func(error)) {
|
||||
var cli zrpc.Client
|
||||
if s.dialer != nil {
|
||||
cli = s.dialer(up.Grpc)
|
||||
} else {
|
||||
cli = zrpc.MustNewClient(up.Grpc)
|
||||
}
|
||||
|
||||
source, err := s.createDescriptorSource(cli, up)
|
||||
if err != nil {
|
||||
cancel(fmt.Errorf("%s: %w", up.Name, err))
|
||||
return
|
||||
}
|
||||
|
||||
methods, err := internal.GetMethods(source)
|
||||
if err != nil {
|
||||
cancel(fmt.Errorf("%s: %w", up.Name, err))
|
||||
return
|
||||
}
|
||||
|
||||
resolver := grpcurl.AnyResolverFromDescriptorSource(source)
|
||||
for _, m := range methods {
|
||||
if len(m.HttpMethod) > 0 && len(m.HttpPath) > 0 {
|
||||
writer.Write(rest.Route{
|
||||
Method: m.HttpMethod,
|
||||
Path: m.HttpPath,
|
||||
Handler: s.buildHandler(source, resolver, cli, m.RpcPath),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
methodSet := make(map[string]struct{})
|
||||
for _, m := range methods {
|
||||
methodSet[m.RpcPath] = struct{}{}
|
||||
}
|
||||
for _, m := range up.Mappings {
|
||||
if _, ok := methodSet[m.RpcPath]; !ok {
|
||||
cancel(fmt.Errorf("%s: rpc method %s not found", up.Name, m.RpcPath))
|
||||
return
|
||||
}
|
||||
|
||||
writer.Write(rest.Route{
|
||||
Method: strings.ToUpper(m.Method),
|
||||
Path: m.Path,
|
||||
Handler: s.buildHandler(source, resolver, cli, m.RpcPath),
|
||||
})
|
||||
// up.Grpc and up.Http are exclusive
|
||||
if up.Grpc != nil {
|
||||
s.buildGrpcRoute(up, writer, cancel)
|
||||
} else if up.Http != nil {
|
||||
s.buildHttpRoute(up, writer)
|
||||
}
|
||||
}, func(pipe <-chan rest.Route, cancel func(error)) {
|
||||
for route := range pipe {
|
||||
@@ -118,7 +103,7 @@ func (s *Server) build() error {
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) buildHandler(source grpcurl.DescriptorSource, resolver jsonpb.AnyResolver,
|
||||
func (s *Server) buildGrpcHandler(source grpcurl.DescriptorSource, resolver jsonpb.AnyResolver,
|
||||
cli zrpc.Client, rpcPath string) func(http.ResponseWriter, *http.Request) {
|
||||
return func(w http.ResponseWriter, r *http.Request) {
|
||||
parser, err := internal.NewRequestParser(r, resolver)
|
||||
@@ -141,31 +126,121 @@ func (s *Server) buildHandler(source grpcurl.DescriptorSource, resolver jsonpb.A
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) createDescriptorSource(cli zrpc.Client, up Upstream) (grpcurl.DescriptorSource, error) {
|
||||
var source grpcurl.DescriptorSource
|
||||
var err error
|
||||
|
||||
if len(up.ProtoSets) > 0 {
|
||||
source, err = grpcurl.DescriptorSourceFromProtoSets(up.ProtoSets...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
func (s *Server) buildGrpcRoute(up Upstream, writer mr.Writer[rest.Route], cancel func(error)) {
|
||||
var cli zrpc.Client
|
||||
if s.dialer != nil {
|
||||
cli = s.dialer(*up.Grpc)
|
||||
} else {
|
||||
client := grpcreflect.NewClientAuto(context.Background(), cli.Conn())
|
||||
source = grpcurl.DescriptorSourceFromServer(context.Background(), client)
|
||||
cli = zrpc.MustNewClient(*up.Grpc)
|
||||
}
|
||||
s.conns = append(s.conns, cli)
|
||||
|
||||
source, err := createDescriptorSource(cli, up)
|
||||
if err != nil {
|
||||
cancel(fmt.Errorf("%s: %w", up.Name, err))
|
||||
return
|
||||
}
|
||||
|
||||
return source, nil
|
||||
methods, err := internal.GetMethods(source)
|
||||
if err != nil {
|
||||
cancel(fmt.Errorf("%s: %w", up.Name, err))
|
||||
return
|
||||
}
|
||||
|
||||
resolver := grpcurl.AnyResolverFromDescriptorSource(source)
|
||||
for _, m := range methods {
|
||||
if len(m.HttpMethod) > 0 && len(m.HttpPath) > 0 {
|
||||
writer.Write(rest.Route{
|
||||
Method: m.HttpMethod,
|
||||
Path: m.HttpPath,
|
||||
Handler: s.buildGrpcHandler(source, resolver, cli, m.RpcPath),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
methodSet := make(map[string]struct{})
|
||||
for _, m := range methods {
|
||||
methodSet[m.RpcPath] = struct{}{}
|
||||
}
|
||||
for _, m := range up.Mappings {
|
||||
if _, ok := methodSet[m.RpcPath]; !ok {
|
||||
cancel(fmt.Errorf("%s: rpc method %s not found", up.Name, m.RpcPath))
|
||||
return
|
||||
}
|
||||
|
||||
writer.Write(rest.Route{
|
||||
Method: strings.ToUpper(m.Method),
|
||||
Path: m.Path,
|
||||
Handler: s.buildGrpcHandler(source, resolver, cli, m.RpcPath),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) buildHttpHandler(target *HttpClientConf) http.HandlerFunc {
|
||||
return func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set(httpx.ContentType, httpx.JsonContentType)
|
||||
req, err := buildRequestWithNewTarget(r, target)
|
||||
if err != nil {
|
||||
httpx.ErrorCtx(r.Context(), w, err)
|
||||
return
|
||||
}
|
||||
|
||||
// set the timeout if it's configured, take effect only if it's greater than 0
|
||||
// and less than the deadline of the original request
|
||||
if target.Timeout > 0 {
|
||||
timeout := time.Duration(target.Timeout) * time.Millisecond
|
||||
ctx, cancel := context.WithTimeout(r.Context(), timeout)
|
||||
defer cancel()
|
||||
req = req.WithContext(ctx)
|
||||
}
|
||||
|
||||
resp, err := httpc.DoRequest(req)
|
||||
if err != nil {
|
||||
httpx.ErrorCtx(r.Context(), w, err)
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
for key, values := range resp.Header {
|
||||
for _, value := range values {
|
||||
w.Header().Add(key, value)
|
||||
}
|
||||
}
|
||||
|
||||
w.WriteHeader(resp.StatusCode)
|
||||
if _, err = io.Copy(w, resp.Body); err != nil {
|
||||
// log the error with original request info
|
||||
logc.Error(r.Context(), err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) buildHttpRoute(up Upstream, writer mr.Writer[rest.Route]) {
|
||||
for _, m := range up.Mappings {
|
||||
writer.Write(rest.Route{
|
||||
Method: strings.ToUpper(m.Method),
|
||||
Path: m.Path,
|
||||
Handler: s.buildHttpHandler(up.Http),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) ensureUpstreamNames() error {
|
||||
for i := 0; i < len(s.upstreams); i++ {
|
||||
target, err := s.upstreams[i].Grpc.BuildTarget()
|
||||
if err != nil {
|
||||
return err
|
||||
if len(s.upstreams[i].Name) > 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
s.upstreams[i].Name = target
|
||||
if s.upstreams[i].Grpc != nil {
|
||||
target, err := s.upstreams[i].Grpc.BuildTarget()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
s.upstreams[i].Name = target
|
||||
} else if s.upstreams[i].Http != nil {
|
||||
s.upstreams[i].Name = s.upstreams[i].Http.Target
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
@@ -188,6 +263,53 @@ func WithHeaderProcessor(processHeader func(http.Header) []string) func(*Server)
|
||||
}
|
||||
}
|
||||
|
||||
func buildRequestWithNewTarget(r *http.Request, target *HttpClientConf) (*http.Request, error) {
|
||||
u := *r.URL
|
||||
u.Host = target.Target
|
||||
if len(u.Scheme) == 0 {
|
||||
u.Scheme = defaultHttpScheme
|
||||
}
|
||||
|
||||
if len(target.Prefix) > 0 {
|
||||
var err error
|
||||
u.Path, err = url.JoinPath(target.Prefix, u.Path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
newReq := &http.Request{
|
||||
Method: r.Method,
|
||||
URL: &u,
|
||||
Header: r.Header.Clone(),
|
||||
Proto: r.Proto,
|
||||
ProtoMajor: r.ProtoMajor,
|
||||
ProtoMinor: r.ProtoMinor,
|
||||
ContentLength: r.ContentLength,
|
||||
Body: io.NopCloser(r.Body),
|
||||
}
|
||||
|
||||
// make sure the context is passed to the new request
|
||||
return newReq.WithContext(r.Context()), nil
|
||||
}
|
||||
|
||||
func createDescriptorSource(cli zrpc.Client, up Upstream) (grpcurl.DescriptorSource, error) {
|
||||
var source grpcurl.DescriptorSource
|
||||
var err error
|
||||
|
||||
if len(up.ProtoSets) > 0 {
|
||||
source, err = grpcurl.DescriptorSourceFromProtoSets(up.ProtoSets...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
} else {
|
||||
client := grpcreflect.NewClientAuto(context.Background(), cli.Conn())
|
||||
source = grpcurl.DescriptorSourceFromServer(context.Background(), client)
|
||||
}
|
||||
|
||||
return source, nil
|
||||
}
|
||||
|
||||
// withDialer sets a dialer to create a gRPC client.
|
||||
func withDialer(dialer func(conf zrpc.RpcClientConf) zrpc.Client) func(*Server) {
|
||||
return func(s *Server) {
|
||||
|
||||
@@ -2,9 +2,12 @@ package gateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"log"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -46,7 +49,7 @@ func dialer() func(context.Context, string) (net.Conn, error) {
|
||||
func TestMustNewServer(t *testing.T) {
|
||||
var c GatewayConf
|
||||
assert.NoError(t, conf.FillDefault(&c))
|
||||
// avoid popup alert on macos for asking permissions
|
||||
// avoid popup alert on MacOS for asking permissions
|
||||
c.DevServer.Host = "localhost"
|
||||
c.Host = "localhost"
|
||||
c.Port = 18881
|
||||
@@ -65,7 +68,7 @@ func TestMustNewServer(t *testing.T) {
|
||||
RpcPath: "mock.DepositService/Deposit",
|
||||
},
|
||||
},
|
||||
Grpc: zrpc.RpcClientConf{
|
||||
Grpc: &zrpc.RpcClientConf{
|
||||
Endpoints: []string{"foo"},
|
||||
Timeout: 1000,
|
||||
Middlewares: zrpc.ClientMiddlewaresConf{
|
||||
@@ -98,7 +101,7 @@ func TestServer_ensureUpstreamNames(t *testing.T) {
|
||||
var s = Server{
|
||||
upstreams: []Upstream{
|
||||
{
|
||||
Grpc: zrpc.RpcClientConf{
|
||||
Grpc: &zrpc.RpcClientConf{
|
||||
Target: "target",
|
||||
},
|
||||
},
|
||||
@@ -113,7 +116,7 @@ func TestServer_ensureUpstreamNames_badEtcd(t *testing.T) {
|
||||
var s = Server{
|
||||
upstreams: []Upstream{
|
||||
{
|
||||
Grpc: zrpc.RpcClientConf{
|
||||
Grpc: &zrpc.RpcClientConf{
|
||||
Etcd: discov.EtcdConf{},
|
||||
},
|
||||
},
|
||||
@@ -125,3 +128,193 @@ func TestServer_ensureUpstreamNames_badEtcd(t *testing.T) {
|
||||
s.Start()
|
||||
})
|
||||
}
|
||||
|
||||
func TestHttpToHttp(t *testing.T) {
|
||||
server := startTestServer(t)
|
||||
defer server.Close()
|
||||
|
||||
var c GatewayConf
|
||||
assert.NoError(t, conf.FillDefault(&c))
|
||||
c.DevServer.Host = "localhost"
|
||||
c.Host = "localhost"
|
||||
c.Port = 18882
|
||||
|
||||
s := MustNewServer(c)
|
||||
s.upstreams = []Upstream{
|
||||
{
|
||||
Name: "test",
|
||||
Mappings: []RouteMapping{
|
||||
{
|
||||
Method: "get",
|
||||
Path: "/api/ping",
|
||||
},
|
||||
},
|
||||
Http: &HttpClientConf{
|
||||
Target: "localhost:45678",
|
||||
Timeout: 3000,
|
||||
},
|
||||
},
|
||||
{
|
||||
Mappings: []RouteMapping{
|
||||
{
|
||||
Method: "get",
|
||||
Path: "/ping",
|
||||
},
|
||||
},
|
||||
Http: &HttpClientConf{
|
||||
Target: "localhost:45678",
|
||||
Prefix: "/api",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
go s.Start()
|
||||
defer s.Stop()
|
||||
|
||||
time.Sleep(time.Millisecond * 200)
|
||||
|
||||
t.Run("/api/ping", func(t *testing.T) {
|
||||
resp, err := httpc.Do(context.Background(), http.MethodGet,
|
||||
"http://localhost:18882/api/ping", nil)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, http.StatusOK, resp.StatusCode)
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if assert.NoError(t, err) {
|
||||
assert.Equal(t, "pong", string(body))
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("/ping", func(t *testing.T) {
|
||||
resp, err := httpc.Do(context.Background(), http.MethodGet,
|
||||
"http://localhost:18882/ping", nil)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, http.StatusOK, resp.StatusCode)
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if assert.NoError(t, err) {
|
||||
assert.Equal(t, "pong", string(body))
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("no upstream", func(t *testing.T) {
|
||||
resp, err := httpc.Do(context.Background(), http.MethodGet,
|
||||
"http://localhost:18882/ping/bad", nil)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, http.StatusNotFound, resp.StatusCode)
|
||||
})
|
||||
}
|
||||
|
||||
func TestHttpToHttpBadUpstream(t *testing.T) {
|
||||
var c GatewayConf
|
||||
assert.NoError(t, conf.FillDefault(&c))
|
||||
c.DevServer.Host = "localhost"
|
||||
c.Host = "localhost"
|
||||
c.Port = 18883
|
||||
|
||||
s := MustNewServer(c)
|
||||
s.upstreams = []Upstream{
|
||||
{
|
||||
Mappings: []RouteMapping{
|
||||
{
|
||||
Method: "get",
|
||||
Path: "/api/ping",
|
||||
},
|
||||
},
|
||||
Http: &HttpClientConf{
|
||||
Target: "localhost:45678",
|
||||
Prefix: "\x7f/api",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
go s.Start()
|
||||
defer s.Stop()
|
||||
|
||||
time.Sleep(time.Millisecond * 200)
|
||||
|
||||
t.Run("/api/ping", func(t *testing.T) {
|
||||
resp, err := httpc.Do(context.Background(), http.MethodGet,
|
||||
"http://localhost:18883/api/ping", nil)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, http.StatusBadRequest, resp.StatusCode)
|
||||
})
|
||||
}
|
||||
|
||||
func TestHttpToHttpBadWriter(t *testing.T) {
|
||||
t.Run("bad url", func(t *testing.T) {
|
||||
handler := new(Server).buildHttpHandler(&HttpClientConf{
|
||||
Target: "http://example.com",
|
||||
Timeout: 3000,
|
||||
})
|
||||
w := httptest.NewRecorder()
|
||||
handler.ServeHTTP(&badResponseWriter{w},
|
||||
httptest.NewRequest(http.MethodGet, "http://localhost:18884", nil))
|
||||
assert.Equal(t, http.StatusBadRequest, w.Code)
|
||||
})
|
||||
|
||||
t.Run("bad url", func(t *testing.T) {
|
||||
var c GatewayConf
|
||||
assert.NoError(t, conf.FillDefault(&c))
|
||||
c.DevServer.Host = "localhost"
|
||||
c.Host = "localhost"
|
||||
c.Port = 18884
|
||||
|
||||
s := MustNewServer(c)
|
||||
s.upstreams = []Upstream{
|
||||
{
|
||||
Mappings: []RouteMapping{
|
||||
{
|
||||
Method: "get",
|
||||
Path: "/api/ping",
|
||||
},
|
||||
},
|
||||
Http: &HttpClientConf{
|
||||
Target: "localhost:45678",
|
||||
Prefix: "\x7f/api",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
go s.Start()
|
||||
defer s.Stop()
|
||||
|
||||
handler := new(Server).buildHttpHandler(&HttpClientConf{
|
||||
Target: "localhost:18884",
|
||||
Timeout: 3000,
|
||||
})
|
||||
w := httptest.NewRecorder()
|
||||
handler.ServeHTTP(&badResponseWriter{w},
|
||||
httptest.NewRequest(http.MethodGet, "http://localhost:18884/api/ping", nil))
|
||||
assert.Equal(t, http.StatusBadRequest, w.Code)
|
||||
})
|
||||
}
|
||||
|
||||
// Handler function for the root route
|
||||
func pingHandler(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write([]byte("pong"))
|
||||
}
|
||||
|
||||
func startTestServer(t *testing.T) *http.Server {
|
||||
http.HandleFunc("/api/ping", pingHandler)
|
||||
|
||||
server := &http.Server{
|
||||
Addr: ":45678",
|
||||
Handler: http.DefaultServeMux,
|
||||
}
|
||||
|
||||
go func() {
|
||||
if err := server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
|
||||
t.Errorf("failed to start server: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
return server
|
||||
}
|
||||
|
||||
type badResponseWriter struct {
|
||||
http.ResponseWriter
|
||||
}
|
||||
|
||||
func (w *badResponseWriter) Write([]byte) (int, error) {
|
||||
return 0, errors.New("bad writer")
|
||||
}
|
||||
|
||||
44
go.mod
44
go.mod
@@ -1,28 +1,28 @@
|
||||
module github.com/zeromicro/go-zero
|
||||
|
||||
go 1.20
|
||||
go 1.21
|
||||
|
||||
require (
|
||||
github.com/DATA-DOG/go-sqlmock v1.5.2
|
||||
github.com/alicebob/miniredis/v2 v2.33.0
|
||||
github.com/alicebob/miniredis/v2 v2.34.0
|
||||
github.com/fatih/color v1.18.0
|
||||
github.com/fullstorydev/grpcurl v1.9.1
|
||||
github.com/go-sql-driver/mysql v1.8.1
|
||||
github.com/golang-jwt/jwt/v4 v4.5.1
|
||||
github.com/fullstorydev/grpcurl v1.9.3
|
||||
github.com/go-sql-driver/mysql v1.9.0
|
||||
github.com/golang-jwt/jwt/v4 v4.5.2
|
||||
github.com/golang/mock v1.6.0
|
||||
github.com/golang/protobuf v1.5.4
|
||||
github.com/google/uuid v1.6.0
|
||||
github.com/jackc/pgx/v5 v5.6.0
|
||||
github.com/jackc/pgx/v5 v5.7.4
|
||||
github.com/jhump/protoreflect v1.17.0
|
||||
github.com/olekukonko/tablewriter v0.0.5
|
||||
github.com/pelletier/go-toml/v2 v2.2.2
|
||||
github.com/prometheus/client_golang v1.20.5
|
||||
github.com/redis/go-redis/v9 v9.7.0
|
||||
github.com/prometheus/client_golang v1.21.1
|
||||
github.com/redis/go-redis/v9 v9.7.3
|
||||
github.com/spaolacci/murmur3 v1.1.0
|
||||
github.com/stretchr/testify v1.9.0
|
||||
github.com/stretchr/testify v1.10.0
|
||||
go.etcd.io/etcd/api/v3 v3.5.15
|
||||
go.etcd.io/etcd/client/v3 v3.5.15
|
||||
go.mongodb.org/mongo-driver v1.17.1
|
||||
go.mongodb.org/mongo-driver v1.17.3
|
||||
go.opentelemetry.io/otel v1.24.0
|
||||
go.opentelemetry.io/otel/exporters/jaeger v1.17.0
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.24.0
|
||||
@@ -33,12 +33,12 @@ require (
|
||||
go.opentelemetry.io/otel/trace v1.24.0
|
||||
go.uber.org/automaxprocs v1.6.0
|
||||
go.uber.org/goleak v1.3.0
|
||||
golang.org/x/net v0.31.0
|
||||
golang.org/x/sys v0.27.0
|
||||
golang.org/x/time v0.8.0
|
||||
golang.org/x/net v0.35.0
|
||||
golang.org/x/sys v0.30.0
|
||||
golang.org/x/time v0.10.0
|
||||
google.golang.org/genproto/googleapis/api v0.0.0-20240711142825-46eb208f015d
|
||||
google.golang.org/grpc v1.65.0
|
||||
google.golang.org/protobuf v1.35.2
|
||||
google.golang.org/protobuf v1.36.5
|
||||
gopkg.in/cheggaaa/pb.v1 v1.0.28
|
||||
gopkg.in/h2non/gock.v1 v1.1.2
|
||||
gopkg.in/yaml.v2 v2.4.0
|
||||
@@ -77,10 +77,10 @@ require (
|
||||
github.com/h2non/parth v0.0.0-20190131123155-b4df798d6542 // indirect
|
||||
github.com/jackc/pgpassfile v1.0.0 // indirect
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
|
||||
github.com/jackc/puddle/v2 v2.2.1 // indirect
|
||||
github.com/jackc/puddle/v2 v2.2.2 // indirect
|
||||
github.com/josharian/intern v1.0.0 // indirect
|
||||
github.com/json-iterator/go v1.1.12 // indirect
|
||||
github.com/klauspost/compress v1.17.9 // indirect
|
||||
github.com/klauspost/compress v1.17.11 // indirect
|
||||
github.com/kylelemons/godebug v1.1.0 // indirect
|
||||
github.com/mailru/easyjson v0.7.7 // indirect
|
||||
github.com/mattn/go-colorable v0.1.13 // indirect
|
||||
@@ -93,7 +93,7 @@ require (
|
||||
github.com/openzipkin/zipkin-go v0.4.3 // indirect
|
||||
github.com/pmezard/go-difflib v1.0.0 // indirect
|
||||
github.com/prometheus/client_model v0.6.1 // indirect
|
||||
github.com/prometheus/common v0.55.0 // indirect
|
||||
github.com/prometheus/common v0.62.0 // indirect
|
||||
github.com/prometheus/procfs v0.15.1 // indirect
|
||||
github.com/rivo/uniseg v0.2.0 // indirect
|
||||
github.com/stretchr/objx v0.5.2 // indirect
|
||||
@@ -109,11 +109,11 @@ require (
|
||||
go.uber.org/atomic v1.10.0 // indirect
|
||||
go.uber.org/multierr v1.9.0 // indirect
|
||||
go.uber.org/zap v1.24.0 // indirect
|
||||
golang.org/x/crypto v0.29.0 // indirect
|
||||
golang.org/x/oauth2 v0.21.0 // indirect
|
||||
golang.org/x/sync v0.9.0 // indirect
|
||||
golang.org/x/term v0.26.0 // indirect
|
||||
golang.org/x/text v0.20.0 // indirect
|
||||
golang.org/x/crypto v0.33.0 // indirect
|
||||
golang.org/x/oauth2 v0.24.0 // indirect
|
||||
golang.org/x/sync v0.11.0 // indirect
|
||||
golang.org/x/term v0.29.0 // indirect
|
||||
golang.org/x/text v0.22.0 // indirect
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20240701130421-f6361c86f094 // indirect
|
||||
gopkg.in/inf.v0 v0.9.1 // indirect
|
||||
gopkg.in/yaml.v3 v3.0.1 // indirect
|
||||
|
||||
96
go.sum
96
go.sum
@@ -4,13 +4,16 @@ github.com/DATA-DOG/go-sqlmock v1.5.2 h1:OcvFkGmslmlZibjAjaHm3L//6LiuBgolP7Oputl
|
||||
github.com/DATA-DOG/go-sqlmock v1.5.2/go.mod h1:88MAG/4G7SMwSE3CeA0ZKzrT5CiOU3OJ+JlNzwDqpNU=
|
||||
github.com/alicebob/gopher-json v0.0.0-20230218143504-906a9b012302 h1:uvdUDbHQHO85qeSydJtItA4T55Pw6BtAejd0APRJOCE=
|
||||
github.com/alicebob/gopher-json v0.0.0-20230218143504-906a9b012302/go.mod h1:SGnFV6hVsYE877CKEZ6tDNTjaSXYUk6QqoIK6PrAtcc=
|
||||
github.com/alicebob/miniredis/v2 v2.33.0 h1:uvTF0EDeu9RLnUEG27Db5I68ESoIxTiXbNUiji6lZrA=
|
||||
github.com/alicebob/miniredis/v2 v2.33.0/go.mod h1:MhP4a3EU7aENRi9aO+tHfTBZicLqQevyi/DJpoj6mi0=
|
||||
github.com/alicebob/miniredis/v2 v2.34.0 h1:mBFWMaJSNL9RwdGRyEDoAAv8OQc5UlEhLDQggTglU/0=
|
||||
github.com/alicebob/miniredis/v2 v2.34.0/go.mod h1:kWShP4b58T1CW0Y5dViCd5ztzrDqRWqM3nksiyXk5s8=
|
||||
github.com/benbjohnson/clock v1.1.0 h1:Q92kusRqC1XV2MjkWETPvjJVqKetz1OzxZB7mHJLju8=
|
||||
github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
|
||||
github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
|
||||
github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
|
||||
github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs=
|
||||
github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c=
|
||||
github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA=
|
||||
github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0=
|
||||
github.com/bufbuild/protocompile v0.14.1 h1:iA73zAf/fyljNjQKwYzUHD6AD4R8KMasmwa/FBatYVw=
|
||||
github.com/bufbuild/protocompile v0.14.1/go.mod h1:ppVdAIhbr2H8asPk6k4pY7t9zB1OU5DoEw9xY/FUi1c=
|
||||
github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK343L8=
|
||||
@@ -37,8 +40,8 @@ github.com/envoyproxy/protoc-gen-validate v1.0.4 h1:gVPz/FMfvh57HdSJQyvBtF00j8JU
|
||||
github.com/envoyproxy/protoc-gen-validate v1.0.4/go.mod h1:qys6tmnRsYrQqIhm2bvKZH4Blx/1gTIZ2UKVY1M+Yew=
|
||||
github.com/fatih/color v1.18.0 h1:S8gINlzdQ840/4pfAwic/ZE0djQEH3wM94VfqLTZcOM=
|
||||
github.com/fatih/color v1.18.0/go.mod h1:4FelSpRwEGDpQ12mAdzqdOukCy4u8WUtOY6lkT/6HfU=
|
||||
github.com/fullstorydev/grpcurl v1.9.1 h1:YxX1aCcCc4SDBQfj9uoWcTLe8t4NWrZe1y+mk83BQgo=
|
||||
github.com/fullstorydev/grpcurl v1.9.1/go.mod h1:i8gKLIC6s93WdU3LSmkE5vtsCxyRmihUj5FK1cNW5EM=
|
||||
github.com/fullstorydev/grpcurl v1.9.3 h1:PC1Xi3w+JAvEE2Tg2Gf2RfVgPbf9+tbuQr1ZkyVU3jk=
|
||||
github.com/fullstorydev/grpcurl v1.9.3/go.mod h1:/b4Wxe8bG6ndAjlfSUjwseQReUDUvBJiFEB7UllOlUE=
|
||||
github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
|
||||
github.com/go-logr/logr v1.3.0/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
|
||||
github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY=
|
||||
@@ -52,14 +55,15 @@ github.com/go-openapi/jsonreference v0.20.2/go.mod h1:Bl1zwGIM8/wsvqjsOQLJ/SH+En
|
||||
github.com/go-openapi/swag v0.22.3/go.mod h1:UzaqsxGiab7freDnrUUra0MwWfN/q7tE4j+VcZ0yl14=
|
||||
github.com/go-openapi/swag v0.22.4 h1:QLMzNJnMGPRNDCbySlcj1x01tzU8/9LTTL9hZZZogBU=
|
||||
github.com/go-openapi/swag v0.22.4/go.mod h1:UzaqsxGiab7freDnrUUra0MwWfN/q7tE4j+VcZ0yl14=
|
||||
github.com/go-sql-driver/mysql v1.8.1 h1:LedoTUt/eveggdHS9qUFC1EFSa8bU2+1pZjSRpvNJ1Y=
|
||||
github.com/go-sql-driver/mysql v1.8.1/go.mod h1:wEBSXgmK//2ZFJyE+qWnIsVGmvmEKlqwuVSjsCm7DZg=
|
||||
github.com/go-sql-driver/mysql v1.9.0 h1:Y0zIbQXhQKmQgTp44Y1dp3wTXcn804QoTptLZT1vtvo=
|
||||
github.com/go-sql-driver/mysql v1.9.0/go.mod h1:pDetrLJeA3oMujJuvXc8RJoasr589B6A9fwzD3QMrqw=
|
||||
github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 h1:tfuBGBXKqDEevZMzYi5KSi8KkcZtzBcTgAUUtapy0OI=
|
||||
github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572/go.mod h1:9Pwr4B2jHnOSGXyyzV8ROjYa2ojvAY6HCGYYfMoC3Ls=
|
||||
github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
|
||||
github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
|
||||
github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
|
||||
github.com/golang-jwt/jwt/v4 v4.5.1 h1:JdqV9zKUdtaa9gdPlywC3aeoEsR681PlKC+4F5gQgeo=
|
||||
github.com/golang-jwt/jwt/v4 v4.5.1/go.mod h1:m21LjoU+eqJr34lmDMbreY2eSTRJ1cv77w39/MY0Ch0=
|
||||
github.com/golang-jwt/jwt/v4 v4.5.2 h1:YtQM7lnr8iZ+j5q71MGKkNw9Mn7AjHM68uc9g5fXeUI=
|
||||
github.com/golang-jwt/jwt/v4 v4.5.2/go.mod h1:m21LjoU+eqJr34lmDMbreY2eSTRJ1cv77w39/MY0Ch0=
|
||||
github.com/golang/mock v1.6.0 h1:ErTB+efbowRARo13NNdxyJji2egdxLGQhRaY+DUumQc=
|
||||
github.com/golang/mock v1.6.0/go.mod h1:p6yTPP+5HYm5mzsMV8JkE6ZKdX+/wYM6Hr+LicevLPs=
|
||||
github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
|
||||
@@ -75,6 +79,7 @@ github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/
|
||||
github.com/google/gofuzz v1.2.0 h1:xRy4A+RhZaiKjJ1bPfwQ8sedCA+YS2YcCHW6ec7JMi0=
|
||||
github.com/google/gofuzz v1.2.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
|
||||
github.com/google/pprof v0.0.0-20210720184732-4bb14d4b1be1 h1:K6RDEckDVWvDI9JAJYCmNdQXq6neHJOYx3V6jnqNEec=
|
||||
github.com/google/pprof v0.0.0-20210720184732-4bb14d4b1be1/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE=
|
||||
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/grpc-ecosystem/grpc-gateway/v2 v2.20.0 h1:bkypFPDjIYGfCYD5mRBvpqxfYX1YCS1PXdKYWi8FsN0=
|
||||
@@ -85,10 +90,10 @@ github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsI
|
||||
github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
|
||||
github.com/jackc/pgx/v5 v5.6.0 h1:SWJzexBzPL5jb0GEsrPMLIsi/3jOo7RHlzTjcAeDrPY=
|
||||
github.com/jackc/pgx/v5 v5.6.0/go.mod h1:DNZ/vlrUnhWCoFGxHAG8U2ljioxukquj7utPDgtQdTw=
|
||||
github.com/jackc/puddle/v2 v2.2.1 h1:RhxXJtFG022u4ibrCSMSiu5aOq1i77R3OHKNJj77OAk=
|
||||
github.com/jackc/puddle/v2 v2.2.1/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
|
||||
github.com/jackc/pgx/v5 v5.7.4 h1:9wKznZrhWa2QiHL+NjTSPP6yjl3451BX3imWDnokYlg=
|
||||
github.com/jackc/pgx/v5 v5.7.4/go.mod h1:ncY89UGWxg82EykZUwSpUKEfccBGGYq1xjrOpsbsfGQ=
|
||||
github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo=
|
||||
github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
|
||||
github.com/jhump/protoreflect v1.17.0 h1:qOEr613fac2lOuTgWN4tPAtLL7fUSbuJL5X5XumQh94=
|
||||
github.com/jhump/protoreflect v1.17.0/go.mod h1:h9+vUUL38jiBzck8ck+6G/aeMX8Z4QUY/NiJPwPNi+8=
|
||||
github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY=
|
||||
@@ -98,10 +103,11 @@ github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHm
|
||||
github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8=
|
||||
github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
|
||||
github.com/kisielk/sqlstruct v0.0.0-20201105191214-5f3e10d3ab46/go.mod h1:yyMNCyc/Ib3bDTKd379tNMpB/7/H5TjM2Y9QJ5THLbE=
|
||||
github.com/klauspost/compress v1.17.9 h1:6KIumPrER1LHsvBVuDa0r5xaG0Es51mhhB9BQB2qeMA=
|
||||
github.com/klauspost/compress v1.17.9/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=
|
||||
github.com/klauspost/compress v1.17.11 h1:In6xLpyWOi1+C7tXUUWv2ot1QvBjxevKAaI6IXrJmUc=
|
||||
github.com/klauspost/compress v1.17.11/go.mod h1:pMDklpSncoRMuLFrf1W9Ss9KT+0rH90U12bZKk7uwG0=
|
||||
github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
|
||||
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
|
||||
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
|
||||
github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
|
||||
github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
|
||||
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
|
||||
@@ -132,31 +138,37 @@ github.com/nbio/st v0.0.0-20140626010706-e9e8d9816f32/go.mod h1:9wM+0iRr9ahx58uY
|
||||
github.com/olekukonko/tablewriter v0.0.5 h1:P2Ga83D34wi1o9J6Wh1mRuqd4mF/x/lgBS7N7AbDhec=
|
||||
github.com/olekukonko/tablewriter v0.0.5/go.mod h1:hPp6KlRPjbx+hW8ykQs1w3UBbZlj6HuIJcUGPhkA7kY=
|
||||
github.com/onsi/ginkgo/v2 v2.13.0 h1:0jY9lJquiL8fcf3M4LAXN5aMlS/b2BV86HFFPCPMgE4=
|
||||
github.com/onsi/ginkgo/v2 v2.13.0/go.mod h1:TE309ZR8s5FsKKpuB1YAQYBzCaAfUgatB/xlT/ETL/o=
|
||||
github.com/onsi/gomega v1.29.0 h1:KIA/t2t5UBzoirT4H9tsML45GEbo3ouUnBHsCfD2tVg=
|
||||
github.com/onsi/gomega v1.29.0/go.mod h1:9sxs+SwGrKI0+PWe4Fxa9tFQQBG5xSsSbMXOI8PPpoQ=
|
||||
github.com/openzipkin/zipkin-go v0.4.3 h1:9EGwpqkgnwdEIJ+Od7QVSEIH+ocmm5nPat0G7sjsSdg=
|
||||
github.com/openzipkin/zipkin-go v0.4.3/go.mod h1:M9wCJZFWCo2RiY+o1eBCEMe0Dp2S5LDHcMZmk3RmK7c=
|
||||
github.com/pelletier/go-toml/v2 v2.2.2 h1:aYUidT7k73Pcl9nb2gScu7NSrKCSHIDE89b3+6Wq+LM=
|
||||
github.com/pelletier/go-toml/v2 v2.2.2/go.mod h1:1t835xjRzz80PqgE6HHgN2JOsmgYu/h4qDAS4n929Rs=
|
||||
github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
|
||||
github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
|
||||
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/prashantv/gostub v1.1.0 h1:BTyx3RfQjRHnUWaGF9oQos79AlQ5k8WNktv7VGvVH4g=
|
||||
github.com/prometheus/client_golang v1.20.5 h1:cxppBPuYhUnsO6yo/aoRol4L7q7UFfdm+bR9r+8l63Y=
|
||||
github.com/prometheus/client_golang v1.20.5/go.mod h1:PIEt8X02hGcP8JWbeHyeZ53Y/jReSnHgO035n//V5WE=
|
||||
github.com/prashantv/gostub v1.1.0/go.mod h1:A5zLQHz7ieHGG7is6LLXLz7I8+3LZzsrV0P1IAHhP5U=
|
||||
github.com/prometheus/client_golang v1.21.1 h1:DOvXXTqVzvkIewV/CDPFdejpMCGeMcbGCQ8YOmu+Ibk=
|
||||
github.com/prometheus/client_golang v1.21.1/go.mod h1:U9NM32ykUErtVBxdvD3zfi+EuFkkaBvMb09mIfe0Zgg=
|
||||
github.com/prometheus/client_model v0.6.1 h1:ZKSh/rekM+n3CeS952MLRAdFwIKqeY8b62p8ais2e9E=
|
||||
github.com/prometheus/client_model v0.6.1/go.mod h1:OrxVMOVHjw3lKMa8+x6HeMGkHMQyHDk9E3jmP2AmGiY=
|
||||
github.com/prometheus/common v0.55.0 h1:KEi6DK7lXW/m7Ig5i47x0vRzuBsHuvJdi5ee6Y3G1dc=
|
||||
github.com/prometheus/common v0.55.0/go.mod h1:2SECS4xJG1kd8XF9IcM1gMX6510RAEL65zxzNImwdc8=
|
||||
github.com/prometheus/common v0.62.0 h1:xasJaQlnWAeyHdUBeGjXmutelfJHWMRr+Fg4QszZ2Io=
|
||||
github.com/prometheus/common v0.62.0/go.mod h1:vyBcEuLSvWos9B1+CyL7JZ2up+uFzXhkqml0W5zIY1I=
|
||||
github.com/prometheus/procfs v0.15.1 h1:YagwOFzUgYfKKHX6Dr+sHT7km/hxC76UB0learggepc=
|
||||
github.com/prometheus/procfs v0.15.1/go.mod h1:fB45yRUv8NstnjriLhBQLuOUt+WW4BsoGhij/e3PBqk=
|
||||
github.com/redis/go-redis/v9 v9.7.0 h1:HhLSs+B6O021gwzl+locl0zEDnyNkxMtf/Z3NNBMa9E=
|
||||
github.com/redis/go-redis/v9 v9.7.0/go.mod h1:f6zhXITC7JUJIlPEiBOTXxJgPLdZcA93GewI7inzyWw=
|
||||
github.com/redis/go-redis/v9 v9.7.3 h1:YpPyAayJV+XErNsatSElgRZZVCwXX9QzkKYNvO7x0wM=
|
||||
github.com/redis/go-redis/v9 v9.7.3/go.mod h1:bGUrSggJ9X9GUmZpZNEOQKaANxSGgOEBRltRTZHSvrA=
|
||||
github.com/rivo/uniseg v0.2.0 h1:S1pD9weZBuJdFmowNwbpi7BJ8TNftyUImj/0WQi72jY=
|
||||
github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
|
||||
github.com/rogpeppe/go-internal v1.10.0 h1:TMyTOH3F/DB16zRVcYyreMH6GnZZrwQVAoYjRBZyWFQ=
|
||||
github.com/rogpeppe/go-internal v1.10.0/go.mod h1:UQnix2H7Ngw/k4C5ijL5+65zddjncjaFoBhdsK/akog=
|
||||
github.com/spaolacci/murmur3 v1.1.0 h1:7c1g84S4BPRrfL5Xrdp6fOJ206sU9y293DDHaoy0bLI=
|
||||
github.com/spaolacci/murmur3 v1.1.0/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA=
|
||||
github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
|
||||
github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
|
||||
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
|
||||
github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
|
||||
github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
|
||||
@@ -168,8 +180,9 @@ github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/
|
||||
github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
|
||||
github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
|
||||
github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
|
||||
github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
|
||||
github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
|
||||
github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
|
||||
github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
|
||||
github.com/xdg-go/pbkdf2 v1.0.0 h1:Su7DPu48wXMwC3bs7MCNG+z4FhcyEuz5dlvchbq0B0c=
|
||||
github.com/xdg-go/pbkdf2 v1.0.0/go.mod h1:jrpuAogTd400dnrH08LKmI/xc1MbPOebTwRqcT5RDeI=
|
||||
github.com/xdg-go/scram v1.1.2 h1:FHX5I5B4i4hKRVRBCFRxq1iQRej7WO3hhBuJf+UUySY=
|
||||
@@ -190,8 +203,8 @@ go.etcd.io/etcd/client/pkg/v3 v3.5.15 h1:fo0HpWz/KlHGMCC+YejpiCmyWDEuIpnTDzpJLB5
|
||||
go.etcd.io/etcd/client/pkg/v3 v3.5.15/go.mod h1:mXDI4NAOwEiszrHCb0aqfAYNCrZP4e9hRca3d1YK8EU=
|
||||
go.etcd.io/etcd/client/v3 v3.5.15 h1:23M0eY4Fd/inNv1ZfU3AxrbbOdW79r9V9Rl62Nm6ip4=
|
||||
go.etcd.io/etcd/client/v3 v3.5.15/go.mod h1:CLSJxrYjvLtHsrPKsy7LmZEE+DK2ktfd2bN4RhBMwlU=
|
||||
go.mongodb.org/mongo-driver v1.17.1 h1:Wic5cJIwJgSpBhe3lx3+/RybR5PiYRMpVFgO7cOHyIM=
|
||||
go.mongodb.org/mongo-driver v1.17.1/go.mod h1:wwWm/+BuOddhcq3n68LKRmgk2wXzmF6s0SFOa0GINL4=
|
||||
go.mongodb.org/mongo-driver v1.17.3 h1:TQyXhnsWfWtgAhMtOgtYHMTkZIfBTpMTsMnd9ZBeHxQ=
|
||||
go.mongodb.org/mongo-driver v1.17.3/go.mod h1:Hy04i7O2kC4RS06ZrhPRqj/u4DTYkFDAAccj+rVKqgQ=
|
||||
go.opentelemetry.io/otel v1.24.0 h1:0LAOdjNmQeSTzGBzduGe/rU4tZhMwL5rWgtp9Ku5Jfo=
|
||||
go.opentelemetry.io/otel v1.24.0/go.mod h1:W7b9Ozg4nkF5tWI5zsXkaKKDjdVjpD4oAt9Qi/MArHo=
|
||||
go.opentelemetry.io/otel/exporters/jaeger v1.17.0 h1:D7UpUy2Xc2wsi1Ras6V40q806WM07rqoCWzXu7Sqy+4=
|
||||
@@ -228,8 +241,8 @@ golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACk
|
||||
golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
|
||||
golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
|
||||
golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
|
||||
golang.org/x/crypto v0.29.0 h1:L5SG1JTTXupVV3n6sUqMTeWbjAyfPwoda2DLX8J8FrQ=
|
||||
golang.org/x/crypto v0.29.0/go.mod h1:+F4F4N5hv6v38hfeYwTdx20oUvLLc+QfrE9Ax9HtgRg=
|
||||
golang.org/x/crypto v0.33.0 h1:IOBPskki6Lysi0lo9qQvbxiQ+FvsCC/YWOecCHAixus=
|
||||
golang.org/x/crypto v0.33.0/go.mod h1:bVdXmD7IV/4GdElGPozy6U7lWdRXA4qyRVGJV57uQ5M=
|
||||
golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
|
||||
golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
|
||||
golang.org/x/mod v0.4.2/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
|
||||
@@ -241,17 +254,17 @@ golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwY
|
||||
golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
|
||||
golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4/go.mod h1:p54w0d4576C0XHj96bSt6lcn1PtDYWL6XObtHCRCNQM=
|
||||
golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
|
||||
golang.org/x/net v0.31.0 h1:68CPQngjLL0r2AlUKiSxtQFKvzRVbnzLwMUn5SzcLHo=
|
||||
golang.org/x/net v0.31.0/go.mod h1:P4fl1q7dY2hnZFxEk4pPSkDHF+QqjitcnDjUQyMM+pM=
|
||||
golang.org/x/oauth2 v0.21.0 h1:tsimM75w1tF/uws5rbeHzIWxEqElMehnc+iW793zsZs=
|
||||
golang.org/x/oauth2 v0.21.0/go.mod h1:XYTD2NtWslqkgxebSiOHnXEap4TF09sJSc7H1sXbhtI=
|
||||
golang.org/x/net v0.35.0 h1:T5GQRQb2y08kTAByq9L4/bz8cipCdA8FbRTXewonqY8=
|
||||
golang.org/x/net v0.35.0/go.mod h1:EglIi67kWsHKlRzzVMUD93VMSWGFOMSZgxFjparz1Qk=
|
||||
golang.org/x/oauth2 v0.24.0 h1:KTBBxWqUa0ykRPLtV69rRto9TLXcqYkeswu48x/gvNE=
|
||||
golang.org/x/oauth2 v0.24.0/go.mod h1:XYTD2NtWslqkgxebSiOHnXEap4TF09sJSc7H1sXbhtI=
|
||||
golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
golang.org/x/sync v0.9.0 h1:fEo0HyrW1GIgZdpbhCRO0PkJajUS5H9IFUztCgEo2jQ=
|
||||
golang.org/x/sync v0.9.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
|
||||
golang.org/x/sync v0.11.0 h1:GGz8+XQP4FvTTrjZPzNKTMFtSXH80RAzG+5ghFPgK9w=
|
||||
golang.org/x/sync v0.11.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
|
||||
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
@@ -263,20 +276,20 @@ golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBc
|
||||
golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.27.0 h1:wBqf8DvsY9Y/2P8gAfPDEYNuS30J4lPHJxXSb/nJZ+s=
|
||||
golang.org/x/sys v0.27.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
|
||||
golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc=
|
||||
golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
|
||||
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
|
||||
golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
|
||||
golang.org/x/term v0.26.0 h1:WEQa6V3Gja/BhNxg540hBip/kkaYtRg3cxg4oXSw4AU=
|
||||
golang.org/x/term v0.26.0/go.mod h1:Si5m1o57C5nBNQo5z1iq+XDijt21BDBDp2bK0QI8e3E=
|
||||
golang.org/x/term v0.29.0 h1:L6pJp37ocefwRRtYPKSWOWzOtWSxVajvz2ldH/xi3iU=
|
||||
golang.org/x/term v0.29.0/go.mod h1:6bl4lRlvVuDgSf3179VpIxBF0o10JUpXWOnI7nErv7s=
|
||||
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
||||
golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
|
||||
golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
|
||||
golang.org/x/text v0.3.8/go.mod h1:E6s5w1FMmriuDzIBO73fBruAKo1PCIq6d2Q6DHfQ8WQ=
|
||||
golang.org/x/text v0.20.0 h1:gK/Kv2otX8gz+wn7Rmb3vT96ZwuoxnQlY+HlJVj7Qug=
|
||||
golang.org/x/text v0.20.0/go.mod h1:D4IsuqiFMhST5bX19pQ9ikHC2GsaKyk/oF+pn3ducp4=
|
||||
golang.org/x/time v0.8.0 h1:9i3RxcPv3PZnitoVGMPDKZSq1xW1gK1Xy3ArNOGZfEg=
|
||||
golang.org/x/time v0.8.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
|
||||
golang.org/x/text v0.22.0 h1:bofq7m3/HAFvbF51jz3Q9wLg3jkvSPuiZu/pD1XwgtM=
|
||||
golang.org/x/text v0.22.0/go.mod h1:YRoo4H8PVmsu+E3Ou7cqLVH8oXWIHVoX0jqUWALQhfY=
|
||||
golang.org/x/time v0.10.0 h1:3usCWA8tQn0L8+hFJQNgzpWbd89begxN66o1Ojdn5L4=
|
||||
golang.org/x/time v0.10.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
|
||||
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
|
||||
golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
|
||||
golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=
|
||||
@@ -284,6 +297,7 @@ golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4f
|
||||
golang.org/x/tools v0.1.1/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
|
||||
golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
|
||||
golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d h1:vU5i/LfpvrRCpgM/VPfJLg5KjxD3E+hfT1SH+d9zLwg=
|
||||
golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d/go.mod h1:aiJjzUbINMkxbQROHiO6hDPo2LHcIPhhQsa9DLh0yGk=
|
||||
golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
@@ -294,8 +308,8 @@ google.golang.org/genproto/googleapis/rpc v0.0.0-20240701130421-f6361c86f094 h1:
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20240701130421-f6361c86f094/go.mod h1:Ue6ibwXGpU+dqIcODieyLOcgj7z8+IcskoNIgZxtrFY=
|
||||
google.golang.org/grpc v1.65.0 h1:bs/cUb4lp1G5iImFFd3u5ixQzweKizoZJAwBNLR42lc=
|
||||
google.golang.org/grpc v1.65.0/go.mod h1:WgYC2ypjlB0EiQi6wdKixMqukr6lBc0Vo+oOgjrM5ZQ=
|
||||
google.golang.org/protobuf v1.35.2 h1:8Ar7bF+apOIoThw1EdZl0p1oWvMqTHmpA2fRTyZO8io=
|
||||
google.golang.org/protobuf v1.35.2/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE=
|
||||
google.golang.org/protobuf v1.36.5 h1:tPhr+woSbjfYvY6/GPufUoYizxw1cF/yFoxJ2fmpwlM=
|
||||
google.golang.org/protobuf v1.36.5/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
|
||||
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
|
||||
|
||||
@@ -111,6 +111,10 @@ func (p *comboHealthManager) IsReady() bool {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
|
||||
if len(p.probes) == 0 {
|
||||
return false
|
||||
}
|
||||
|
||||
for _, probe := range p.probes {
|
||||
if !probe.IsReady() {
|
||||
return false
|
||||
|
||||
@@ -43,7 +43,7 @@ func TestComboHealthManager(t *testing.T) {
|
||||
hm1 := NewHealthManager(probeName)
|
||||
hm2 := NewHealthManager(probeName + "2")
|
||||
|
||||
assert.True(t, chm.IsReady())
|
||||
assert.False(t, chm.IsReady())
|
||||
chm.addProbe(hm1)
|
||||
chm.addProbe(hm2)
|
||||
assert.False(t, chm.IsReady())
|
||||
@@ -57,7 +57,7 @@ func TestComboHealthManager(t *testing.T) {
|
||||
chm := newComboHealthManager()
|
||||
hm := NewHealthManager(probeName)
|
||||
|
||||
assert.True(t, chm.IsReady())
|
||||
assert.False(t, chm.IsReady())
|
||||
chm.addProbe(hm)
|
||||
assert.False(t, chm.IsReady())
|
||||
hm.MarkReady()
|
||||
@@ -127,7 +127,7 @@ func TestCreateHttpHandler(t *testing.T) {
|
||||
resp, err := http.Get(srv.URL)
|
||||
assert.Nil(t, err)
|
||||
_ = resp.Body.Close()
|
||||
assert.Equal(t, http.StatusOK, resp.StatusCode)
|
||||
assert.Equal(t, http.StatusServiceUnavailable, resp.StatusCode)
|
||||
|
||||
hm := NewHealthManager(probeName)
|
||||
defaultHealthManager.addProbe(hm)
|
||||
|
||||
40
mcp/config.go
Normal file
40
mcp/config.go
Normal file
@@ -0,0 +1,40 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/zeromicro/go-zero/rest"
|
||||
)
|
||||
|
||||
// McpConf defines the configuration for an MCP server.
|
||||
// It embeds rest.RestConf for HTTP server settings
|
||||
// and adds MCP-specific configuration options.
|
||||
type McpConf struct {
|
||||
rest.RestConf
|
||||
Mcp struct {
|
||||
// Name is the server name reported in initialize responses
|
||||
Name string `json:",optional"`
|
||||
|
||||
// Version is the server version reported in initialize responses
|
||||
Version string `json:",default=1.0.0"`
|
||||
|
||||
// ProtocolVersion is the MCP protocol version implemented
|
||||
ProtocolVersion string `json:",default=2024-11-05"`
|
||||
|
||||
// BaseUrl is the base URL for the server, used in SSE endpoint messages
|
||||
// If not set, defaults to http://localhost:{Port}
|
||||
BaseUrl string `json:",optional"`
|
||||
|
||||
// SseEndpoint is the path for Server-Sent Events connections
|
||||
SseEndpoint string `json:",default=/sse"`
|
||||
|
||||
// MessageEndpoint is the path for JSON-RPC requests
|
||||
MessageEndpoint string `json:",default=/message"`
|
||||
|
||||
// Cors contains allowed CORS origins
|
||||
Cors []string `json:",optional"`
|
||||
|
||||
// ToolTimeout is the maximum time allowed for tool execution
|
||||
ToolTimeout time.Duration `json:",default=30s"`
|
||||
}
|
||||
}
|
||||
63
mcp/config_test.go
Normal file
63
mcp/config_test.go
Normal file
@@ -0,0 +1,63 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/zeromicro/go-zero/core/conf"
|
||||
)
|
||||
|
||||
func TestMcpConfDefaults(t *testing.T) {
|
||||
// Test default values are set correctly when unmarshalled from JSON
|
||||
jsonConfig := `name: test-service
|
||||
port: 8080
|
||||
mcp:
|
||||
name: test-mcp-server
|
||||
version: 1.0.0
|
||||
`
|
||||
|
||||
var c McpConf
|
||||
err := conf.LoadFromYamlBytes([]byte(jsonConfig), &c)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Check default values
|
||||
assert.Equal(t, "test-mcp-server", c.Mcp.Name)
|
||||
assert.Equal(t, "1.0.0", c.Mcp.Version, "Default version should be 1.0.0")
|
||||
assert.Equal(t, "2024-11-05", c.Mcp.ProtocolVersion, "Default protocol version should be 2024-11-05")
|
||||
assert.Equal(t, "/sse", c.Mcp.SseEndpoint, "Default SSE endpoint should be /sse")
|
||||
assert.Equal(t, "/message", c.Mcp.MessageEndpoint, "Default message endpoint should be /message")
|
||||
assert.Equal(t, 30*time.Second, c.Mcp.ToolTimeout, "Default tool timeout should be 30s")
|
||||
}
|
||||
|
||||
func TestMcpConfCustomValues(t *testing.T) {
|
||||
// Test custom values can be set
|
||||
jsonConfig := `{
|
||||
"Name": "test-service",
|
||||
"Port": 8080,
|
||||
"Mcp": {
|
||||
"Name": "test-mcp-server",
|
||||
"Version": "2.0.0",
|
||||
"ProtocolVersion": "2025-01-01",
|
||||
"BaseUrl": "http://example.com",
|
||||
"SseEndpoint": "/custom-sse",
|
||||
"MessageEndpoint": "/custom-message",
|
||||
"Cors": ["http://localhost:3000", "http://example.com"],
|
||||
"ToolTimeout": "60s"
|
||||
}
|
||||
}`
|
||||
|
||||
var c McpConf
|
||||
err := conf.LoadFromJsonBytes([]byte(jsonConfig), &c)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Check custom values
|
||||
assert.Equal(t, "test-mcp-server", c.Mcp.Name, "Name should be inherited from RestConf")
|
||||
assert.Equal(t, "2.0.0", c.Mcp.Version, "Version should be customizable")
|
||||
assert.Equal(t, "2025-01-01", c.Mcp.ProtocolVersion, "Protocol version should be customizable")
|
||||
assert.Equal(t, "http://example.com", c.Mcp.BaseUrl, "BaseUrl should be customizable")
|
||||
assert.Equal(t, "/custom-sse", c.Mcp.SseEndpoint, "SSE endpoint should be customizable")
|
||||
assert.Equal(t, "/custom-message", c.Mcp.MessageEndpoint, "Message endpoint should be customizable")
|
||||
assert.Equal(t, []string{"http://localhost:3000", "http://example.com"}, c.Mcp.Cors, "CORS settings should be customizable")
|
||||
assert.Equal(t, 60*time.Second, c.Mcp.ToolTimeout, "Tool timeout should be customizable")
|
||||
}
|
||||
443
mcp/integration_test.go
Normal file
443
mcp/integration_test.go
Normal file
@@ -0,0 +1,443 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// syncResponseRecorder is a thread-safe wrapper around httptest.ResponseRecorder
|
||||
type syncResponseRecorder struct {
|
||||
*httptest.ResponseRecorder
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
// Create a new synchronized response recorder
|
||||
func newSyncResponseRecorder() *syncResponseRecorder {
|
||||
return &syncResponseRecorder{
|
||||
ResponseRecorder: httptest.NewRecorder(),
|
||||
}
|
||||
}
|
||||
|
||||
// Override Write method to synchronize access
|
||||
func (srr *syncResponseRecorder) Write(p []byte) (int, error) {
|
||||
srr.mu.Lock()
|
||||
defer srr.mu.Unlock()
|
||||
return srr.ResponseRecorder.Write(p)
|
||||
}
|
||||
|
||||
// Override WriteHeader method to synchronize access
|
||||
func (srr *syncResponseRecorder) WriteHeader(statusCode int) {
|
||||
srr.mu.Lock()
|
||||
defer srr.mu.Unlock()
|
||||
srr.ResponseRecorder.WriteHeader(statusCode)
|
||||
}
|
||||
|
||||
// Override Result method to synchronize access
|
||||
func (srr *syncResponseRecorder) Result() *http.Response {
|
||||
srr.mu.Lock()
|
||||
defer srr.mu.Unlock()
|
||||
return srr.ResponseRecorder.Result()
|
||||
}
|
||||
|
||||
// TestHTTPHandlerIntegration tests the HTTP handlers with a real server instance
|
||||
func TestHTTPHandlerIntegration(t *testing.T) {
|
||||
// Skip in short test mode
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test in short mode")
|
||||
}
|
||||
|
||||
// Create a test configuration
|
||||
conf := McpConf{}
|
||||
conf.Mcp.Name = "test-integration"
|
||||
conf.Mcp.Version = "1.0.0-test"
|
||||
conf.Mcp.ToolTimeout = 1 * time.Second
|
||||
|
||||
// Create a mock server directly
|
||||
server := &sseMcpServer{
|
||||
conf: conf,
|
||||
clients: make(map[string]*mcpClient),
|
||||
tools: make(map[string]Tool),
|
||||
prompts: make(map[string]Prompt),
|
||||
resources: make(map[string]Resource),
|
||||
}
|
||||
|
||||
// Register a test tool
|
||||
err := server.RegisterTool(Tool{
|
||||
Name: "echo",
|
||||
Description: "Echo tool for testing",
|
||||
InputSchema: InputSchema{
|
||||
Type: "object",
|
||||
Properties: map[string]any{
|
||||
"message": map[string]any{
|
||||
"type": "string",
|
||||
"description": "Message to echo",
|
||||
},
|
||||
},
|
||||
},
|
||||
Handler: func(params map[string]any) (any, error) {
|
||||
if msg, ok := params["message"].(string); ok {
|
||||
return fmt.Sprintf("Echo: %s", msg), nil
|
||||
}
|
||||
return "Echo: no message provided", nil
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create a test HTTP request to the SSE endpoint
|
||||
req := httptest.NewRequest("GET", "/sse", nil)
|
||||
w := newSyncResponseRecorder()
|
||||
|
||||
// Create a done channel to signal completion of test
|
||||
done := make(chan bool)
|
||||
|
||||
// Start the SSE handler in a goroutine
|
||||
go func() {
|
||||
// lock.Lock()
|
||||
server.handleSSE(w, req)
|
||||
// lock.Unlock()
|
||||
done <- true
|
||||
}()
|
||||
|
||||
// Allow time for the handler to process
|
||||
select {
|
||||
case <-time.After(100 * time.Millisecond):
|
||||
// Expected - handler would normally block indefinitely
|
||||
case <-done:
|
||||
// This shouldn't happen immediately - the handler should block
|
||||
t.Error("SSE handler returned unexpectedly")
|
||||
}
|
||||
|
||||
// Check the initial headers
|
||||
resp := w.Result()
|
||||
assert.Equal(t, "chunked", resp.Header.Get("Transfer-Encoding"))
|
||||
resp.Body.Close()
|
||||
|
||||
// The handler creates a client and sends the endpoint message
|
||||
var sessionId string
|
||||
|
||||
// Give the handler time to set up the client
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
// Check that a client was created
|
||||
server.clientsLock.Lock()
|
||||
assert.Equal(t, 1, len(server.clients))
|
||||
for id := range server.clients {
|
||||
sessionId = id
|
||||
}
|
||||
server.clientsLock.Unlock()
|
||||
|
||||
require.NotEmpty(t, sessionId, "Expected a session ID to be created")
|
||||
|
||||
// Now that we have a session ID, we can test the message endpoint
|
||||
messageBody, _ := json.Marshal(Request{
|
||||
JsonRpc: "2.0",
|
||||
ID: 1,
|
||||
Method: methodInitialize,
|
||||
Params: json.RawMessage(`{}`),
|
||||
})
|
||||
|
||||
// Create a message request
|
||||
reqURL := fmt.Sprintf("/message?%s=%s", sessionIdKey, sessionId)
|
||||
msgReq := httptest.NewRequest("POST", reqURL, bytes.NewReader(messageBody))
|
||||
msgW := newSyncResponseRecorder()
|
||||
|
||||
// Process the message
|
||||
server.handleRequest(msgW, msgReq)
|
||||
|
||||
// Check the response
|
||||
msgResp := msgW.Result()
|
||||
assert.Equal(t, http.StatusAccepted, msgResp.StatusCode)
|
||||
msgResp.Body.Close() // Ensure response body is closed
|
||||
}
|
||||
|
||||
// TestHandlerResponseFlow tests the flow of a full request/response cycle
|
||||
func TestHandlerResponseFlow(t *testing.T) {
|
||||
// Create a mock server for testing
|
||||
server := &sseMcpServer{
|
||||
conf: McpConf{},
|
||||
clients: map[string]*mcpClient{
|
||||
"test-session": {
|
||||
id: "test-session",
|
||||
channel: make(chan string, 10),
|
||||
initialized: true,
|
||||
},
|
||||
},
|
||||
tools: make(map[string]Tool),
|
||||
prompts: make(map[string]Prompt),
|
||||
resources: make(map[string]Resource),
|
||||
}
|
||||
|
||||
// Register test resources
|
||||
server.RegisterTool(Tool{
|
||||
Name: "test.tool",
|
||||
Description: "Test tool",
|
||||
InputSchema: InputSchema{Type: "object"},
|
||||
Handler: func(params map[string]any) (any, error) {
|
||||
return "tool result", nil
|
||||
},
|
||||
})
|
||||
|
||||
server.RegisterPrompt(Prompt{
|
||||
Name: "test.prompt",
|
||||
Description: "Test prompt",
|
||||
})
|
||||
|
||||
server.RegisterResource(Resource{
|
||||
Name: "test.resource",
|
||||
URI: "http://example.com",
|
||||
Description: "Test resource",
|
||||
})
|
||||
|
||||
// Create a request with session ID parameter
|
||||
reqURL := fmt.Sprintf("/message?%s=%s", sessionIdKey, "test-session")
|
||||
|
||||
// Test tools/list request
|
||||
toolsListBody, _ := json.Marshal(Request{
|
||||
JsonRpc: "2.0",
|
||||
ID: 1,
|
||||
Method: methodToolsList,
|
||||
Params: json.RawMessage(`{}`),
|
||||
})
|
||||
|
||||
toolsReq := httptest.NewRequest("POST", reqURL, bytes.NewReader(toolsListBody))
|
||||
toolsW := newSyncResponseRecorder()
|
||||
|
||||
// Process the request
|
||||
server.handleRequest(toolsW, toolsReq)
|
||||
|
||||
// Check the response code
|
||||
toolsResp := toolsW.Result()
|
||||
assert.Equal(t, http.StatusAccepted, toolsResp.StatusCode)
|
||||
toolsResp.Body.Close()
|
||||
|
||||
// Check the channel message
|
||||
client := server.clients["test-session"]
|
||||
select {
|
||||
case message := <-client.channel:
|
||||
assert.Contains(t, message, `"tools":[{"name":"test.tool"`)
|
||||
case <-time.After(100 * time.Millisecond):
|
||||
t.Fatal("Timed out waiting for tools/list response")
|
||||
}
|
||||
|
||||
// Test prompts/list request
|
||||
promptsListBody, _ := json.Marshal(Request{
|
||||
JsonRpc: "2.0",
|
||||
ID: 2,
|
||||
Method: methodPromptsList,
|
||||
Params: json.RawMessage(`{}`),
|
||||
})
|
||||
|
||||
promptsReq := httptest.NewRequest("POST", reqURL, bytes.NewReader(promptsListBody))
|
||||
promptsW := newSyncResponseRecorder()
|
||||
|
||||
// Process the request
|
||||
server.handleRequest(promptsW, promptsReq)
|
||||
|
||||
// Check the response code
|
||||
promptsResp := promptsW.Result()
|
||||
assert.Equal(t, http.StatusAccepted, promptsResp.StatusCode)
|
||||
promptsResp.Body.Close()
|
||||
|
||||
// Check the channel message
|
||||
select {
|
||||
case message := <-client.channel:
|
||||
assert.Contains(t, message, `"prompts":[{"name":"test.prompt"`)
|
||||
case <-time.After(100 * time.Millisecond):
|
||||
t.Fatal("Timed out waiting for prompts/list response")
|
||||
}
|
||||
|
||||
// Test resources/list request
|
||||
resourcesListBody, _ := json.Marshal(Request{
|
||||
JsonRpc: "2.0",
|
||||
ID: 3,
|
||||
Method: methodResourcesList,
|
||||
Params: json.RawMessage(`{}`),
|
||||
})
|
||||
|
||||
resourcesReq := httptest.NewRequest("POST", reqURL, bytes.NewReader(resourcesListBody))
|
||||
resourcesW := newSyncResponseRecorder()
|
||||
|
||||
// Process the request
|
||||
server.handleRequest(resourcesW, resourcesReq)
|
||||
|
||||
// Check the response code
|
||||
resourcesResp := resourcesW.Result()
|
||||
assert.Equal(t, http.StatusAccepted, resourcesResp.StatusCode)
|
||||
resourcesResp.Body.Close()
|
||||
|
||||
// Check the channel message
|
||||
select {
|
||||
case message := <-client.channel:
|
||||
assert.Contains(t, message, `"name":"test.resource"`)
|
||||
case <-time.After(100 * time.Millisecond):
|
||||
t.Fatal("Timed out waiting for resources/list response")
|
||||
}
|
||||
}
|
||||
|
||||
// TestProcessListMethods tests the list processing methods with pagination
|
||||
func TestProcessListMethods(t *testing.T) {
|
||||
server := &sseMcpServer{
|
||||
tools: make(map[string]Tool),
|
||||
prompts: make(map[string]Prompt),
|
||||
resources: make(map[string]Resource),
|
||||
}
|
||||
|
||||
// Add some test data
|
||||
for i := 1; i <= 5; i++ {
|
||||
tool := Tool{
|
||||
Name: fmt.Sprintf("tool%d", i),
|
||||
Description: fmt.Sprintf("Tool %d", i),
|
||||
InputSchema: InputSchema{Type: "object"},
|
||||
}
|
||||
server.tools[tool.Name] = tool
|
||||
|
||||
prompt := Prompt{
|
||||
Name: fmt.Sprintf("prompt%d", i),
|
||||
Description: fmt.Sprintf("Prompt %d", i),
|
||||
}
|
||||
server.prompts[prompt.Name] = prompt
|
||||
|
||||
resource := Resource{
|
||||
Name: fmt.Sprintf("resource%d", i),
|
||||
URI: fmt.Sprintf("http://example.com/%d", i),
|
||||
Description: fmt.Sprintf("Resource %d", i),
|
||||
}
|
||||
server.resources[resource.Name] = resource
|
||||
}
|
||||
|
||||
// Create a test client
|
||||
client := &mcpClient{
|
||||
id: "test-client",
|
||||
channel: make(chan string, 10),
|
||||
initialized: true,
|
||||
}
|
||||
|
||||
// Test processListTools
|
||||
req := Request{
|
||||
JsonRpc: "2.0",
|
||||
ID: 1,
|
||||
Method: methodToolsList,
|
||||
Params: json.RawMessage(`{"cursor": "", "_meta": {"progressToken": "token1"}}`),
|
||||
}
|
||||
|
||||
server.processListTools(client, req)
|
||||
|
||||
// Read response
|
||||
select {
|
||||
case response := <-client.channel:
|
||||
assert.Contains(t, response, `"tools":`)
|
||||
assert.Contains(t, response, `"progressToken":"token1"`)
|
||||
case <-time.After(100 * time.Millisecond):
|
||||
t.Fatal("Timed out waiting for tools/list response")
|
||||
}
|
||||
|
||||
// Test processListPrompts
|
||||
req.ID = 2
|
||||
req.Method = methodPromptsList
|
||||
req.Params = json.RawMessage(`{"cursor": "next"}`)
|
||||
server.processListPrompts(client, req)
|
||||
|
||||
// Read response
|
||||
select {
|
||||
case response := <-client.channel:
|
||||
assert.Contains(t, response, `"prompts":`)
|
||||
case <-time.After(100 * time.Millisecond):
|
||||
t.Fatal("Timed out waiting for prompts/list response")
|
||||
}
|
||||
|
||||
// Test processListResources
|
||||
req.ID = 3
|
||||
req.Method = methodResourcesList
|
||||
req.Params = json.RawMessage(`{"cursor": "next"}`)
|
||||
server.processListResources(client, req)
|
||||
|
||||
// Read response
|
||||
select {
|
||||
case response := <-client.channel:
|
||||
assert.Contains(t, response, `"resources":`)
|
||||
case <-time.After(100 * time.Millisecond):
|
||||
t.Fatal("Timed out waiting for resources/list response")
|
||||
}
|
||||
}
|
||||
|
||||
// TestErrorResponseHandling tests error handling in the server
|
||||
func TestErrorResponseHandling(t *testing.T) {
|
||||
server := &sseMcpServer{
|
||||
tools: make(map[string]Tool),
|
||||
prompts: make(map[string]Prompt),
|
||||
resources: make(map[string]Resource),
|
||||
}
|
||||
|
||||
// Create a test client
|
||||
client := &mcpClient{
|
||||
id: "test-client",
|
||||
channel: make(chan string, 10),
|
||||
initialized: true,
|
||||
}
|
||||
|
||||
// Test invalid method
|
||||
req := Request{
|
||||
JsonRpc: "2.0",
|
||||
ID: 1,
|
||||
Method: "invalid_method",
|
||||
Params: json.RawMessage(`{}`),
|
||||
}
|
||||
|
||||
// Mock handleRequest by directly calling error handler
|
||||
server.sendErrorResponse(client, req.ID, "Method not found", errCodeMethodNotFound)
|
||||
|
||||
// Check response
|
||||
select {
|
||||
case response := <-client.channel:
|
||||
assert.Contains(t, response, `"error":{"code":-32601,"message":"Method not found"}`)
|
||||
case <-time.After(100 * time.Millisecond):
|
||||
t.Fatal("Timed out waiting for error response")
|
||||
}
|
||||
|
||||
// Test invalid tool
|
||||
toolReq := Request{
|
||||
JsonRpc: "2.0",
|
||||
ID: 2,
|
||||
Method: methodToolsCall,
|
||||
Params: json.RawMessage(`{"name":"non_existent_tool"}`),
|
||||
}
|
||||
|
||||
// Call process method directly
|
||||
server.processToolCall(client, toolReq)
|
||||
|
||||
// Check response
|
||||
select {
|
||||
case response := <-client.channel:
|
||||
assert.Contains(t, response, `"error":{"code":-32602,"message":"Tool 'non_existent_tool' not found"}`)
|
||||
case <-time.After(100 * time.Millisecond):
|
||||
t.Fatal("Timed out waiting for error response")
|
||||
}
|
||||
|
||||
// Test invalid prompt
|
||||
promptReq := Request{
|
||||
JsonRpc: "2.0",
|
||||
ID: 3,
|
||||
Method: methodPromptsGet,
|
||||
Params: json.RawMessage(`{"name":"non_existent_prompt"}`),
|
||||
}
|
||||
|
||||
// Call process method directly
|
||||
server.processGetPrompt(client, promptReq)
|
||||
|
||||
// Check response
|
||||
select {
|
||||
case response := <-client.channel:
|
||||
assert.Contains(t, response, `"error":{"code":-32602,"message":"Prompt 'non_existent_prompt' not found"}`)
|
||||
case <-time.After(100 * time.Millisecond):
|
||||
t.Fatal("Timed out waiting for error response")
|
||||
}
|
||||
}
|
||||
62
mcp/readme.md
Normal file
62
mcp/readme.md
Normal file
@@ -0,0 +1,62 @@
|
||||
# Model Context Protocol (MCP) SDK Implementation
|
||||
|
||||
## Overview
|
||||
This package implements a Model Context Protocol (MCP) server in Go that facilitates real-time communication between AI models and clients using Server-Sent Events (SSE). The implementation provides a framework for building AI-assisted applications with bidirectional communication capabilities.
|
||||
|
||||
## Core Components
|
||||
|
||||
### Server-Sent Events (SSE) Communication
|
||||
- **Real-time Communication**: Robust SSE-based communication system that maintains persistent connections with clients
|
||||
- **Connection Management**: Client registration, message broadcasting, and client cleanup mechanisms
|
||||
- **Event Handling**: Event types for tools, prompts, and resources changes
|
||||
|
||||
### JSON-RPC Implementation
|
||||
- **Request Processing**: Complete JSON-RPC request processor for handling MCP protocol methods
|
||||
- **Response Formatting**: Proper response formatting according to JSON-RPC specifications
|
||||
- **Error Handling**: Comprehensive error handling with appropriate error codes
|
||||
|
||||
### Tool Management
|
||||
- **Tool Registration**: System to register custom tools with handlers
|
||||
- **Tool Execution**: Mechanism to execute tool functions with proper timeout handling
|
||||
- **Result Handling**: Flexible result handling supporting various return types (string, JSON, images)
|
||||
|
||||
### Prompt System
|
||||
- **Prompt Registration**: System for registering both static and dynamic prompts
|
||||
- **Argument Validation**: Validation for required arguments and default values for optional ones
|
||||
- **Message Generation**: Handlers that generate properly formatted conversation messages
|
||||
|
||||
### Resource Management
|
||||
- **Resource Registration**: System for managing and accessing external resources
|
||||
- **Content Delivery**: Handlers for delivering resource content to clients on demand
|
||||
- **Resource Subscription**: Mechanisms for clients to subscribe to resource updates
|
||||
|
||||
### Protocol Features
|
||||
- **Initialization Sequence**: Proper handshaking with capability negotiation
|
||||
- **Notification Handling**: Support for both standard and client-specific notifications
|
||||
- **Message Routing**: Intelligent routing of requests to appropriate handlers
|
||||
|
||||
## Technical Highlights
|
||||
|
||||
### Configuration System
|
||||
- **Flexible Configuration**: Configuration system with sensible defaults and customization options
|
||||
- **CORS Support**: Configurable CORS settings for cross-origin requests
|
||||
- **Server Information**: Proper server identification and versioning
|
||||
|
||||
### Client Session Management
|
||||
- **Session Tracking**: Client session tracking with unique identifiers
|
||||
- **Connection Health**: Ping/pong mechanism to maintain connection health
|
||||
- **Initialization State**: Client initialization state tracking
|
||||
|
||||
### Content Handling
|
||||
- **Multi-format Content**: Support for text, code, and binary content
|
||||
- **MIME Type Support**: Proper MIME type identification for various content types
|
||||
- **Audience Annotations**: Content audience annotations for user/assistant targeting
|
||||
|
||||
## Usage
|
||||
|
||||
To create and use an MCP server, see the examples directory for practical implementation examples including:
|
||||
- Tool registration and execution
|
||||
- Static and dynamic prompt creation
|
||||
- Resource handling with proper URI identification
|
||||
- Embedded resources in prompt responses
|
||||
- Client connection management
|
||||
974
mcp/server.go
Normal file
974
mcp/server.go
Normal file
@@ -0,0 +1,974 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/zeromicro/go-zero/core/logx"
|
||||
"github.com/zeromicro/go-zero/rest"
|
||||
)
|
||||
|
||||
func NewMcpServer(c McpConf) McpServer {
|
||||
var server *rest.Server
|
||||
if len(c.Mcp.Cors) == 0 {
|
||||
server = rest.MustNewServer(c.RestConf)
|
||||
} else {
|
||||
server = rest.MustNewServer(c.RestConf, rest.WithCors(c.Mcp.Cors...))
|
||||
}
|
||||
|
||||
if len(c.Mcp.Name) == 0 {
|
||||
c.Mcp.Name = c.Name
|
||||
}
|
||||
if len(c.Mcp.BaseUrl) == 0 {
|
||||
c.Mcp.BaseUrl = fmt.Sprintf("http://localhost:%d", c.Port)
|
||||
}
|
||||
|
||||
s := &sseMcpServer{
|
||||
conf: c,
|
||||
server: server,
|
||||
clients: make(map[string]*mcpClient),
|
||||
tools: make(map[string]Tool),
|
||||
prompts: make(map[string]Prompt),
|
||||
resources: make(map[string]Resource),
|
||||
}
|
||||
|
||||
// SSE endpoint for real-time updates
|
||||
s.server.AddRoute(rest.Route{
|
||||
Method: http.MethodGet,
|
||||
Path: s.conf.Mcp.SseEndpoint,
|
||||
Handler: s.handleSSE,
|
||||
}, rest.WithSSE())
|
||||
|
||||
// JSON-RPC message endpoint for regular requests
|
||||
s.server.AddRoute(rest.Route{
|
||||
Method: http.MethodPost,
|
||||
Path: s.conf.Mcp.MessageEndpoint,
|
||||
Handler: s.handleRequest,
|
||||
})
|
||||
|
||||
return s
|
||||
}
|
||||
|
||||
// RegisterPrompt registers a new prompt with the server
|
||||
func (s *sseMcpServer) RegisterPrompt(prompt Prompt) {
|
||||
s.promptsLock.Lock()
|
||||
s.prompts[prompt.Name] = prompt
|
||||
s.promptsLock.Unlock()
|
||||
// Notify clients about the new prompt
|
||||
s.broadcast(eventPromptsListChanged, map[string][]Prompt{keyPrompts: {prompt}})
|
||||
}
|
||||
|
||||
// RegisterResource registers a new resource with the server
|
||||
func (s *sseMcpServer) RegisterResource(resource Resource) {
|
||||
s.resourcesLock.Lock()
|
||||
s.resources[resource.URI] = resource
|
||||
s.resourcesLock.Unlock()
|
||||
// Notify clients about the new resource
|
||||
s.broadcast(eventResourcesListChanged, map[string][]Resource{keyResources: {resource}})
|
||||
}
|
||||
|
||||
// RegisterTool registers a new tool with the server
|
||||
func (s *sseMcpServer) RegisterTool(tool Tool) error {
|
||||
if tool.Handler == nil {
|
||||
return fmt.Errorf("tool '%s' has no handler function", tool.Name)
|
||||
}
|
||||
|
||||
s.toolsLock.Lock()
|
||||
s.tools[tool.Name] = tool
|
||||
s.toolsLock.Unlock()
|
||||
// Notify clients about the new tool
|
||||
s.broadcast(eventToolsListChanged, map[string][]Tool{keyTools: {tool}})
|
||||
return nil
|
||||
}
|
||||
|
||||
// Start implements McpServer.
|
||||
func (s *sseMcpServer) Start() {
|
||||
s.server.Start()
|
||||
}
|
||||
|
||||
func (s *sseMcpServer) Stop() {
|
||||
s.server.Stop()
|
||||
}
|
||||
|
||||
// broadcast sends a message to all connected clients
|
||||
// It uses Server-Sent Events (SSE) format for real-time communication
|
||||
func (s *sseMcpServer) broadcast(event string, data any) {
|
||||
jsonData, err := json.Marshal(data)
|
||||
if err != nil {
|
||||
logx.Errorf("Failed to marshal broadcast data: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
// Lock only while reading the clients map
|
||||
s.clientsLock.Lock()
|
||||
clients := make([]*mcpClient, 0, len(s.clients))
|
||||
for _, client := range s.clients {
|
||||
clients = append(clients, client)
|
||||
}
|
||||
s.clientsLock.Unlock()
|
||||
|
||||
clientCount := len(clients)
|
||||
if clientCount == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
logx.Infof("Broadcasting event '%s' to %d clients", event, clientCount)
|
||||
|
||||
// Use CRLF line endings as per SSE specification
|
||||
message := fmt.Sprintf("event: %s\r\ndata: %s\r\n\r\n", event, string(jsonData))
|
||||
|
||||
// Send messages without holding the lock
|
||||
for _, client := range clients {
|
||||
select {
|
||||
case client.channel <- message:
|
||||
// Message sent successfully
|
||||
default:
|
||||
// Channel buffer is full, log warning and continue
|
||||
logx.Errorf("Client channel buffer full, dropping message for client %s", client.id)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// cleanupClient removes a client from the active clients map
|
||||
func (s *sseMcpServer) cleanupClient(sessionId string) {
|
||||
s.clientsLock.Lock()
|
||||
defer s.clientsLock.Unlock()
|
||||
|
||||
if client, exists := s.clients[sessionId]; exists {
|
||||
// Close the channel to signal any goroutines waiting on it
|
||||
close(client.channel)
|
||||
// Remove from active clients
|
||||
delete(s.clients, sessionId)
|
||||
logx.Infof("Cleaned up client %s (remaining clients: %d)", sessionId, len(s.clients))
|
||||
}
|
||||
}
|
||||
|
||||
// handleRequest handles MCP JSON-RPC requests
|
||||
func (s *sseMcpServer) handleRequest(w http.ResponseWriter, r *http.Request) {
|
||||
// Extract sessionId from query parameters
|
||||
sessionId := r.URL.Query().Get(sessionIdKey)
|
||||
if len(sessionId) == 0 {
|
||||
http.Error(w, fmt.Sprintf("Missing %s", sessionIdKey), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
// Check if the client with this sessionId exists
|
||||
s.clientsLock.Lock()
|
||||
client, exists := s.clients[sessionId]
|
||||
s.clientsLock.Unlock()
|
||||
|
||||
if !exists {
|
||||
http.Error(w, fmt.Sprintf("Invalid or expired %s", sessionIdKey), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
var req Request
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, "Invalid request", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
w.WriteHeader(http.StatusAccepted)
|
||||
|
||||
// For notification methods (no ID), we don't send a response
|
||||
isNotification := req.ID == 0
|
||||
|
||||
// Special handling for initialization sequence
|
||||
// Always allow initialize and notifications/initialized regardless of client state
|
||||
if req.Method == methodInitialize {
|
||||
logx.Infof("Processing initialize request with ID: %d", req.ID)
|
||||
s.processInitialize(client, req)
|
||||
logx.Infof("Sent initialize response for ID: %d, waiting for notifications/initialized", req.ID)
|
||||
return
|
||||
} else if req.Method == methodNotificationsInitialized {
|
||||
// Handle initialized notification
|
||||
logx.Info("Received notifications/initialized notification")
|
||||
if !isNotification {
|
||||
s.sendErrorResponse(client, req.ID, "Method should be used as a notification", errCodeInvalidRequest)
|
||||
return
|
||||
}
|
||||
s.processNotificationInitialized(client)
|
||||
return
|
||||
} else if !client.initialized && req.Method != methodNotificationsCancelled {
|
||||
// Block most requests until client is initialized (except for cancellations)
|
||||
s.sendErrorResponse(client, req.ID, "Client not fully initialized, waiting for notifications/initialized",
|
||||
errCodeClientNotInitialized)
|
||||
return
|
||||
}
|
||||
|
||||
// Process normal requests only after initialization
|
||||
switch req.Method {
|
||||
case methodToolsCall:
|
||||
logx.Infof("Received tools call request with ID: %d", req.ID)
|
||||
s.processToolCall(client, req)
|
||||
logx.Infof("Sent tools call response for ID: %d", req.ID)
|
||||
case methodToolsList:
|
||||
logx.Infof("Processing tools/list request with ID: %d", req.ID)
|
||||
s.processListTools(client, req)
|
||||
logx.Infof("Sent tools/list response for ID: %d", req.ID)
|
||||
case methodPromptsList:
|
||||
logx.Infof("Processing prompts/list request with ID: %d", req.ID)
|
||||
s.processListPrompts(client, req)
|
||||
logx.Infof("Sent prompts/list response for ID: %d", req.ID)
|
||||
case methodPromptsGet:
|
||||
logx.Infof("Processing prompts/get request with ID: %d", req.ID)
|
||||
s.processGetPrompt(client, req)
|
||||
logx.Infof("Sent prompts/get response for ID: %d", req.ID)
|
||||
case methodResourcesList:
|
||||
logx.Infof("Processing resources/list request with ID: %d", req.ID)
|
||||
s.processListResources(client, req)
|
||||
logx.Infof("Sent resources/list response for ID: %d", req.ID)
|
||||
case methodResourcesRead:
|
||||
logx.Infof("Processing resources/read request with ID: %d", req.ID)
|
||||
s.processResourcesRead(client, req)
|
||||
logx.Infof("Sent resources/read response for ID: %d", req.ID)
|
||||
case methodResourcesSubscribe:
|
||||
logx.Infof("Processing resources/subscribe request with ID: %d", req.ID)
|
||||
s.processResourceSubscribe(client, req)
|
||||
logx.Infof("Sent resources/subscribe response for ID: %d", req.ID)
|
||||
case methodPing:
|
||||
logx.Infof("Processing ping request with ID: %d", req.ID)
|
||||
s.processPing(client, req)
|
||||
case methodNotificationsCancelled:
|
||||
logx.Infof("Received notifications/cancelled notification: %v", req.Params)
|
||||
s.processNotificationCancelled(client, req)
|
||||
default:
|
||||
logx.Infof("Unknown method: %s", req.Method)
|
||||
s.sendErrorResponse(client, req.ID, "Method not found", errCodeMethodNotFound)
|
||||
}
|
||||
}
|
||||
|
||||
// handleSSE handles Server-Sent Events connections
|
||||
func (s *sseMcpServer) handleSSE(w http.ResponseWriter, r *http.Request) {
|
||||
// Generate a unique session ID for this client
|
||||
sessionId := uuid.New().String()
|
||||
|
||||
// Create new client with buffered channel to prevent blocking
|
||||
client := &mcpClient{
|
||||
id: sessionId,
|
||||
channel: make(chan string, eventChanSize),
|
||||
}
|
||||
|
||||
// Add client to active clients map
|
||||
s.clientsLock.Lock()
|
||||
s.clients[sessionId] = client
|
||||
activeClients := len(s.clients)
|
||||
s.clientsLock.Unlock()
|
||||
|
||||
logx.Infof("New SSE connection established for client %s (active clients: %d)",
|
||||
sessionId, activeClients)
|
||||
|
||||
// Set proper SSE headers
|
||||
w.Header().Set("Transfer-Encoding", "chunked")
|
||||
|
||||
// Enable streaming
|
||||
flusher, ok := w.(http.Flusher)
|
||||
if !ok {
|
||||
logx.Error("Streaming not supported by the underlying http.ResponseWriter")
|
||||
http.Error(w, "Streaming not supported", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
// Send the message endpoint URL to the client
|
||||
endpoint := fmt.Sprintf("%s%s?%s=%s",
|
||||
s.conf.Mcp.BaseUrl, s.conf.Mcp.MessageEndpoint, sessionIdKey, sessionId)
|
||||
|
||||
// Format and send the endpoint message
|
||||
endpointMsg := formatSSEMessage(eventEndpoint, []byte(endpoint))
|
||||
if _, err := fmt.Fprint(w, endpointMsg); err != nil {
|
||||
logx.Errorf("Failed to send endpoint message to client %s: %v", sessionId, err)
|
||||
s.cleanupClient(sessionId)
|
||||
return
|
||||
}
|
||||
flusher.Flush()
|
||||
|
||||
// Set up keep-alive ping and client cleanup
|
||||
ticker := time.NewTicker(pingInterval.Load())
|
||||
defer func() {
|
||||
ticker.Stop()
|
||||
s.cleanupClient(sessionId)
|
||||
logx.Infof("SSE connection closed for client %s", sessionId)
|
||||
}()
|
||||
|
||||
// Message processing loop
|
||||
for {
|
||||
select {
|
||||
case message, ok := <-client.channel:
|
||||
if !ok {
|
||||
// Channel was closed, end connection
|
||||
logx.Infof("Client channel was closed for %s", sessionId)
|
||||
return
|
||||
}
|
||||
|
||||
// Write message to the response
|
||||
if _, err := fmt.Fprint(w, message); err != nil {
|
||||
logx.Infof("Failed to write message to client %s: %v", sessionId, err)
|
||||
return
|
||||
}
|
||||
flusher.Flush()
|
||||
case <-ticker.C:
|
||||
// Send keep-alive ping to maintain connection
|
||||
ping := fmt.Sprintf(`{"type":"ping","timestamp":"%s"}`, time.Now().String())
|
||||
pingMsg := formatSSEMessage("ping", []byte(ping))
|
||||
if _, err := fmt.Fprint(w, pingMsg); err != nil {
|
||||
logx.Errorf("Failed to send ping to client %s, closing connection: %v", sessionId, err)
|
||||
return
|
||||
}
|
||||
flusher.Flush()
|
||||
case <-r.Context().Done():
|
||||
// Client disconnected or request was canceled
|
||||
logx.Infof("Client %s disconnected: context done", sessionId)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// processInitialize processes the initialize request
|
||||
func (s *sseMcpServer) processInitialize(client *mcpClient, req Request) {
|
||||
// Create a proper JSON-RPC response that preserves the client's request ID
|
||||
result := initializationResponse{
|
||||
ProtocolVersion: s.conf.Mcp.ProtocolVersion,
|
||||
Capabilities: capabilities{
|
||||
Prompts: struct {
|
||||
ListChanged bool `json:"listChanged"`
|
||||
}{
|
||||
ListChanged: true,
|
||||
},
|
||||
Resources: struct {
|
||||
Subscribe bool `json:"subscribe"`
|
||||
ListChanged bool `json:"listChanged"`
|
||||
}{
|
||||
Subscribe: true,
|
||||
ListChanged: true,
|
||||
},
|
||||
Tools: struct {
|
||||
ListChanged bool `json:"listChanged"`
|
||||
}{
|
||||
ListChanged: true,
|
||||
},
|
||||
},
|
||||
ServerInfo: serverInfo{
|
||||
Name: s.conf.Mcp.Name,
|
||||
Version: s.conf.Mcp.Version,
|
||||
},
|
||||
}
|
||||
|
||||
// Mark client as initialized
|
||||
client.initialized = true
|
||||
|
||||
// Send response with client's original request ID
|
||||
s.sendResponse(client, req.ID, result)
|
||||
}
|
||||
|
||||
// processListTools processes the tools/list request
|
||||
func (s *sseMcpServer) processListTools(client *mcpClient, req Request) {
|
||||
// Extract pagination params if any
|
||||
var nextCursor string
|
||||
var progressToken any
|
||||
|
||||
// Extract meta data including progress token
|
||||
if req.Params != nil {
|
||||
var metaParams struct {
|
||||
Cursor string `json:"cursor"`
|
||||
Meta struct {
|
||||
ProgressToken any `json:"progressToken"`
|
||||
} `json:"_meta"`
|
||||
}
|
||||
if err := json.Unmarshal(req.Params, &metaParams); err == nil {
|
||||
if len(metaParams.Cursor) > 0 {
|
||||
nextCursor = metaParams.Cursor
|
||||
}
|
||||
progressToken = metaParams.Meta.ProgressToken
|
||||
}
|
||||
}
|
||||
|
||||
var toolsList []Tool
|
||||
s.toolsLock.Lock()
|
||||
for _, tool := range s.tools {
|
||||
toolsList = append(toolsList, tool)
|
||||
}
|
||||
s.toolsLock.Unlock()
|
||||
|
||||
result := ListToolsResult{
|
||||
PaginatedResult: PaginatedResult{
|
||||
Result: Result{},
|
||||
NextCursor: Cursor(nextCursor),
|
||||
},
|
||||
Tools: toolsList,
|
||||
}
|
||||
|
||||
// Add meta information if progress token was provided
|
||||
if progressToken != nil {
|
||||
result.Result.Meta = map[string]any{
|
||||
"progressToken": progressToken,
|
||||
}
|
||||
}
|
||||
|
||||
s.sendResponse(client, req.ID, result)
|
||||
}
|
||||
|
||||
// processListPrompts processes the prompts/list request
|
||||
func (s *sseMcpServer) processListPrompts(client *mcpClient, req Request) {
|
||||
// Extract pagination params if any
|
||||
var nextCursor string
|
||||
if req.Params != nil {
|
||||
var cursorParams struct {
|
||||
Cursor string `json:"cursor"`
|
||||
}
|
||||
if err := json.Unmarshal(req.Params, &cursorParams); err == nil && cursorParams.Cursor != "" {
|
||||
// If we have a valid cursor, we could use it for pagination
|
||||
// For now, we're not actually implementing pagination, so this is just
|
||||
// to show how it would be extracted from the request
|
||||
_ = cursorParams.Cursor
|
||||
}
|
||||
}
|
||||
|
||||
// Prepare prompt list
|
||||
var promptsList []Prompt
|
||||
s.promptsLock.Lock()
|
||||
for _, prompt := range s.prompts {
|
||||
promptsList = append(promptsList, prompt)
|
||||
}
|
||||
s.promptsLock.Unlock()
|
||||
|
||||
// In a real implementation, you'd handle pagination here
|
||||
// For now, we'll return all prompts at once
|
||||
result := struct {
|
||||
Prompts []Prompt `json:"prompts"`
|
||||
NextCursor string `json:"nextCursor,omitempty"`
|
||||
Meta *struct{} `json:"_meta,omitempty"`
|
||||
}{
|
||||
Prompts: promptsList,
|
||||
NextCursor: nextCursor,
|
||||
}
|
||||
|
||||
s.sendResponse(client, req.ID, result)
|
||||
}
|
||||
|
||||
// processListResources processes the resources/list request
|
||||
func (s *sseMcpServer) processListResources(client *mcpClient, req Request) {
|
||||
// Extract pagination params if any
|
||||
var nextCursor string
|
||||
var progressToken any
|
||||
|
||||
// Extract meta information including progress token if available
|
||||
if req.Params != nil {
|
||||
var metaParams PaginatedParams
|
||||
if err := json.Unmarshal(req.Params, &metaParams); err == nil {
|
||||
if len(metaParams.Cursor) > 0 {
|
||||
nextCursor = metaParams.Cursor
|
||||
}
|
||||
progressToken = metaParams.Meta.ProgressToken
|
||||
}
|
||||
}
|
||||
|
||||
var resourcesList []Resource
|
||||
s.resourcesLock.Lock()
|
||||
for _, resource := range s.resources {
|
||||
// Create a copy without the handler function which shouldn't be sent to clients
|
||||
resourceCopy := Resource{
|
||||
URI: resource.URI,
|
||||
Name: resource.Name,
|
||||
Description: resource.Description,
|
||||
MimeType: resource.MimeType,
|
||||
}
|
||||
resourcesList = append(resourcesList, resourceCopy)
|
||||
}
|
||||
s.resourcesLock.Unlock()
|
||||
|
||||
// Create proper ResourcesListResult according to MCP specification
|
||||
result := ResourcesListResult{
|
||||
PaginatedResult: PaginatedResult{
|
||||
Result: Result{},
|
||||
NextCursor: Cursor(nextCursor),
|
||||
},
|
||||
Resources: resourcesList,
|
||||
}
|
||||
|
||||
// Add meta information if progress token was provided
|
||||
if progressToken != nil {
|
||||
result.Result.Meta = map[string]any{
|
||||
"progressToken": progressToken,
|
||||
}
|
||||
}
|
||||
|
||||
s.sendResponse(client, req.ID, result)
|
||||
}
|
||||
|
||||
// processGetPrompt processes the prompts/get request
|
||||
func (s *sseMcpServer) processGetPrompt(client *mcpClient, req Request) {
|
||||
type GetPromptParams struct {
|
||||
Name string `json:"name"`
|
||||
Arguments map[string]string `json:"arguments,omitempty"`
|
||||
}
|
||||
|
||||
var params GetPromptParams
|
||||
if err := json.Unmarshal(req.Params, ¶ms); err != nil {
|
||||
s.sendErrorResponse(client, req.ID, "Invalid parameters", errCodeInvalidParams)
|
||||
return
|
||||
}
|
||||
|
||||
// Check if prompt exists
|
||||
s.promptsLock.Lock()
|
||||
prompt, exists := s.prompts[params.Name]
|
||||
s.promptsLock.Unlock()
|
||||
if !exists {
|
||||
message := fmt.Sprintf("Prompt '%s' not found", params.Name)
|
||||
s.sendErrorResponse(client, req.ID, message, errCodeInvalidParams)
|
||||
return
|
||||
}
|
||||
|
||||
logx.Infof("Processing prompt request: %s with %d arguments", prompt.Name, len(params.Arguments))
|
||||
|
||||
// Validate required arguments
|
||||
missingArgs := validatePromptArguments(prompt, params.Arguments)
|
||||
if len(missingArgs) > 0 {
|
||||
message := fmt.Sprintf("Missing required arguments: %s", strings.Join(missingArgs, ", "))
|
||||
s.sendErrorResponse(client, req.ID, message, errCodeInvalidParams)
|
||||
return
|
||||
}
|
||||
|
||||
// Apply default values for missing optional arguments
|
||||
args := applyDefaultArguments(prompt, params.Arguments)
|
||||
|
||||
// Generate messages using handler or static content
|
||||
var messages []PromptMessage
|
||||
var err error
|
||||
|
||||
if prompt.Handler != nil {
|
||||
// Use dynamic handler to generate messages
|
||||
logx.Info("Using prompt handler to generate content")
|
||||
messages, err = prompt.Handler(args)
|
||||
if err != nil {
|
||||
logx.Errorf("Error from prompt handler: %v", err)
|
||||
s.sendErrorResponse(client, req.ID, fmt.Sprintf("Error generating prompt content: %v", err), errCodeInternalError)
|
||||
return
|
||||
}
|
||||
} else {
|
||||
// No handler, generate messages from static content
|
||||
var messageText string
|
||||
if prompt.Content != "" {
|
||||
messageText = prompt.Content
|
||||
|
||||
// Apply argument substitutions to static content
|
||||
for key, value := range args {
|
||||
placeholder := fmt.Sprintf("{{%s}}", key)
|
||||
messageText = strings.Replace(messageText, placeholder, value, -1)
|
||||
}
|
||||
} else {
|
||||
// No content, use a default fallback
|
||||
topic := "this topic"
|
||||
if t, ok := args["topic"]; ok && t != "" {
|
||||
topic = t
|
||||
}
|
||||
messageText = fmt.Sprintf("Tell me about %s", topic)
|
||||
}
|
||||
|
||||
// Create a single user message with the content
|
||||
messages = []PromptMessage{
|
||||
{
|
||||
Role: roleUser,
|
||||
Content: TextContent{
|
||||
Type: contentTypeText,
|
||||
Text: messageText,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Construct the response according to MCP spec
|
||||
result := struct {
|
||||
Description string `json:"description,omitempty"`
|
||||
Messages []PromptMessage `json:"messages"`
|
||||
}{
|
||||
Description: prompt.Description,
|
||||
Messages: messages,
|
||||
}
|
||||
|
||||
s.sendResponse(client, req.ID, result)
|
||||
}
|
||||
|
||||
// validatePromptArguments checks if all required arguments are provided
|
||||
// Returns a list of missing required arguments
|
||||
func validatePromptArguments(prompt Prompt, providedArgs map[string]string) []string {
|
||||
var missingArgs []string
|
||||
|
||||
for _, arg := range prompt.Arguments {
|
||||
if arg.Required {
|
||||
if value, exists := providedArgs[arg.Name]; !exists || len(value) == 0 {
|
||||
missingArgs = append(missingArgs, arg.Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return missingArgs
|
||||
}
|
||||
|
||||
// applyDefaultArguments adds default values for missing optional arguments
|
||||
func applyDefaultArguments(prompt Prompt, providedArgs map[string]string) map[string]string {
|
||||
result := make(map[string]string)
|
||||
|
||||
// Copy all provided arguments
|
||||
for k, v := range providedArgs {
|
||||
result[k] = v
|
||||
}
|
||||
|
||||
// Add defaults for missing arguments
|
||||
for _, arg := range prompt.Arguments {
|
||||
if _, exists := result[arg.Name]; !exists && arg.Default != "" {
|
||||
result[arg.Name] = arg.Default
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// processToolCall processes the tools/call request
|
||||
func (s *sseMcpServer) processToolCall(client *mcpClient, req Request) {
|
||||
var toolCallParams struct {
|
||||
Name string `json:"name"`
|
||||
Arguments map[string]any `json:"arguments,omitempty"`
|
||||
Meta struct {
|
||||
ProgressToken any `json:"progressToken"`
|
||||
} `json:"_meta,omitempty"`
|
||||
}
|
||||
|
||||
// Handle different types of req.Params
|
||||
// If it's a RawMessage (JSON), unmarshal it
|
||||
if err := json.Unmarshal(req.Params, &toolCallParams); err != nil {
|
||||
logx.Errorf("Failed to unmarshal tool call params: %v", err)
|
||||
s.sendErrorResponse(client, req.ID, "Invalid tool call parameters", errCodeInvalidParams)
|
||||
return
|
||||
}
|
||||
|
||||
// Extract progress token if available
|
||||
progressToken := toolCallParams.Meta.ProgressToken
|
||||
|
||||
// Find the requested tool
|
||||
s.toolsLock.Lock()
|
||||
tool, exists := s.tools[toolCallParams.Name]
|
||||
s.toolsLock.Unlock()
|
||||
if !exists {
|
||||
s.sendErrorResponse(client, req.ID, fmt.Sprintf("Tool '%s' not found",
|
||||
toolCallParams.Name), errCodeInvalidParams)
|
||||
return
|
||||
}
|
||||
|
||||
// Create a context with the configured timeout
|
||||
ctx, cancel := context.WithTimeout(context.Background(), s.conf.Mcp.ToolTimeout)
|
||||
defer cancel()
|
||||
|
||||
// Log parameters before execution
|
||||
logx.Infof("Executing tool '%s' with arguments: %#v", toolCallParams.Name, toolCallParams.Arguments)
|
||||
|
||||
// Execute the tool handler with timeout handling
|
||||
var result any
|
||||
var err error
|
||||
|
||||
// Create a channel to receive the result
|
||||
resultCh := make(chan struct {
|
||||
result any
|
||||
err error
|
||||
}, 1)
|
||||
|
||||
// Execute the tool handler in a goroutine
|
||||
go func() {
|
||||
toolResult, toolErr := tool.Handler(toolCallParams.Arguments)
|
||||
resultCh <- struct {
|
||||
result any
|
||||
err error
|
||||
}{
|
||||
result: toolResult,
|
||||
err: toolErr,
|
||||
}
|
||||
}()
|
||||
|
||||
// Wait for either the result or a timeout
|
||||
select {
|
||||
case res := <-resultCh:
|
||||
result = res.result
|
||||
err = res.err
|
||||
case <-ctx.Done():
|
||||
// Handle timeout
|
||||
logx.Errorf("Tool execution timed out after %v: %s", s.conf.Mcp.ToolTimeout, toolCallParams.Name)
|
||||
s.sendErrorResponse(client, req.ID, "Tool execution timed out", errCodeTimeout)
|
||||
return
|
||||
}
|
||||
|
||||
// Create the base result structure with metadata
|
||||
callToolResult := CallToolResult{
|
||||
Result: Result{},
|
||||
Content: []any{},
|
||||
IsError: false,
|
||||
}
|
||||
|
||||
// Add meta information if progress token was provided
|
||||
if progressToken != nil {
|
||||
callToolResult.Result.Meta = map[string]any{
|
||||
"progressToken": progressToken,
|
||||
}
|
||||
}
|
||||
|
||||
// Check if there was an error during tool execution
|
||||
if err != nil {
|
||||
// According to the spec, for tool-level errors (as opposed to protocol-level errors),
|
||||
// we should report them inside the result with isError=true
|
||||
logx.Errorf("Tool execution reported error: %v", err)
|
||||
|
||||
callToolResult.Content = []any{
|
||||
TextContent{
|
||||
Type: contentTypeText,
|
||||
Text: fmt.Sprintf("Error: %v", err),
|
||||
},
|
||||
}
|
||||
callToolResult.IsError = true
|
||||
s.sendResponse(client, req.ID, callToolResult)
|
||||
return
|
||||
}
|
||||
|
||||
// Format the response according to the CallToolResult schema
|
||||
switch v := result.(type) {
|
||||
case string:
|
||||
// Simple string becomes text content
|
||||
callToolResult.Content = append(callToolResult.Content, TextContent{
|
||||
Type: contentTypeText,
|
||||
Text: v,
|
||||
Annotations: &Annotations{
|
||||
Audience: []roleType{roleUser, roleAssistant},
|
||||
},
|
||||
})
|
||||
case map[string]any:
|
||||
// JSON-like object becomes formatted JSON text
|
||||
jsonStr, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
jsonStr = []byte(err.Error())
|
||||
}
|
||||
callToolResult.Content = append(callToolResult.Content, TextContent{
|
||||
Type: contentTypeText,
|
||||
Text: string(jsonStr),
|
||||
Annotations: &Annotations{
|
||||
Audience: []roleType{roleUser, roleAssistant},
|
||||
},
|
||||
})
|
||||
case TextContent:
|
||||
// Direct TextContent object
|
||||
callToolResult.Content = append(callToolResult.Content, v)
|
||||
case ImageContent:
|
||||
// Direct ImageContent object
|
||||
callToolResult.Content = append(callToolResult.Content, v)
|
||||
case []any:
|
||||
// Array of content items
|
||||
callToolResult.Content = v
|
||||
case ToolResult:
|
||||
// Handle legacy ToolResult type
|
||||
switch v.Type {
|
||||
case contentTypeText:
|
||||
callToolResult.Content = append(callToolResult.Content, TextContent{
|
||||
Type: contentTypeText,
|
||||
Text: fmt.Sprintf("%v", v.Content),
|
||||
Annotations: &Annotations{
|
||||
Audience: []roleType{roleUser, roleAssistant},
|
||||
},
|
||||
})
|
||||
case contentTypeImage:
|
||||
if imgData, ok := v.Content.(map[string]any); ok {
|
||||
callToolResult.Content = append(callToolResult.Content, ImageContent{
|
||||
Type: contentTypeImage,
|
||||
Data: fmt.Sprintf("%v", imgData["data"]),
|
||||
MimeType: fmt.Sprintf("%v", imgData["mimeType"]),
|
||||
})
|
||||
}
|
||||
default:
|
||||
callToolResult.Content = append(callToolResult.Content, TextContent{
|
||||
Type: contentTypeText,
|
||||
Text: fmt.Sprintf("%v", v.Content),
|
||||
Annotations: &Annotations{
|
||||
Audience: []roleType{roleUser, roleAssistant},
|
||||
},
|
||||
})
|
||||
}
|
||||
default:
|
||||
// For any other type, convert to string
|
||||
callToolResult.Content = append(callToolResult.Content, TextContent{
|
||||
Type: contentTypeText,
|
||||
Text: fmt.Sprintf("%v", v),
|
||||
Annotations: &Annotations{
|
||||
Audience: []roleType{roleUser, roleAssistant},
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
logx.Infof("Tool call result: %#v", callToolResult)
|
||||
s.sendResponse(client, req.ID, callToolResult)
|
||||
}
|
||||
|
||||
// processResourcesRead processes the resources/read request
|
||||
func (s *sseMcpServer) processResourcesRead(client *mcpClient, req Request) {
|
||||
var params ResourceReadParams
|
||||
if err := json.Unmarshal(req.Params, ¶ms); err != nil {
|
||||
s.sendErrorResponse(client, req.ID, "Invalid parameters", errCodeInvalidParams)
|
||||
return
|
||||
}
|
||||
|
||||
// Find resource that matches the URI
|
||||
s.resourcesLock.Lock()
|
||||
resource, exists := s.resources[params.URI]
|
||||
s.resourcesLock.Unlock()
|
||||
|
||||
if !exists {
|
||||
s.sendErrorResponse(client, req.ID, fmt.Sprintf("Resource with URI '%s' not found",
|
||||
params.URI), errCodeResourceNotFound)
|
||||
return
|
||||
}
|
||||
|
||||
// If no handler is provided, return an empty content array
|
||||
if resource.Handler == nil {
|
||||
result := ResourceReadResult{
|
||||
Contents: []ResourceContent{
|
||||
{
|
||||
URI: params.URI,
|
||||
MimeType: resource.MimeType,
|
||||
Text: "",
|
||||
},
|
||||
},
|
||||
}
|
||||
s.sendResponse(client, req.ID, result)
|
||||
return
|
||||
}
|
||||
|
||||
// Execute the resource handler
|
||||
content, err := resource.Handler()
|
||||
if err != nil {
|
||||
s.sendErrorResponse(client, req.ID, fmt.Sprintf("Error reading resource: %v", err),
|
||||
errCodeInternalError)
|
||||
return
|
||||
}
|
||||
|
||||
// Ensure the URI is set if not already provided by the handler
|
||||
if len(content.URI) == 0 {
|
||||
content.URI = params.URI
|
||||
}
|
||||
|
||||
// Ensure MimeType is set if available from the resource definition
|
||||
if len(content.MimeType) == 0 && resource.MimeType != "" {
|
||||
content.MimeType = resource.MimeType
|
||||
}
|
||||
|
||||
// Create response with contents from the handler
|
||||
// The MCP specification requires a contents array
|
||||
result := ResourceReadResult{
|
||||
Contents: []ResourceContent{content},
|
||||
}
|
||||
|
||||
s.sendResponse(client, req.ID, result)
|
||||
}
|
||||
|
||||
// processResourceSubscribe processes the resources/subscribe request
|
||||
func (s *sseMcpServer) processResourceSubscribe(client *mcpClient, req Request) {
|
||||
var params ResourceSubscribeParams
|
||||
if err := json.Unmarshal(req.Params, ¶ms); err != nil {
|
||||
s.sendErrorResponse(client, req.ID, "Invalid parameters", errCodeInvalidParams)
|
||||
return
|
||||
}
|
||||
|
||||
// Check if the resource exists
|
||||
s.resourcesLock.Lock()
|
||||
_, exists := s.resources[params.URI]
|
||||
s.resourcesLock.Unlock()
|
||||
|
||||
if !exists {
|
||||
s.sendErrorResponse(client, req.ID, fmt.Sprintf("Resource with URI '%s' not found",
|
||||
params.URI), errCodeResourceNotFound)
|
||||
return
|
||||
}
|
||||
|
||||
// Send success response for the subscription
|
||||
s.sendResponse(client, req.ID, struct{}{})
|
||||
|
||||
logx.Infof("Client %s subscribed to resource '%s'", client.id, params.URI)
|
||||
}
|
||||
|
||||
// processNotificationCancelled processes the notifications/cancelled notification
|
||||
func (s *sseMcpServer) processNotificationCancelled(client *mcpClient, req Request) {
|
||||
// Extract the requestId that was canceled
|
||||
type CancelParams struct {
|
||||
RequestId int64 `json:"requestId"`
|
||||
Reason string `json:"reason"`
|
||||
}
|
||||
|
||||
var params CancelParams
|
||||
if err := json.Unmarshal(req.Params, ¶ms); err != nil {
|
||||
logx.Errorf("Failed to parse cancellation params: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
logx.Infof("Request %d was cancelled by client. Reason: %s", params.RequestId, params.Reason)
|
||||
}
|
||||
|
||||
// processNotificationInitialized processes the notifications/initialized notification
|
||||
func (s *sseMcpServer) processNotificationInitialized(client *mcpClient) {
|
||||
// Mark the client as properly initialized
|
||||
client.initialized = true
|
||||
logx.Infof("Client %s is now fully initialized and ready for normal operations", client.id)
|
||||
}
|
||||
|
||||
// processPing processes the ping request and responds immediately
|
||||
func (s *sseMcpServer) processPing(client *mcpClient, req Request) {
|
||||
// A ping request should simply respond with an empty result to confirm the server is alive
|
||||
logx.Infof("Received ping request with ID: %d", req.ID)
|
||||
|
||||
// Send an empty response with client's original request ID
|
||||
s.sendResponse(client, req.ID, struct{}{})
|
||||
|
||||
logx.Infof("Sent ping response for ID: %d", req.ID)
|
||||
}
|
||||
|
||||
// sendErrorResponse sends an error response via the SSE channel
|
||||
func (s *sseMcpServer) sendErrorResponse(client *mcpClient, id int64, message string, code int) {
|
||||
errorResponse := struct {
|
||||
JsonRpc string `json:"jsonrpc"`
|
||||
ID int64 `json:"id"`
|
||||
Error errorMessage `json:"error"`
|
||||
}{
|
||||
JsonRpc: jsonRpcVersion,
|
||||
ID: id,
|
||||
Error: errorMessage{
|
||||
Code: code,
|
||||
Message: message,
|
||||
},
|
||||
}
|
||||
|
||||
// all fields are primitive types, impossible to fail
|
||||
jsonData, _ := json.Marshal(errorResponse)
|
||||
// Use CRLF line endings as requested
|
||||
sseMessage := fmt.Sprintf("event: %s\r\ndata: %s\r\n\r\n", eventMessage, string(jsonData))
|
||||
logx.Infof("Sending error for ID %d: %s", id, sseMessage)
|
||||
|
||||
client.channel <- sseMessage
|
||||
}
|
||||
|
||||
// sendResponse sends a success response via the SSE channel
|
||||
func (s *sseMcpServer) sendResponse(client *mcpClient, id int64, result any) {
|
||||
response := Response{
|
||||
JsonRpc: jsonRpcVersion,
|
||||
ID: id,
|
||||
Result: result,
|
||||
}
|
||||
|
||||
jsonData, err := json.Marshal(response)
|
||||
if err != nil {
|
||||
s.sendErrorResponse(client, id, "Failed to marshal response", errCodeInternalError)
|
||||
return
|
||||
}
|
||||
|
||||
// Use CRLF line endings as requested
|
||||
sseMessage := fmt.Sprintf("event: %s\r\ndata: %s\r\n\r\n", eventMessage, string(jsonData))
|
||||
logx.Infof("Sending response for ID %d: %s", id, sseMessage)
|
||||
|
||||
client.channel <- sseMessage
|
||||
}
|
||||
3418
mcp/server_test.go
Normal file
3418
mcp/server_test.go
Normal file
File diff suppressed because it is too large
Load Diff
294
mcp/types.go
Normal file
294
mcp/types.go
Normal file
@@ -0,0 +1,294 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"sync"
|
||||
|
||||
"github.com/zeromicro/go-zero/rest"
|
||||
)
|
||||
|
||||
// Cursor is an opaque token used for pagination
|
||||
type Cursor string
|
||||
|
||||
// Request represents a generic MCP request following JSON-RPC 2.0 specification
|
||||
type Request struct {
|
||||
SessionId string `form:"session_id"` // Session identifier for client tracking
|
||||
JsonRpc string `json:"jsonrpc"` // Must be "2.0" per JSON-RPC spec
|
||||
ID int64 `json:"id"` // Request identifier for matching responses
|
||||
Method string `json:"method"` // Method name to invoke
|
||||
Params json.RawMessage `json:"params"` // Parameters for the method
|
||||
}
|
||||
|
||||
type PaginatedParams struct {
|
||||
Cursor string `json:"cursor"`
|
||||
Meta struct {
|
||||
ProgressToken any `json:"progressToken"`
|
||||
} `json:"_meta"`
|
||||
}
|
||||
|
||||
// Result is the base interface for all results
|
||||
type Result struct {
|
||||
Meta map[string]any `json:"_meta,omitempty"` // Optional metadata
|
||||
}
|
||||
|
||||
// PaginatedResult is a base for results that support pagination
|
||||
type PaginatedResult struct {
|
||||
Result
|
||||
NextCursor Cursor `json:"nextCursor,omitempty"` // Opaque token for fetching next page
|
||||
}
|
||||
|
||||
// ListToolsResult represents the response to a tools/list request
|
||||
type ListToolsResult struct {
|
||||
PaginatedResult
|
||||
Tools []Tool `json:"tools"` // List of available tools
|
||||
}
|
||||
|
||||
// Message Content Types
|
||||
|
||||
// roleType represents the sender or recipient of messages in a conversation
|
||||
type roleType string
|
||||
|
||||
// PromptArgument defines a single argument that can be passed to a prompt
|
||||
type PromptArgument struct {
|
||||
Name string `json:"name"` // Argument name
|
||||
Description string `json:"description,omitempty"` // Human-readable description
|
||||
Required bool `json:"required,omitempty"` // Whether this argument is required
|
||||
Default string `json:"default,omitempty"` // Default value if not provided
|
||||
}
|
||||
|
||||
// PromptHandler is a function that dynamically generates prompt content
|
||||
type PromptHandler func(args map[string]string) ([]PromptMessage, error)
|
||||
|
||||
// Prompt represents an MCP Prompt definition
|
||||
type Prompt struct {
|
||||
Name string `json:"name"` // Unique identifier for the prompt
|
||||
Description string `json:"description,omitempty"` // Human-readable description
|
||||
Arguments []PromptArgument `json:"arguments,omitempty"` // Arguments for customization
|
||||
Content string `json:"-"` // Static content (internal use only)
|
||||
Handler PromptHandler `json:"-"` // Handler for dynamic content generation
|
||||
}
|
||||
|
||||
// PromptMessage represents a message in a conversation
|
||||
type PromptMessage struct {
|
||||
Role roleType `json:"role"` // Message sender role
|
||||
Content any `json:"content"` // Message content (TextContent, ImageContent, etc.)
|
||||
}
|
||||
|
||||
// TextContent represents text content in a message
|
||||
type TextContent struct {
|
||||
Type string `json:"type"` // Always "text"
|
||||
Text string `json:"text"` // The text content
|
||||
Annotations *Annotations `json:"annotations,omitempty"` // Optional annotations
|
||||
}
|
||||
|
||||
// ImageContent represents image data in a message
|
||||
type ImageContent struct {
|
||||
Type string `json:"type"` // Always "image"
|
||||
Data string `json:"data"` // Base64-encoded image data
|
||||
MimeType string `json:"mimeType"` // MIME type (e.g., "image/png")
|
||||
}
|
||||
|
||||
// AudioContent represents audio data in a message
|
||||
type AudioContent struct {
|
||||
Type string `json:"type"` // Always "audio"
|
||||
Data string `json:"data"` // Base64-encoded audio data
|
||||
MimeType string `json:"mimeType"` // MIME type (e.g., "audio/mp3")
|
||||
}
|
||||
|
||||
// FileContent represents file content
|
||||
type FileContent struct {
|
||||
URI string `json:"uri"` // URI identifying the file
|
||||
MimeType string `json:"mimeType"` // MIME type of the file
|
||||
Text string `json:"text"` // File content as text
|
||||
}
|
||||
|
||||
// EmbeddedResource represents a resource embedded in a message
|
||||
type EmbeddedResource struct {
|
||||
Type string `json:"type"` // Always "resource"
|
||||
Resource struct {
|
||||
URI string `json:"uri"` // Resource URI
|
||||
MimeType string `json:"mimeType"` // MIME type of the resource
|
||||
Text string `json:"text,omitempty"` // Text content (if available)
|
||||
Blob string `json:"blob,omitempty"` // Base64 encoded blob data (if available)
|
||||
} `json:"resource"` // The resource data
|
||||
}
|
||||
|
||||
// Annotations provides additional metadata for content
|
||||
type Annotations struct {
|
||||
Audience []roleType `json:"audience,omitempty"` // Who should see this content
|
||||
Priority *float64 `json:"priority,omitempty"` // Optional priority (0-1)
|
||||
}
|
||||
|
||||
// Tool-related Types
|
||||
|
||||
// Tool Definition Types
|
||||
|
||||
// ToolHandler is a function that handles tool calls
|
||||
type ToolHandler func(params map[string]any) (any, error)
|
||||
|
||||
// Tool represents a Model Context Protocol Tool definition
|
||||
type Tool struct {
|
||||
Name string `json:"name"` // Unique identifier for the tool
|
||||
Description string `json:"description"` // Human-readable description
|
||||
InputSchema InputSchema `json:"inputSchema"` // JSON Schema for parameters
|
||||
Handler ToolHandler `json:"-"` // Not sent to clients
|
||||
}
|
||||
|
||||
// InputSchema represents tool's input schema in JSON Schema format
|
||||
type InputSchema struct {
|
||||
Type string `json:"type"` // Always "object" for tool inputs
|
||||
Properties map[string]any `json:"properties"` // Property definitions
|
||||
Required []string `json:"required,omitempty"` // List of required properties
|
||||
}
|
||||
|
||||
// CallToolResult represents a tool call result that conforms to the MCP schema
|
||||
type CallToolResult struct {
|
||||
Result
|
||||
Content []interface{} `json:"content"` // Content items (text, images, etc.)
|
||||
IsError bool `json:"isError,omitempty"` // True if tool execution failed
|
||||
}
|
||||
|
||||
// Resource represents a Model Context Protocol Resource definition
|
||||
type Resource struct {
|
||||
URI string `json:"uri"` // Unique resource identifier (RFC3986)
|
||||
Name string `json:"name"` // Human-readable name
|
||||
Description string `json:"description,omitempty"` // Optional description
|
||||
MimeType string `json:"mimeType,omitempty"` // Optional MIME type
|
||||
Handler ResourceHandler `json:"-"` // Internal handler not sent to clients
|
||||
}
|
||||
|
||||
// ResourceHandler is a function that handles resource read requests
|
||||
type ResourceHandler func() (ResourceContent, error)
|
||||
|
||||
// ResourceContent represents the content of a resource
|
||||
type ResourceContent struct {
|
||||
URI string `json:"uri"` // Resource URI (required)
|
||||
MimeType string `json:"mimeType,omitempty"` // MIME type of the resource
|
||||
Text string `json:"text,omitempty"` // Text content (if available)
|
||||
Blob string `json:"blob,omitempty"` // Base64 encoded blob data (if available)
|
||||
}
|
||||
|
||||
// ResourcesListResult represents the response to a resources/list request
|
||||
type ResourcesListResult struct {
|
||||
PaginatedResult
|
||||
Resources []Resource `json:"resources"` // List of available resources
|
||||
}
|
||||
|
||||
// ResourceReadParams contains parameters for a resources/read request
|
||||
type ResourceReadParams struct {
|
||||
URI string `json:"uri"` // URI of the resource to read
|
||||
}
|
||||
|
||||
// ResourceReadResult contains the result of a resources/read request
|
||||
type ResourceReadResult struct {
|
||||
Result
|
||||
Contents []ResourceContent `json:"contents"` // Array of resource content
|
||||
}
|
||||
|
||||
// ResourceSubscribeParams contains parameters for a resources/subscribe request
|
||||
type ResourceSubscribeParams struct {
|
||||
URI string `json:"uri"` // URI of the resource to subscribe to
|
||||
}
|
||||
|
||||
// ResourceUpdateNotification represents a notification about a resource update
|
||||
type ResourceUpdateNotification struct {
|
||||
URI string `json:"uri"` // URI of the updated resource
|
||||
Content ResourceContent `json:"content"` // New resource content
|
||||
}
|
||||
|
||||
// Client and Server Types
|
||||
|
||||
// mcpClient represents an SSE client connection
|
||||
type mcpClient struct {
|
||||
id string // Unique client identifier
|
||||
channel chan string // Channel for sending SSE messages
|
||||
initialized bool // Tracks if client has sent notifications/initialized
|
||||
}
|
||||
|
||||
// McpServer defines the interface for Model Context Protocol servers
|
||||
type McpServer interface {
|
||||
Start()
|
||||
Stop()
|
||||
RegisterTool(tool Tool) error
|
||||
RegisterPrompt(prompt Prompt)
|
||||
RegisterResource(resource Resource)
|
||||
}
|
||||
|
||||
// sseMcpServer implements the McpServer interface using SSE
|
||||
type sseMcpServer struct {
|
||||
conf McpConf
|
||||
server *rest.Server
|
||||
clients map[string]*mcpClient
|
||||
clientsLock sync.Mutex
|
||||
tools map[string]Tool
|
||||
toolsLock sync.Mutex
|
||||
prompts map[string]Prompt
|
||||
promptsLock sync.Mutex
|
||||
resources map[string]Resource
|
||||
resourcesLock sync.Mutex
|
||||
}
|
||||
|
||||
// Response Types
|
||||
|
||||
// errorObj represents a JSON-RPC error object
|
||||
type errorObj struct {
|
||||
Code int `json:"code"` // Error code
|
||||
Message string `json:"message"` // Error message
|
||||
}
|
||||
|
||||
// Response represents a JSON-RPC response
|
||||
type Response struct {
|
||||
JsonRpc string `json:"jsonrpc"` // Always "2.0"
|
||||
ID int64 `json:"id"` // Same as request ID
|
||||
Result any `json:"result"` // Result object (null if error)
|
||||
Error *errorObj `json:"error,omitempty"` // Error object (null if success)
|
||||
}
|
||||
|
||||
// Server Information Types
|
||||
|
||||
// serverInfo provides information about the server
|
||||
type serverInfo struct {
|
||||
Name string `json:"name"` // Server name
|
||||
Version string `json:"version"` // Server version
|
||||
}
|
||||
|
||||
// capabilities describes the server's capabilities
|
||||
type capabilities struct {
|
||||
Logging struct{} `json:"logging"`
|
||||
Prompts struct {
|
||||
ListChanged bool `json:"listChanged"` // Server will notify on prompt changes
|
||||
} `json:"prompts"`
|
||||
Resources struct {
|
||||
Subscribe bool `json:"subscribe"` // Server supports resource subscriptions
|
||||
ListChanged bool `json:"listChanged"` // Server will notify on resource changes
|
||||
} `json:"resources"`
|
||||
Tools struct {
|
||||
ListChanged bool `json:"listChanged"` // Server will notify on tool changes
|
||||
} `json:"tools"`
|
||||
}
|
||||
|
||||
// initializationResponse is sent in response to an initialize request
|
||||
type initializationResponse struct {
|
||||
ProtocolVersion string `json:"protocolVersion"` // Protocol version
|
||||
Capabilities capabilities `json:"capabilities"` // Server capabilities
|
||||
ServerInfo serverInfo `json:"serverInfo"` // Server information
|
||||
}
|
||||
|
||||
// ToolCallParams contains the parameters for a tool call
|
||||
type ToolCallParams struct {
|
||||
Name string `json:"name"` // Tool name
|
||||
Parameters map[string]any `json:"parameters"` // Tool parameters
|
||||
}
|
||||
|
||||
// ToolResult contains the result of a tool execution
|
||||
type ToolResult struct {
|
||||
Type string `json:"type"` // Content type (text, image, etc.)
|
||||
Content any `json:"content"` // Result content
|
||||
}
|
||||
|
||||
// errorMessage represents a detailed error message
|
||||
type errorMessage struct {
|
||||
Code int `json:"code"` // Error code
|
||||
Message string `json:"message"` // Error message
|
||||
Data any `json:",omitempty"` // Additional error data
|
||||
}
|
||||
212
mcp/types_test.go
Normal file
212
mcp/types_test.go
Normal file
@@ -0,0 +1,212 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestResponseMarshaling(t *testing.T) {
|
||||
// Test that the Response struct marshals correctly
|
||||
resp := Response{
|
||||
JsonRpc: "2.0",
|
||||
ID: 123,
|
||||
Result: map[string]string{
|
||||
"key": "value",
|
||||
},
|
||||
}
|
||||
|
||||
data, err := json.Marshal(resp)
|
||||
assert.NoError(t, err)
|
||||
assert.Contains(t, string(data), `"jsonrpc":"2.0"`)
|
||||
assert.Contains(t, string(data), `"id":123`)
|
||||
assert.Contains(t, string(data), `"result":{"key":"value"}`)
|
||||
|
||||
// Test response with error
|
||||
respWithError := Response{
|
||||
JsonRpc: "2.0",
|
||||
ID: 456,
|
||||
Error: &errorObj{
|
||||
Code: errCodeInvalidRequest,
|
||||
Message: "Invalid Request",
|
||||
},
|
||||
}
|
||||
|
||||
data, err = json.Marshal(respWithError)
|
||||
assert.NoError(t, err)
|
||||
assert.Contains(t, string(data), `"jsonrpc":"2.0"`)
|
||||
assert.Contains(t, string(data), `"id":456`)
|
||||
assert.Contains(t, string(data), `"error":{"code":-32600,"message":"Invalid Request"}`)
|
||||
}
|
||||
|
||||
func TestRequestUnmarshaling(t *testing.T) {
|
||||
// Test that the Request struct unmarshals correctly
|
||||
jsonStr := `{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 789,
|
||||
"method": "test_method",
|
||||
"params": {"key": "value"}
|
||||
}`
|
||||
|
||||
var req Request
|
||||
err := json.Unmarshal([]byte(jsonStr), &req)
|
||||
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, "2.0", req.JsonRpc)
|
||||
assert.Equal(t, int64(789), req.ID)
|
||||
assert.Equal(t, "test_method", req.Method)
|
||||
|
||||
// Check params unmarshaled correctly
|
||||
var params map[string]string
|
||||
err = json.Unmarshal(req.Params, ¶ms)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, "value", params["key"])
|
||||
}
|
||||
|
||||
func TestToolStructs(t *testing.T) {
|
||||
// Test Tool struct
|
||||
tool := Tool{
|
||||
Name: "test.tool",
|
||||
Description: "A test tool",
|
||||
InputSchema: InputSchema{
|
||||
Type: "object",
|
||||
Properties: map[string]any{
|
||||
"input": map[string]any{
|
||||
"type": "string",
|
||||
"description": "Input parameter",
|
||||
},
|
||||
},
|
||||
Required: []string{"input"},
|
||||
},
|
||||
Handler: func(params map[string]any) (any, error) {
|
||||
return "result", nil
|
||||
},
|
||||
}
|
||||
|
||||
// Verify fields are correct
|
||||
assert.Equal(t, "test.tool", tool.Name)
|
||||
assert.Equal(t, "A test tool", tool.Description)
|
||||
assert.Equal(t, "object", tool.InputSchema.Type)
|
||||
assert.Contains(t, tool.InputSchema.Properties, "input")
|
||||
propMap, ok := tool.InputSchema.Properties["input"].(map[string]any)
|
||||
assert.True(t, ok, "Property should be a map")
|
||||
assert.Equal(t, "string", propMap["type"])
|
||||
assert.NotNil(t, tool.Handler)
|
||||
|
||||
// Verify JSON marshalling (which should exclude Handler function)
|
||||
data, err := json.Marshal(tool)
|
||||
assert.NoError(t, err)
|
||||
assert.Contains(t, string(data), `"name":"test.tool"`)
|
||||
assert.Contains(t, string(data), `"description":"A test tool"`)
|
||||
assert.Contains(t, string(data), `"inputSchema":`)
|
||||
assert.NotContains(t, string(data), `"Handler":`)
|
||||
}
|
||||
|
||||
func TestPromptStructs(t *testing.T) {
|
||||
// Test Prompt struct
|
||||
prompt := Prompt{
|
||||
Name: "test.prompt",
|
||||
Description: "A test prompt description",
|
||||
}
|
||||
|
||||
// Verify fields are correct
|
||||
assert.Equal(t, "test.prompt", prompt.Name)
|
||||
assert.Equal(t, "A test prompt description", prompt.Description)
|
||||
|
||||
// Verify JSON marshalling
|
||||
data, err := json.Marshal(prompt)
|
||||
assert.NoError(t, err)
|
||||
assert.Contains(t, string(data), `"name":"test.prompt"`)
|
||||
assert.Contains(t, string(data), `"description":"A test prompt description"`)
|
||||
}
|
||||
|
||||
func TestResourceStructs(t *testing.T) {
|
||||
// Test Resource struct
|
||||
resource := Resource{
|
||||
Name: "test.resource",
|
||||
URI: "http://example.com/resource",
|
||||
Description: "A test resource",
|
||||
}
|
||||
|
||||
// Verify fields are correct
|
||||
assert.Equal(t, "test.resource", resource.Name)
|
||||
assert.Equal(t, "http://example.com/resource", resource.URI)
|
||||
assert.Equal(t, "A test resource", resource.Description)
|
||||
|
||||
// Verify JSON marshalling
|
||||
data, err := json.Marshal(resource)
|
||||
assert.NoError(t, err)
|
||||
assert.Contains(t, string(data), `"name":"test.resource"`)
|
||||
assert.Contains(t, string(data), `"uri":"http://example.com/resource"`)
|
||||
assert.Contains(t, string(data), `"description":"A test resource"`)
|
||||
}
|
||||
|
||||
func TestContentTypes(t *testing.T) {
|
||||
// Test TextContent
|
||||
textContent := TextContent{
|
||||
Type: "text",
|
||||
Text: "Sample text",
|
||||
Annotations: &Annotations{
|
||||
Audience: []roleType{roleUser, roleAssistant},
|
||||
Priority: ptr(1.0),
|
||||
},
|
||||
}
|
||||
|
||||
data, err := json.Marshal(textContent)
|
||||
assert.NoError(t, err)
|
||||
assert.Contains(t, string(data), `"type":"text"`)
|
||||
assert.Contains(t, string(data), `"text":"Sample text"`)
|
||||
assert.Contains(t, string(data), `"audience":["user","assistant"]`)
|
||||
assert.Contains(t, string(data), `"priority":1`)
|
||||
|
||||
// Test ImageContent
|
||||
imageContent := ImageContent{
|
||||
Type: "image",
|
||||
Data: "base64data",
|
||||
MimeType: "image/png",
|
||||
}
|
||||
|
||||
data, err = json.Marshal(imageContent)
|
||||
assert.NoError(t, err)
|
||||
assert.Contains(t, string(data), `"type":"image"`)
|
||||
assert.Contains(t, string(data), `"data":"base64data"`)
|
||||
assert.Contains(t, string(data), `"mimeType":"image/png"`)
|
||||
|
||||
// Test AudioContent
|
||||
audioContent := AudioContent{
|
||||
Type: "audio",
|
||||
Data: "base64audio",
|
||||
MimeType: "audio/mp3",
|
||||
}
|
||||
|
||||
data, err = json.Marshal(audioContent)
|
||||
assert.NoError(t, err)
|
||||
assert.Contains(t, string(data), `"type":"audio"`)
|
||||
assert.Contains(t, string(data), `"data":"base64audio"`)
|
||||
assert.Contains(t, string(data), `"mimeType":"audio/mp3"`)
|
||||
}
|
||||
|
||||
func TestCallToolResult(t *testing.T) {
|
||||
// Test CallToolResult
|
||||
result := CallToolResult{
|
||||
Result: Result{
|
||||
Meta: map[string]any{
|
||||
"progressToken": "token123",
|
||||
},
|
||||
},
|
||||
Content: []interface{}{
|
||||
TextContent{
|
||||
Type: "text",
|
||||
Text: "Sample result",
|
||||
},
|
||||
},
|
||||
IsError: false,
|
||||
}
|
||||
|
||||
data, err := json.Marshal(result)
|
||||
assert.NoError(t, err)
|
||||
assert.Contains(t, string(data), `"_meta":{"progressToken":"token123"}`)
|
||||
assert.Contains(t, string(data), `"content":[{"type":"text","text":"Sample result"}]`)
|
||||
assert.NotContains(t, string(data), `"isError":`)
|
||||
}
|
||||
15
mcp/util.go
Normal file
15
mcp/util.go
Normal file
@@ -0,0 +1,15 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
)
|
||||
|
||||
// ptr is a helper function to get a pointer to a value
|
||||
func ptr[T any](v T) *T {
|
||||
return &v
|
||||
}
|
||||
|
||||
// formatSSEMessage formats a Server-Sent Event message with proper CRLF line endings
|
||||
func formatSSEMessage(event string, data []byte) string {
|
||||
return fmt.Sprintf("event: %s\r\ndata: %s\r\n\r\n", event, string(data))
|
||||
}
|
||||
63
mcp/util_test.go
Normal file
63
mcp/util_test.go
Normal file
@@ -0,0 +1,63 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestPtr(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
v interface{}
|
||||
}{
|
||||
{"string", "test"},
|
||||
{"int", 42},
|
||||
{"bool", true},
|
||||
{"float", 3.14},
|
||||
{"struct", struct{ Name string }{"test"}},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := ptr(tt.v)
|
||||
assert.NotNil(t, got, "ptr() should not return nil")
|
||||
assert.Equal(t, tt.v, *got, "dereferenced pointer should equal input value")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
type Event struct {
|
||||
Type string
|
||||
Data map[string]any
|
||||
}
|
||||
|
||||
func parseEvent(input string) (*Event, error) {
|
||||
var evt Event
|
||||
var dataStr string
|
||||
|
||||
scanner := bufio.NewScanner(strings.NewReader(input))
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
if strings.HasPrefix(line, "event:") {
|
||||
evt.Type = strings.TrimSpace(strings.TrimPrefix(line, "event:"))
|
||||
} else if strings.HasPrefix(line, "data:") {
|
||||
dataStr = strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
||||
}
|
||||
}
|
||||
if err := scanner.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if len(dataStr) > 0 {
|
||||
if err := json.Unmarshal([]byte(dataStr), &evt.Data); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse data: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
return &evt, nil
|
||||
}
|
||||
137
mcp/vars.go
Normal file
137
mcp/vars.go
Normal file
@@ -0,0 +1,137 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/zeromicro/go-zero/core/syncx"
|
||||
)
|
||||
|
||||
// Protocol constants
|
||||
const (
|
||||
// JSON-RPC version as defined in the specification
|
||||
jsonRpcVersion = "2.0"
|
||||
|
||||
// Session identifier key used in request URLs
|
||||
sessionIdKey = "session_id"
|
||||
)
|
||||
|
||||
// Server-Sent Events (SSE) event types
|
||||
const (
|
||||
// Standard message event for JSON-RPC responses
|
||||
eventMessage = "message"
|
||||
|
||||
// Endpoint event for sending endpoint URL to clients
|
||||
eventEndpoint = "endpoint"
|
||||
)
|
||||
|
||||
// Content type identifiers
|
||||
const (
|
||||
// Text content type
|
||||
contentTypeText = "text"
|
||||
|
||||
// Image content type
|
||||
contentTypeImage = "image"
|
||||
)
|
||||
|
||||
// Collection keys for broadcast events
|
||||
const (
|
||||
// Key for prompts collection
|
||||
keyPrompts = "prompts"
|
||||
|
||||
// Key for resources collection
|
||||
keyResources = "resources"
|
||||
|
||||
// Key for tools collection
|
||||
keyTools = "tools"
|
||||
)
|
||||
|
||||
// JSON-RPC error codes
|
||||
// Standard error codes from JSON-RPC 2.0 spec
|
||||
const (
|
||||
// Invalid JSON was received by the server
|
||||
errCodeInvalidRequest = -32600
|
||||
|
||||
// The method does not exist / is not available
|
||||
errCodeMethodNotFound = -32601
|
||||
|
||||
// Invalid method parameter(s)
|
||||
errCodeInvalidParams = -32602
|
||||
|
||||
// Internal JSON-RPC error
|
||||
errCodeInternalError = -32603
|
||||
|
||||
// Tool execution timed out
|
||||
errCodeTimeout = -32001
|
||||
|
||||
// Resource not found error
|
||||
errCodeResourceNotFound = -32002
|
||||
|
||||
// Client hasn't completed initialization
|
||||
errCodeClientNotInitialized = -32800
|
||||
)
|
||||
|
||||
// User and assistant role definitions
|
||||
const (
|
||||
// The "user" role - the entity asking questions
|
||||
roleUser roleType = "user"
|
||||
|
||||
// The "assistant" role - the entity providing responses
|
||||
roleAssistant roleType = "assistant"
|
||||
)
|
||||
|
||||
// Method names as defined in the MCP specification
|
||||
const (
|
||||
// Initialize the connection between client and server
|
||||
methodInitialize = "initialize"
|
||||
|
||||
// List available tools
|
||||
methodToolsList = "tools/list"
|
||||
|
||||
// Call a specific tool
|
||||
methodToolsCall = "tools/call"
|
||||
|
||||
// List available prompts
|
||||
methodPromptsList = "prompts/list"
|
||||
|
||||
// Get a specific prompt
|
||||
methodPromptsGet = "prompts/get"
|
||||
|
||||
// List available resources
|
||||
methodResourcesList = "resources/list"
|
||||
|
||||
// Read a specific resource
|
||||
methodResourcesRead = "resources/read"
|
||||
|
||||
// Subscribe to resource updates
|
||||
methodResourcesSubscribe = "resources/subscribe"
|
||||
|
||||
// Simple ping to check server availability
|
||||
methodPing = "ping"
|
||||
|
||||
// Notification that client is fully initialized
|
||||
methodNotificationsInitialized = "notifications/initialized"
|
||||
|
||||
// Notification that a request was canceled
|
||||
methodNotificationsCancelled = "notifications/cancelled"
|
||||
)
|
||||
|
||||
// Event names for Server-Sent Events (SSE)
|
||||
const (
|
||||
// Notification of tool list changes
|
||||
eventToolsListChanged = "tools/list_changed"
|
||||
|
||||
// Notification of prompt list changes
|
||||
eventPromptsListChanged = "prompts/list_changed"
|
||||
|
||||
// Notification of resource list changes
|
||||
eventResourcesListChanged = "resources/list_changed"
|
||||
)
|
||||
|
||||
var (
|
||||
// Default channel size for events
|
||||
eventChanSize = 10
|
||||
|
||||
// Default ping interval for checking connection availability
|
||||
// use syncx.ForAtomicDuration to ensure atomicity in test race
|
||||
pingInterval = syncx.ForAtomicDuration(30 * time.Second)
|
||||
)
|
||||
214
mcp/vars_test.go
Normal file
214
mcp/vars_test.go
Normal file
@@ -0,0 +1,214 @@
|
||||
// filepath: /Users/kevin/Develop/go/opensource/go-zero/mcp/vars_test.go
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
// TestErrorCodes ensures error codes are applied correctly in error responses
|
||||
func TestErrorCodes(t *testing.T) {
|
||||
testCases := []struct {
|
||||
name string
|
||||
code int
|
||||
message string
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "invalid request error",
|
||||
code: errCodeInvalidRequest,
|
||||
message: "Invalid request",
|
||||
expected: `"code":-32600`,
|
||||
},
|
||||
{
|
||||
name: "method not found error",
|
||||
code: errCodeMethodNotFound,
|
||||
message: "Method not found",
|
||||
expected: `"code":-32601`,
|
||||
},
|
||||
{
|
||||
name: "invalid params error",
|
||||
code: errCodeInvalidParams,
|
||||
message: "Invalid parameters",
|
||||
expected: `"code":-32602`,
|
||||
},
|
||||
{
|
||||
name: "internal error",
|
||||
code: errCodeInternalError,
|
||||
message: "Internal server error",
|
||||
expected: `"code":-32603`,
|
||||
},
|
||||
{
|
||||
name: "timeout error",
|
||||
code: errCodeTimeout,
|
||||
message: "Operation timed out",
|
||||
expected: `"code":-32001`,
|
||||
},
|
||||
{
|
||||
name: "resource not found error",
|
||||
code: errCodeResourceNotFound,
|
||||
message: "Resource not found",
|
||||
expected: `"code":-32002`,
|
||||
},
|
||||
{
|
||||
name: "client not initialized error",
|
||||
code: errCodeClientNotInitialized,
|
||||
message: "Client not initialized",
|
||||
expected: `"code":-32800`,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
resp := Response{
|
||||
JsonRpc: jsonRpcVersion,
|
||||
ID: int64(1),
|
||||
Error: &errorObj{
|
||||
Code: tc.code,
|
||||
Message: tc.message,
|
||||
},
|
||||
}
|
||||
data, err := json.Marshal(resp)
|
||||
assert.NoError(t, err)
|
||||
assert.Contains(t, string(data), tc.expected, "Error code should match expected value")
|
||||
assert.Contains(t, string(data), tc.message, "Error message should be included")
|
||||
assert.Contains(t, string(data), jsonRpcVersion, "JSON-RPC version should be included")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestJsonRpcVersion ensures the correct JSON-RPC version is used
|
||||
func TestJsonRpcVersion(t *testing.T) {
|
||||
assert.Equal(t, "2.0", jsonRpcVersion, "JSON-RPC version should be 2.0")
|
||||
|
||||
// Test that it's used in responses
|
||||
resp := Response{
|
||||
JsonRpc: jsonRpcVersion,
|
||||
ID: int64(1),
|
||||
Result: "test",
|
||||
}
|
||||
data, err := json.Marshal(resp)
|
||||
assert.NoError(t, err)
|
||||
assert.Contains(t, string(data), `"jsonrpc":"2.0"`, "Response should use correct JSON-RPC version")
|
||||
|
||||
// Test that it's expected in requests
|
||||
reqStr := `{"jsonrpc":"2.0","id":1,"method":"test"}`
|
||||
var req Request
|
||||
err = json.Unmarshal([]byte(reqStr), &req)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, jsonRpcVersion, req.JsonRpc, "Request should parse correct JSON-RPC version")
|
||||
}
|
||||
|
||||
// TestSessionIdKey ensures session ID extraction works correctly
|
||||
func TestSessionIdKey(t *testing.T) {
|
||||
// Create a mock server implementation
|
||||
mock := newMockMcpServer(t)
|
||||
defer mock.shutdown()
|
||||
|
||||
// Verify the key constant
|
||||
assert.Equal(t, "session_id", sessionIdKey, "Session ID key should be 'session_id'")
|
||||
|
||||
// Test that session ID is extracted correctly
|
||||
mockR := httptest.NewRequest("GET", "/?"+sessionIdKey+"=test-session", nil)
|
||||
|
||||
// Since the mock server is using the same session key logic,
|
||||
// we can test this by accessing the request query parameters directly
|
||||
sessionID := mockR.URL.Query().Get(sessionIdKey)
|
||||
assert.Equal(t, "test-session", sessionID, "Session ID should be extracted correctly")
|
||||
}
|
||||
|
||||
// TestEventTypes ensures event types are set correctly in SSE responses
|
||||
func TestEventTypes(t *testing.T) {
|
||||
// Test message event
|
||||
assert.Equal(t, "message", eventMessage, "Message event should be 'message'")
|
||||
|
||||
// Test endpoint event
|
||||
assert.Equal(t, "endpoint", eventEndpoint, "Endpoint event should be 'endpoint'")
|
||||
|
||||
// Verify them in an actual SSE format string
|
||||
messageEvent := "event: " + eventMessage + "\ndata: test\n\n"
|
||||
assert.Contains(t, messageEvent, "event: message", "Message event should format correctly")
|
||||
|
||||
endpointEvent := "event: " + eventEndpoint + "\ndata: test\n\n"
|
||||
assert.Contains(t, endpointEvent, "event: endpoint", "Endpoint event should format correctly")
|
||||
}
|
||||
|
||||
// TestCollectionKeys checks that collection keys are used correctly
|
||||
func TestCollectionKeys(t *testing.T) {
|
||||
// Verify collection key constants
|
||||
assert.Equal(t, "prompts", keyPrompts, "Prompts key should be 'prompts'")
|
||||
assert.Equal(t, "resources", keyResources, "Resources key should be 'resources'")
|
||||
assert.Equal(t, "tools", keyTools, "Tools key should be 'tools'")
|
||||
}
|
||||
|
||||
// TestRoleTypes checks that role types are used correctly
|
||||
func TestRoleTypes(t *testing.T) {
|
||||
// Verify role type constants
|
||||
assert.Equal(t, "user", string(roleUser), "User role should be 'user'")
|
||||
assert.Equal(t, "assistant", string(roleAssistant), "Assistant role should be 'assistant'")
|
||||
|
||||
// Test in annotations
|
||||
annotations := Annotations{
|
||||
Audience: []roleType{roleUser, roleAssistant},
|
||||
}
|
||||
data, err := json.Marshal(annotations)
|
||||
assert.NoError(t, err)
|
||||
assert.Contains(t, string(data), `"audience":["user","assistant"]`, "Role types should marshal correctly")
|
||||
}
|
||||
|
||||
// TestMethodNames checks that method names are used correctly
|
||||
func TestMethodNames(t *testing.T) {
|
||||
// Verify method name constants
|
||||
methods := map[string]string{
|
||||
"initialize": methodInitialize,
|
||||
"tools/list": methodToolsList,
|
||||
"tools/call": methodToolsCall,
|
||||
"prompts/list": methodPromptsList,
|
||||
"prompts/get": methodPromptsGet,
|
||||
"resources/list": methodResourcesList,
|
||||
"resources/read": methodResourcesRead,
|
||||
"resources/subscribe": methodResourcesSubscribe,
|
||||
"ping": methodPing,
|
||||
"notifications/initialized": methodNotificationsInitialized,
|
||||
"notifications/cancelled": methodNotificationsCancelled,
|
||||
}
|
||||
|
||||
for expected, actual := range methods {
|
||||
assert.Equal(t, expected, actual, "Method name should be "+expected)
|
||||
}
|
||||
|
||||
// Test in a request
|
||||
for methodName := range methods {
|
||||
req := Request{
|
||||
JsonRpc: jsonRpcVersion,
|
||||
ID: int64(1),
|
||||
Method: methodName,
|
||||
}
|
||||
data, err := json.Marshal(req)
|
||||
assert.NoError(t, err)
|
||||
assert.Contains(t, string(data), `"method":"`+methodName+`"`, "Method name should be used in requests")
|
||||
}
|
||||
}
|
||||
|
||||
// TestEventNames checks that event names are used correctly
|
||||
func TestEventNames(t *testing.T) {
|
||||
// Verify event name constants
|
||||
events := map[string]string{
|
||||
"tools/list_changed": eventToolsListChanged,
|
||||
"prompts/list_changed": eventPromptsListChanged,
|
||||
"resources/list_changed": eventResourcesListChanged,
|
||||
}
|
||||
|
||||
for expected, actual := range events {
|
||||
assert.Equal(t, expected, actual, "Event name should be "+expected)
|
||||
}
|
||||
|
||||
// Test event names in SSE format
|
||||
for _, eventName := range events {
|
||||
sseEvent := "event: " + eventName + "\ndata: test\n\n"
|
||||
assert.Contains(t, sseEvent, "event: "+eventName, "Event name should format correctly in SSE")
|
||||
}
|
||||
}
|
||||
@@ -299,6 +299,9 @@ go-zero 已被许多公司用于生产部署,接入场景如在线教育、电
|
||||
>100. 上海邑脉科技有限公司
|
||||
>101. 上海巨瓴科技有限公司
|
||||
>102. 深圳市兴海物联科技有限公司
|
||||
>103. 爱芯元智半导体股份有限公司
|
||||
>104. 杭州升恒科技有限公司
|
||||
>105. 昆仑万维科技股份有限公司
|
||||
|
||||
如果贵公司也已使用 go-zero,欢迎在 [登记地址](https://github.com/zeromicro/go-zero/issues/602) 登记,仅仅为了推广,不做其它用途。
|
||||
|
||||
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
"github.com/zeromicro/go-zero/rest/handler"
|
||||
"github.com/zeromicro/go-zero/rest/httpx"
|
||||
"github.com/zeromicro/go-zero/rest/internal"
|
||||
"github.com/zeromicro/go-zero/rest/internal/header"
|
||||
"github.com/zeromicro/go-zero/rest/internal/response"
|
||||
)
|
||||
|
||||
@@ -54,6 +55,9 @@ func newEngine(c RestConf) *engine {
|
||||
}
|
||||
|
||||
func (ng *engine) addRoutes(r featuredRoutes) {
|
||||
if r.sse {
|
||||
r.routes = buildSSERoutes(r.routes)
|
||||
}
|
||||
ng.routes = append(ng.routes, r)
|
||||
|
||||
// need to guarantee the timeout is the max of all routes
|
||||
@@ -63,6 +67,20 @@ func (ng *engine) addRoutes(r featuredRoutes) {
|
||||
}
|
||||
}
|
||||
|
||||
func buildSSERoutes(routes []Route) []Route {
|
||||
for i, route := range routes {
|
||||
h := route.Handler
|
||||
routes[i].Handler = func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set(header.ContentType, header.ContentTypeEventStream)
|
||||
w.Header().Set(header.CacheControl, header.CacheControlNoCache)
|
||||
w.Header().Set(header.Connection, header.ConnectionKeepAlive)
|
||||
h(w, r)
|
||||
}
|
||||
}
|
||||
|
||||
return routes
|
||||
}
|
||||
|
||||
func (ng *engine) appendAuthHandler(fr featuredRoutes, chn chain.Chain,
|
||||
verifier func(chain.Chain) chain.Chain) chain.Chain {
|
||||
if fr.jwt.enabled {
|
||||
|
||||
@@ -7,7 +7,7 @@ import (
|
||||
"net/http/httputil"
|
||||
|
||||
"github.com/golang-jwt/jwt/v4"
|
||||
"github.com/zeromicro/go-zero/core/logx"
|
||||
"github.com/zeromicro/go-zero/core/logc"
|
||||
"github.com/zeromicro/go-zero/rest/internal/response"
|
||||
"github.com/zeromicro/go-zero/rest/token"
|
||||
)
|
||||
@@ -100,7 +100,7 @@ func WithUnauthorizedCallback(callback UnauthorizedCallback) AuthorizeOption {
|
||||
func detailAuthLog(r *http.Request, reason string) {
|
||||
// discard dump error, only for debug purpose
|
||||
details, _ := httputil.DumpRequest(r, true)
|
||||
logx.Errorf("authorize failed: %s\n=> %+v", reason, string(details))
|
||||
logc.Errorf(r.Context(), "authorize failed: %s\n=> %+v", reason, string(details))
|
||||
}
|
||||
|
||||
func unauthorized(w http.ResponseWriter, r *http.Request, err error, callback UnauthorizedCallback) {
|
||||
|
||||
@@ -6,7 +6,7 @@ import (
|
||||
"strings"
|
||||
|
||||
"github.com/zeromicro/go-zero/core/breaker"
|
||||
"github.com/zeromicro/go-zero/core/logx"
|
||||
"github.com/zeromicro/go-zero/core/logc"
|
||||
"github.com/zeromicro/go-zero/core/stat"
|
||||
"github.com/zeromicro/go-zero/rest/httpx"
|
||||
"github.com/zeromicro/go-zero/rest/internal/response"
|
||||
@@ -22,7 +22,7 @@ func BreakerHandler(method, path string, metrics *stat.Metrics) func(http.Handle
|
||||
promise, err := brk.Allow()
|
||||
if err != nil {
|
||||
metrics.AddDrop()
|
||||
logx.Errorf("[http] dropped, %s - %s - %s",
|
||||
logc.Errorf(r.Context(), "[http] dropped, %s - %s - %s",
|
||||
r.RequestURI, httpx.GetRemoteAddr(r), r.UserAgent())
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
return
|
||||
|
||||
@@ -5,7 +5,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/zeromicro/go-zero/core/codec"
|
||||
"github.com/zeromicro/go-zero/core/logx"
|
||||
"github.com/zeromicro/go-zero/core/logc"
|
||||
"github.com/zeromicro/go-zero/rest/httpx"
|
||||
"github.com/zeromicro/go-zero/rest/internal/security"
|
||||
)
|
||||
@@ -34,11 +34,11 @@ func LimitContentSecurityHandler(limitBytes int64, decrypters map[string]codec.R
|
||||
case http.MethodDelete, http.MethodGet, http.MethodPost, http.MethodPut:
|
||||
header, err := security.ParseContentSecurity(decrypters, r)
|
||||
if err != nil {
|
||||
logx.Errorf("Signature parse failed, X-Content-Security: %s, error: %s",
|
||||
logc.Errorf(r.Context(), "Signature parse failed, X-Content-Security: %s, error: %s",
|
||||
r.Header.Get(contentSecurity), err.Error())
|
||||
executeCallbacks(w, r, next, strict, httpx.CodeSignatureInvalidHeader, callbacks)
|
||||
} else if code := security.VerifySignature(r, header, tolerance); code != httpx.CodeSignaturePass {
|
||||
logx.Errorf("Signature verification failed, X-Content-Security: %s",
|
||||
logc.Errorf(r.Context(), "Signature verification failed, X-Content-Security: %s",
|
||||
r.Header.Get(contentSecurity))
|
||||
executeCallbacks(w, r, next, strict, code, callbacks)
|
||||
} else if r.ContentLength > 0 && header.Encrypted() {
|
||||
|
||||
@@ -3,6 +3,7 @@ package handler
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"io"
|
||||
@@ -10,7 +11,7 @@ import (
|
||||
"net/http"
|
||||
|
||||
"github.com/zeromicro/go-zero/core/codec"
|
||||
"github.com/zeromicro/go-zero/core/logx"
|
||||
"github.com/zeromicro/go-zero/core/logc"
|
||||
)
|
||||
|
||||
const maxBytes = 1 << 20 // 1 MiB
|
||||
@@ -27,7 +28,7 @@ func LimitCryptionHandler(limitBytes int64, key []byte) func(http.Handler) http.
|
||||
return func(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
cw := newCryptionResponseWriter(w)
|
||||
defer cw.flush(key)
|
||||
defer cw.flush(r.Context(), key)
|
||||
|
||||
if r.ContentLength <= 0 {
|
||||
next.ServeHTTP(cw, r)
|
||||
@@ -118,7 +119,7 @@ func (w *cryptionResponseWriter) WriteHeader(statusCode int) {
|
||||
w.ResponseWriter.WriteHeader(statusCode)
|
||||
}
|
||||
|
||||
func (w *cryptionResponseWriter) flush(key []byte) {
|
||||
func (w *cryptionResponseWriter) flush(ctx context.Context, key []byte) {
|
||||
if w.buf.Len() == 0 {
|
||||
return
|
||||
}
|
||||
@@ -131,8 +132,8 @@ func (w *cryptionResponseWriter) flush(key []byte) {
|
||||
|
||||
body := base64.StdEncoding.EncodeToString(content)
|
||||
if n, err := io.WriteString(w.ResponseWriter, body); err != nil {
|
||||
logx.Errorf("write response failed, error: %s", err)
|
||||
logc.Errorf(ctx, "write response failed, error: %s", err)
|
||||
} else if n < len(body) {
|
||||
logx.Errorf("actual bytes: %d, written bytes: %d", len(body), n)
|
||||
logc.Errorf(ctx, "actual bytes: %d, written bytes: %d", len(body), n)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,6 +2,7 @@ package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/base64"
|
||||
"io"
|
||||
@@ -174,7 +175,7 @@ func TestCryptionResponseWriter_Flush(t *testing.T) {
|
||||
w := newCryptionResponseWriter(f)
|
||||
_, err := w.Write(body)
|
||||
assert.NoError(t, err)
|
||||
w.flush(aesKey)
|
||||
w.flush(context.Background(), aesKey)
|
||||
b, err := io.ReadAll(recorder.Body)
|
||||
assert.NoError(t, err)
|
||||
expected, err := codec.EcbEncrypt(aesKey, body)
|
||||
@@ -191,7 +192,7 @@ func TestCryptionResponseWriter_Flush(t *testing.T) {
|
||||
w := newCryptionResponseWriter(f)
|
||||
_, err := w.Write(body)
|
||||
assert.NoError(t, err)
|
||||
w.flush(aesKey)
|
||||
w.flush(context.Background(), aesKey)
|
||||
b, err := io.ReadAll(recorder.Body)
|
||||
assert.NoError(t, err)
|
||||
expected, err := codec.EcbEncrypt(aesKey, body)
|
||||
@@ -207,7 +208,7 @@ func TestCryptionResponseWriter_Flush(t *testing.T) {
|
||||
w := newCryptionResponseWriter(f)
|
||||
_, err := w.Write(body)
|
||||
assert.NoError(t, err)
|
||||
w.flush(aesKey)
|
||||
w.flush(context.Background(), aesKey)
|
||||
assert.True(t, strings.Contains(buf.Content(), io.ErrClosedPipe.Error()))
|
||||
})
|
||||
}
|
||||
|
||||
@@ -158,9 +158,9 @@ func logDetails(r *http.Request, response *detailLoggedResponseWriter, timer *ut
|
||||
logger := logx.WithContext(r.Context())
|
||||
buf.WriteString(fmt.Sprintf("[HTTP] %s - %d - %s - %s\n=> %s\n",
|
||||
r.Method, code, r.RemoteAddr, timex.ReprOfDuration(duration), dumpRequest(r)))
|
||||
if duration > defaultSlowThreshold {
|
||||
if duration > slowThreshold.Load() {
|
||||
logger.Slowf("[HTTP] %s - %d - %s - slowcall(%s)\n=> %s\n", r.Method, code, r.RemoteAddr,
|
||||
fmt.Sprintf("slowcall(%s)", timex.ReprOfDuration(duration)), dumpRequest(r))
|
||||
timex.ReprOfDuration(duration), dumpRequest(r))
|
||||
}
|
||||
|
||||
body := logs.Flush()
|
||||
|
||||
@@ -5,7 +5,7 @@ import (
|
||||
"sync"
|
||||
|
||||
"github.com/zeromicro/go-zero/core/load"
|
||||
"github.com/zeromicro/go-zero/core/logx"
|
||||
"github.com/zeromicro/go-zero/core/logc"
|
||||
"github.com/zeromicro/go-zero/core/stat"
|
||||
"github.com/zeromicro/go-zero/rest/httpx"
|
||||
"github.com/zeromicro/go-zero/rest/internal/response"
|
||||
@@ -35,7 +35,7 @@ func SheddingHandler(shedder load.Shedder, metrics *stat.Metrics) func(http.Hand
|
||||
if err != nil {
|
||||
metrics.AddDrop()
|
||||
sheddingStat.IncrementDrop()
|
||||
logx.Errorf("[http] dropped, %s - %s - %s",
|
||||
logc.Errorf(r.Context(), "[http] dropped, %s - %s - %s",
|
||||
r.RequestURI, httpx.GetRemoteAddr(r), r.UserAgent())
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
return
|
||||
|
||||
@@ -6,15 +6,11 @@ import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/golang/mock/gomock"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/zeromicro/go-zero/core/logx"
|
||||
)
|
||||
|
||||
func TestMetricsInterceptor(t *testing.T) {
|
||||
c := gomock.NewController(t)
|
||||
defer c.Finish()
|
||||
|
||||
logx.Disable()
|
||||
|
||||
svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
@@ -105,7 +105,7 @@ func buildRequest(ctx context.Context, method, url string, data any) (*http.Requ
|
||||
req.URL.RawQuery = buildFormQuery(u, val[formKey])
|
||||
fillHeader(req, val[headerKey])
|
||||
if hasJsonBody {
|
||||
req.Header.Set(header.ContentType, header.JsonContentType)
|
||||
req.Header.Set(header.ContentType, header.ContentTypeJson)
|
||||
}
|
||||
|
||||
return req, nil
|
||||
@@ -183,7 +183,6 @@ func request(r *http.Request, cli client) (*http.Response, error) {
|
||||
for i := len(respHandlers) - 1; i >= 0; i-- {
|
||||
respHandlers[i](resp, err)
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
span.RecordError(err)
|
||||
span.SetStatus(codes.Error, err.Error())
|
||||
|
||||
@@ -45,7 +45,7 @@ func TestDoRequest_NotFound(t *testing.T) {
|
||||
defer svr.Close()
|
||||
req, err := http.NewRequest(http.MethodPost, svr.URL, nil)
|
||||
assert.Nil(t, err)
|
||||
req.Header.Set(header.ContentType, header.JsonContentType)
|
||||
req.Header.Set(header.ContentType, header.ContentTypeJson)
|
||||
resp, err := DoRequest(req)
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, http.StatusNotFound, resp.StatusCode)
|
||||
|
||||
@@ -18,7 +18,7 @@ func TestParse(t *testing.T) {
|
||||
}
|
||||
svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("foo", "bar")
|
||||
w.Header().Set(header.ContentType, header.JsonContentType)
|
||||
w.Header().Set(header.ContentType, header.ContentTypeJson)
|
||||
w.Write([]byte(`{"name":"kevin","value":100}`))
|
||||
}))
|
||||
defer svr.Close()
|
||||
@@ -38,7 +38,7 @@ func TestParseHeaderError(t *testing.T) {
|
||||
}
|
||||
svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("foo", "bar")
|
||||
w.Header().Set(header.ContentType, header.JsonContentType)
|
||||
w.Header().Set(header.ContentType, header.ContentTypeJson)
|
||||
}))
|
||||
defer svr.Close()
|
||||
req, err := http.NewRequest(http.MethodGet, svr.URL, nil)
|
||||
@@ -54,7 +54,7 @@ func TestParseNoBody(t *testing.T) {
|
||||
}
|
||||
svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("foo", "bar")
|
||||
w.Header().Set(header.ContentType, header.JsonContentType)
|
||||
w.Header().Set(header.ContentType, header.ContentTypeJson)
|
||||
}))
|
||||
defer svr.Close()
|
||||
req, err := http.NewRequest(http.MethodGet, svr.URL, nil)
|
||||
@@ -72,7 +72,7 @@ func TestParseWithZeroValue(t *testing.T) {
|
||||
}
|
||||
svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("foo", "0")
|
||||
w.Header().Set(header.ContentType, header.JsonContentType)
|
||||
w.Header().Set(header.ContentType, header.ContentTypeJson)
|
||||
w.Write([]byte(`{"bar":0}`))
|
||||
}))
|
||||
defer svr.Close()
|
||||
@@ -90,7 +90,7 @@ func TestParseWithNegativeContentLength(t *testing.T) {
|
||||
Bar int `json:"bar"`
|
||||
}
|
||||
svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set(header.ContentType, header.JsonContentType)
|
||||
w.Header().Set(header.ContentType, header.ContentTypeJson)
|
||||
w.Write([]byte(`{"bar":0}`))
|
||||
}))
|
||||
defer svr.Close()
|
||||
@@ -124,7 +124,7 @@ func TestParseWithNegativeContentLength(t *testing.T) {
|
||||
func TestParseWithNegativeContentLengthNoBody(t *testing.T) {
|
||||
var val struct{}
|
||||
svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set(header.ContentType, header.JsonContentType)
|
||||
w.Header().Set(header.ContentType, header.ContentTypeJson)
|
||||
}))
|
||||
defer svr.Close()
|
||||
req, err := http.NewRequest(http.MethodGet, svr.URL, nil)
|
||||
@@ -156,7 +156,7 @@ func TestParseWithNegativeContentLengthNoBody(t *testing.T) {
|
||||
func TestParseJsonBody_BodyError(t *testing.T) {
|
||||
var val struct{}
|
||||
svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set(header.ContentType, header.JsonContentType)
|
||||
w.Header().Set(header.ContentType, header.ContentTypeJson)
|
||||
}))
|
||||
defer svr.Close()
|
||||
req, err := http.NewRequest(http.MethodGet, svr.URL, nil)
|
||||
|
||||
@@ -44,7 +44,7 @@ func TestNamedService_DoRequestPost(t *testing.T) {
|
||||
service := NewService("foo")
|
||||
req, err := http.NewRequest(http.MethodPost, svr.URL, nil)
|
||||
assert.Nil(t, err)
|
||||
req.Header.Set(header.ContentType, header.JsonContentType)
|
||||
req.Header.Set(header.ContentType, header.ContentTypeJson)
|
||||
resp, err := service.DoRequest(req)
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, http.StatusNotFound, resp.StatusCode)
|
||||
|
||||
@@ -5,7 +5,7 @@ import (
|
||||
"net/http"
|
||||
"reflect"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"sync"
|
||||
|
||||
"github.com/zeromicro/go-zero/core/mapping"
|
||||
"github.com/zeromicro/go-zero/core/validation"
|
||||
@@ -33,7 +33,11 @@ var (
|
||||
pathKey,
|
||||
mapping.WithStringValues(),
|
||||
mapping.WithOpaqueKeys())
|
||||
validator atomic.Value
|
||||
|
||||
// panic: sync/atomic: store of inconsistently typed value into Value
|
||||
// don't use atomic.Value to store the validator, different concrete types still panic
|
||||
validator Validator
|
||||
validatorLock sync.RWMutex
|
||||
)
|
||||
|
||||
// Validator defines the interface for validating the request.
|
||||
@@ -65,8 +69,8 @@ func Parse(r *http.Request, v any) error {
|
||||
|
||||
if valid, ok := v.(validation.Validator); ok {
|
||||
return valid.Validate()
|
||||
} else if val := validator.Load(); val != nil {
|
||||
return val.(Validator).Validate(r, v)
|
||||
} else if val := getValidator(); val != nil {
|
||||
return val.Validate(r, v)
|
||||
}
|
||||
|
||||
return nil
|
||||
@@ -135,7 +139,15 @@ func ParsePath(r *http.Request, v any) error {
|
||||
// The validator is used to validate the request, only called in Parse,
|
||||
// not in ParseHeaders, ParseForm, ParseHeader, ParseJsonBody, ParsePath.
|
||||
func SetValidator(val Validator) {
|
||||
validator.Store(val)
|
||||
validatorLock.Lock()
|
||||
defer validatorLock.Unlock()
|
||||
validator = val
|
||||
}
|
||||
|
||||
func getValidator() Validator {
|
||||
validatorLock.RLock()
|
||||
defer validatorLock.RUnlock()
|
||||
return validator
|
||||
}
|
||||
|
||||
func withJsonBody(r *http.Request) bool {
|
||||
|
||||
@@ -2,6 +2,7 @@ package httpx
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
@@ -47,6 +48,21 @@ func TestParseForm(t *testing.T) {
|
||||
assert.Nil(t, Parse(r, &v))
|
||||
assert.Equal(t, 0, len(v.NoValue))
|
||||
})
|
||||
|
||||
t.Run("slice with one value on array format", func(t *testing.T) {
|
||||
var v struct {
|
||||
Names string `form:"names"`
|
||||
}
|
||||
|
||||
r, err := http.NewRequest(
|
||||
http.MethodGet,
|
||||
"/a?names=1,2,3",
|
||||
http.NoBody)
|
||||
assert.NoError(t, err)
|
||||
if assert.NoError(t, Parse(r, &v)) {
|
||||
assert.Equal(t, "1,2,3", v.Names)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestParseFormArray(t *testing.T) {
|
||||
@@ -88,6 +104,36 @@ func TestParseFormArray(t *testing.T) {
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("slice with empty", func(t *testing.T) {
|
||||
var v struct {
|
||||
Name []string `form:"name,optional"`
|
||||
}
|
||||
|
||||
r, err := http.NewRequest(
|
||||
http.MethodGet,
|
||||
"/a",
|
||||
http.NoBody)
|
||||
assert.NoError(t, err)
|
||||
if assert.NoError(t, Parse(r, &v)) {
|
||||
assert.ElementsMatch(t, []string{}, v.Name)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("slice with empty", func(t *testing.T) {
|
||||
var v struct {
|
||||
Name []string `form:"name,optional"`
|
||||
}
|
||||
|
||||
r, err := http.NewRequest(
|
||||
http.MethodGet,
|
||||
"/a?name=",
|
||||
http.NoBody)
|
||||
assert.NoError(t, err)
|
||||
if assert.NoError(t, Parse(r, &v)) {
|
||||
assert.Empty(t, v.Name)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("slice with empty and non-empty", func(t *testing.T) {
|
||||
var v struct {
|
||||
Name []string `form:"name"`
|
||||
@@ -102,6 +148,154 @@ func TestParseFormArray(t *testing.T) {
|
||||
assert.ElementsMatch(t, []string{"1"}, v.Name)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("slice with one value on array format", func(t *testing.T) {
|
||||
var v struct {
|
||||
Names []string `form:"names"`
|
||||
}
|
||||
|
||||
r, err := http.NewRequest(
|
||||
http.MethodGet,
|
||||
"/a?names=1,2,3",
|
||||
http.NoBody)
|
||||
assert.NoError(t, err)
|
||||
if assert.NoError(t, Parse(r, &v)) {
|
||||
assert.ElementsMatch(t, []string{"1,2,3"}, v.Names)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("slice with one value on combined array format", func(t *testing.T) {
|
||||
var v struct {
|
||||
Names []string `form:"names"`
|
||||
}
|
||||
|
||||
r, err := http.NewRequest(
|
||||
http.MethodGet,
|
||||
"/a?names=[1,2,3]&names=4",
|
||||
http.NoBody)
|
||||
assert.NoError(t, err)
|
||||
if assert.NoError(t, Parse(r, &v)) {
|
||||
assert.ElementsMatch(t, []string{"[1,2,3]", "4"}, v.Names)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("slice with one value on integer array format", func(t *testing.T) {
|
||||
var v struct {
|
||||
Numbers []int `form:"numbers"`
|
||||
}
|
||||
|
||||
r, err := http.NewRequest(
|
||||
http.MethodGet,
|
||||
"/a?numbers=1,2,3",
|
||||
http.NoBody)
|
||||
assert.NoError(t, err)
|
||||
assert.Error(t, Parse(r, &v))
|
||||
})
|
||||
|
||||
t.Run("slice with one value on array format brackets", func(t *testing.T) {
|
||||
var v struct {
|
||||
Names []string `form:"names"`
|
||||
}
|
||||
|
||||
r, err := http.NewRequest(
|
||||
http.MethodGet,
|
||||
"/a?names[]=1&names[]=2&names[]=3",
|
||||
http.NoBody)
|
||||
assert.NoError(t, err)
|
||||
if assert.NoError(t, Parse(r, &v)) {
|
||||
assert.ElementsMatch(t, []string{"1", "2", "3"}, v.Names)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("slice with one empty value on integer array format", func(t *testing.T) {
|
||||
var v struct {
|
||||
Numbers []int `form:"numbers,optional"`
|
||||
}
|
||||
|
||||
r, err := http.NewRequest(
|
||||
http.MethodGet,
|
||||
"/a?numbers=",
|
||||
http.NoBody)
|
||||
assert.NoError(t, err)
|
||||
if assert.NoError(t, Parse(r, &v)) {
|
||||
assert.Empty(t, v.Numbers)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("slice with one value on integer array format", func(t *testing.T) {
|
||||
var v struct {
|
||||
Numbers []int `form:"numbers,optional"`
|
||||
}
|
||||
|
||||
r, err := http.NewRequest(
|
||||
http.MethodGet,
|
||||
"/a?numbers=&numbers=2",
|
||||
http.NoBody)
|
||||
assert.NoError(t, err)
|
||||
if assert.NoError(t, Parse(r, &v)) {
|
||||
assert.ElementsMatch(t, []int{2}, v.Numbers)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("slice with one empty value on float64 array format", func(t *testing.T) {
|
||||
var v struct {
|
||||
Numbers []float64 `form:"numbers,optional"`
|
||||
}
|
||||
|
||||
r, err := http.NewRequest(
|
||||
http.MethodGet,
|
||||
"/a?numbers=",
|
||||
http.NoBody)
|
||||
assert.NoError(t, err)
|
||||
if assert.NoError(t, Parse(r, &v)) {
|
||||
assert.Empty(t, v.Numbers)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("slice with one value on float64 array format", func(t *testing.T) {
|
||||
var v struct {
|
||||
Numbers []float64 `form:"numbers,optional"`
|
||||
}
|
||||
|
||||
r, err := http.NewRequest(
|
||||
http.MethodGet,
|
||||
"/a?numbers=&numbers=2",
|
||||
http.NoBody)
|
||||
assert.NoError(t, err)
|
||||
if assert.NoError(t, Parse(r, &v)) {
|
||||
assert.ElementsMatch(t, []float64{2}, v.Numbers)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("slice with one value", func(t *testing.T) {
|
||||
var v struct {
|
||||
Codes []string `form:"codes"`
|
||||
}
|
||||
|
||||
r, err := http.NewRequest(
|
||||
http.MethodGet,
|
||||
"/a?codes=aaa,bbb,ccc",
|
||||
http.NoBody)
|
||||
assert.NoError(t, err)
|
||||
if assert.NoError(t, Parse(r, &v)) {
|
||||
assert.ElementsMatch(t, []string{"aaa,bbb,ccc"}, v.Codes)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("slice with multiple values", func(t *testing.T) {
|
||||
var v struct {
|
||||
Codes []string `form:"codes,arrayComma=false"`
|
||||
}
|
||||
|
||||
r, err := http.NewRequest(
|
||||
http.MethodGet,
|
||||
"/a?codes=aaa,bbb,ccc&codes=ccc,ddd,eee",
|
||||
http.NoBody)
|
||||
assert.NoError(t, err)
|
||||
if assert.NoError(t, Parse(r, &v)) {
|
||||
assert.ElementsMatch(t, []string{"aaa,bbb,ccc", "ccc,ddd,eee"}, v.Codes)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestParseForm_Error(t *testing.T) {
|
||||
@@ -282,7 +476,7 @@ func TestParseJsonBody(t *testing.T) {
|
||||
|
||||
body := `{"name":"kevin", "age": 18}`
|
||||
r := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body))
|
||||
r.Header.Set(ContentType, header.JsonContentType)
|
||||
r.Header.Set(ContentType, header.ContentTypeJson)
|
||||
|
||||
if assert.NoError(t, Parse(r, &v)) {
|
||||
assert.Equal(t, "kevin", v.Name)
|
||||
@@ -298,7 +492,7 @@ func TestParseJsonBody(t *testing.T) {
|
||||
|
||||
body := `{"name":"kevin", "ag": 18}`
|
||||
r := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body))
|
||||
r.Header.Set(ContentType, header.JsonContentType)
|
||||
r.Header.Set(ContentType, header.ContentTypeJson)
|
||||
|
||||
assert.Error(t, Parse(r, &v))
|
||||
})
|
||||
@@ -323,7 +517,7 @@ func TestParseJsonBody(t *testing.T) {
|
||||
|
||||
body := `[{"name":"kevin", "age": 18}]`
|
||||
r := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body))
|
||||
r.Header.Set(ContentType, header.JsonContentType)
|
||||
r.Header.Set(ContentType, header.ContentTypeJson)
|
||||
|
||||
assert.NoError(t, Parse(r, &v))
|
||||
assert.Equal(t, 1, len(v))
|
||||
@@ -343,13 +537,33 @@ func TestParseJsonBody(t *testing.T) {
|
||||
|
||||
body := `[{"name":"apple", "age": 18}]`
|
||||
r := httptest.NewRequest(http.MethodPost, "/a?product=tree", strings.NewReader(body))
|
||||
r.Header.Set(ContentType, header.JsonContentType)
|
||||
r.Header.Set(ContentType, header.ContentTypeJson)
|
||||
|
||||
assert.NoError(t, Parse(r, &v))
|
||||
assert.Equal(t, 1, len(v))
|
||||
assert.Equal(t, "apple", v[0].Name)
|
||||
assert.Equal(t, 18, v[0].Age)
|
||||
})
|
||||
|
||||
t.Run("bytes field", func(t *testing.T) {
|
||||
type v struct {
|
||||
Signature []byte `json:"signature,optional"`
|
||||
}
|
||||
v1 := v{
|
||||
Signature: []byte{0x01, 0xff, 0x00},
|
||||
}
|
||||
body, _ := json.Marshal(v1)
|
||||
t.Logf("body:%s", string(body))
|
||||
r := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(string(body)))
|
||||
r.Header.Set(ContentType, header.ContentTypeJson)
|
||||
var v2 v
|
||||
err := ParseJsonBody(r, &v2)
|
||||
if assert.NoError(t, err) {
|
||||
assert.Greater(t, len(v2.Signature), 0)
|
||||
}
|
||||
t.Logf("%x", v2.Signature)
|
||||
assert.EqualValues(t, v1.Signature, v2.Signature)
|
||||
})
|
||||
}
|
||||
|
||||
func TestParseRequired(t *testing.T) {
|
||||
@@ -395,7 +609,7 @@ func TestParseHeaders(t *testing.T) {
|
||||
request.Header.Add("addrs", "addr2")
|
||||
request.Header.Add("X-Forwarded-For", "10.0.10.11")
|
||||
request.Header.Add("x-real-ip", "10.0.11.10")
|
||||
request.Header.Add("Accept", header.JsonContentType)
|
||||
request.Header.Add("Accept", header.ContentTypeJson)
|
||||
err = ParseHeaders(request, &v)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
@@ -405,7 +619,7 @@ func TestParseHeaders(t *testing.T) {
|
||||
assert.Equal(t, []string{"addr1", "addr2"}, v.Addrs)
|
||||
assert.Equal(t, "10.0.10.11", v.XForwardedFor)
|
||||
assert.Equal(t, "10.0.11.10", v.XRealIP)
|
||||
assert.Equal(t, header.JsonContentType, v.Accept)
|
||||
assert.Equal(t, header.ContentTypeJson, v.Accept)
|
||||
}
|
||||
|
||||
func TestParseHeaders_Error(t *testing.T) {
|
||||
@@ -497,7 +711,7 @@ func TestParseWithFloatPtr(t *testing.T) {
|
||||
}
|
||||
body := `{"weightFloat32": 3.2}`
|
||||
r := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body))
|
||||
r.Header.Set(ContentType, header.JsonContentType)
|
||||
r.Header.Set(ContentType, header.ContentTypeJson)
|
||||
|
||||
if assert.NoError(t, Parse(r, &v)) {
|
||||
assert.Equal(t, float32(3.2), *v.WeightFloat32)
|
||||
@@ -528,6 +742,42 @@ func TestCustomUnmarshalerStructRequest(t *testing.T) {
|
||||
assert.Equal(t, "hello", v.Foo.Name)
|
||||
}
|
||||
|
||||
func TestParseJsonStringRequest(t *testing.T) {
|
||||
type GoodsInfo struct {
|
||||
Sku int64 `json:"sku,optional"`
|
||||
}
|
||||
|
||||
type GetReq struct {
|
||||
GoodsList []*GoodsInfo `json:"goods_list"`
|
||||
}
|
||||
|
||||
input := `{"goods_list":"[{\"sku\":11},{\"sku\":22}]"}`
|
||||
r := httptest.NewRequest(http.MethodPost, "/a", strings.NewReader(input))
|
||||
r.Header.Set(ContentType, JsonContentType)
|
||||
var v GetReq
|
||||
assert.NotPanics(t, func() {
|
||||
assert.NoError(t, Parse(r, &v))
|
||||
assert.Equal(t, 2, len(v.GoodsList))
|
||||
assert.ElementsMatch(t, []int64{11, 22}, []int64{v.GoodsList[0].Sku, v.GoodsList[1].Sku})
|
||||
})
|
||||
}
|
||||
|
||||
type valid1 struct{}
|
||||
|
||||
func (v valid1) Validate(*http.Request, any) error { return nil }
|
||||
|
||||
type valid2 struct{}
|
||||
|
||||
func (v valid2) Validate(*http.Request, any) error { return nil }
|
||||
|
||||
func TestSetValidatorTwice(t *testing.T) {
|
||||
// panic: sync/atomic: store of inconsistently typed value into Value
|
||||
assert.NotPanics(t, func() {
|
||||
SetValidator(valid1{})
|
||||
SetValidator(valid2{})
|
||||
})
|
||||
}
|
||||
|
||||
func BenchmarkParseRaw(b *testing.B) {
|
||||
r, err := http.NewRequest(http.MethodGet, "http://hello.com/a?name=hello&age=18&percent=3.4", http.NoBody)
|
||||
if err != nil {
|
||||
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
"net/http"
|
||||
"sync"
|
||||
|
||||
"github.com/zeromicro/go-zero/core/logc"
|
||||
"github.com/zeromicro/go-zero/core/logx"
|
||||
"github.com/zeromicro/go-zero/rest/internal/errcode"
|
||||
"github.com/zeromicro/go-zero/rest/internal/header"
|
||||
@@ -119,7 +120,7 @@ func WriteJson(w http.ResponseWriter, code int, v any) {
|
||||
// WriteJsonCtx writes v as json string into w with code.
|
||||
func WriteJsonCtx(ctx context.Context, w http.ResponseWriter, code int, v any) {
|
||||
if err := doWriteJson(w, code, v); err != nil {
|
||||
logx.WithContext(ctx).Error(err)
|
||||
logc.Error(ctx, err)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -178,7 +179,7 @@ func doWriteJson(w http.ResponseWriter, code int, v any) error {
|
||||
return fmt.Errorf("marshal json failed, error: %w", err)
|
||||
}
|
||||
|
||||
w.Header().Set(ContentType, header.JsonContentType)
|
||||
w.Header().Set(ContentType, header.ContentTypeJson)
|
||||
w.WriteHeader(code)
|
||||
|
||||
if n, err := w.Write(bs); err != nil {
|
||||
|
||||
@@ -2,12 +2,23 @@ package httpx
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
)
|
||||
|
||||
const xForwardedFor = "X-Forwarded-For"
|
||||
const (
|
||||
xForwardedFor = "X-Forwarded-For"
|
||||
arraySuffix = "[]"
|
||||
// most servers and clients have a limit of 8192 bytes (8 KB)
|
||||
// one parameter at least take 4 chars, for example `?a=b&c=d`
|
||||
maxFormParamCount = 2048
|
||||
)
|
||||
|
||||
// GetFormValues returns the form values.
|
||||
// GetFormValues returns the form values supporting three array notation formats:
|
||||
// 1. Standard notation: /api?names=alice&names=bob
|
||||
// 2. Comma notation: /api?names=alice,bob
|
||||
// 3. Bracket notation: /api?names[]=alice&names[]=bob
|
||||
func GetFormValues(r *http.Request) (map[string]any, error) {
|
||||
if err := r.ParseForm(); err != nil {
|
||||
return nil, err
|
||||
@@ -19,16 +30,33 @@ func GetFormValues(r *http.Request) (map[string]any, error) {
|
||||
}
|
||||
}
|
||||
|
||||
var n int
|
||||
params := make(map[string]any, len(r.Form))
|
||||
for name, values := range r.Form {
|
||||
filtered := make([]string, 0, len(values))
|
||||
for _, v := range values {
|
||||
if len(v) > 0 {
|
||||
// ignore empty values, especially for optional int parameters
|
||||
// e.g. /api?ids=
|
||||
// e.g. /api
|
||||
// type Req struct {
|
||||
// IDs []int `form:"ids,optional"`
|
||||
// }
|
||||
if len(v) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
if n < maxFormParamCount {
|
||||
filtered = append(filtered, v)
|
||||
n++
|
||||
} else {
|
||||
return nil, fmt.Errorf("too many form values, error: %s", r.Form.Encode())
|
||||
}
|
||||
}
|
||||
|
||||
if len(filtered) > 0 {
|
||||
if strings.HasSuffix(name, arraySuffix) {
|
||||
name = name[:len(name)-2]
|
||||
}
|
||||
params[name] = filtered
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
package httpx
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
@@ -23,3 +25,23 @@ func TestGetRemoteAddrNoHeader(t *testing.T) {
|
||||
|
||||
assert.True(t, len(GetRemoteAddr(r)) == 0)
|
||||
}
|
||||
|
||||
func TestGetFormValues_TooManyValues(t *testing.T) {
|
||||
form := url.Values{}
|
||||
|
||||
// Add more values than the limit
|
||||
for i := 0; i < maxFormParamCount+10; i++ {
|
||||
form.Add("param", fmt.Sprintf("value%d", i))
|
||||
}
|
||||
|
||||
// Create a new request with the form data
|
||||
req, err := http.NewRequest("POST", "/test", strings.NewReader(form.Encode()))
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Set the content type for form data
|
||||
req.Header.Set(ContentType, "application/x-www-form-urlencoded")
|
||||
|
||||
_, err = GetFormValues(req)
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "too many form values")
|
||||
}
|
||||
|
||||
@@ -10,7 +10,7 @@ const (
|
||||
// ContentType means Content-Type.
|
||||
ContentType = header.ContentType
|
||||
// JsonContentType means application/json.
|
||||
JsonContentType = header.JsonContentType
|
||||
JsonContentType = header.ContentTypeJson
|
||||
// KeyField means key.
|
||||
KeyField = "key"
|
||||
// SecretField means secret.
|
||||
|
||||
@@ -3,8 +3,18 @@ package header
|
||||
const (
|
||||
// ApplicationJson stands for application/json.
|
||||
ApplicationJson = "application/json"
|
||||
// CacheControl is the header key for Cache-Control.
|
||||
CacheControl = "Cache-Control"
|
||||
// CacheControlNoCache is the value for Cache-Control: no-cache.
|
||||
CacheControlNoCache = "no-cache"
|
||||
// Connection is the header key for Connection.
|
||||
Connection = "Connection"
|
||||
// ConnectionKeepAlive is the value for Connection: keep-alive.
|
||||
ConnectionKeepAlive = "keep-alive"
|
||||
// ContentType is the header key for Content-Type.
|
||||
ContentType = "Content-Type"
|
||||
// JsonContentType is the content type for JSON.
|
||||
JsonContentType = "application/json; charset=utf-8"
|
||||
// ContentTypeJson is the content type for JSON.
|
||||
ContentTypeJson = "application/json; charset=utf-8"
|
||||
// ContentTypeEventStream is the content type for event stream.
|
||||
ContentTypeEventStream = "text/event-stream"
|
||||
)
|
||||
|
||||
@@ -14,7 +14,7 @@ import (
|
||||
|
||||
"github.com/zeromicro/go-zero/core/codec"
|
||||
"github.com/zeromicro/go-zero/core/iox"
|
||||
"github.com/zeromicro/go-zero/core/logx"
|
||||
"github.com/zeromicro/go-zero/core/logc"
|
||||
"github.com/zeromicro/go-zero/rest/httpx"
|
||||
)
|
||||
|
||||
@@ -123,7 +123,7 @@ func VerifySignature(r *http.Request, securityHeader *ContentSecurityHeader, tol
|
||||
return httpx.CodeSignaturePass
|
||||
}
|
||||
|
||||
logx.Infof("signature different, expect: %s, actual: %s",
|
||||
logc.Infof(r.Context(), "signature different, expect: %s, actual: %s",
|
||||
securityHeader.Signature, actualSignature)
|
||||
|
||||
return httpx.CodeSignatureInvalidToken
|
||||
|
||||
@@ -516,28 +516,55 @@ func TestParsePtrInRequestEmpty(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestParseQueryOptional(t *testing.T) {
|
||||
r, err := http.NewRequest(http.MethodGet, "http://hello.com/kevin/2017?nickname=whatever&zipcode=", nil)
|
||||
assert.Nil(t, err)
|
||||
t.Run("optional with string", func(t *testing.T) {
|
||||
r, err := http.NewRequest(http.MethodGet, "http://hello.com/kevin/2017?nickname=whatever&zipcode=", nil)
|
||||
assert.Nil(t, err)
|
||||
|
||||
router := NewRouter()
|
||||
err = router.Handle(http.MethodGet, "/:name/:year", http.HandlerFunc(
|
||||
func(w http.ResponseWriter, r *http.Request) {
|
||||
v := struct {
|
||||
Nickname string `form:"nickname"`
|
||||
Zipcode int64 `form:"zipcode,optional"`
|
||||
}{}
|
||||
router := NewRouter()
|
||||
err = router.Handle(http.MethodGet, "/:name/:year", http.HandlerFunc(
|
||||
func(w http.ResponseWriter, r *http.Request) {
|
||||
v := struct {
|
||||
Nickname string `form:"nickname"`
|
||||
Zipcode string `form:"zipcode,optional"`
|
||||
}{}
|
||||
|
||||
err = httpx.Parse(r, &v)
|
||||
assert.Nil(t, err)
|
||||
_, err = io.WriteString(w, fmt.Sprintf("%s:%d", v.Nickname, v.Zipcode))
|
||||
assert.Nil(t, err)
|
||||
}))
|
||||
assert.Nil(t, err)
|
||||
err = httpx.Parse(r, &v)
|
||||
assert.Nil(t, err)
|
||||
_, err = io.WriteString(w, fmt.Sprintf("%s:%s", v.Nickname, v.Zipcode))
|
||||
assert.Nil(t, err)
|
||||
}))
|
||||
assert.Nil(t, err)
|
||||
|
||||
rr := httptest.NewRecorder()
|
||||
router.ServeHTTP(rr, r)
|
||||
rr := httptest.NewRecorder()
|
||||
router.ServeHTTP(rr, r)
|
||||
|
||||
assert.Equal(t, "whatever:0", rr.Body.String())
|
||||
assert.Equal(t, "whatever:", rr.Body.String())
|
||||
})
|
||||
|
||||
t.Run("optional with int", func(t *testing.T) {
|
||||
r, err := http.NewRequest(http.MethodGet, "http://hello.com/kevin/2017?nickname=whatever", nil)
|
||||
assert.Nil(t, err)
|
||||
|
||||
router := NewRouter()
|
||||
err = router.Handle(http.MethodGet, "/:name/:year", http.HandlerFunc(
|
||||
func(w http.ResponseWriter, r *http.Request) {
|
||||
v := struct {
|
||||
Nickname string `form:"nickname"`
|
||||
Zipcode int `form:"zipcode,optional"`
|
||||
}{}
|
||||
|
||||
err = httpx.Parse(r, &v)
|
||||
assert.Nil(t, err)
|
||||
_, err = io.WriteString(w, fmt.Sprintf("%s:%d", v.Nickname, v.Zipcode))
|
||||
assert.Nil(t, err)
|
||||
}))
|
||||
assert.Nil(t, err)
|
||||
|
||||
rr := httptest.NewRecorder()
|
||||
router.ServeHTTP(rr, r)
|
||||
|
||||
assert.Equal(t, "whatever:0", rr.Body.String())
|
||||
})
|
||||
}
|
||||
|
||||
func TestParse(t *testing.T) {
|
||||
@@ -601,7 +628,7 @@ func TestParseWrappedRequest(t *testing.T) {
|
||||
func TestParseWrappedGetRequestWithJsonHeader(t *testing.T) {
|
||||
r, err := http.NewRequest(http.MethodGet, "http://hello.com/kevin/2017", bytes.NewReader(nil))
|
||||
assert.Nil(t, err)
|
||||
r.Header.Set(httpx.ContentType, header.JsonContentType)
|
||||
r.Header.Set(httpx.ContentType, header.ContentTypeJson)
|
||||
|
||||
type (
|
||||
Request struct {
|
||||
@@ -634,7 +661,7 @@ func TestParseWrappedGetRequestWithJsonHeader(t *testing.T) {
|
||||
func TestParseWrappedHeadRequestWithJsonHeader(t *testing.T) {
|
||||
r, err := http.NewRequest(http.MethodHead, "http://hello.com/kevin/2017", bytes.NewReader(nil))
|
||||
assert.Nil(t, err)
|
||||
r.Header.Set(httpx.ContentType, header.JsonContentType)
|
||||
r.Header.Set(httpx.ContentType, header.ContentTypeJson)
|
||||
|
||||
type (
|
||||
Request struct {
|
||||
@@ -731,7 +758,7 @@ func TestParseWithAllUtf8(t *testing.T) {
|
||||
r, err := http.NewRequest(http.MethodPost, "http://hello.com/kevin/2017?nickname=whatever&zipcode=200000",
|
||||
bytes.NewBufferString(`{"location": "shanghai", "time": 20170912}`))
|
||||
assert.Nil(t, err)
|
||||
r.Header.Set(httpx.ContentType, header.JsonContentType)
|
||||
r.Header.Set(httpx.ContentType, header.ContentTypeJson)
|
||||
|
||||
router := NewRouter()
|
||||
err = router.Handle(http.MethodPost, "/:name/:year", http.HandlerFunc(
|
||||
@@ -921,7 +948,7 @@ func TestParseWithMissingAllPaths(t *testing.T) {
|
||||
func TestParseGetWithContentLengthHeader(t *testing.T) {
|
||||
r, err := http.NewRequest(http.MethodGet, "http://hello.com/kevin/2017?nickname=whatever&zipcode=200000", nil)
|
||||
assert.Nil(t, err)
|
||||
r.Header.Set(httpx.ContentType, header.JsonContentType)
|
||||
r.Header.Set(httpx.ContentType, header.ContentTypeJson)
|
||||
r.Header.Set(contentLength, "1024")
|
||||
|
||||
router := NewRouter()
|
||||
@@ -949,7 +976,7 @@ func TestParseJsonPostWithTypeMismatch(t *testing.T) {
|
||||
r, err := http.NewRequest(http.MethodPost, "http://hello.com/kevin/2017?nickname=whatever&zipcode=200000",
|
||||
bytes.NewBufferString(`{"time": "20170912"}`))
|
||||
assert.Nil(t, err)
|
||||
r.Header.Set(httpx.ContentType, header.JsonContentType)
|
||||
r.Header.Set(httpx.ContentType, header.ContentTypeJson)
|
||||
|
||||
router := NewRouter()
|
||||
err = router.Handle(http.MethodPost, "/:name/:year", http.HandlerFunc(
|
||||
@@ -975,7 +1002,7 @@ func TestParseJsonPostWithInt2String(t *testing.T) {
|
||||
r, err := http.NewRequest(http.MethodPost, "http://hello.com/kevin/2017",
|
||||
bytes.NewBufferString(`{"time": 20170912}`))
|
||||
assert.Nil(t, err)
|
||||
r.Header.Set(httpx.ContentType, header.JsonContentType)
|
||||
r.Header.Set(httpx.ContentType, header.ContentTypeJson)
|
||||
|
||||
router := NewRouter()
|
||||
err = router.Handle(http.MethodPost, "/:name/:year", http.HandlerFunc(
|
||||
|
||||
@@ -63,6 +63,11 @@ func NewServer(c RestConf, opts ...RunOption) (*Server, error) {
|
||||
return server, nil
|
||||
}
|
||||
|
||||
// AddRoute adds given route into the Server.
|
||||
func (s *Server) AddRoute(r Route, opts ...RouteOption) {
|
||||
s.AddRoutes([]Route{r}, opts...)
|
||||
}
|
||||
|
||||
// AddRoutes add given routes into the Server.
|
||||
func (s *Server) AddRoutes(rs []Route, opts ...RouteOption) {
|
||||
r := featuredRoutes{
|
||||
@@ -74,11 +79,6 @@ func (s *Server) AddRoutes(rs []Route, opts ...RouteOption) {
|
||||
s.ngin.addRoutes(r)
|
||||
}
|
||||
|
||||
// AddRoute adds given route into the Server.
|
||||
func (s *Server) AddRoute(r Route, opts ...RouteOption) {
|
||||
s.AddRoutes([]Route{r}, opts...)
|
||||
}
|
||||
|
||||
// PrintRoutes prints the added routes to stdout.
|
||||
func (s *Server) PrintRoutes() {
|
||||
s.ngin.print()
|
||||
@@ -95,25 +95,6 @@ func (s *Server) Routes() []Route {
|
||||
return routes
|
||||
}
|
||||
|
||||
// ServeHTTP is for test purpose, allow developer to do a unit test with
|
||||
// all defined router without starting an HTTP Server.
|
||||
//
|
||||
// For example:
|
||||
//
|
||||
// server := MustNewServer(...)
|
||||
// server.addRoute(...) // router a
|
||||
// server.addRoute(...) // router b
|
||||
// server.addRoute(...) // router c
|
||||
//
|
||||
// r, _ := http.NewRequest(...)
|
||||
// w := httptest.NewRecorder(...)
|
||||
// server.ServeHTTP(w, r)
|
||||
// // verify the response
|
||||
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
s.ngin.bindRoutes(s.router)
|
||||
s.router.ServeHTTP(w, r)
|
||||
}
|
||||
|
||||
// Start starts the Server.
|
||||
// Graceful shutdown is enabled by default.
|
||||
// Use proc.SetTimeToForceQuit to customize the graceful shutdown period.
|
||||
@@ -298,6 +279,14 @@ func WithSignature(signature SignatureConf) RouteOption {
|
||||
}
|
||||
}
|
||||
|
||||
// WithSSE returns a RouteOption to enable server-sent events.
|
||||
func WithSSE() RouteOption {
|
||||
return func(r *featuredRoutes) {
|
||||
r.sse = true
|
||||
r.timeout = 0
|
||||
}
|
||||
}
|
||||
|
||||
// WithTimeout returns a RouteOption to set timeout with given value.
|
||||
func WithTimeout(timeout time.Duration) RouteOption {
|
||||
return func(r *featuredRoutes) {
|
||||
|
||||
@@ -20,6 +20,7 @@ import (
|
||||
"github.com/zeromicro/go-zero/rest/chain"
|
||||
"github.com/zeromicro/go-zero/rest/httpx"
|
||||
"github.com/zeromicro/go-zero/rest/internal/cors"
|
||||
"github.com/zeromicro/go-zero/rest/internal/header"
|
||||
"github.com/zeromicro/go-zero/rest/router"
|
||||
)
|
||||
|
||||
@@ -231,7 +232,7 @@ func TestWithFileServerMiddleware(t *testing.T) {
|
||||
req := httptest.NewRequest(http.MethodGet, tt.requestPath, nil)
|
||||
rr := httptest.NewRecorder()
|
||||
|
||||
server.ServeHTTP(rr, req)
|
||||
serve(server, rr, req)
|
||||
|
||||
assert.Equal(t, tt.expectedStatus, rr.Code)
|
||||
if len(tt.expectedContent) > 0 {
|
||||
@@ -458,7 +459,7 @@ Port: 54321
|
||||
// we would need to verify the behavior here. Since we don't have
|
||||
// direct access to headers, we'll mock newCorsRouter to capture it.
|
||||
w := httptest.NewRecorder()
|
||||
svr.ServeHTTP(w, httptest.NewRequest(http.MethodOptions, "/", nil))
|
||||
serve(svr, w, httptest.NewRequest(http.MethodOptions, "/", nil))
|
||||
|
||||
vals := w.Header().Values("Access-Control-Allow-Headers")
|
||||
respHeaders := make(map[string]struct{})
|
||||
@@ -748,12 +749,46 @@ Port: 54321
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest("GET", test.path, nil)
|
||||
svr.ServeHTTP(w, req)
|
||||
serve(svr, w, req)
|
||||
assert.Equal(t, test.code, w.Code)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerEventStream(t *testing.T) {
|
||||
server := MustNewServer(RestConf{})
|
||||
server.AddRoutes([]Route{
|
||||
{
|
||||
Method: http.MethodGet,
|
||||
Path: "/foo",
|
||||
Handler: func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Write([]byte("foo"))
|
||||
},
|
||||
},
|
||||
{
|
||||
Method: http.MethodGet,
|
||||
Path: "/bar",
|
||||
Handler: func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Write([]byte("bar"))
|
||||
},
|
||||
},
|
||||
}, WithSSE())
|
||||
|
||||
check := func(val string) {
|
||||
req, err := http.NewRequest(http.MethodGet, fmt.Sprintf("/%s", val), http.NoBody)
|
||||
assert.Nil(t, err)
|
||||
rr := httptest.NewRecorder()
|
||||
serve(server, rr, req)
|
||||
assert.Equal(t, http.StatusOK, rr.Code)
|
||||
assert.Equal(t, header.ContentTypeEventStream, rr.Header().Get(header.ContentType))
|
||||
assert.Equal(t, header.CacheControlNoCache, rr.Header().Get(header.CacheControl))
|
||||
assert.Equal(t, header.ConnectionKeepAlive, rr.Header().Get(header.Connection))
|
||||
assert.Equal(t, val, rr.Body.String())
|
||||
}
|
||||
check("foo")
|
||||
check("bar")
|
||||
}
|
||||
|
||||
//go:embed testdata
|
||||
var content embed.FS
|
||||
|
||||
@@ -765,6 +800,25 @@ func TestServerEmbedFileSystem(t *testing.T) {
|
||||
req, err := http.NewRequest(http.MethodGet, "/assets/sample.txt", http.NoBody)
|
||||
assert.Nil(t, err)
|
||||
rr := httptest.NewRecorder()
|
||||
server.ServeHTTP(rr, req)
|
||||
serve(server, rr, req)
|
||||
assert.Equal(t, sampleContent, rr.Body.String())
|
||||
}
|
||||
|
||||
// serve is for test purpose, allow developer to do a unit test with
|
||||
// all defined routes without starting an HTTP Server.
|
||||
//
|
||||
// For example:
|
||||
//
|
||||
// server := MustNewServer(...)
|
||||
// server.addRoute(...) // router a
|
||||
// server.addRoute(...) // router b
|
||||
// server.addRoute(...) // router c
|
||||
//
|
||||
// r, _ := http.NewRequest(...)
|
||||
// w := httptest.NewRecorder(...)
|
||||
// serve(server, w, r)
|
||||
// // verify the response
|
||||
func serve(s *Server, w http.ResponseWriter, r *http.Request) {
|
||||
s.ngin.bindRoutes(s.router)
|
||||
s.router.ServeHTTP(w, r)
|
||||
}
|
||||
|
||||
@@ -35,6 +35,7 @@ type (
|
||||
priority bool
|
||||
jwt jwtSetting
|
||||
signature signatureSetting
|
||||
sse bool
|
||||
routes []Route
|
||||
maxBytes int64
|
||||
}
|
||||
|
||||
@@ -3,7 +3,8 @@ FROM golang:alpine AS builder
|
||||
LABEL stage=gobuilder
|
||||
|
||||
ENV CGO_ENABLED=0
|
||||
ENV GOPROXY=https://goproxy.cn,direct
|
||||
# if you are in China, you can use the following command to speed up the download
|
||||
# ENV GOPROXY=https://goproxy.cn,direct
|
||||
|
||||
RUN apk update --no-cache && apk add --no-cache tzdata
|
||||
RUN go install google.golang.org/protobuf/cmd/protoc-gen-go@latest
|
||||
|
||||
@@ -2,6 +2,7 @@ package api
|
||||
|
||||
import (
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"github.com/zeromicro/go-zero/tools/goctl/api/apigen"
|
||||
"github.com/zeromicro/go-zero/tools/goctl/api/dartgen"
|
||||
"github.com/zeromicro/go-zero/tools/goctl/api/docgen"
|
||||
@@ -10,10 +11,12 @@ import (
|
||||
"github.com/zeromicro/go-zero/tools/goctl/api/javagen"
|
||||
"github.com/zeromicro/go-zero/tools/goctl/api/ktgen"
|
||||
"github.com/zeromicro/go-zero/tools/goctl/api/new"
|
||||
"github.com/zeromicro/go-zero/tools/goctl/api/swagger"
|
||||
"github.com/zeromicro/go-zero/tools/goctl/api/tsgen"
|
||||
"github.com/zeromicro/go-zero/tools/goctl/api/validate"
|
||||
"github.com/zeromicro/go-zero/tools/goctl/config"
|
||||
"github.com/zeromicro/go-zero/tools/goctl/internal/cobrax"
|
||||
"github.com/zeromicro/go-zero/tools/goctl/pkg/env"
|
||||
"github.com/zeromicro/go-zero/tools/goctl/plugin"
|
||||
)
|
||||
|
||||
@@ -31,6 +34,7 @@ var (
|
||||
ktCmd = cobrax.NewCommand("kt", cobrax.WithRunE(ktgen.KtCommand))
|
||||
pluginCmd = cobrax.NewCommand("plugin", cobrax.WithRunE(plugin.PluginCommand))
|
||||
tsCmd = cobrax.NewCommand("ts", cobrax.WithRunE(tsgen.TsCommand))
|
||||
swaggerCmd = cobrax.NewCommand("swagger", cobrax.WithRunE(swagger.Command))
|
||||
)
|
||||
|
||||
func init() {
|
||||
@@ -46,6 +50,7 @@ func init() {
|
||||
pluginCmdFlags = pluginCmd.Flags()
|
||||
tsCmdFlags = tsCmd.Flags()
|
||||
validateCmdFlags = validateCmd.Flags()
|
||||
swaggerCmdFlags = swaggerCmd.Flags()
|
||||
)
|
||||
|
||||
apiCmdFlags.StringVar(&apigen.VarStringOutput, "o")
|
||||
@@ -72,6 +77,7 @@ func init() {
|
||||
goCmdFlags.StringVar(&gogen.VarStringHome, "home")
|
||||
goCmdFlags.StringVar(&gogen.VarStringRemote, "remote")
|
||||
goCmdFlags.StringVar(&gogen.VarStringBranch, "branch")
|
||||
goCmdFlags.BoolVar(&gogen.VarBoolWithTest, "test")
|
||||
goCmdFlags.StringVarWithDefaultValue(&gogen.VarStringStyle, "style", config.DefaultFormat)
|
||||
|
||||
javaCmdFlags.StringVar(&javagen.VarStringDir, "dir")
|
||||
@@ -96,8 +102,16 @@ func init() {
|
||||
tsCmdFlags.StringVar(&tsgen.VarStringCaller, "caller")
|
||||
tsCmdFlags.BoolVar(&tsgen.VarBoolUnWrap, "unwrap")
|
||||
|
||||
swaggerCmdFlags.StringVar(&swagger.VarStringAPI, "api")
|
||||
swaggerCmdFlags.StringVar(&swagger.VarStringDir, "dir")
|
||||
swaggerCmdFlags.StringVar(&swagger.VarStringFilename, "filename")
|
||||
swaggerCmdFlags.BoolVar(&swagger.VarBoolYaml, "yaml")
|
||||
|
||||
validateCmdFlags.StringVar(&validate.VarStringAPI, "api")
|
||||
|
||||
// Add sub-commands
|
||||
Cmd.AddCommand(dartCmd, docCmd, formatCmd, goCmd, javaCmd, ktCmd, newCmd, pluginCmd, tsCmd, validateCmd)
|
||||
if env.UseExperimental() {
|
||||
Cmd.AddCommand(swaggerCmd)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,10 +10,10 @@ const (
|
||||
import 'package:shared_preferences/shared_preferences.dart';
|
||||
import '../data/tokens.dart';
|
||||
|
||||
/// 保存tokens到本地
|
||||
/// store tokens to local
|
||||
///
|
||||
/// 传入null则删除本地tokens
|
||||
/// 返回:true:设置成功 false:设置失败
|
||||
/// pass null will clean local stored tokens
|
||||
/// returns true if success, otherwise false
|
||||
Future<bool> setTokens(Tokens tokens) async {
|
||||
var sp = await SharedPreferences.getInstance();
|
||||
if (tokens == null) {
|
||||
@@ -23,9 +23,9 @@ Future<bool> setTokens(Tokens tokens) async {
|
||||
return await sp.setString('tokens', jsonEncode(tokens.toJson()));
|
||||
}
|
||||
|
||||
/// 获取本地存储的tokens
|
||||
/// get local stored tokens
|
||||
///
|
||||
/// 如果没有,则返回null
|
||||
/// if no, returns null
|
||||
Future<Tokens> getTokens() async {
|
||||
try {
|
||||
var sp = await SharedPreferences.getInstance();
|
||||
@@ -82,7 +82,8 @@ func genVars(dir string, isLegacy bool, scheme string, hostname string) error {
|
||||
}
|
||||
|
||||
if !fileExists(dir + "vars.dart") {
|
||||
err = os.WriteFile(dir+"vars.dart", []byte(fmt.Sprintf(`const serverHost='%s://%s';`, scheme, hostname)), 0o644)
|
||||
err = os.WriteFile(dir+"vars.dart", []byte(fmt.Sprintf(`const serverHost='%s://%s';`,
|
||||
scheme, hostname)), 0o644)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -38,7 +38,8 @@ var (
|
||||
// VarStringBranch describes the branch.
|
||||
VarStringBranch string
|
||||
// VarStringStyle describes the style of output files.
|
||||
VarStringStyle string
|
||||
VarStringStyle string
|
||||
VarBoolWithTest bool
|
||||
)
|
||||
|
||||
// GoCommand gen go project files from command line
|
||||
@@ -49,6 +50,7 @@ func GoCommand(_ *cobra.Command, _ []string) error {
|
||||
home := VarStringHome
|
||||
remote := VarStringRemote
|
||||
branch := VarStringBranch
|
||||
withTest := VarBoolWithTest
|
||||
if len(remote) > 0 {
|
||||
repo, _ := util.CloneIntoGitHome(remote, branch)
|
||||
if len(repo) > 0 {
|
||||
@@ -66,11 +68,11 @@ func GoCommand(_ *cobra.Command, _ []string) error {
|
||||
return errors.New("missing -dir")
|
||||
}
|
||||
|
||||
return DoGenProject(apiFile, dir, namingStyle)
|
||||
return DoGenProject(apiFile, dir, namingStyle, withTest)
|
||||
}
|
||||
|
||||
// DoGenProject gen go project files with api file
|
||||
func DoGenProject(apiFile, dir, style string) error {
|
||||
func DoGenProject(apiFile, dir, style string, withTest bool) error {
|
||||
api, err := parser.Parse(apiFile)
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -100,6 +102,10 @@ func DoGenProject(apiFile, dir, style string) error {
|
||||
logx.Must(genHandlers(dir, rootPkg, cfg, api))
|
||||
logx.Must(genLogic(dir, rootPkg, cfg, api))
|
||||
logx.Must(genMiddleware(dir, cfg, api))
|
||||
if withTest {
|
||||
logx.Must(genHandlersTest(dir, rootPkg, cfg, api))
|
||||
logx.Must(genLogicTest(dir, rootPkg, cfg, api))
|
||||
}
|
||||
|
||||
if err := backupAndSweep(apiFile); err != nil {
|
||||
return err
|
||||
|
||||
@@ -348,7 +348,7 @@ func validateWithCamel(t *testing.T, api, camel string) {
|
||||
assert.Nil(t, err)
|
||||
err = initMod(dir)
|
||||
assert.Nil(t, err)
|
||||
err = DoGenProject(api, dir, camel)
|
||||
err = DoGenProject(api, dir, camel, true)
|
||||
assert.Nil(t, err)
|
||||
filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
|
||||
if strings.HasSuffix(path, ".go") {
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user