forked from zph/moresql
-
Notifications
You must be signed in to change notification settings - Fork 0
/
db.go
102 lines (94 loc) · 2.34 KB
/
db.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
102
package moresql
import (
"bytes"
"crypto/tls"
"crypto/x509"
"encoding/binary"
"errors"
"io/ioutil"
"math"
"net"
"time"
log "github.com/Sirupsen/logrus"
"github.com/jmoiron/sqlx"
_ "github.com/lib/pq"
mgo "gopkg.in/mgo.v2"
"gopkg.in/mgo.v2/bson"
)
func GetMongoConnection(env Env) (session *mgo.Session) {
if env.UseSSL() {
var tlsConfig *tls.Config
if env.SSLCert != "" {
// Certificate Provided
clientCert, err := ioutil.ReadFile(env.SSLCert)
if err != nil {
log.Fatalln("Unable to read ssl certificate")
}
roots := x509.NewCertPool()
ok := roots.AppendCertsFromPEM([]byte(clientCert))
if !ok {
log.Fatalln("failed to parse root certificate")
}
tlsConfig = &tls.Config{RootCAs: roots}
} else if env.SSLInsecureSkipVerify {
tlsConfig = &tls.Config{InsecureSkipVerify: true}
}
c, err := mgo.ParseURL(env.urls.mongo)
if err != nil {
log.Fatalf("Unable to parse mongo url")
}
dialInfo := &mgo.DialInfo{
Addrs: c.Addrs,
Database: c.Database,
Source: c.Source,
Username: c.Username,
Password: c.Password,
DialServer: func(addr *mgo.ServerAddr) (net.Conn, error) {
return tls.Dial("tcp", addr.String(), tlsConfig)
},
Timeout: time.Second * 10,
}
session, err = mgo.DialWithInfo(dialInfo)
if err != nil {
log.Fatal(err)
}
session.SetMode(mgo.Primary, true)
} else {
var err error
session, err = mgo.Dial(env.urls.mongo)
if err != nil {
log.Fatal(err)
}
session.SetMode(mgo.Primary, true)
}
return
}
func GetPostgresConnection(env Env) (pg *sqlx.DB) {
var err error
pg, err = sqlx.Connect("postgres", env.urls.postgres)
if err != nil {
log.Fatal(err)
}
setupPgDefaults(pg)
return
}
// setupPgDefaults: Set safe cap so workers do not overwhelm server
func setupPgDefaults(pg *sqlx.DB) {
pg.SetMaxIdleConns(50)
pg.SetMaxOpenConns(50)
}
// Credit: https://github.com/go-mgo/mgo/pull/202/files#diff-b47d6566744e81abad9312022bdc8896R374
// From @mwmahlberg
func NewMongoTimestamp(t time.Time, c uint32) (bson.MongoTimestamp, error) {
var tv uint32
u := t.Unix()
if u < 0 || u > math.MaxUint32 {
return -1, errors.New("invalid value for time")
}
tv = uint32(u)
buf := bytes.Buffer{}
binary.Write(&buf, binary.BigEndian, tv)
binary.Write(&buf, binary.BigEndian, c)
i := int64(binary.BigEndian.Uint64(buf.Bytes()))
return bson.MongoTimestamp(i), nil
}