-
Notifications
You must be signed in to change notification settings - Fork 9
/
driver.go
101 lines (89 loc) · 2.6 KB
/
driver.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
package driver
import (
"context"
"fmt"
"github.com/dtm-labs/dtmdriver"
consul "github.com/go-kratos/kratos/contrib/registry/consul/v2"
etcd "github.com/go-kratos/kratos/contrib/registry/etcd/v2"
"github.com/go-kratos/kratos/v2/registry"
_ "github.com/go-kratos/kratos/v2/transport/grpc/resolver/direct"
"github.com/go-kratos/kratos/v2/transport/grpc/resolver/discovery"
consulAPI "github.com/hashicorp/consul/api"
etcdAPI "go.etcd.io/etcd/client/v3"
"google.golang.org/grpc/resolver"
"net/url"
"strings"
)
const (
DriverName = "dtm-driver-kratos"
DefaultScheme = "discovery"
EtcdScheme = "etcd"
ConsulScheme = "consul"
)
type kratosDriver struct{}
func (k *kratosDriver) GetName() string {
return DriverName
}
func (k *kratosDriver) RegisterAddrResolver() {
}
func (k *kratosDriver) RegisterService(target string, endpoint string) error {
if target == "" {
return nil
}
u, err := url.Parse(target)
if err != nil {
return err
}
switch u.Scheme {
case DefaultScheme:
fallthrough
case EtcdScheme:
registerInstance := ®istry.ServiceInstance{
Name: strings.TrimPrefix(u.Path, "/"),
Endpoints: strings.Split(endpoint, ","),
}
client, err := etcdAPI.New(etcdAPI.Config{
Endpoints: strings.Split(u.Host, ","),
})
if err != nil {
return err
}
registry := etcd.New(client)
//add resolver so that dtm can handle discovery://
resolver.Register(discovery.NewBuilder(registry, discovery.WithInsecure(true)))
return registry.Register(context.Background(), registerInstance)
case ConsulScheme:
registerInstance := ®istry.ServiceInstance{
Name: strings.TrimPrefix(u.Path, "/"),
Endpoints: strings.Split(endpoint, ","),
}
client, err := consulAPI.NewClient(&consulAPI.Config{Address: u.Host})
if err != nil {
return err
}
registry := consul.New(client)
//add resolver so that dtm can handle discovery://
resolver.Register(discovery.NewBuilder(registry, discovery.WithInsecure(true)))
return registry.Register(context.Background(), registerInstance)
default:
return fmt.Errorf("unknown scheme: %s", u.Scheme)
}
}
func (k *kratosDriver) ParseServerMethod(uri string) (server string, method string, err error) {
if !strings.Contains(uri, "//") {
sep := strings.IndexByte(uri, '/')
if sep == -1 {
return "", "", fmt.Errorf("bad url: '%s'. no '/' found", uri)
}
return uri[:sep], uri[sep:], nil
}
u, err := url.Parse(uri)
if err != nil {
return "", "", nil
}
index := strings.IndexByte(u.Path[1:], '/') + 1
return u.Scheme + "://" + u.Host + u.Path[:index], u.Path[index:], nil
}
func init() {
dtmdriver.Register(&kratosDriver{})
}