pgproxy.go

package pgproxy

import (
	"log"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

	pg "github.com/go-pg/pg/v9"
)

// DBProxyCallback pg数据库操作的回调函数
type dbProxyCallback func(dbCli *pg.DB) error

// DBProxy 数据库客户端的代理
type dbProxy struct {
	proxyLock sync.RWMutex //代理的锁
	Options   *pg.Options
	conn      *pg.DB
	callBacks []dbProxyCallback
}

// New 创建一个新的数据库客户端代理
func New() *dbProxy {
	proxy := new(dbProxy)
	proxy.proxyLock = sync.RWMutex{}
	return proxy
}

// IsOk 检查代理是否已经可用
func (proxy *dbProxy) IsOk() bool {
	if proxy.conn == nil {
		return false
	}
	return true
}

//GetConn 获取被代理的连接
func (proxy *dbProxy) GetConn() (*pg.DB, error) {
	if !proxy.IsOk() {
		return proxy.conn, ErrProxyNotInited
	}
	proxy.proxyLock.RLock()
	defer proxy.proxyLock.RUnlock()
	return proxy.conn, nil
}

// Close 关闭pg
func (proxy *dbProxy) Close() {
	if proxy.IsOk() {
		proxy.conn.Close()
		proxy.proxyLock.Lock()
		proxy.conn = nil
		proxy.proxyLock.Unlock()
	}
}

//SetConnect 设置连接的客户端
func (proxy *dbProxy) SetConnect(cli *pg.DB) {
	proxy.proxyLock.Lock()
	proxy.conn = cli
	proxy.proxyLock.Unlock()
	for _, cb := range proxy.callBacks {
		err := cb(proxy.conn)
		if err != nil {
			log.Println("regist callback get error", err)
		} else {
			log.Println("regist callback done")
		}
	}
}

// 将url解析为pg的初始化参数
func parseDBURL(address string) (*pg.Options, error) {
	result := &pg.Options{}
	u, err := url.Parse(address)
	if err != nil {
		return result, err
	}
	if u.Scheme != "postgres" {
		return result, ErrURLSchemaWrong
	}

	user := u.User.Username()
	if user == "" {
		result.User = "postgres"
	} else {
		result.User = user
	}
	password, has := u.User.Password()
	if has == false {
		result.Password = "postgres"
	} else {
		result.Password = password
	}
	result.Addr = u.Host
	result.Database = u.Path[1:]
	if u.RawQuery != "" {
		v, err := url.ParseQuery(u.RawQuery)
		if err != nil {
			log.Fatal(err)
			return result, nil
		}
		for key, value := range v {
			switch strings.ToLower(key) {
			case "maxretries":
				{
					maxretries, err := strconv.Atoi(value[0])
					if err != nil {
						log.Fatal(err)
					} else {
						result.MaxRetries = maxretries
					}
				}
			case "retrystatementtimeout":
				{
					bv := strings.ToLower(value[0])
					switch bv {
					case "true":
						result.RetryStatementTimeout = true
					case "false":
						result.RetryStatementTimeout = false
					default:
						log.Fatal("unknown value for RetryStatementTimeout")
					}
				}
			case "minretrybackoff":
				{
					number, err := strconv.Atoi(value[0])
					if err != nil {
						log.Fatal(err)
					} else {
						result.MinRetryBackoff = time.Duration(number) * time.Second
					}
				}
			case "maxretrybackoff":
				{
					number, err := strconv.Atoi(value[0])
					if err != nil {
						log.Fatal(err)
					} else {
						result.MaxRetryBackoff = time.Duration(number) * time.Second
					}
				}
			case "dialtimeout":
				{
					number, err := strconv.Atoi(value[0])
					if err != nil {
						log.Fatal(err)
					} else {
						result.DialTimeout = time.Duration(number) * time.Second
					}
				}
			case "readtimeout":
				{
					number, err := strconv.Atoi(value[0])
					if err != nil {
						log.Fatal(err)
					} else {
						result.ReadTimeout = time.Duration(number) * time.Second
					}
				}
			case "writetimeout":
				{
					number, err := strconv.Atoi(value[0])
					if err != nil {
						log.Fatal(err)
					} else {
						result.WriteTimeout = time.Duration(number) * time.Second
					}
				}
			case "maxconnage":
				{
					number, err := strconv.Atoi(value[0])
					if err != nil {
						log.Fatal(err)
					} else {
						result.MaxConnAge = time.Duration(number) * time.Second
					}
				}
			case "pooltimeout":
				{
					number, err := strconv.Atoi(value[0])
					if err != nil {
						log.Fatal(err)
					} else {
						result.PoolTimeout = time.Duration(number) * time.Second
					}
				}
			case "idletimeout":
				{
					number, err := strconv.Atoi(value[0])
					if err != nil {
						log.Fatal(err)
					} else {
						result.IdleTimeout = time.Duration(number) * time.Second
					}
				}
			case "idlecheckfrequency":
				{
					number, err := strconv.Atoi(value[0])
					if err != nil {
						log.Fatal(err)
					} else {
						result.IdleCheckFrequency = time.Duration(number) * time.Second
					}
				}
			case "PoolSize":
				{
					number, err := strconv.Atoi(value[0])
					if err != nil {
						log.Fatal(err)
					} else {
						result.PoolSize = number
					}
				}
			case "MinIdleConns":
				{
					number, err := strconv.Atoi(value[0])
					if err != nil {
						log.Fatal(err)
					} else {
						result.MinIdleConns = number
					}
				}
			}
		}
	}
	return result, nil
}

// Init 给代理赋值客户端实例
func (proxy *dbProxy) Init(cli *pg.DB) error {
	if proxy.IsOk() {
		return ErrProxyAlreadyInited
	}
	proxy.SetConnect(cli)
	return nil
}

// InitFromOptions 从配置条件初始化代理对象
func (proxy *dbProxy) InitFromOptions(options *pg.Options) error {
	proxy.Options = options
	db := pg.Connect(options)
	return proxy.Init(db)
}

// InitFromURL 使用配置给代理赋值客户端实例
func (proxy *dbProxy) InitFromURL(address string) error {
	options, err := parseDBURL(address)
	if err != nil {
		return err
	}
	err = proxy.InitFromOptions(options)
	return err
}

func (proxy *dbProxy) Exec(query interface{}, params ...interface{}) (res pg.Result, err error) {
	if !proxy.IsOk() {
		return nil, ErrProxyNotInited
	}
	return proxy.conn.Exec(query, params)
}

// Regist 注册回调函数,在init执行后执行回调函数
func (proxy *dbProxy) Regist(cb dbProxyCallback) {
	proxy.callBacks = append(proxy.callBacks, cb)
}

//Proxy 默认的pg代理对象
var Proxy = New()