GORMのソースコードを読む1 – Open

Go言語でDBを操作するライブラリはいくつか存在するが、手軽さとAutoMigrationの便利さから、個人的にはGORMが気に入っている。ただ、ある程度使っていると、バッチInsertができないなど、いくつか気になる点が出てきた。せっかくなので勉強を兼ねてコードを読んでいくことにした。今回は長くなるので、何回かに分けて書いていく。

パッケージ構成

SQLダイアレクト以外の全てのファイルがルートディレクトリに配置されている。全てgormパッケージ。

gormにはmain関数はないので、実際にライブラリとして使うときの関数呼び出しの順番でコードを読んで行こう。

Open

db, err := gorm.Open("sqlite3", "test.db")

まずは、DBへの接続と切断を行うOpenを見ていく。Openはmain.goに定義されている。

func Open(dialect string, args ...interface{}) (db *DB, err error) {

Openは可変長引数argsをとる。通常はここで接続先のDBの情報を指定する。読み進めやすくするために、argsを文字列で指定した場合の実装部分のみを見ていく。

var driver = dialect
if len(args) == 1 {
	source = value
} else if len(args) >= 2 {
	driver = value
	source = args[1].(string)
}

args[0]が文字列の場合は、配列の長さが1の場合はdialectがdriver、その文字列をsourceとして、長さが2以上の場合はargs[0]をdriver、args[1]をsourceとして利用する。どのような場合に必要となるのかわからないが、dialectとdriverを別々に指定できるようになっているようだ。

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

ここで取得したdriverとsourceを元に、Go言語デフォルトのsqlパッケージのOpen関数を使ってDB接続をオープンしている。接続に失敗すると、errにエラー情報が含まれた状態で処理が返ってくる。errが空でなければすぐにreturnしてしまえば良いような気がするのだが、なぜかその処理はgorm.Open関数返り値dbをセットした後で行うようになっている。

db = &DB{
	db:        dbSQL,
	logger:    defaultLogger,
	values:    map[string]interface{}{},
	callbacks: DefaultCallback,
	dialect:   newDialect(dialect, dbSQL),
}
db.parent = db
if err != nil {
	return
}

SQLの接続に失敗していたら、このdbオブジェクトも使い道がないので、あまり意味がないような気がするので、この実装は不思議だ。テストの書きやすさなどの点で、実際には使えないオブジェクトだったとしても、dbをnilで返したくない事情があるのだろうか?

var DefaultCallback = &Callback{}

DefaultCallbackの定義はcallback.goに記述されている。このDefaultCallbackの中身は、callback_create.goなどの各ファイルのinit関数内で追加されているのだが、長くなるのでこのタイミングでは一旦省略する。

var dialectsMap = map[string]Dialect{}

func newDialect(name string, db SQLCommon) Dialect {
	if value, ok := dialectsMap[name]; ok {
		dialect := reflect.New(reflect.TypeOf(value).Elem()).Interface().(Dialect)
		dialect.SetDB(db)
		return dialect
	}

	fmt.Printf("`%v` is not officially supported, running under compatibility mode.\n", name)
	commontDialect := &commonDialect{}
	commontDialect.SetDB(db)
	return commontDialect
}

newDialect関数はdialect.goに定義されている。dialectsMapはmapになっており、そこに対応する名前のdialectが定義されていればそれを使い、なければcommonDialectを使うようになっている。

func init() {
	RegisterDialect("mysql", &mysql{})
}

dialectsMapの中身は、dialect_mysql.goなどのinit関数内で、RegisterDialect関数を通じて登録されている。

// 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

DBインスタンスを作ったら、最後にpingを送って生存確認を行う。pingに失敗した場合はDBをCloseしてreturnするが、この場合も第一引数のdbは空ではなくインスタンスが含まれた状態になる。

dbSQLの中身がsql.DBではない場合に処理をスキップしているのは、テスト用と思われる。説明を省略したが、引数argsの型がstringではなく、かつSQLCommonインターフェイスを継承したオブジェクトの場合は、dbSQLにはそのオブジェクトが含まれるようになっている。

テスト

Openに関するテストの部分も見てみよう。

var (
	DB                 *gorm.DB
	t1, t2, t3, t4, t5 time.Time
)

func init() {
	var err error

	if DB, err = OpenTestConnection(); err != nil {
		panic(fmt.Sprintf("No error should happen when connecting to test database, but got err=%+v", err))
	}

	runMigration()
}

main_test.goにはinit関数が定義されており、テストの実行前にDBの接続設定とマイグレーションが行われる。接続に失敗した場合はpanicにより強制終了するようになっている。

func OpenTestConnection() (db *gorm.DB, err error) {
	dbDSN := os.Getenv("GORM_DSN")
	switch os.Getenv("GORM_DIALECT") {
	case "mysql":
		fmt.Println("testing mysql...")
		if dbDSN == "" {
			dbDSN = "gorm:gorm@tcp(localhost:9910)/gorm?charset=utf8&parseTime=True"
		}
		db, err = gorm.Open("mysql", dbDSN)
	case "postgres":
		...
	default:
		fmt.Println("testing sqlite3...")
		db, err = gorm.Open("sqlite3", filepath.Join(os.TempDir(), "gorm.db"))
	}

	// db.SetLogger(Logger{log.New(os.Stdout, "\r\n", 0)})
	// db.SetLogger(log.New(os.Stdout, "\r\n", 0))
	if debug := os.Getenv("DEBUG"); debug == "true" {
		db.LogMode(true)
	} else if debug == "false" {
		db.LogMode(false)
	}

	db.DB().SetMaxIdleConns(10)

	return
}

OpenTestConnection関数内ではDBの接続処理が行われる。環境変数でGORM_DSNとGORM_DIALECTを指定することで、任意のDBを使ってテストを実行できるようになっている。省略された場合はsqlite3データベースをTempディレクトリ以下に新規で作成して、それを使うようになっているようだ。

func runMigration() {
	if err := DB.DropTableIfExists(&User{}).Error; err != nil {
		fmt.Printf("Got error when try to delete table users, %+v\n", err)
	}

	for _, table := range []string{"animals", "user_languages"} {
		DB.Exec(fmt.Sprintf("drop table %v;", table))
	}

	values := []interface{}{&Short{}, &ReallyLongThingThatReferencesShort{}, &ReallyLongTableNameToTestMySQLNameLengthLimit{}, &NotSoLongTableName{}, &Product{}, &Email{}, &Address{}, &CreditCard{}, &Company{}, &Role{}, &Language{}, &HNPost{}, &EngadgetPost{}, &Animal{}, &User{}, &JoinTable{}, &Post{}, &Category{}, &Comment{}, &Cat{}, &Dog{}, &Hamster{}, &Toy{}, &ElementWithIgnoredField{}}
	for _, value := range values {
		DB.DropTable(value)
	}
	if err := DB.AutoMigrate(values...).Error; err != nil {
		panic(fmt.Sprintf("No error should happen when create table, but got %+v", err))
	}
}

runMigration関数の中では、テストに使うテーブルの再生成が行われているようだ。特別なことはしていないので、省略。

Open関数自体に関しては下記のようなテストが書かれている。

func TestOpen_ReturnsError_WithBadArgs(t *testing.T) {
	stringRef := "foo"
	testCases := []interface{}{42, time.Now(), &stringRef}
	for _, tc := range testCases {
		t.Run(fmt.Sprintf("%v", tc), func(t *testing.T) {
			_, err := gorm.Open("postgresql", tc)
			if err == nil {
				t.Error("Should got error with invalid database source")
			}
			if !strings.HasPrefix(err.Error(), "invalid database source:") {
				t.Errorf("Should got error starting with \"invalid database source:\", but got %q", err.Error())
			}
		})
	}
}

TestOpen_ReturnsError_WithBadArgsでは、argsに無効な値を指定して初期化した場合のテストが書かれている。Table Driven Testの形を取っているが、テストケースは一つしかない。エラー内容はstrings.HasPrefixを使って、文字列での比較となっているが、これは適切なのだろうか。

func TestOpenExistingDB(t *testing.T) {
	DB.Save(&User{Name: "jnfeinstein"})
	dialect := os.Getenv("GORM_DIALECT")

	db, err := gorm.Open(dialect, DB.DB())
	if err != nil {
		t.Errorf("Should have wrapped the existing DB connection")
	}

	var user User
	if db.Where("name = ?", "jnfeinstein").First(&user).Error == gorm.ErrRecordNotFound {
		t.Errorf("Should have found existing record")
	}
}

TestOpenExistingDBでは、OpenTestConnectionで接続したDBにデータを保存した上で、Open関数を使って新規のgorm.DBインスタンスを作成、保存したデータを読み込めるかをテストしている。DBインスタンスはテスト全体を通じて存在し続けるので、ここでSaveしたデータが他の箇所に影響を及ぼさないか、気になる。

func TestOpenWithOneParameter(t *testing.T) {
	db, err := gorm.Open("dialect")
	if db != nil {
		t.Error("Open with one parameter returned non nil for db")
	}
	if err == nil {
		t.Error("Open with one parameter returned err as nil")
	}
}

TestOpenWithOneParameterは、可変長引数argsを全く指定しなかった場合に問題が起きないかどうかをチェックしている。