expose sql.DB to let orm operate on it (#1015)

* expose sql.DB to let orm operate on it

* add missing RawDB methods

* add NewSqlConnFromDB for cooperate with dtm

core/stores/sqlx/sqlconn.go

@@ -6,6 +6,9 @@ import (
 	"github.com/tal-tech/go-zero/core/breaker"
 )
 
+// datasource placeholder for logging error.
+const rawDB = "sql.DB"
+
 // ErrNotFound is an alias of sql.ErrNoRows
 var ErrNotFound = sql.ErrNoRows
 
@@ -23,6 +26,7 @@ type (
 	// SqlConn only stands for raw connections, so Transact method can be called.
 	SqlConn interface {
 		Session
+		RawDB() (*sql.DB, error)
 		Transact(func(session Session) error) error
 	}
 
@@ -43,13 +47,15 @@ type (
 	// Because CORBA doesn't support PREPARE, so we need to combine the
 	// query arguments into one string and do underlying query without arguments
 	commonSqlConn struct {
-		driverName string
 		datasource string
+		connProv   connProvider
 		beginTx    beginnable
 		brk        breaker.Breaker
 		accept     func(error) bool
 	}
 
+	connProvider func() (*sql.DB, error)
+
 	sessionConn interface {
 		Exec(query string, args ...interface{}) (sql.Result, error)
 		Query(query string, args ...interface{}) (*sql.Rows, error)
@@ -69,10 +75,30 @@ type (
 // NewSqlConn returns a SqlConn with given driver name and datasource.
 func NewSqlConn(driverName, datasource string, opts ...SqlOption) SqlConn {
 	conn := &commonSqlConn{
-		driverName: driverName,
 		datasource: datasource,
-		beginTx:    begin,
-		brk:        breaker.NewBreaker(),
+		connProv: func() (*sql.DB, error) {
+			return getSqlConn(driverName, datasource)
+		},
+		beginTx: begin,
+		brk:     breaker.NewBreaker(),
+	}
+	for _, opt := range opts {
+		opt(conn)
+	}
+
+	return conn
+}
+
+// NewSqlConnFromDB returns a SqlConn with the given sql.DB.
+// Use it with caution, it's provided for other ORM to interact with.
+func NewSqlConnFromDB(db *sql.DB, opts ...SqlOption) SqlConn {
+	conn := &commonSqlConn{
+		datasource: rawDB,
+		connProv: func() (*sql.DB, error) {
+			return db, nil
+		},
+		beginTx: begin,
+		brk:     breaker.NewBreaker(),
 	}
 	for _, opt := range opts {
 		opt(conn)
@@ -84,7 +110,7 @@ func NewSqlConn(driverName, datasource string, opts ...SqlOption) SqlConn {
 func (db *commonSqlConn) Exec(q string, args ...interface{}) (result sql.Result, err error) {
 	err = db.brk.DoWithAcceptable(func() error {
 		var conn *sql.DB
-		conn, err = getSqlConn(db.driverName, db.datasource)
+		conn, err = db.connProv()
 		if err != nil {
 			logInstanceError(db.datasource, err)
 			return err
@@ -100,7 +126,7 @@ func (db *commonSqlConn) Exec(q string, args ...interface{}) (result sql.Result,
 func (db *commonSqlConn) Prepare(query string) (stmt StmtSession, err error) {
 	err = db.brk.DoWithAcceptable(func() error {
 		var conn *sql.DB
-		conn, err = getSqlConn(db.driverName, db.datasource)
+		conn, err = db.connProv()
 		if err != nil {
 			logInstanceError(db.datasource, err)
 			return err
@@ -145,6 +171,10 @@ func (db *commonSqlConn) QueryRowsPartial(v interface{}, q string, args ...inter
 	}, q, args...)
 }
 
+func (db *commonSqlConn) RawDB() (*sql.DB, error) {
+	return db.connProv()
+}
+
 func (db *commonSqlConn) Transact(fn func(Session) error) error {
 	return db.brk.DoWithAcceptable(func() error {
 		return transact(db, db.beginTx, fn)
@@ -163,7 +193,7 @@ func (db *commonSqlConn) acceptable(err error) bool {
 func (db *commonSqlConn) queryRows(scanner func(*sql.Rows) error, q string, args ...interface{}) error {
 	var qerr error
 	return db.brk.DoWithAcceptable(func() error {
-		conn, err := getSqlConn(db.driverName, db.datasource)
+		conn, err := db.connProv()
 		if err != nil {
 			logInstanceError(db.datasource, err)
 			return err