GoのORM sqlcのType Overridesを自動生成する(MySQL用)

前回はsqlcのType Overridesの紹介をしましたが、Type Overridesで自前の型にマッピングすると便利だけど、手でやるのはちょっとしんどいかもしれません。

はい、そうですね、自動生成すればいいですね。

おさらい: Type Overrides の何が嬉しいのか？

その前に、何が嬉しいのか、もう一度、おさらいしましょう(お腹いっぱいの人は次へGo)。例えば、category_id というものがあったとします。ですが、categoryは別にテーブル定義されておらず、カテゴリ名はハードコーディングされていたとします。残念ですね。DBで引っ張ってから、Category型でおきかえるかー...。いや、Type Overridesしましょう。

type CategoryID int

var categoryMap map[CategoryID]string = map[CategoryID]string {
  1: "Sports",
  2: "Music",
  3: "Outdoor",
}

func (c *CategoryID) Name() string {
    return categoryMap[c]
}

このような型が存在しており、Type Overridesで、category_idがCategoryIDにマッピングされているならば、Name()と呼ぶだけで、名前が取れます。おや、便利ですね。

もう少しマシな例でしたら、例えば、user_statusというカラムがあったとします。

type UserStatus int

func (u UserStatus) IsTemporary() bool {
	return u == 0
}

func (u UserStatus) IsActive() bool {
	return u == 1
}

func (u UserStatus) IsSuspend() bool {
	return u == 2
}

func (u UserStatus) IsWithdraw() bool {
	return u == 3
}

このような型にマッピングしておけば、使い勝手が良いですね。
ですが、user_statusが他のテーブルにある(のは、考えにくいですが)とか、category_idが他のテーブルにある(のは、ありそう)場合、全部を列挙しないといけません。だるい。

MySQLのinformation_schemaからカラム名を取得してType Overridesを自動生成する

そこで、information_schemaを使います。information_schemaには、DBのメタ情報が整理されて入っています。テーブルのメタ情報も当然。プロセスなんかも入ってます(余談ですが、show processlist見るよりPROCESSES見たほうが10倍くらいは便利です)。

今回はCOLUMNSテーブルを使います。

SELECT TABLE_SCHEMA, TABLE_NAME, COLUMN_NAME FROM information_schema.COLUMNS;

とすれば、どのDB(TABLE_SCHEMA)のどのテーブル(TABLE_NAME)にどのカラム(COLUMN_NAME)が入ってるかわかります。以下のようなクエリで、category_idがどこで使われてるかがわかるわけです。

SELECT TABLE_SCHEMA, TABLE_NAME, COLUMN_NAME FROM information_schema.COLUMNS
WHERE COLUMN_NAME = "category_id";

これがわかれば、sqlc.yml にType Overridesを追加することも容易です。

自動生成するコードは以下のようになりますが、いくつか前提があります。

• 引数は、DSNとYAMLファイルのパスとDB名
• *.yml.origin を編集して、*.yml を上書きする
• overridesの設定は、typeOverrides の変数にハードコードする
• const yourPackageに指定しているpackageに独自の型が定義されており、基本的にそれをimportに指定する
• Option[から始まる型を指定した場合のみ、go-optionalをimportに指定する

package main

import (
	"database/sql"
	"log"
	"os"
	"sort"
	"strings"

	"gopkg.in/yaml.v3"
)

type def struct {
	typeName string
	nullable bool
}

// 独自の型を定義しているpackageを書く
const yourPackage = "your-project/valueobject"

// Type Overrides するものは、ここに書く
var typeOverrides = map[string]def{
	// 変換したいカラム名: {型, nullable(bool)}
	"*.category_id":    {"CategoryID", true},
	"user.user_status": {"UserStatus", false},
}

// sqlc の YAMLファイルの型定義始まり
type Database struct {
	URI string `yaml:"uri"`
}

type GoType struct {
	Import string `yaml:"import"`
	Type   string `yaml:"type"`
}

type Overrides struct {
	DBType      *string `yaml:"db_type,omitempty"`
	Nullable    *bool   `yaml:"nullable,omitempty"`
	Unsigned    *bool   `yaml:"unsigned,omitempty"`
	Column      *string `yaml:"column,omitempty"`
	GoType      *GoType `yaml:"go_type"`
	GoStructTag *string `yaml:"go_struct_tag,omitempty"`
}

type Go struct {
	Pacakge             string       `yaml:"package"`
	SQLPackage          string       `yaml:"sql_package,omitempty"`
	SQLDriver           string       `yaml:"sql_driver,omitempty"`
	Out                 string       `yaml:"out"`
	EmitJsonTags        bool         `yaml:"emit_json_tags,omitempty"`
	EmitDBTags          bool         `yaml:"emit_db_tags,omitempty"`
	EmitExactTableNames bool         `yaml:"emit_exact_table_names,omitempty"`
	EmitEmptySlices     bool         `yaml:"emit_empty_slices,omitempty"`
	Overrides           []*Overrides `yaml:"overrides"`
}

type Gen struct {
	Go *Go `yaml:"go"`
}

type SQL struct {
	Engine   string    `yaml:"engine"`
	Queries  string    `yaml:"queries"`
	Schema   string    `yaml:"schema"`
	Database *Database `yaml:"database"`
	Gen      *Gen      `yaml:"gen"`
}

type SqlcYaml struct {
	Version string `yaml:"version"`
	SQL     []*SQL `yaml:"sql"`
}

// sqlc の YAMLファイルの型定義終わり

func main() {
	args := os.Args
	if len(args) < 3 {
		log.Fatal("pass dsn and yaml file")
	}
	dsn := os.Args[1]
	file := os.Args[2]
	tableSchema := os.Args[3]

	bytes, err := os.ReadFile(file + ".origin")
	if err != nil {
		log.Fatal(err)
	}
	r := strings.NewReader(string(bytes))
	d := yaml.NewDecoder(r)
	yamlData := &SqlcYaml{}
	d.Decode(yamlData)

	db, err := sql.Open("mysql", dsn)
	if err != nil {
		log.Fatal(err)
	}

	sortedKeys := make([]string, 0, len(typeOverrides))
	for key := range typeOverrides {
		sortedKeys = append(sortedKeys, key)
	}
	sort.Strings(sortedKeys)

	overrides := make([]*Overrides, 0)
	for _, fullColumnName := range sortedKeys {
		fullTableColumn := strings.Split(fullColumnName, ".")
		def := typeOverrides[fullColumnName]

		tableName := fullTableColumn[0]
		columnName := fullTableColumn[1]

		if tableName != "*" {
			nullable := &def.nullable
			if !*nullable {
				nullable = nil
			}
			importPackage := yourPackage
			if strings.HasPrefix(def.typeName, "Option[") {
				importPackage = "github.com/moznion/go-optional"
			}
			overrides = append(overrides, &Overrides{
				Nullable: nullable,
				Column:   &fullColumnName,
				GoType: &GoType{
					Import: importPackage,
					Type:   def.typeName,
				},
			})
			continue
		}

		result, err := db.Query(`
	     SELECT TABLE_NAME FROM information_schema.COLUMNS
	      WHERE COLUMN_NAME = ? AND TABLE_SCHEMA = ?`, columnName, tableSchema)
		if err != nil {
			log.Fatal(err)
		}

		for result.Next() {
			var tableName string

			err := result.Scan(&tableName)
			if err != nil {
				log.Fatal(err)
			}
			importPackage := yourPackage
			if strings.HasPrefix(def.typeName, "Option[") {
				importPackage = "github.com/moznion/go-optional"
			}
			fullColumnName := tableName + "." + columnName
			nullable := &def.nullable
			if !*nullable {
				nullable = nil
			}
			overrides = append(overrides, &Overrides{
				Column:   &fullColumnName,
				Nullable: nullable,

				GoType: &GoType{
					Import: importPackage,
					Type:   def.typeName,
				},
			})
		}
	}

	yamlData.SQL[0].Gen.Go.Overrides = append(yamlData.SQL[0].Gen.Go.Overrides, overrides...)

	f, err := os.OpenFile(file, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0664)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	enc := yaml.NewEncoder(f)
	err = enc.Encode(yamlData)
	if err != nil {
		log.Fatal(err)
	}

	enc.Close()
}

終わり

改めて見直すと、nullableかどうかは、information_schemaでわかるな...とか、修正したい点がいくつかみつかりましたが...まぁ、必要に応じて改善していただければ良いのではないかなと思います。

これで、YAMLファイルを自動生成できるので、お気軽にカラムを特定の型でoverridesできるようになりました。overridesをちまちま書くのから開放されましたね。ぜひ、overridesを活用しましょう。

ちなみに、今やってるプロジェクトだと、63くらいoverridesの設定がありました。多いですが、管理は全然苦労していません。