01GORM源码解读

• 简介
• 起步
• 数据库连接
• gorm.DB
• 事务实现
• 总结

简介

GORM 源码解读, 基于 v1.9.11 版本.

起步

官方文档上入门的例子如下:

package main

import (
  "github.com/jinzhu/gorm"
  _ "github.com/jinzhu/gorm/dialects/sqlite"
)

type Product struct {
  gorm.Model
  Code string
  Price uint
}

func main() {
  db, err := gorm.Open("sqlite3", "test.db")
  if err != nil {
    panic("failed to connect database")
  }
  defer db.Close()

  // Migrate the schema
  db.AutoMigrate(&Product{})

  // 创建
  db.Create(&Product{Code: "L1212", Price: 1000})

  // 读取
  var product Product
  db.First(&product, 1) // 查询id为1的product
  db.First(&product, "code = ?", "L1212") // 查询code为l1212的product

  // 更新 - 更新product的price为2000
  db.Model(&product).Update("Price", 2000)

  // 删除 - 删除product
  db.Delete(&product)
}

数据库连接

从 gorm.Open 开始看起吧, 看数据库是怎么连接的:

// Open initialize a new db connection, need to import driver first, e.g:
//
//     import _ "github.com/go-sql-driver/mysql"
//     func main() {
//       db, err := gorm.Open("mysql", "user:password@/dbname?charset=utf8&parseTime=True&loc=Local")
//     }
// GORM has wrapped some drivers, for easier to remember driver's import path, so you could import the mysql driver with
//    import _ "github.com/jinzhu/gorm/dialects/mysql"
//    // import _ "github.com/jinzhu/gorm/dialects/postgres"
//    // import _ "github.com/jinzhu/gorm/dialects/sqlite"
//    // import _ "github.com/jinzhu/gorm/dialects/mssql"
func Open(dialect string, args ...interface{}) (db *DB, err error) {
	if len(args) == 0 {
		err = errors.New("invalid database source")
		return nil, err
	}
	var source string
	var dbSQL SQLCommon
	var ownDbSQL bool

	switch value := args[0].(type) {
	case string:
		var driver = dialect
		if len(args) == 1 {
			source = value
		} else if len(args) >= 2 {
			driver = value
			source = args[1].(string)
		}
		dbSQL, err = sql.Open(driver, source)
		ownDbSQL = true
	case SQLCommon:
		dbSQL = value
		ownDbSQL = false
	default:
		return nil, fmt.Errorf("invalid database source: %v is not a valid type", value)
	}

	db = &DB{
		db:        dbSQL,
		logger:    defaultLogger,
		callbacks: DefaultCallback,
		dialect:   newDialect(dialect, dbSQL),
	}
	db.parent = db
	if err != nil {
		return
	}
	// Send a ping to make sure the database connection is alive.
	if d, ok := dbSQL.(*sql.DB); ok {
		if err = d.Ping(); err != nil && ownDbSQL {
			d.Close()
		}
	}
	return
}

gorm.Open 有两个参数, 一个是数据库名称, 其余是连接参数.

从 switch 语句中, 可以发现如果第一个参数是 string 类型, 实际上是通过 Golang 中的 sql 模块连接:

dbSQL, err = sql.Open(driver, source)

也可以直接传递一个实现了 SQLCommon 接口的实例.

然后初始化了一个 gorm.DB 实例, 并在最后执行了一次 ping 请求, 测试数据库连接是否正常.

看一下 gorm.DB 结构体:

// DB contains information for current db connection
type DB struct {
	sync.RWMutex
	Value        interface{}
	Error        error
	RowsAffected int64

	// single db
	db                SQLCommon
	blockGlobalUpdate bool
	logMode           logModeValue
	logger            logger
	search            *search
	values            sync.Map

	// global db
	parent        *DB
	callbacks     *Callback
	dialect       Dialect
	singularTable bool

	// function to be used to override the creating of a new timestamp
	nowFuncOverride func() time.Time
}

gorm.DB 扩展自 sync.RWMutex 读写互斥锁.

gorm.DB

上面已经看过了 gorm.DB 结构体的定义了, 从入门的示例代码中可以看出, 所有的操作都是围绕它来进行的, 所以 gorm.DB 是核心的结构体. 看下它具体实现了哪些方法.

// New clone a new db connection without search conditions
func (s *DB) New() *DB {
	clone := s.clone()
	clone.search = nil
	clone.Value = nil
	return clone
}

type closer interface {
	Close() error
}

// Close close current db connection.  If database connection is not an io.Closer, returns an error.
func (s *DB) Close() error {
	if db, ok := s.parent.db.(closer); ok {
		return db.Close()
	}
	return errors.New("can't close current db")
}

克隆数据库的连接和关闭数据库连接. New 方法内部使用到了 s.clone(),

func (s *DB) clone() *DB {
	db := &DB{
		db:                s.db,
		parent:            s.parent,
		logger:            s.logger,
		logMode:           s.logMode,
		Value:             s.Value,
		Error:             s.Error,
		blockGlobalUpdate: s.blockGlobalUpdate,
		dialect:           newDialect(s.dialect.GetName(), s.db),
		nowFuncOverride:   s.nowFuncOverride,
	}

	s.values.Range(func(k, v interface{}) bool {
		db.values.Store(k, v)
		return true
	})

	if s.search == nil {
		db.search = &search{limit: -1, offset: -1}
	} else {
		db.search = s.search.clone()
	}

	db.search.db = db
	return db
}

略过一些简单的 get/set 方法, 接着看

// NewScope create a scope for current operation
func (s *DB) NewScope(value interface{}) *Scope {
	dbClone := s.clone()
	dbClone.Value = value
	scope := &Scope{db: dbClone, Value: value}
	if s.search != nil {
		scope.Search = s.search.clone()
	} else {
		scope.Search = &search{}
	}
	return scope
}

NewScope 会为当前的操作创建一个新的 scope (作用域).

// QueryExpr returns the query as expr object
func (s *DB) QueryExpr() *expr {
	scope := s.NewScope(s.Value)
	scope.InstanceSet("skip_bindvar", true)
	scope.prepareQuerySQL()

	return Expr(scope.SQL, scope.SQLVars...)
}

// SubQuery returns the query as sub query
func (s *DB) SubQuery() *expr {
	scope := s.NewScope(s.Value)
	scope.InstanceSet("skip_bindvar", true)
	scope.prepareQuerySQL()

	return Expr(fmt.Sprintf("(%v)", scope.SQL), scope.SQLVars...)
}

QueryExpr 和 SubQuery 都用到了 NewScope, 在当前的作用域下获取查询表达式和进行子查询.

接着是很多查询方法, 类似 Where,

// Where return a new relation, filter records with given conditions, accepts `map`, `struct` or `string` as conditions, refer http://jinzhu.github.io/gorm/crud.html#query
func (s *DB) Where(query interface{}, args ...interface{}) *DB {
	return s.clone().search.Where(query, args...).db
}

跳过这些方法, 等后面探究查询表达式的时候再详细研究.

// Scopes pass current database connection to arguments `func(*DB) *DB`, which could be used to add conditions dynamically
//     func AmountGreaterThan1000(db *gorm.DB) *gorm.DB {
//         return db.Where("amount > ?", 1000)
//     }
//
//     func OrderStatus(status []string) func (db *gorm.DB) *gorm.DB {
//         return func (db *gorm.DB) *gorm.DB {
//             return db.Scopes(AmountGreaterThan1000).Where("status in (?)", status)
//         }
//     }
//
//     db.Scopes(AmountGreaterThan1000, OrderStatus([]string{"paid", "shipped"})).Find(&orders)
// Refer https://jinzhu.github.io/gorm/crud.html#scopes
func (s *DB) Scopes(funcs ...func(*DB) *DB) *DB {
	for _, f := range funcs {
		s = f(s)
	}
	return s
}

Scopes 是一个钩子函数, 用于动态添加查询条件, 这在函数是一等公民的语言里是一个常见的模式.

事务实现

看一下事务是如何实现的:

// Begin begins a transaction
func (s *DB) Begin() *DB {
	return s.BeginTx(context.Background(), &sql.TxOptions{})
}

// BeginTx begins a transaction with options
func (s *DB) BeginTx(ctx context.Context, opts *sql.TxOptions) *DB {
	c := s.clone()
	if db, ok := c.db.(sqlDb); ok && db != nil {
		tx, err := db.BeginTx(ctx, opts)
		c.db = interface{}(tx).(SQLCommon)

		c.dialect.SetDB(c.db)
		c.AddError(err)
	} else {
		c.AddError(ErrCantStartTransaction)
	}
	return c
}

这一部分是开始事务时的操作, 实际上是 c.db 实现了 sqlDb 接口, 调用了 BeginTx 方法.

type sqlDb interface {
	Begin() (*sql.Tx, error)
	BeginTx(ctx context.Context, opts *sql.TxOptions) (*sql.Tx, error)
}

接着看如何提交事务:

// Commit commit a transaction
func (s *DB) Commit() *DB {
	var emptySQLTx *sql.Tx
	if db, ok := s.db.(sqlTx); ok && db != nil && db != emptySQLTx {
		s.AddError(db.Commit())
	} else {
		s.AddError(ErrInvalidTransaction)
	}
	return s
}

和开始事务类似, s.db 实现了 sqlTx 接口, 调用了 Commit 方法.

type sqlTx interface {
	Commit() error
	Rollback() error
}

sqlTx 接口里还有个 Rollback 方法, 所以回滚操作也是类似的:

// Rollback rollback a transaction
func (s *DB) Rollback() *DB {
	var emptySQLTx *sql.Tx
	if db, ok := s.db.(sqlTx); ok && db != nil && db != emptySQLTx {
		if err := db.Rollback(); err != nil && err != sql.ErrTxDone {
			s.AddError(err)
		}
	} else {
		s.AddError(ErrInvalidTransaction)
	}
	return s
}

// RollbackUnlessCommitted rollback a transaction if it has not yet been
// committed.
func (s *DB) RollbackUnlessCommitted() *DB {
	var emptySQLTx *sql.Tx
	if db, ok := s.db.(sqlTx); ok && db != nil && db != emptySQLTx {
		err := db.Rollback()
		// Ignore the error indicating that the transaction has already
		// been committed.
		if err != sql.ErrTxDone {
			s.AddError(err)
		}
	} else {
		s.AddError(ErrInvalidTransaction)
	}
	return s
}

RollbackUnlessCommitted 和 Rollback 的区别在于前者少了一个 err != nil 的判断, 看了半天还是难以理解这有什么差别.

RollbackUnlessCommitted 作者给出的例子如下:

func doTransaction(DB *gorm.DB) error {
  tx := DB.Begin()
  defer tx.RollbackUnlessCommitted()

  u := User{Name: "test"}
  if err != tx.Save(&User).Error; err != nil {
    return err
  }
  return tx.Commit().Error
}

相比较而言, 官方文档上事务的例子如下:

func CreateAnimals(db *gorm.DB) error {
  // 请注意，事务一旦开始，你就应该使用 tx 作为数据库句柄
  tx := db.Begin()
  defer func() {
    if r := recover(); r != nil {
      tx.Rollback()
    }
  }()

  if err := tx.Error; err != nil {
    return err
  }

  if err := tx.Create(&Animal{Name: "Giraffe"}).Error; err != nil {
    tx.Rollback()
    return err
  }

  if err := tx.Create(&Animal{Name: "Lion"}).Error; err != nil {
    tx.Rollback()
    return err
  }

  return tx.Commit().Error
}

总结

暂时就看到这里吧, gorm.DB 还有很多方法等待后续发掘.

主要看了数据库的连接过程, 基本是通过 dbSQL, err = sql.Open(driver, source) 实现的.

也看了事务部分, 主要是要实现两个接口:

type sqlDb interface {
	Begin() (*sql.Tx, error)
	BeginTx(ctx context.Context, opts *sql.TxOptions) (*sql.Tx, error)
}

type sqlTx interface {
	Commit() error
	Rollback() error
}

当然, 对于其中的 RollbackUnlessCommitted 和 Rollback 有点疑惑, 因为我想不明白到底有什么不同.

既然是 ORM, 模型定义应该是重中之重, 后续将探索 Model 实现.