1
2 package mongodb
3
4
5
6
7 import (
8 "context"
9 "crypto/tls"
10 "fmt"
11 "net"
12 "strings"
13
14 "github.com/Azure/azure-sdk-for-go/services/cosmos-db/mgmt/2015-04-08/documentdb"
15 "github.com/Azure/go-autorest/autorest"
16 "github.com/Azure/go-autorest/autorest/adal"
17 "github.com/Azure/go-autorest/autorest/azure"
18 "github.com/globalsign/mgo"
19 )
20
21 const (
22 cosmosDbConnectionPort = 10255
23 )
24
25
26 func NewMongoDBClientWithConnectionString(connectionString string) (*mgo.Session, error) {
27
28
29 connectionString = strings.Replace(connectionString, "ssl=true", "", -1)
30 dialInfo, err := mgo.ParseURL(connectionString)
31
32 if err != nil {
33 return nil, err
34 }
35
36 return NewMongoDBClient(dialInfo)
37 }
38
39
40 func NewMongoDBClientWithCredentials(username, password, host string) (*mgo.Session, error) {
41
42 dialInfo := &mgo.DialInfo{
43 Addrs: []string{fmt.Sprintf("%s:%d", host, cosmosDbConnectionPort)},
44 Username: username,
45 Password: password,
46 }
47
48 return NewMongoDBClient(dialInfo)
49 }
50
51
52 func NewMongoDBClientWithSPToken(spToken *adal.ServicePrincipalToken, subscriptionID, resourceGroup, account string, environment azure.Environment) (*mgo.Session, error) {
53
54 authorizer := autorest.NewBearerAuthorizer(spToken)
55
56 cosmosDbClient := documentdb.NewDatabaseAccountsClientWithBaseURI(environment.ResourceManagerEndpoint, subscriptionID)
57 cosmosDbClient.Authorizer = authorizer
58 cosmosDbClient.AddToUserAgent("dataplane mongodb")
59
60 result, err := cosmosDbClient.ListConnectionStrings(context.Background(), resourceGroup, account)
61
62 if err != nil {
63 return nil, err
64 }
65
66 connectionStrings := *result.ConnectionStrings
67
68 for _, connectionString := range connectionStrings {
69 session, err := NewMongoDBClientWithConnectionString(*connectionString.ConnectionString)
70
71 if session != nil && err == nil {
72 return session, nil
73 }
74 }
75
76 return nil, err
77 }
78
79
80 func NewMongoDBClientWithMSI(subscriptionID, resourceGroup, account string, environment azure.Environment) (*mgo.Session, error) {
81
82 msiEndpoint, err := adal.GetMSIVMEndpoint()
83 spToken, err := adal.NewServicePrincipalTokenFromMSI(msiEndpoint, environment.ResourceManagerEndpoint)
84
85 if err != nil {
86 return nil, err
87 }
88
89 return NewMongoDBClientWithSPToken(spToken, subscriptionID, resourceGroup, account, environment)
90 }
91
92
93 func NewMongoDBClient(dialInfo *mgo.DialInfo) (*mgo.Session, error) {
94
95 dialInfo.DialServer = func(addr *mgo.ServerAddr) (net.Conn, error) {
96 return tls.Dial("tcp", addr.String(), &tls.Config{})
97 }
98
99 session, err := mgo.DialWithInfo(dialInfo)
100
101 if err != nil {
102 return nil, err
103 }
104
105 return session, nil
106 }
107
View as plain text