Initial commit

This commit is contained in:
Donny
2019-04-22 20:46:32 +08:00
commit 49ab8aadd1
25441 changed files with 4055000 additions and 0 deletions

View File

@@ -0,0 +1,59 @@
load(
"@io_bazel_rules_go//go:def.bzl",
"go_library",
"go_test",
)
package(default_visibility = ["//visibility:public"])
filegroup(
name = "package-srcs",
srcs = glob(["**"]),
tags = ["automanaged"],
visibility = ["//visibility:private"],
)
filegroup(
name = "all-srcs",
srcs = [":package-srcs"],
tags = ["automanaged"],
)
go_library(
name = "go_default_library",
srcs = [
"discovery.go",
"node_proc.go",
"sql.go",
"tidb.go",
],
importpath = "go-common/library/database/tidb",
tags = ["automanaged"],
deps = [
"//library/conf/env:go_default_library",
"//library/log:go_default_library",
"//library/naming:go_default_library",
"//library/naming/discovery:go_default_library",
"//library/net/netutil/breaker:go_default_library",
"//library/net/trace:go_default_library",
"//library/stat:go_default_library",
"//library/time:go_default_library",
"//vendor/github.com/go-sql-driver/mysql:go_default_library",
"//vendor/github.com/pkg/errors:go_default_library",
],
)
go_test(
name = "go_default_test",
srcs = [
"sql_test.go",
"tidb_test.go",
],
embed = [":go_default_library"],
rundir = ".",
tags = ["automanaged"],
deps = [
"//library/net/netutil/breaker:go_default_library",
"//library/time:go_default_library",
],
)

View File

@@ -0,0 +1,4 @@
### database/tidb
##### Version 1.0.0
> 1. 接入discovery

View File

@@ -0,0 +1,7 @@
# Author
wangxu01
# Reviewer
maojian
haoguanwei
zhapuyu

View File

@@ -0,0 +1,14 @@
# See the OWNERS docs at https://go.k8s.io/owners
approvers:
- maojian
- wangxu01
labels:
- library
- library/database/tidb
reviewers:
- chenzhihui
- haoguanwei
- maojian
- wangxu01
- zhapuyu

View File

@@ -0,0 +1,17 @@
#### database/tidb
##### 项目简介
TiDB数据库驱动 对mysql驱动进行封装
##### 功能
1. 支持discovery服务发现 多节点直连
2. 支持通过lvs单一地址连接
3. 支持prepare绑定多个节点
4. 支持动态增减节点负载均衡
5. 日志区分运行节点
##### 编译环境
> 请只用golang v1.8.x以上版本编译执行。
##### 依赖包
> 1.[Go-MySQL-Driver](https://github.com/go-sql-driver/mysql)

View File

@@ -0,0 +1,56 @@
package tidb
import (
"context"
"fmt"
"strings"
"time"
"go-common/library/conf/env"
"go-common/library/log"
"go-common/library/naming"
"go-common/library/naming/discovery"
)
var _schema = "tidb://"
func (db *DB) nodeList() (nodes []string) {
var (
insMap map[string][]*naming.Instance
ins []*naming.Instance
ok bool
)
if insMap, ok = db.dis.Fetch(context.Background()); !ok {
return
}
if ins, ok = insMap[env.Zone]; !ok || len(ins) == 0 {
return
}
for _, in := range ins {
for _, addr := range in.Addrs {
if strings.HasPrefix(addr, _schema) {
addr = strings.Replace(addr, _schema, "", -1)
nodes = append(nodes, addr)
}
}
}
log.Info("tidb get %s instances(%v)", db.appid, nodes)
return
}
func (db *DB) disc() (nodes []string) {
db.dis = discovery.Build(db.appid)
e := db.dis.Watch()
select {
case <-e:
nodes = db.nodeList()
case <-time.After(10 * time.Second):
panic("tidb init discovery err")
}
if len(nodes) == 0 {
panic(fmt.Sprintf("tidb %s no instance", db.appid))
}
go db.nodeproc(e)
log.Info("init tidb discvoery info successfully")
return
}

View File

@@ -0,0 +1,82 @@
package tidb
import (
"time"
"go-common/library/log"
)
func (db *DB) nodeproc(e <-chan struct{}) {
if db.dis == nil {
return
}
for {
<-e
nodes := db.nodeList()
if len(nodes) == 0 {
continue
}
cm := make(map[string]*conn)
var conns []*conn
for _, conn := range db.conns {
cm[conn.addr] = conn
}
for _, node := range nodes {
if cm[node] != nil {
conns = append(conns, cm[node])
continue
}
c, err := db.connectDSN(genDSN(db.conf.DSN, node))
if err == nil {
conns = append(conns, c)
} else {
log.Error("tidb: connect addr: %s err: %+v", node, err)
}
}
if len(conns) == 0 {
log.Error("tidb: no nodes ignore event")
continue
}
oldConns := db.conns
db.mutex.Lock()
db.conns = conns
db.mutex.Unlock()
log.Info("tidb: new nodes: %v", nodes)
var removedConn []*conn
for _, conn := range oldConns {
var exist bool
for _, c := range conns {
if c.addr == conn.addr {
exist = true
break
}
}
if !exist {
removedConn = append(removedConn, conn)
}
}
go db.closeConns(removedConn)
}
}
func (db *DB) closeConns(conns []*conn) {
if len(conns) == 0 {
return
}
du := db.conf.QueryTimeout
if db.conf.ExecTimeout > du {
du = db.conf.ExecTimeout
}
if db.conf.TranTimeout > du {
du = db.conf.TranTimeout
}
time.Sleep(time.Duration(du))
for _, conn := range conns {
err := conn.Close()
if err != nil {
log.Error("tidb: close removed conn: %s err: %v", conn.addr, err)
} else {
log.Info("tidb: close removed conn: %s", conn.addr)
}
}
}

View File

@@ -0,0 +1,715 @@
package tidb
import (
"context"
"database/sql"
"fmt"
"sync"
"sync/atomic"
"time"
"go-common/library/naming"
"go-common/library/net/netutil/breaker"
"go-common/library/net/trace"
"github.com/go-sql-driver/mysql"
"github.com/pkg/errors"
)
const (
_family = "tidb_client"
)
var (
// ErrStmtNil prepared stmt error
ErrStmtNil = errors.New("sql: prepare failed and stmt nil")
// ErrNoRows is returned by Scan when QueryRow doesn't return a row.
// In such a case, QueryRow returns a placeholder *Row value that defers
// this error until a Scan.
ErrNoRows = sql.ErrNoRows
// ErrTxDone transaction done.
ErrTxDone = sql.ErrTxDone
)
// DB database.
type DB struct {
conf *Config
conns []*conn
idx int64
dis naming.Resolver
appid string
mutex sync.RWMutex
breakerGroup *breaker.Group
}
// conn database connection
type conn struct {
*sql.DB
breaker breaker.Breaker
conf *Config
addr string
}
// Tx transaction.
type Tx struct {
db *conn
tx *sql.Tx
t trace.Trace
c context.Context
cancel func()
}
// Row row.
type Row struct {
err error
*sql.Row
db *conn
query string
args []interface{}
t trace.Trace
cancel func()
}
// Scan copies the columns from the matched row into the values pointed at by dest.
func (r *Row) Scan(dest ...interface{}) (err error) {
if r.t != nil {
defer r.t.Finish(&err)
}
if r.err != nil {
err = r.err
} else if r.Row == nil {
err = ErrStmtNil
}
if err != nil {
return
}
err = r.Row.Scan(dest...)
if r.cancel != nil {
r.cancel()
}
r.db.onBreaker(&err)
if err != ErrNoRows {
err = errors.Wrapf(err, "addr: %s, query %s args %+v", r.db.addr, r.query, r.args)
}
return
}
// Rows rows.
type Rows struct {
*sql.Rows
cancel func()
}
// Close closes the Rows, preventing further enumeration. If Next is called
// and returns false and there are no further result sets,
// the Rows are closed automatically and it will suffice to check the
// result of Err. Close is idempotent and does not affect the result of Err.
func (rs *Rows) Close() (err error) {
err = errors.WithStack(rs.Rows.Close())
if rs.cancel != nil {
rs.cancel()
}
return
}
// Stmt prepared stmt.
type Stmt struct {
db *conn
tx bool
query string
stmt atomic.Value
t trace.Trace
}
// Stmts random prepared stmt.
type Stmts struct {
query string
sts map[string]*Stmt
mu sync.RWMutex
db *DB
}
// Open opens a database specified by its database driver name and a
// driver-specific data source name, usually consisting of at least a database
// name and connection information.
func Open(c *Config) (db *DB, err error) {
db = &DB{conf: c, breakerGroup: breaker.NewGroup(c.Breaker)}
cfg, err := mysql.ParseDSN(c.DSN)
if err != nil {
return
}
var dsns []string
if cfg.Net == "discovery" {
db.appid = cfg.Addr
for _, addr := range db.disc() {
dsns = append(dsns, genDSN(c.DSN, addr))
}
} else {
dsns = append(dsns, c.DSN)
}
cs := make([]*conn, 0, len(dsns))
for _, dsn := range dsns {
r, err := db.connectDSN(dsn)
if err != nil {
return db, err
}
cs = append(cs, r)
}
db.conns = cs
return
}
func (db *DB) connectDSN(dsn string) (c *conn, err error) {
d, err := connect(db.conf, dsn)
if err != nil {
return
}
addr := parseDSNAddr(dsn)
brk := db.breakerGroup.Get(addr)
c = &conn{DB: d, breaker: brk, conf: db.conf, addr: addr}
return
}
func connect(c *Config, dataSourceName string) (*sql.DB, error) {
d, err := sql.Open("mysql", dataSourceName)
if err != nil {
err = errors.WithStack(err)
return nil, err
}
d.SetMaxOpenConns(c.Active)
d.SetMaxIdleConns(c.Idle)
d.SetConnMaxLifetime(time.Duration(c.IdleTimeout))
return d, nil
}
func (db *DB) conn() (c *conn) {
db.mutex.RLock()
c = db.conns[db.index()]
db.mutex.RUnlock()
return
}
// Begin starts a transaction. The isolation level is dependent on the driver.
func (db *DB) Begin(c context.Context) (tx *Tx, err error) {
return db.conn().begin(c)
}
// Exec executes a query without returning any rows.
// The args are for any placeholder parameters in the query.
func (db *DB) Exec(c context.Context, query string, args ...interface{}) (res sql.Result, err error) {
return db.conn().exec(c, query, args...)
}
// Prepare creates a prepared statement for later queries or executions.
// Multiple queries or executions may be run concurrently from the returned
// statement. The caller must call the statement's Close method when the
// statement is no longer needed.
func (db *DB) Prepare(query string) (*Stmt, error) {
return db.conn().prepare(query)
}
// Prepared creates a prepared statement for later queries or executions.
// Multiple queries or executions may be run concurrently from the returned
// statement. The caller must call the statement's Close method when the
// statement is no longer needed.
func (db *DB) Prepared(query string) (s *Stmts) {
s = &Stmts{query: query, sts: make(map[string]*Stmt), db: db}
for _, c := range db.conns {
st := c.prepared(query)
s.mu.Lock()
s.sts[c.addr] = st
s.mu.Unlock()
}
return
}
// Query executes a query that returns rows, typically a SELECT. The args are
// for any placeholder parameters in the query.
func (db *DB) Query(c context.Context, query string, args ...interface{}) (rows *Rows, err error) {
return db.conn().query(c, query, args...)
}
// QueryRow executes a query that is expected to return at most one row.
// QueryRow always returns a non-nil value. Errors are deferred until Row's
// Scan method is called.
func (db *DB) QueryRow(c context.Context, query string, args ...interface{}) *Row {
return db.conn().queryRow(c, query, args...)
}
func (db *DB) index() int {
if len(db.conns) == 1 {
return 0
}
v := atomic.AddInt64(&db.idx, 1)
return int(v) % len(db.conns)
}
// Close closes the databases, releasing any open resources.
func (db *DB) Close() (err error) {
db.mutex.RLock()
defer db.mutex.RUnlock()
for _, d := range db.conns {
if e := d.Close(); e != nil {
err = errors.WithStack(e)
}
}
return
}
// Ping verifies a connection to the database is still alive, establishing a
// connection if necessary.
func (db *DB) Ping(c context.Context) (err error) {
if err = db.conn().ping(c); err != nil {
return
}
return
}
func (db *conn) onBreaker(err *error) {
if err != nil && *err != nil && *err != sql.ErrNoRows && *err != sql.ErrTxDone {
db.breaker.MarkFailed()
} else {
db.breaker.MarkSuccess()
}
}
func (db *conn) begin(c context.Context) (tx *Tx, err error) {
now := time.Now()
t, ok := trace.FromContext(c)
if ok {
t = t.Fork(_family, "begin")
t.SetTag(trace.String(trace.TagAddress, db.addr), trace.String(trace.TagComment, ""))
defer func() {
if err != nil {
t.Finish(&err)
}
}()
}
if err = db.breaker.Allow(); err != nil {
stats.Incr("tidb:begin", "breaker")
return
}
_, c, cancel := db.conf.TranTimeout.Shrink(c)
rtx, err := db.BeginTx(c, nil)
stats.Timing("tidb:begin", int64(time.Since(now)/time.Millisecond))
if err != nil {
err = errors.WithStack(err)
cancel()
return
}
tx = &Tx{tx: rtx, t: t, db: db, c: c, cancel: cancel}
return
}
func (db *conn) exec(c context.Context, query string, args ...interface{}) (res sql.Result, err error) {
now := time.Now()
if t, ok := trace.FromContext(c); ok {
t = t.Fork(_family, "exec")
t.SetTag(trace.String(trace.TagAddress, db.addr), trace.String(trace.TagComment, query))
defer t.Finish(&err)
}
if err = db.breaker.Allow(); err != nil {
stats.Incr("tidb:exec", "breaker")
return
}
_, c, cancel := db.conf.ExecTimeout.Shrink(c)
res, err = db.ExecContext(c, query, args...)
cancel()
db.onBreaker(&err)
stats.Timing("tidb:exec", int64(time.Since(now)/time.Millisecond))
if err != nil {
err = errors.Wrapf(err, "addr: %s exec:%s, args:%+v", db.addr, query, args)
}
return
}
func (db *conn) ping(c context.Context) (err error) {
now := time.Now()
if t, ok := trace.FromContext(c); ok {
t = t.Fork(_family, "ping")
t.SetTag(trace.String(trace.TagAddress, db.addr), trace.String(trace.TagComment, ""))
defer t.Finish(&err)
}
if err = db.breaker.Allow(); err != nil {
stats.Incr("tidb:ping", "breaker")
return
}
_, c, cancel := db.conf.ExecTimeout.Shrink(c)
err = db.PingContext(c)
cancel()
db.onBreaker(&err)
stats.Timing("tidb:ping", int64(time.Since(now)/time.Millisecond))
if err != nil {
err = errors.WithStack(err)
}
return
}
func (db *conn) prepare(query string) (*Stmt, error) {
stmt, err := db.Prepare(query)
if err != nil {
err = errors.Wrapf(err, "addr: %s prepare %s", db.addr, query)
return nil, err
}
st := &Stmt{query: query, db: db}
st.stmt.Store(stmt)
return st, nil
}
func (db *conn) prepared(query string) (stmt *Stmt) {
stmt = &Stmt{query: query, db: db}
s, err := db.Prepare(query)
if err == nil {
stmt.stmt.Store(s)
return
}
return
}
func (db *conn) query(c context.Context, query string, args ...interface{}) (rows *Rows, err error) {
now := time.Now()
if t, ok := trace.FromContext(c); ok {
t = t.Fork(_family, "query")
t.SetTag(trace.String(trace.TagAddress, db.addr), trace.String(trace.TagComment, query))
defer t.Finish(&err)
}
if err = db.breaker.Allow(); err != nil {
stats.Incr("tidb:query", "breaker")
return
}
_, c, cancel := db.conf.QueryTimeout.Shrink(c)
rs, err := db.DB.QueryContext(c, query, args...)
db.onBreaker(&err)
stats.Timing("tidb:query", int64(time.Since(now)/time.Millisecond))
if err != nil {
err = errors.Wrapf(err, "addr: %s, query:%s, args:%+v", db.addr, query, args)
cancel()
return
}
rows = &Rows{Rows: rs, cancel: cancel}
return
}
func (db *conn) queryRow(c context.Context, query string, args ...interface{}) *Row {
now := time.Now()
t, ok := trace.FromContext(c)
if ok {
t = t.Fork(_family, "queryrow")
t.SetTag(trace.String(trace.TagAddress, db.addr), trace.String(trace.TagComment, query))
}
if err := db.breaker.Allow(); err != nil {
stats.Incr("tidb:queryrow", "breaker")
return &Row{db: db, t: t, err: err}
}
_, c, cancel := db.conf.QueryTimeout.Shrink(c)
r := db.DB.QueryRowContext(c, query, args...)
stats.Timing("tidb:queryrow", int64(time.Since(now)/time.Millisecond))
return &Row{db: db, Row: r, query: query, args: args, t: t, cancel: cancel}
}
// Close closes the statement.
func (s *Stmt) Close() (err error) {
stmt, ok := s.stmt.Load().(*sql.Stmt)
if ok {
err = errors.WithStack(stmt.Close())
}
return
}
func (s *Stmt) prepare() (st *sql.Stmt) {
var ok bool
if st, ok = s.stmt.Load().(*sql.Stmt); ok {
return
}
var err error
if st, err = s.db.Prepare(s.query); err == nil {
s.stmt.Store(st)
}
return
}
// Exec executes a prepared statement with the given arguments and returns a
// Result summarizing the effect of the statement.
func (s *Stmt) Exec(c context.Context, args ...interface{}) (res sql.Result, err error) {
now := time.Now()
if s.tx {
if s.t != nil {
s.t.SetTag(trace.String(trace.TagAnnotation, s.query))
}
} else if t, ok := trace.FromContext(c); ok {
t = t.Fork(_family, "exec")
t.SetTag(trace.String(trace.TagAddress, s.db.addr), trace.String(trace.TagComment, s.query))
defer t.Finish(&err)
}
if err = s.db.breaker.Allow(); err != nil {
stats.Incr("tidb:stmt:exec", "breaker")
return
}
stmt := s.prepare()
if stmt == nil {
err = ErrStmtNil
return
}
_, c, cancel := s.db.conf.ExecTimeout.Shrink(c)
res, err = stmt.ExecContext(c, args...)
cancel()
s.db.onBreaker(&err)
stats.Timing("tidb:stmt:exec", int64(time.Since(now)/time.Millisecond))
if err != nil {
err = errors.Wrapf(err, "addr: %s exec:%s, args:%+v", s.db.addr, s.query, args)
}
return
}
// Query executes a prepared query statement with the given arguments and
// returns the query results as a *Rows.
func (s *Stmt) Query(c context.Context, args ...interface{}) (rows *Rows, err error) {
if s.tx {
if s.t != nil {
s.t.SetTag(trace.String(trace.TagAnnotation, s.query))
}
} else if t, ok := trace.FromContext(c); ok {
t = t.Fork(_family, "query")
t.SetTag(trace.String(trace.TagAddress, s.db.addr), trace.String(trace.TagComment, s.query))
defer t.Finish(&err)
}
if err = s.db.breaker.Allow(); err != nil {
stats.Incr("tidb:stmt:query", "breaker")
return
}
stmt := s.prepare()
if stmt == nil {
err = ErrStmtNil
return
}
now := time.Now()
_, c, cancel := s.db.conf.QueryTimeout.Shrink(c)
rs, err := stmt.QueryContext(c, args...)
s.db.onBreaker(&err)
stats.Timing("tidb:stmt:query", int64(time.Since(now)/time.Millisecond))
if err != nil {
err = errors.Wrapf(err, "addr: %s, query:%s, args:%+v", s.db.addr, s.query, args)
cancel()
return
}
rows = &Rows{Rows: rs, cancel: cancel}
return
}
// QueryRow executes a prepared query statement with the given arguments.
// If an error occurs during the execution of the statement, that error will
// be returned by a call to Scan on the returned *Row, which is always non-nil.
// If the query selects no rows, the *Row's Scan will return ErrNoRows.
// Otherwise, the *Row's Scan scans the first selected row and discards the rest.
func (s *Stmt) QueryRow(c context.Context, args ...interface{}) (row *Row) {
now := time.Now()
row = &Row{db: s.db, query: s.query, args: args}
if s.tx {
if s.t != nil {
s.t.SetTag(trace.String(trace.TagAnnotation, s.query))
}
} else if t, ok := trace.FromContext(c); ok {
t = t.Fork(_family, "queryrow")
t.SetTag(trace.String(trace.TagAddress, s.db.addr), trace.String(trace.TagComment, s.query))
row.t = t
}
if row.err = s.db.breaker.Allow(); row.err != nil {
stats.Incr("tidb:stmt:queryrow", "breaker")
return
}
stmt := s.prepare()
if stmt == nil {
return
}
_, c, cancel := s.db.conf.QueryTimeout.Shrink(c)
row.Row = stmt.QueryRowContext(c, args...)
row.cancel = cancel
stats.Timing("tidb:stmt:queryrow", int64(time.Since(now)/time.Millisecond))
return
}
func (s *Stmts) prepare(conn *conn) (st *Stmt) {
if conn == nil {
conn = s.db.conn()
}
s.mu.RLock()
st = s.sts[conn.addr]
s.mu.RUnlock()
if st == nil {
st = conn.prepared(s.query)
s.mu.Lock()
s.sts[conn.addr] = st
s.mu.Unlock()
}
return
}
// Exec executes a prepared statement with the given arguments and returns a
// Result summarizing the effect of the statement.
func (s *Stmts) Exec(c context.Context, args ...interface{}) (res sql.Result, err error) {
return s.prepare(nil).Exec(c, args...)
}
// Query executes a prepared query statement with the given arguments and
// returns the query results as a *Rows.
func (s *Stmts) Query(c context.Context, args ...interface{}) (rows *Rows, err error) {
return s.prepare(nil).Query(c, args...)
}
// QueryRow executes a prepared query statement with the given arguments.
// If an error occurs during the execution of the statement, that error will
// be returned by a call to Scan on the returned *Row, which is always non-nil.
// If the query selects no rows, the *Row's Scan will return ErrNoRows.
// Otherwise, the *Row's Scan scans the first selected row and discards the rest.
func (s *Stmts) QueryRow(c context.Context, args ...interface{}) (row *Row) {
return s.prepare(nil).QueryRow(c, args...)
}
// Close closes the statement.
func (s *Stmts) Close() (err error) {
for _, st := range s.sts {
if err = errors.WithStack(st.Close()); err != nil {
return
}
}
return
}
// Commit commits the transaction.
func (tx *Tx) Commit() (err error) {
err = tx.tx.Commit()
tx.cancel()
tx.db.onBreaker(&err)
if tx.t != nil {
tx.t.Finish(&err)
}
if err != nil {
err = errors.WithStack(err)
}
return
}
// Rollback aborts the transaction.
func (tx *Tx) Rollback() (err error) {
err = tx.tx.Rollback()
tx.cancel()
tx.db.onBreaker(&err)
if tx.t != nil {
tx.t.Finish(&err)
}
if err != nil {
err = errors.WithStack(err)
}
return
}
// Exec executes a query that doesn't return rows. For example: an INSERT and
// UPDATE.
func (tx *Tx) Exec(query string, args ...interface{}) (res sql.Result, err error) {
now := time.Now()
if tx.t != nil {
tx.t.SetTag(trace.String(trace.TagAnnotation, fmt.Sprintf("exec %s", query)))
}
res, err = tx.tx.ExecContext(tx.c, query, args...)
stats.Timing("tidb:tx:exec", int64(time.Since(now)/time.Millisecond))
if err != nil {
err = errors.Wrapf(err, "addr: %s exec:%s, args:%+v", tx.db.addr, query, args)
}
return
}
// Query executes a query that returns rows, typically a SELECT.
func (tx *Tx) Query(query string, args ...interface{}) (rows *Rows, err error) {
if tx.t != nil {
tx.t.SetTag(trace.String(trace.TagAnnotation, fmt.Sprintf("query %s", query)))
}
now := time.Now()
defer func() {
stats.Timing("tidb:tx:query", int64(time.Since(now)/time.Millisecond))
}()
rs, err := tx.tx.QueryContext(tx.c, query, args...)
if err == nil {
rows = &Rows{Rows: rs}
} else {
err = errors.Wrapf(err, "addr: %s, query:%s, args:%+v", tx.db.addr, query, args)
}
return
}
// QueryRow executes a query that is expected to return at most one row.
// QueryRow always returns a non-nil value. Errors are deferred until Row's
// Scan method is called.
func (tx *Tx) QueryRow(query string, args ...interface{}) *Row {
if tx.t != nil {
tx.t.SetTag(trace.String(trace.TagAnnotation, fmt.Sprintf("queryrow %s", query)))
}
now := time.Now()
defer func() {
stats.Timing("tidb:tx:queryrow", int64(time.Since(now)/time.Millisecond))
}()
r := tx.tx.QueryRowContext(tx.c, query, args...)
return &Row{Row: r, db: tx.db, query: query, args: args}
}
// Stmt returns a transaction-specific prepared statement from an existing statement.
func (tx *Tx) Stmt(stmt *Stmt) *Stmt {
if stmt == nil {
return nil
}
as, ok := stmt.stmt.Load().(*sql.Stmt)
if !ok {
return nil
}
ts := tx.tx.StmtContext(tx.c, as)
st := &Stmt{query: stmt.query, tx: true, t: tx.t, db: tx.db}
st.stmt.Store(ts)
return st
}
// Stmts returns a transaction-specific prepared statement from an existing statement.
func (tx *Tx) Stmts(stmt *Stmts) *Stmt {
return tx.Stmt(stmt.prepare(tx.db))
}
// Prepare creates a prepared statement for use within a transaction.
// The returned statement operates within the transaction and can no longer be
// used once the transaction has been committed or rolled back.
// To use an existing prepared statement on this transaction, see Tx.Stmt.
func (tx *Tx) Prepare(query string) (*Stmt, error) {
if tx.t != nil {
tx.t.SetTag(trace.String(trace.TagAnnotation, fmt.Sprintf("prepare %s", query)))
}
stmt, err := tx.tx.Prepare(query)
if err != nil {
err = errors.Wrapf(err, "addr: %s prepare %s", tx.db.addr, query)
return nil, err
}
st := &Stmt{query: query, tx: true, t: tx.t, db: tx.db}
st.stmt.Store(stmt)
return st, nil
}
// parseDSNAddr parse dsn name and return addr.
func parseDSNAddr(dsn string) (addr string) {
if dsn == "" {
return
}
cfg, err := mysql.ParseDSN(dsn)
if err != nil {
return
}
addr = cfg.Addr
return
}
func genDSN(dsn, addr string) (res string) {
cfg, err := mysql.ParseDSN(dsn)
if err != nil {
return
}
cfg.Addr = addr
cfg.Net = "tcp"
res = cfg.FormatDSN()
return
}

View File

@@ -0,0 +1,21 @@
package tidb
import "testing"
func TestParseDSNAddrr(t *testing.T) {
dsn := "u:p@tcp(127.0.0.1:3306)/db?timeout=5s&readTimeout=5s&writeTimeout=5s&parseTime=true&loc=Local&charset=utf8,utf8mb4"
addr := parseDSNAddr(dsn)
if addr != "127.0.0.1:3306" {
t.Errorf("expect 127.0.0.1:3306 got: %s", addr)
}
}
func TestGenDSN(t *testing.T) {
dsn := "u:p@tcp(127.0.0.1:3306)/db?loc=Local&parseTime=true&readTimeout=5s&timeout=5s&writeTimeout=5s&charset=utf8mb4"
addr := "127.0.0.2:3308"
expect := "u:p@tcp(127.0.0.2:3308)/db?loc=Local&parseTime=true&readTimeout=5s&timeout=5s&writeTimeout=5s&charset=utf8mb4"
newdsn := genDSN(dsn, addr)
if newdsn != expect {
t.Errorf("expect %s got: %s", expect, newdsn)
}
}

View File

@@ -0,0 +1,38 @@
package tidb
import (
"go-common/library/log"
"go-common/library/net/netutil/breaker"
"go-common/library/stat"
"go-common/library/time"
// database driver
_ "github.com/go-sql-driver/mysql"
)
var stats = stat.DB
// Config mysql config.
type Config struct {
DSN string // dsn
Active int // pool
Idle int // pool
IdleTimeout time.Duration // connect max life time.
QueryTimeout time.Duration // query sql timeout
ExecTimeout time.Duration // execute sql timeout
TranTimeout time.Duration // transaction sql timeout
Breaker *breaker.Config // breaker
}
// NewTiDB new db and retry connection when has error.
func NewTiDB(c *Config) (db *DB) {
if c.QueryTimeout == 0 || c.ExecTimeout == 0 || c.TranTimeout == 0 {
panic("tidb must be set query/execute/transction timeout")
}
db, err := Open(c)
if err != nil {
log.Error("open tidb error(%v)", err)
panic(err)
}
return
}

View File

@@ -0,0 +1,234 @@
package tidb
import (
"context"
"database/sql"
"os"
"testing"
"time"
"go-common/library/net/netutil/breaker"
xtime "go-common/library/time"
)
func TestMySQL(t *testing.T) {
bc := &breaker.Config{
Window: xtime.Duration(10 * time.Second),
Sleep: xtime.Duration(10 * time.Second),
Bucket: 10,
Ratio: 0.5,
Request: 100,
}
dsn := os.Getenv("TEST_MYSQL_DSN")
if dsn == "" {
t.Skipf("TEST_MYSQL_DSN is empty, sql test skipped")
}
dsn = dsn + "?timeout=5s&readTimeout=5s&writeTimeout=5s&parseTime=true&loc=Local&charset=utf8"
c := &Config{
DSN: dsn,
Active: 10,
Idle: 5,
IdleTimeout: xtime.Duration(time.Minute),
QueryTimeout: xtime.Duration(time.Minute),
ExecTimeout: xtime.Duration(time.Minute),
TranTimeout: xtime.Duration(time.Minute),
Breaker: bc,
}
db := NewTiDB(c)
defer db.Close()
testPing(t, db)
testTable(t, db)
testExec(t, db)
testQuery(t, db)
testQueryRow(t, db)
testPrepare(t, db)
testPrepared(t, db)
testTransaction(t, db)
}
func testTransaction(t *testing.T, db *DB) {
var (
tx *Tx
err error
execSQL = "INSERT INTO test(name) VALUES(?)"
selSQL = "SELECT name FROM test WHERE name=?"
txstmt *Stmt
)
if tx, err = db.Begin(context.TODO()); err != nil {
t.Errorf("MySQL: db transaction Begin err(%v)", err)
tx.Rollback()
return
}
t.Log("MySQL: db transaction begin")
if txstmt, err = tx.Prepare(execSQL); err != nil {
t.Errorf("MySQL: tx.Prepare err(%v)", err)
}
if stmt := tx.Stmt(txstmt); stmt == nil {
t.Errorf("MySQL:tx.Stmt err(%v)", err)
}
// exec
if _, err = tx.Exec(execSQL, "tx1"); err != nil {
t.Errorf("MySQL: tx.Exec err(%v)", err)
tx.Rollback()
return
}
t.Logf("MySQL:tx.Exec tx1")
if _, err = tx.Exec(execSQL, "tx1"); err != nil {
t.Errorf("MySQL: tx.Exec err(%v)", err)
tx.Rollback()
return
}
t.Logf("MySQL:tx.Exec tx1")
// query
rows, err := tx.Query(selSQL, "tx2")
if err != nil {
t.Errorf("MySQL:tx.Query err(%v)", err)
tx.Rollback()
return
}
rows.Close()
t.Log("MySQL: tx.Query tx2")
// queryrow
var name string
row := tx.QueryRow(selSQL, "noexist")
if err = row.Scan(&name); err != sql.ErrNoRows {
t.Errorf("MySQL: queryRow name: noexist")
}
if err = tx.Commit(); err != nil {
t.Errorf("MySQL:tx.Commit err(%v)", err)
}
if err = tx.Commit(); err != nil {
t.Logf("MySQL:tx.Commit err(%v)", err)
}
if err = tx.Rollback(); err != nil {
t.Logf("MySQL:tx Rollback err(%v)", err)
}
}
func testPing(t *testing.T, db *DB) {
if err := db.Ping(context.TODO()); err != nil {
t.Errorf("MySQL: ping error(%v)", err)
t.FailNow()
} else {
t.Log("MySQL: ping ok")
}
}
func testTable(t *testing.T, db *DB) {
table := "CREATE TABLE IF NOT EXISTS `test` (`id` int(11) NOT NULL AUTO_INCREMENT COMMENT '自增ID', `name` varchar(16) NOT NULL DEFAULT '' COMMENT '名称', PRIMARY KEY (`id`)) ENGINE=InnoDB DEFAULT CHARSET=utf8"
if _, err := db.Exec(context.TODO(), table); err != nil {
t.Errorf("MySQL: create table error(%v)", err)
} else {
t.Log("MySQL: create table ok")
}
}
func testExec(t *testing.T, db *DB) {
sql := "INSERT INTO test(name) VALUES(?)"
if _, err := db.Exec(context.TODO(), sql, "test"); err != nil {
t.Errorf("MySQL: insert error(%v)", err)
} else {
t.Log("MySQL: insert ok")
}
}
func testQuery(t *testing.T, db *DB) {
sql := "SELECT name FROM test WHERE name=?"
rows, err := db.Query(context.TODO(), sql, "test")
if err != nil {
t.Errorf("MySQL: query error(%v)", err)
}
defer rows.Close()
for rows.Next() {
name := ""
if err := rows.Scan(&name); err != nil {
t.Errorf("MySQL: query scan error(%v)", err)
} else {
t.Logf("MySQL: query name: %s", name)
}
}
}
func testQueryRow(t *testing.T, db *DB) {
sql := "SELECT name FROM test WHERE name=?"
name := ""
row := db.QueryRow(context.TODO(), sql, "test")
if err := row.Scan(&name); err != nil {
t.Errorf("MySQL: queryRow error(%v)", err)
} else {
t.Logf("MySQL: queryRow name: %s", name)
}
}
func testPrepared(t *testing.T, db *DB) {
sql := "SELECT name FROM test WHERE name=?"
name := ""
stmt := db.Prepared(sql)
row := stmt.QueryRow(context.TODO(), "test")
if err := row.Scan(&name); err != nil {
t.Errorf("MySQL: prepared query error(%v)", err)
} else {
t.Logf("MySQL: prepared query name: %s", name)
}
if err := stmt.Close(); err != nil {
t.Errorf("MySQL:stmt.Close err(%v)", err)
}
}
func testPrepare(t *testing.T, db *DB) {
var (
selsql = "SELECT name FROM test WHERE name=?"
execsql = "INSERT INTO test(name) VALUES(?)"
name = ""
)
selstmt, err := db.Prepare(selsql)
if err != nil {
t.Errorf("MySQL:Prepare err(%v)", err)
return
}
row := selstmt.QueryRow(context.TODO(), "noexit")
if err = row.Scan(&name); err == sql.ErrNoRows {
t.Logf("MySQL: prepare query error(%v)", err)
} else {
t.Errorf("MySQL: prepared query name: noexist")
}
rows, err := selstmt.Query(context.TODO(), "test")
if err != nil {
t.Errorf("MySQL:stmt.Query err(%v)", err)
}
rows.Close()
execstmt, err := db.Prepare(execsql)
if err != nil {
t.Errorf("MySQL:Prepare err(%v)", err)
return
}
if _, err := execstmt.Exec(context.TODO(), "test"); err != nil {
t.Errorf("MySQL: stmt.Exec(%v)", err)
}
}
func BenchmarkMySQL(b *testing.B) {
c := &Config{
DSN: "test:test@tcp(172.16.0.148:3306)/test?timeout=5s&readTimeout=5s&writeTimeout=5s&parseTime=true&loc=Local&charset=utf8",
Active: 10,
Idle: 5,
IdleTimeout: xtime.Duration(time.Minute),
}
db := NewTiDB(c)
defer db.Close()
b.RunParallel(func(pb *testing.PB) {
for pb.Next() {
sql := "SELECT name FROM test WHERE name=?"
rows, err := db.Query(context.TODO(), sql, "test")
if err == nil {
for rows.Next() {
var name string
if err = rows.Scan(&name); err != nil {
break
}
}
rows.Close()
}
}
})
}