🚧

GoのORM「ent」で意図しないUPDATE/DELETE ALLを阻止する

2022/06/05に公開

ORM (もっと言うと query builder) あるあるだが、WHERE 句の組立時に何らかのバグで条件が空になってしまって、意図せず UPDATE/DELETE ALL を発行してしまうことがある。

gorm なんかだとデフォルトで gorm.ErrMissingWhereClause が返るようになっており、もし UPDATE/DELETE ALL したいのであれば AllowGlobalUpdate: true と設定する必要がある。

https://gorm.io/docs/gorm_config.html#AllowGlobalUpdate

ちなみに MySQL では --safe-updates を付けることで防ぐことができるので、ユースケースにも依るだろうが DBMS 側で阻止できるならそれに越したことはない。

https://dev.mysql.com/doc/refman/8.0/ja/mysql-tips.html#safe-updates

さて、Go 製の ORM の1つである ent では、この Global Update を防ぐ機構が (パッと調べた感じ) 存在しない。

今回は WHERE 句の無い UPDATE/DELETE が実行される時にエラーを返すランタイムフックを作成する方法を記す。

Mutation が predicates を返すカスタムテンプレートを作る

WHERE 句の条件に相当する変数 predicates は、生成される各モデルの Mutation の内部に unexported な変数として保持されている。
これの len() を取って0ならばエラーを返せばよいのだが、Getter も存在しないので reflect を使わない限り取得することはできない。

ということで predicates のスライスの長さを返すメソッドをカスタムテンプレートで定義する。

mutation_predicates.tmpl
{{ define "mutation_predicates" }}

{{ $pkg := base $.Config.Package }}
{{ template "header" $ }}

{{ range $n := $.Nodes }}
  func (m *{{ print $n.Name "Mutation" }}) GetPredicatesSize() int {
    return len(m.predicates)
  }
{{ end }}

{{ end }}

これで各モデルの Mutation が内部に持つ predicates の数が返るようになった。

ランタイムフックでエラーを返す

スキーマフックでちまちま Hook を作っても良いが面倒なので、ランタイムフックで全モデルの UPDATE/DELETE をチェックする Hook をテンプレートから生成する。

prevent_global_update_hook.tmpl
{{ define "hook/prevent_global_update" }}

{{ with extend $ "Package" "hook" }}
  {{ template "header" . }}
{{ end }}

func PreventGlobalUpdate() ent.Hook {
  return On(
    func(next ent.Mutator) ent.Mutator {
      return ent.MutateFunc(func(ctx context.Context, m ent.Mutation) (ent.Value, error) {
        switch mut := m.(type) {
          {{ range $n := $.Nodes }}
          case *ent.{{ print $n.Name "Mutation" }}:
            if mut.GetPredicatesSize() == 0 {
              return nil, errors.New("prevented global update")
            }
          {{ end }}
        }
        return next.Mutate(ctx, m)
      })
    },
    ent.OpDelete | ent.OpUpdate,
  )
}

{{ end }}

後は go generate して、これをクライアント作成時に指定するだけ。

client := ent.NewClient()
client.Use(hook.PreventGlobalUpdate())

以上。

GitHubで編集を提案

Discussion