background/db/sqlManager.go

package db

import (
	"background/config"
	"background/logs"
	"crypto/tls"
	"crypto/x509"
	"database/sql"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"time"

	"github.com/go-sql-driver/mysql"
	_ "github.com/go-sql-driver/mysql"
	"github.com/jinzhu/gorm"
	"gopkg.in/mgo.v2"
	"gopkg.in/olivere/elastic.v7"
)

var gEla ElkEngine
var gDb Database
var blogDb Database
var gMongo *mgo.Session
var gElkEngine ElkEngine
var gOrm *gorm.DB

func Init() {
	var e error
	mysqlconf := config.GetMysqlConfig()
	log.Print("api runmode is " + config.ApiConfig().RunMode)
	if config.ApiConfig().RunMode == "debug" {
		gDb = Database{Type: string(""), DB: initMysqlTLS(mysqlconf)}
		rootCertPool := x509.NewCertPool()
		pem, err := ioutil.ReadFile("pem/ca.pem")
		if err != nil {
			log.Fatal(err)
		}
		if ok := rootCertPool.AppendCertsFromPEM(pem); !ok {
			log.Fatal("Failed to append PEM.")
		}
		clientCert := make([]tls.Certificate, 0, 1)
		certs, err := tls.LoadX509KeyPair("pem/client-cert.pem", "pem/client-key.pem")
		if err != nil {
			log.Fatal(err)
		}
		clientCert = append(clientCert, certs)
		mysql.RegisterTLSConfig("custom", &tls.Config{
			RootCAs:            rootCertPool,
			Certificates:       clientCert,
			InsecureSkipVerify: true,
		})
		// try to connect to mysql database.
		cfg := mysql.Config{
			User:                 config.GetMysqlBlogConfig().UserName,
			Passwd:               config.GetMysqlBlogConfig().Password,
			Addr:                 config.GetMysqlBlogConfig().Addr, //IP:PORT
			Net:                  "tcp",
			DBName:               "background",
			Loc:                  time.Local,
			AllowNativePasswords: true,
		}
		cfg.TLSConfig = "custom"
		str := cfg.FormatDSN()
		gOrm, e = gorm.Open("mysql", str)
		if nil != e {
			log.Print(e.Error())
			os.Exit(-1)
		}
	} else {
		rootCertPool := x509.NewCertPool()
		pem, err := ioutil.ReadFile("pem/ca.pem")
		if err != nil {
			log.Fatal(err)
		}
		if ok := rootCertPool.AppendCertsFromPEM(pem); !ok {
			log.Fatal("Failed to append PEM.")
		}
		clientCert := make([]tls.Certificate, 0, 1)
		certs, err := tls.LoadX509KeyPair("pem/client-cert.pem", "pem/client-key.pem")
		if err != nil {
			log.Fatal(err)
		}
		clientCert = append(clientCert, certs)
		log.Print("client cert is ", clientCert)
		mysql.RegisterTLSConfig("custom", &tls.Config{
			RootCAs:      rootCertPool,
			Certificates: clientCert,
		})

		// try to connect to mysql database.
		cfg := mysql.Config{
			User:                 config.GetMysqlBlogConfig().UserName,
			Passwd:               config.GetMysqlBlogConfig().Password,
			Addr:                 config.GetMysqlBlogConfig().Addr, //IP:PORT
			Net:                  "tcp",
			DBName:               "background",
			Loc:                  time.Local,
			AllowNativePasswords: true,
		}
		cfg.TLSConfig = "custom"
		str := cfg.FormatDSN()
		log.Print(str)
		gOrm, e = gorm.Open("mysql", str)
		if nil != e {
			log.Print(e.Error())
			os.Exit(-1)
		}
	}

	InitELK()
}

func InitELK() {
	var e error
	elkconf := config.GetElkConfig()
	gElkEngine.cli, e = elastic.NewClient(
		elastic.SetURL(elkconf.Address),
		// Must turn off sniff in docker
		elastic.SetSniff(false),
		elastic.SetBasicAuth(elkconf.User, elkconf.Password),
	)
	if nil != e {
		logs.Error(e.Error())
		gElkEngine.cli = nil
	}
}


func initMysql(mysqlconf *config.MysqlConfig) *sql.DB {
	cnn := fmt.Sprintf("%s:%s@tcp(%s:3306)/%s?charset=utf8", mysqlconf.UserName, mysqlconf.Password,
		mysqlconf.Addr, mysqlconf.Db)
	log.Print("Connect to mysql " + cnn)
	_db, err := sql.Open("mysql", cnn)
	if err != nil {
		fmt.Println("connect sql server ", err.Error())
		os.Exit(200)
	}
	e := _db.Ping()
	if nil != e {
		fmt.Println("ping mysql server error", e.Error())
		os.Exit(200)
	}
	return _db
}

func initMysqlTLS(mysqlconf *config.MysqlConfig) *sql.DB {
	rootCertPool := x509.NewCertPool()
	pem, err := ioutil.ReadFile("pem/ca.pem")
	if err != nil {
		log.Fatal(err)
	}
	if ok := rootCertPool.AppendCertsFromPEM(pem); !ok {
		log.Fatal("Failed to append PEM.")
	}
	clientCert := make([]tls.Certificate, 0, 1)
	certs, err := tls.LoadX509KeyPair("pem/client-cert.pem", "pem/client-key.pem")
	if err != nil {
		log.Fatal(err)
	}
	clientCert = append(clientCert, certs)
	mysql.RegisterTLSConfig("custom", &tls.Config{
		RootCAs:            rootCertPool,
		Certificates:       clientCert,
		InsecureSkipVerify: true,
	})
	cnn := fmt.Sprintf("%s:%s@tcp(%s:3306)/%s?charset=utf8&tls=custom", mysqlconf.UserName, mysqlconf.Password,
		mysqlconf.Addr, mysqlconf.Db)
	log.Print("Connect to mysql " + cnn)
	_db, err := sql.Open("mysql", cnn)
	if err != nil {
		fmt.Println("connect sql server ", err.Error())
		os.Exit(200)
	}
	e := _db.Ping()
	if nil != e {
		fmt.Println(e.Error())
		os.Exit(200)
	}
	return _db
}
func initMongoDb(conf *config.MongoConfig) *mgo.Session {
	var url string
	if conf.UserName == "" {
		url = fmt.Sprintf("mongodb://%s:%d/%s", conf.Addr, conf.Port, conf.Db)

	} else {
		url = fmt.Sprintf("mongodb://%s:%s@%s:%d/%s", conf.UserName, conf.Password, conf.Addr, conf.Port, conf.Db)

	}
	log.Print("connect to url " + url)
	logs.Debug("connect to url " + url)

	mgo, err := mgo.Dial(url)
	if nil != err {
		logs.Error(err.Error())
		return nil
	}
	return mgo
}

func GetMysqlClient() *Database {
	return &gDb
}
func GetBlogMysql() *Database {
	return &blogDb
}
func GetMongoDb() *mgo.Session {
	return gMongo
}
func InitMongoDb() {
	mongoConf := config.GetMongoDBConfig()
	gMongo = initMongoDb(mongoConf)
	if gMongo == nil {
		log.Print("error mongo initial")
		os.Exit(25)
	}
}
func GetElastic() *ElkEngine {
	return &gElkEngine
}
func GetOrm() *gorm.DB {
	return gOrm
}