-
Notifications
You must be signed in to change notification settings - Fork 0
/
config.go
92 lines (79 loc) · 2.2 KB
/
config.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
package cmd
import (
"context"
"database/sql"
"fmt"
mssql "github.com/denisenkom/go-mssqldb"
"github.com/denisenkom/go-mssqldb/azuread"
"golang.org/x/net/proxy"
"io/ioutil"
"os"
"path"
"strings"
_ "github.com/denisenkom/go-mssqldb/azuread"
"github.com/denisenkom/go-mssqldb/msdsn"
"github.com/pkg/errors"
"github.com/sirupsen/logrus"
"gopkg.in/yaml.v3"
)
type DatabaseConfig struct {
Connection string `yaml:"connection"`
Dsn msdsn.Config
UsePasswordLogin bool
}
func OpenSocks5Sql(dsn string) (*sql.DB, error) {
var err error
var connector *mssql.Connector
if strings.HasPrefix(dsn, "azuresql://") {
connector, err = azuread.NewConnector(dsn)
if err != nil {
return nil, err
}
} else if strings.HasPrefix(dsn, "sqlserver://") {
//dbi, err := sql.Open("sqlserver", dsn)
connector, err = mssql.NewConnector(dsn)
if err != nil {
return nil, err
}
} else {
return nil, errors.New("expected URI-style dsn; sqlserver:// for password login or azuresql:// for AD login")
}
socksProxyAddress := os.Getenv("SQL_SOCKS")
if socksProxyAddress != "" {
dialer, err := proxy.SOCKS5("tcp", socksProxyAddress, nil, nil)
if err != nil {
return nil, errors.Wrap(err, fmt.Sprintf("Could not connect with SOCKS5 to %s", socksProxyAddress))
}
connector.Dialer = dialer.(proxy.ContextDialer)
}
return sql.OpenDB(connector), nil
}
func (dbcfg DatabaseConfig) Open(ctx context.Context, logger logrus.FieldLogger) (*sql.DB, error) {
return OpenSocks5Sql(dbcfg.Connection)
}
type Config struct {
Databases map[string]DatabaseConfig `yaml:"databases"`
ServiceName string `yaml:"servicename"`
}
func LoadConfig() (Config, error) {
var result Config
configFilename := path.Join(directory, "sqlcode.yaml")
if _, err := os.Stat(configFilename); os.IsNotExist(err) {
return Config{}, errors.New("No sqlcode.yaml found in current directory")
}
yamlFile, err := ioutil.ReadFile(configFilename)
if err != nil {
return Config{}, err
}
err = yaml.Unmarshal(yamlFile, &result)
if err != nil {
return Config{}, err
}
for key, dbcfg := range result.Databases {
if err != nil {
return Config{}, err
}
result.Databases[key] = dbcfg
}
return result, nil
}