core/db/db.go

package db

import (
	"database/sql"
	"fmt"
	"log"
	"time"

	"gitea.meta-tech.academy/go/core/util"

	"github.com/go-sql-driver/mysql"
)

const USER_DEFAULT_PRIVILEGES = "CREATE, CREATE VIEW, CREATE TEMPORARY TABLES, SELECT, INSERT, UPDATE, DELETE, EXECUTE, LOCK TABLES, TRIGGER"

var MAX_OPEN_CONNS = 500
var MAX_IDLE_CONNS = 0

type Db struct {
	pool map[string]*sql.DB
}

type DbConfig struct {
	host   string
	dbname string
	user   string
	passwd string
	port   int
}

type PasswordDecode func(in string) (string, bool)

func (dbc *DbConfig) GetHost() string {
	return fmt.Sprintf("%s:%d", dbc.host, dbc.port)
}

func (dbc *DbConfig) GetSqlConfig(passwdGet ...PasswordDecode) *mysql.Config {
	var plain string = dbc.passwd
	var done bool = len(passwdGet) == 0
	if !done {
		plain, done = passwdGet[0](dbc.passwd)
	}
	if done {
		return &mysql.Config{
			User:                 dbc.user,
			Net:                  "tcp",
			Addr:                 dbc.GetHost(),
			DBName:               dbc.dbname,
			Passwd:               plain,
			AllowNativePasswords: true,
		}
	} else {
		fmt.Printf("invalid password unable to connect")
		time.Sleep(3 * time.Second)
	}
	return nil
}

func NewDbConfig(host string, dbname string, user string, passwd string, port int) *DbConfig {
	return &DbConfig{host, dbname, user, passwd, port}
}

func (db *Db) AddDb(dbconf *DbConfig, passwdGet ...PasswordDecode) {
	cnx, err := sql.Open("mysql", dbconf.GetSqlConfig(passwdGet...).FormatDSN())
	if err != nil {
		log.Fatal(err)
	}
	pingErr := cnx.Ping()
	if pingErr != nil {
		log.Fatal(pingErr)
	}
	cnx.SetMaxIdleConns(MAX_IDLE_CONNS)
	cnx.SetMaxOpenConns(MAX_OPEN_CONNS)
	db.pool[dbconf.dbname] = cnx
}

func NewDb() *Db {
	pool := make(map[string]*sql.DB)
	return &Db{pool}
}

func getFullUser(username string, hostname string) string {
	return fmt.Sprintf("'%s'@'%s'", util.EscapeSquote(username), util.EscapeSquote(hostname))
}

func (db *Db) onDbExec(dbname string, instr DbInstruct) bool {
	var done bool
	if sqlDb, ok := db.pool[dbname]; ok {
		return instr(sqlDb)
	} else {
		fmt.Printf("you should add db mysql first\n")
		done = ok
	}
	return done
}

type DbInstruct func(database *sql.DB) bool

func (db *Db) CreateUser(username string, hostname string, passwd string) bool {
	return db.onDbExec("mysql", func(sqlDb *sql.DB) bool {
		query := fmt.Sprintf(
			"CREATE USER IF NOT EXISTS %s IDENTIFIED BY '%s';",
			getFullUser(username, hostname),
			util.EscapeSquote(passwd),
		)
		_, err := sqlDb.Exec(query)
		done := err == nil
		if !done {
			fmt.Printf("unable to add user %s\n%v\n", username, err)
		}
		return done
	})
}

func (db *Db) GrantUserPrivileges(dbname string, username string, hostname string) bool {
	return db.onDbExec("mysql", func(sqlDb *sql.DB) bool {
		query := fmt.Sprintf(
			"GRANT %s ON `%s`.* TO  %s;",
			USER_DEFAULT_PRIVILEGES,
			util.EscapeBt(dbname),
			getFullUser(username, hostname),
		)
		_, err := sqlDb.Exec(query)
		done := err == nil
		if !done {
			fmt.Printf("query : %s\n", query)
			fmt.Printf("unable to grant user privilege %s\n%v\n", username, err)
		}
		return done
	})
}

func (db *Db) CreateDb(dbname string, collate string) bool {
	return db.onDbExec("mysql", func(sqlDb *sql.DB) bool {
		query := fmt.Sprintf(
			"CREATE DATABASE IF NOT EXISTS `%s` CHARACTER SET '%s' COLLATE '%s_unicode_ci'",
			util.EscapeBt(dbname),
			util.EscapeSquote(collate),
			util.EscapeSquote(collate),
		)
		_, err := db.pool["mysql"].Exec(query)
		done := err == nil
		if !done {
			fmt.Printf("unable to create database %s\n%v\n", dbname, err)
		}
		return done
	})
}

func (db *Db) SizeDb(dbname string) float32 {
	type RsSize struct {
		Dbname string
		Size   float32
	}
	rs := RsSize{Size: -1}
	db.onDbExec("information_schema", func(sqlDb *sql.DB) bool {
		query := "SELECT table_schema 'dbname', SUM(data_length + index_length)/1024/1024 'size' FROM TABLES WHERE table_schema = ? GROUP BY table_schema;"
		row := sqlDb.QueryRow(query, dbname)
		err := row.Scan(&rs.Dbname, &rs.Size)
		done := err != nil
		if !done {
			if err == sql.ErrNoRows {
				fmt.Printf("   cannot find db %s or db is empty\n", dbname)
			} else if err != nil {
				fmt.Printf("   what the fuck error : %+v\n", err)
			}
		}
		return done
	})
	return rs.Size
}

type TableSize struct {
	Table string
	Size  float32
}

func (db *Db) SizeTable(dbname string) []TableSize {
	rs := make([]TableSize, 0)
	db.onDbExec("information_schema", func(sqlDb *sql.DB) bool {
		query := "SELECT `TABLE_NAME` AS `table`, ROUND(((`DATA_LENGTH` + `INDEX_LENGTH`) / 1024 / 1024), 2) `size` FROM `TABLES` WHERE `table_schema` = ? AND (`DATA_LENGTH` + `INDEX_LENGTH`) > 0 ORDER BY (`DATA_LENGTH` + `INDEX_LENGTH`) DESC;"
		rows, err := sqlDb.Query(query, dbname)
		if err != nil {
			fmt.Printf("cannot get table size of db %s : %v", dbname, err)
			return false
		}
		defer rows.Close()
		for rows.Next() {
			ts := TableSize{Size: -1}
			if err := rows.Scan(&ts.Table, &ts.Size); err != nil {
				fmt.Printf("cannot get table size of db %s : %v", dbname, err)
				return false
			}
			rs = append(rs, ts)
		}
		return true
	})
	return rs
}