db/DB.go

package db

import (
	"crypto/tls"
	"database/sql"
	"database/sql/driver"
	"encoding/base64"
	"errors"
	"fmt"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"apigo.cc/go/cast"
	"apigo.cc/go/config"
	"apigo.cc/go/crypto"
	"apigo.cc/go/id"
	"apigo.cc/go/log"
	"apigo.cc/go/rand"
	"apigo.cc/go/redis"
	"apigo.cc/go/safe"
)

type Config struct {
	Type          string
	User          string
	Password      string
	pwd           *safe.SafeBuf
	Host          string
	ReadonlyHosts []string
	DB            string
	SSL           string
	tls           *tls.Config
	Args          string
	MaxOpens      int
	MaxIdles      int
	MaxLifeTime   int
	LogSlow       config.Duration
	VersionRedis  string
	logger        *log.Logger
}

type SSL struct {
	Ca       string
	Cert     string
	Key      string
	Insecure bool
}

type ConnectorFunc func(*Config, *safe.SafeBuf, *tls.Config) driver.Connector

var connectors = map[string]ConnectorFunc{}

func RegisterConnector(typ string, conn ConnectorFunc) {
	connectors[typ] = conn
}

var sqlite3PwdMatcher = regexp.MustCompile(`x'\w+'`)

func (dbInfo *Config) Dsn() string {
	args := make([]string, 0)
	if dbInfo.SSL != "" {
		args = append(args, "tls="+dbInfo.SSL)
	}
	if dbInfo.Args != "" {
		args = append(args, dbInfo.Args)
	}
	argsStr := ""
	if len(args) > 0 {
		argsStr = "&" + strings.Join(args, "&")
	}

	if isFileDB(dbInfo.Type) {
		argsStr = sqlite3PwdMatcher.ReplaceAllString(argsStr, "******")
		return fmt.Sprintf("%s://%s?logSlow=%s"+argsStr, dbInfo.Type, dbInfo.Host, dbInfo.LogSlow.TimeDuration())
	} else {
		hosts := []string{dbInfo.Host}
		if dbInfo.ReadonlyHosts != nil {
			hosts = append(hosts, dbInfo.ReadonlyHosts...)
		}
		return fmt.Sprintf("%s://%s:****@%s/%s?logSlow=%s"+argsStr, dbInfo.Type, dbInfo.User, strings.Join(hosts, ","), dbInfo.DB, dbInfo.LogSlow.TimeDuration())
	}
}

var fileDBs = map[string]bool{
	"sqlite":     true,
	"sqlite3":    true,
	"chai":       true,
	"access":     true,
	"mdb":        true,
	"accdb":      true,
	"h2":         true,
	"hsqldb":     true,
	"derby":      true,
	"sqlce":      true,
	"sdf":        true,
	"firebird":   true,
	"fdb":        true,
	"dbase":      true,
	"dbf":        true,
	"berkeleydb": true,
	"bdb":        true,
}

func isFileDB(typ string) bool {
	return fileDBs[typ]
}

func (dbInfo *Config) ConfigureBy(setting string) {
	urlInfo, err := url.Parse(setting)
	if err != nil {
		if dbInfo.logger != nil {
			dbInfo.logger.Error(err.Error(), "url", setting)
		}
		return
	}

	dbInfo.Type = urlInfo.Scheme
	if isFileDB(dbInfo.Type) {
		dbInfo.Host = urlInfo.Host + urlInfo.Path
		dbInfo.DB = strings.SplitN(urlInfo.Host, ".", 2)[0]
	} else {
		if strings.ContainsRune(urlInfo.Host, ',') {
			a := strings.Split(urlInfo.Host, ",")
			dbInfo.Host = a[0]
			dbInfo.ReadonlyHosts = a[1:]
		} else {
			dbInfo.Host = urlInfo.Host
			dbInfo.ReadonlyHosts = nil
		}
		if len(urlInfo.Path) > 1 {
			dbInfo.DB = urlInfo.Path[1:]
		}
	}
	dbInfo.User = urlInfo.User.Username()
	dbInfo.Password, _ = urlInfo.User.Password()

	q := urlInfo.Query()
	dbInfo.MaxIdles = cast.Int(q.Get("maxIdles"))
	dbInfo.MaxLifeTime = cast.Int(q.Get("maxLifeTime"))
	dbInfo.MaxOpens = cast.Int(q.Get("maxOpens"))
	dbInfo.LogSlow = config.Duration(cast.Duration(q.Get("logSlow")))
	dbInfo.VersionRedis = q.Get("versionRedis")
	dbInfo.SSL = q.Get("tls")

	sslCa := q.Get("sslCA")
	sslCert := q.Get("sslCert")
	sslKey := q.Get("sslKey")
	sslSkipVerify := cast.Bool(q.Get("sslSkipVerify"))
	if sslCa == "" || sslCert == "" || sslKey == "" {
		if dbSSL, ok := dbSSLs[dbInfo.SSL]; ok {
			sslCa = dbSSL.Ca
			sslCert = dbSSL.Cert
			sslKey = dbSSL.Key
			sslSkipVerify = dbSSL.Insecure
		}
	}
	if sslCa != "" && sslCert != "" && sslKey != "" {
		sslName := id.MakeID(12)
		dbInfo.SSL = sslName
		decryptedCa, _ := confAes.DecryptBytes([]byte(sslCa))
		decryptedCert, _ := confAes.DecryptBytes([]byte(sslCert))
		decryptedKey, _ := confAes.DecryptBytes([]byte(sslKey))
		tlsConf := BuildTLSConfig(decryptedCa, decryptedCert, decryptedKey, sslSkipVerify)
		if tlsConf != nil {
			dbInfo.tls = tlsConf
		}
		safe.ZeroMemory(decryptedCa)
		safe.ZeroMemory(decryptedCert)
		safe.ZeroMemory(decryptedKey)
	}

	args := make([]string, 0)
	for k := range q {
		if k != "maxIdles" && k != "maxLifeTime" && k != "maxOpens" && k != "logSlow" && k != "tls" && k != "versionRedis" {
			args = append(args, k+"="+q.Get(k))
		}
	}
	if len(args) > 0 {
		dbInfo.Args = strings.Join(args, "&")
	}
}

type DB struct {
	name                string
	conn                *sql.DB
	readonlyConnections []*sql.DB
	Config              *Config
	logger              *dbLogger
	Error               error
	QuoteTag            string
	tables              map[string]*TableStruct
	tablesLock          *sync.RWMutex
}

type TableStruct struct {
	VersionField   string
	HasShadowTable bool
}

var confAes, _ = crypto.NewAESCBCAndEraseKey([]byte("?GQ$0K0GgLdO=f+~L68PLm$uhKr4'=tV"), []byte("VFs7@sK61cj^f?HZ"))
var keysSetted = sync.Once{}

func SetEncryptKeys(key, iv []byte) {
	keysSetted.Do(func() {
		confAes.Close()
		confAes, _ = crypto.NewAESGCMAndEraseKey(key, iv)
	})
}

type dbLogger struct {
	config *Config
	logger *log.Logger
}

func (dl *dbLogger) LogError(errStr string) {
	dl.logger.DBError(errStr, dl.config.Type, dl.config.Dsn(), "", nil, 0)
}

func (dl *dbLogger) LogQuery(query string, args []any, usedTime float32) {
	dl.logger.DB(dl.config.Type, dl.config.Dsn(), query, args, usedTime)
}

func (dl *dbLogger) LogQueryError(errStr string, query string, args []any, usedTime float32) {
	dl.logger.DBError(errStr, dl.config.Type, dl.config.Dsn(), query, args, usedTime)
}

var dbConfigs = make(map[string]*Config)
var dbConfigsLock = sync.RWMutex{}
var dbSSLs = make(map[string]*SSL)
var dbInstances = make(map[string]*DB)
var dbInstancesLock = sync.RWMutex{}
var globalVersionMap = sync.Map{}
var once sync.Once

func (db *DB) NextVersion(key string) int64 {
	if db.Config.VersionRedis != "" {
		r := redis.GetRedis(db.Config.VersionRedis, db.logger.logger)
		if r != nil {
			return r.INCR("db_ver_" + key)
		}
	}

	v, _ := globalVersionMap.LoadOrStore(key, new(int64))
	return atomic.AddInt64(v.(*int64), 1)
}

func GetDBWithoutCache(name string, logger *log.Logger) *DB {
	return getDB(name, logger, false)
}

func GetDB(name string, logger *log.Logger) *DB {
	return getDB(name, logger, true)
}

func getDB(name string, logger *log.Logger, useCache bool) *DB {
	if logger == nil {
		logger = log.DefaultLogger
	}

	if useCache {
		dbInstancesLock.RLock()
		oldConn := dbInstances[name]
		dbInstancesLock.RUnlock()
		if oldConn != nil {
			return oldConn.CopyByLogger(logger)
		}
	}

	var conf *Config
	if strings.Contains(name, "://") {
		conf = new(Config)
		conf.logger = logger
		conf.ConfigureBy(name)
	} else {
		dbConfigsLock.RLock()
		n := len(dbConfigs)
		dbConfigsLock.RUnlock()
		if n == 0 {
			once.Do(func() {
				dbConfigs1 := make(map[string]*Config)
				if err := config.Load("db", &dbConfigs1); err == nil {
					for k, v := range dbConfigs1 {
						if v.Host != "" {
							dbConfigsLock.Lock()
							dbConfigs[k] = v
							dbConfigsLock.Unlock()
						}
					}
				} else {
					logger.Error(err.Error())
				}
				dbConfigs2 := make(map[string]string)
				if err := config.Load("db", &dbConfigs2); err == nil {
					for k, v := range dbConfigs2 {
						if strings.Contains(v, "://") {
							v2 := new(Config)
							v2.ConfigureBy(v)
							if v2.Host != "" {
								v2.logger = logger
								dbConfigsLock.Lock()
								dbConfigs[k] = v2
								dbConfigsLock.Unlock()
							}
						} else {
							dbConfigsLock.Lock()
							v2 := dbConfigs[v]
							if v2 != nil && v2.Host != "" {
								dbConfigs[k] = v2
							}
							dbConfigsLock.Unlock()
						}
					}
				} else {
					logger.Error(err.Error())
				}
			})
		}
		dbConfigsLock.RLock()
		conf = dbConfigs[name]
		dbConfigsLock.RUnlock()
		if conf == nil {
			conf = new(Config)
			dbConfigsLock.Lock()
			dbConfigs[name] = conf
			dbConfigsLock.Unlock()
		}
	}

	if conf.Host == "" {
		if name != "default" {
			logger.Error("db config not exists", "name", name)
		}
		return nil
	}

	if conf.SSL != "" && len(dbSSLs) == 0 {
		_ = config.Load("dbssl", &dbSSLs)
	}

	if conf.SSL != "" && dbSSLs[conf.SSL] == nil {
		logger.Error("dbssl config lost")
	}

	if strings.ContainsRune(conf.Host, ',') {
		a := strings.Split(conf.Host, ",")
		conf.Host = a[0]
		conf.ReadonlyHosts = a[1:]
	} else {
		conf.ReadonlyHosts = nil
	}

	if conf.Password != "" {
		if encryptedPassword, err := base64.URLEncoding.DecodeString(conf.Password); err == nil {
			if pwdSafeBuf, err := confAes.Decrypt(encryptedPassword); err == nil {
				conf.pwd = pwdSafeBuf
			}
		}
		if conf.pwd == nil {
			conf.pwd = safe.NewSafeBuf([]byte(conf.Password))
		}
	} else {
		if !isFileDB(conf.Type) && conf.Host != "127.0.0.1:3306" && conf.User == "root" {
			logger.Warning("password is empty")
		}
	}
	conf.Password = ""

	conn, err := getPool(conf)
	if err != nil {
		logger.DBError(err.Error(), conf.Type, conf.Dsn(), "", nil, 0)
		return &DB{conn: nil, QuoteTag: "\"", Error: err}
	}

	db := new(DB)
	db.QuoteTag = cast.If(conf.Type == "mysql", "`", "\"")
	db.name = name
	db.conn = conn
	db.tables = make(map[string]*TableStruct)
	db.tablesLock = new(sync.RWMutex)

	if conf.ReadonlyHosts != nil {
		readonlyConnections := make([]*sql.DB, 0)
		for _, host := range conf.ReadonlyHosts {
			conn, err := getPoolForHost(conf, host)
			if err != nil {
				logger.DBError(err.Error(), conf.Type, conf.Dsn(), "", nil, 0)
			} else {
				readonlyConnections = append(readonlyConnections, conn)
			}
		}
		if len(readonlyConnections) > 0 {
			db.readonlyConnections = readonlyConnections
		}
	}

	db.Error = nil
	db.Config = conf
	if conf.MaxIdles > 0 {
		conn.SetMaxIdleConns(conf.MaxIdles)
	}
	if conf.MaxOpens > 0 {
		conn.SetMaxOpenConns(conf.MaxOpens)
	}
	if conf.MaxLifeTime > 0 {
		conn.SetConnMaxLifetime(time.Second * time.Duration(conf.MaxLifeTime))
	}
	if conf.LogSlow == 0 {
		conf.LogSlow = config.Duration(1000 * time.Millisecond)
	}
	if useCache {
		dbInstancesLock.Lock()
		dbInstances[name] = db
		dbInstancesLock.Unlock()
	}
	return db.CopyByLogger(logger)
}

func getPool(conf *Config) (*sql.DB, error) {
	return getPoolForHost(conf, "")
}

func getPoolForHost(conf *Config, host string) (*sql.DB, error) {
	connectType := "tcp"
	if host == "" {
		host = conf.Host
	}
	if len(host) > 0 && host[0] == '/' {
		connectType = "unix"
	}

	if connector := connectors[conf.Type]; connector != nil {
		return sql.OpenDB(connector(conf, conf.pwd, conf.tls)), nil
	} else {
		dsn := ""
		args := make([]string, 0)
		if conf.SSL != "" {
			args = append(args, "tls="+conf.SSL)
		}
		if conf.Args != "" {
			args = append(args, conf.Args)
		}
		argsStr := ""
		if len(args) > 0 {
			argsStr = "?" + strings.Join(args, "&")
		}

		if isFileDB(conf.Type) {
			dsn = host + argsStr
		} else {
			pwdBuf := conf.pwd.Open()
			dsn = fmt.Sprintf("%s:%s@%s(%s)/%s"+argsStr, conf.User, pwdBuf.String(), connectType, host, conf.DB)
			pwdBuf.Close()
		}
		return sql.Open(conf.Type, dsn)
	}
}

func (db *DB) CopyByLogger(logger *log.Logger) *DB {
	newDB := new(DB)
	newDB.QuoteTag = db.QuoteTag
	newDB.name = db.name
	newDB.conn = db.conn
	newDB.readonlyConnections = db.readonlyConnections
	newDB.Config = db.Config
	newDB.tables = db.tables
	newDB.tablesLock = db.tablesLock
	if logger == nil {
		logger = log.DefaultLogger
	}
	newDB.logger = &dbLogger{logger: logger, config: db.Config}
	return newDB
}

func (db *DB) SetLogger(logger *log.Logger) {
	db.logger.logger = logger
}

func (db *DB) GetLogger() *log.Logger {
	return db.logger.logger
}

func (db *DB) Destroy() error {
	if db.conn == nil {
		return errors.New("operate on a bad connection")
	}
	err := db.conn.Close()
	if err != nil {
		db.logger.LogError(err.Error())
	}
	dbInstancesLock.Lock()
	delete(dbInstances, db.name)
	dbInstancesLock.Unlock()
	return err
}

func (db *DB) GetOriginDB() *sql.DB {
	return db.conn
}

func (db *DB) Prepare(query string) *Stmt {
	stmt := basePrepare(db.conn, nil, query)
	stmt.logger = db.logger
	if stmt.Error != nil {
		db.logger.LogError(stmt.Error.Error())
	}
	return stmt
}

func (db *DB) Quote(text string) string {
	return quote(db.QuoteTag, text)
}

func (db *DB) Quotes(texts []string) string {
	return quotes(db.QuoteTag, texts)
}

func (db *DB) Begin() *Tx {
	if db.conn == nil {
		return &Tx{db: db, QuoteTag: db.QuoteTag, logSlow: db.Config.LogSlow.TimeDuration(), Error: errors.New("operate on a bad connection"), logger: db.logger}
	}
	sqlTx, err := db.conn.Begin()
	if err != nil {
		db.logger.LogError(err.Error())
		return &Tx{db: db, QuoteTag: db.QuoteTag, logSlow: db.Config.LogSlow.TimeDuration(), Error: err, logger: db.logger}
	}
	return &Tx{db: db, QuoteTag: db.QuoteTag, logSlow: db.Config.LogSlow.TimeDuration(), conn: sqlTx, logger: db.logger}
}

func (db *DB) Exec(query string, args ...any) *ExecResult {
	r := baseExec(db.conn, nil, query, args...)
	r.logger = db.logger
	if r.Error != nil {
		db.logger.LogQueryError(r.Error.Error(), query, args, r.usedTime)
	} else {
		if db.Config.LogSlow > 0 && r.usedTime >= float32(db.Config.LogSlow.TimeDuration()/time.Millisecond) {
			db.logger.LogQuery(query, args, r.usedTime)
		}
	}
	return r
}

func (db *DB) Query(query string, args ...any) *QueryResult {
	conn := db.conn
	if db.readonlyConnections != nil {
		connNum := len(db.readonlyConnections)
		if connNum == 1 {
			conn = db.readonlyConnections[0]
		} else {
			p := rand.Int(0, connNum-1)
			conn = db.readonlyConnections[p]
		}
	}

	r := baseQuery(conn, nil, query, args...)
	r.logger = db.logger
	if r.Error != nil {
		db.logger.LogQueryError(r.Error.Error(), query, args, r.usedTime)
	} else {
		if db.Config.LogSlow > 0 && r.usedTime >= float32(db.Config.LogSlow.TimeDuration()/time.Millisecond) {
			db.logger.LogQuery(query, args, r.usedTime)
		}
	}
	return r
}

func (db *DB) Insert(table string, data any) *ExecResult {
	query, values := db.MakeInsertSql(table, data, false)
	r := baseExec(db.conn, nil, query, values...)
	r.logger = db.logger
	if r.Error != nil {
		db.logger.LogQueryError(r.Error.Error(), query, values, r.usedTime)
	} else {
		if db.Config.LogSlow > 0 && r.usedTime >= float32(db.Config.LogSlow.TimeDuration()/time.Millisecond) {
			db.logger.LogQuery(query, values, r.usedTime)
		}
	}
	return r
}

func (db *DB) Replace(table string, data any) *ExecResult {
	query, values := db.MakeInsertSql(table, data, true)
	r := baseExec(db.conn, nil, query, values...)
	r.logger = db.logger
	if r.Error != nil {
		db.logger.LogQueryError(r.Error.Error(), query, values, r.usedTime)
	} else {
		if db.Config.LogSlow > 0 && r.usedTime >= float32(db.Config.LogSlow.TimeDuration()/time.Millisecond) {
			db.logger.LogQuery(query, values, r.usedTime)
		}
	}
	return r
}

func (db *DB) Update(table string, data any, conditions string, args ...any) *ExecResult {
	query, values := db.MakeUpdateSql(table, data, conditions, args...)
	r := baseExec(db.conn, nil, query, values...)
	r.logger = db.logger
	if r.Error != nil {
		db.logger.LogQueryError(r.Error.Error(), query, values, r.usedTime)
	} else {
		if db.Config.LogSlow > 0 && r.usedTime >= float32(db.Config.LogSlow.TimeDuration()/time.Millisecond) {
			db.logger.LogQuery(query, values, r.usedTime)
		}
	}
	return r
}

func (db *DB) Delete(table string, conditions string, args ...any) *ExecResult {
	ts := db.getTable(table)
	if !ts.HasShadowTable {
		if conditions != "" {
			conditions = " where " + conditions
		}
		query := fmt.Sprintf("delete from %s%s", db.Quote(table), conditions)
		r := baseExec(db.conn, nil, query, args...)
		r.logger = db.logger
		if r.Error != nil {
			db.logger.LogQueryError(r.Error.Error(), query, args, r.usedTime)
		} else {
			if db.Config.LogSlow > 0 && r.usedTime >= float32(db.Config.LogSlow.TimeDuration()/time.Millisecond) {
				db.logger.LogQuery(query, args, r.usedTime)
			}
		}
		return r
	}

	// Shadow delete
	tx := db.Begin()
	defer tx.CheckFinished()
	r := tx.Delete(table, conditions, args...)
	if r.Error == nil {
		tx.Commit()
	}
	return r
}

func (db *DB) getTable(table string) *TableStruct {
	db.tablesLock.RLock()
	ts, ok := db.tables[table]
	db.tablesLock.RUnlock()
	if ok {
		return ts
	}

	db.tablesLock.Lock()
	defer db.tablesLock.Unlock()

	// Double check
	if ts, ok = db.tables[table]; ok {
		return ts
	}

	ts = &TableStruct{}
	// Probe autoVersion
	var query string
	if db.Config.Type == "mysql" {
		query = "SELECT COLUMN_NAME FROM information_schema.COLUMNS WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? AND COLUMN_NAME = 'autoVersion' AND DATA_TYPE = 'bigint' AND COLUMN_TYPE LIKE '%unsigned%'"
		res := db.Query(query, db.Config.DB, table)
		if res.StringOnR1C1() != "" {
			ts.VersionField = "autoVersion"
		}
	} else if isFileDB(db.Config.Type) {
		// For SQLite
		query = fmt.Sprintf("PRAGMA table_info(%s)", db.Quote(table))
		res := db.Query(query)
		rows := res.MapResults()
		for _, row := range rows {
			if cast.String(row["name"]) == "autoVersion" {
				ts.VersionField = "autoVersion"
				break
			}
		}
	}

	// Probe shadow table
	shadowTable := table + "_deleted"
	if db.Config.Type == "mysql" {
		query = "SELECT TABLE_NAME FROM information_schema.TABLES WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ?"
		res := db.Query(query, db.Config.DB, shadowTable)
		if res.StringOnR1C1() != "" {
			ts.HasShadowTable = true
		}
	} else if isFileDB(db.Config.Type) {
		query = "SELECT name FROM sqlite_master WHERE type='table' AND name=?"
		res := db.Query(query, shadowTable)
		if res.StringOnR1C1() != "" {
			ts.HasShadowTable = true
		}
	}

	db.tables[table] = ts
	return ts
}

func (db *DB) InKeys(numArgs int) string {
	return InKeys(numArgs)
}

func InKeys(numArgs int) string {
	a := make([]string, numArgs)
	for i := 0; i < numArgs; i++ {
		a[i] = "?"
	}
	return fmt.Sprintf("(%s)", strings.Join(a, ","))
}