Kong Gatewayのカスタムエンティティを作成する

Kong Gatewayでカスタムプラグインを作成する際、プラグインが扱うデータをどこかしらに保存したい時がある。
例えば鍵認証のようなものの場合、プラグイン設定時に指定した鍵情報をリクエスト送信時に参照して一致するか照合するようなケースである。
このようなケースでは外部のDBを用いる方法も考えられるが、Kong GatewayのDBに格納することも出来る。
この独自に格納するデータをカスタムエンティティと呼ぶが、これを今回はこれの使い方を検証する。

## 基礎知識
Kong GatewayではDBへのアクセスはPDKを使ってDAO(Data Access Object)として抽象化されたもので[アクセスできる](https://docs.konghq.com/gateway/latest/plugin-development/access-the-datastore/)。

カスタムプラグインでこのDAOを実装しようとした場合、DAOを含めたカスタムプラグインの構造は以下となる（[参考](https://docs.konghq.com/gateway/latest/plugin-development/file-structure/#advanced-plugin-modules)）。

```sh
complete-plugin
├── daos.lua
├── handler.lua
├── migrations
│   ├── init.lua
│   └── 000_base_complete_plugin.lua
└── schema.lua
```

`schema.lua`と`hadnler.lua`がカスタムプラグインの実装にあたり、カスタムエンティティの実装は`daos.lua`で行い、DBの初期化は`migrations`以下のコードで行う。
`migrations`以下のファイルは以下のようになる。
- `init.lua`：DBの初期化で利用するファイルの定義
- `000_xxx.lua`：実際に初期化で使うコード

`000_xxx.lua`のファイル名については厳密な命名規則はないが、先頭に3桁の数字をつけ、実行順にあわせて`000_...`、`001_...`と昇順のプリフィックスをつけるようになっている。
また以下の慣習的な命名規則がある。
- 最初のファイルは`000_base_<プラグイン名>`
- 以降のファイルは`00x_<旧バージョン>_to_<新バージョン>`

例えばKey Auth Pluginの場合は
```
000_base_key_auth.lua
002_130_to_140.lua
003_200_to_210.lua
```
のように1.3.0から1.4.0での変更、2.0.0から2.1.0の変更を`130_to_140`や`200_to_210`のようにファイル名で表現している。

カスタムエンティティの実装については[Key Auth Pluginのコード](https://github.com/Kong/kong/tree/master/kong/plugins/key-auth)も参考になるので、必要に応じて参照したい。

## 実装
ここでは以下のような仕様のカスタムプラグインを作成する。
- ヘッダに`"mykey-add: value"`をつけてアクセスするとDBに以下を格納する。
  - UUID (※自動生成）
  - 作成時刻　（※自動生成）
  - value (カラム名は`mykey`とする)
- ヘッダに`"mykey-del: UUID"`をつけてアクセスするとテーブルに指定したUUIDと合致する項目があれば削除する
- バージョンアップを想定し、DBの初期化では最初に利用しないカラム(`col1`)を作成した後、そのカラムを削除して`mykey`を作成する

また、今回はカスタムプラグインのテンプレートである[`kong-plugin`](https://github.com/Kong/kong-plugin)を使って実装し、カスタムプラグインのテストツールである[Pongo](https://github.com/Kong/kong-pongo)を使って動作確認する。
Pongoについては@jirokichi2017さんが書いた「[Kong Gatewayカスタムプラグインをテストする　～Part1: Pongoを実行する～](https://qiita.com/jirokichi2017/items/b6e272ac5fcba113e405)」がよく纏まってるのでそちらを参考にして欲しい。

最初にテンプレートをcloneする。
```sh
git clone https://github.com/Kong/kong-plugin
cd kong-plugin
```
テンプレートでは以下のように`daos.lua`や`migrations`は含まれていないので、それぞれを作っていく。
```
kong/plugins/myplugin/
├── handler.lua
└── schema.lua
```

### テーブルの作成
先にDBのテーブルを作っていく。
まず格納先ディレクトリを作成する。
```sh
mkdir kong/plugins/myplugin/migrations
```
最初にDBの初期化の際にどのファイルを使うかを定義する`init.lua`から作成する。
```lua:kong/plugins/myplugin/migrations/init.lua
cat <<'EOF' > ./kong/plugins/myplugin/migrations/init.lua
return {
  "000_base_my_plugin",
  "001_100_to_110",
}
EOF
```
`init.lua`ではDB初期化時に読み込ませるファイルの一覧を記載する。
検証の冒頭で述べたように、バージョンアップを想定して初期ファイルである`000_base_my_plugin.lua`とバージョンアップで更新する`001_100_to_110.lua`を用意する。

次に`init.lua`に記載した`000_base_my_plugin.lua`と`001_100_to_110.lua`を作成する。
ここでは公式ドキュメントの[Migration file syntax](https://docs.konghq.com/gateway/latest/plugin-development/custom-entities/#migration-file-syntax)にあるコードを参考にした（※公式の記述だと'25/1時点で記載ミスがあり、コピペだと動かないので注意）。
`000_base_my_plugin.lua`では`id`、`created_at`、`col1`というカラムを持つ`my_plugin_table`というテーブルを作成し、`col1`というカラムに基づくインデックス`my_plugin_table_col1`を作成する。
```lua:kong/plugins/myplugin/migrations/000_base_my_plugin.lua
cat <<'EOF' > ./kong/plugins/myplugin/migrations/000_base_my_plugin.lua
return {
 postgres = {
   up = [[
     CREATE TABLE IF NOT EXISTS "my_plugin_table" (
       "id"           UUID                         PRIMARY KEY,
       "created_at"   TIMESTAMP WITHOUT TIME ZONE,
       "col1"         TEXT
     );

     DO $$
     BEGIN
       CREATE INDEX IF NOT EXISTS "my_plugin_table_col1"
                               ON "my_plugin_table" ("col1");
     EXCEPTION WHEN UNDEFINED_COLUMN THEN
       -- Do nothing, accept existing state
     END$$;
   ]],
 }
}
EOF
```
`001_100_to_110.lua`はプラグインのバージョンアップ時に追加されたファイルと想定し、以下の変更をDBに行う。
- 新しいカラム`mykey`を追加
- 古いカラム`col1`を削除
```lua:kong/plugins/myplugin/migrations/001_100_to_110.lua
cat <<'EOF' > ./kong/plugins/myplugin/migrations/001_100_to_110.lua
return {
 postgres = {
   up = [[
     DO $$
     BEGIN
       ALTER TABLE IF EXISTS ONLY "my_plugin_table" ADD "mykey" TEXT UNIQUE;
     EXCEPTION WHEN DUPLICATE_COLUMN THEN
       -- Do nothing, accept existing state
     END$$;
   ]],
   teardown = function(connector, helpers)
     assert(connector:connect_migrations())
     assert(connector:query([[
       DO $$
       BEGIN
         ALTER TABLE IF EXISTS ONLY "my_plugin_table" DROP "col1";
       EXCEPTION WHEN UNDEFINED_COLUMN THEN
         -- Do nothing, accept existing state
       END$$;
     ]]))
   end,
 }
}
EOF
```
`up =`で記載されている箇所は`kong migrations up`で動作し、`teardown =`で記載されている箇所は`kong migrations finish`で動作する。
そのため、以下のような使い分けが推奨されている。
- 非破壊的な操作：`up`で記述
- 破壊的な操作：`teardown`で記述

### カスタムエンティティの作成
カスタムエンティティのためのスキーマは`daos.lua`によって定義できる。
設定できる項目の説明については[Define a schema](https://docs.konghq.com/gateway/latest/plugin-development/custom-entities/#define-a-schema)を参照。

:::note info
なお、上記サイトには載っていないが、`workspaceable`をつけるとWorkspace単位でエンティティが存在するのかどうかも指定できる。（今回はPongoをWorkspaceがないOSS版で起動していたので、有効化はしていない）
`workspaceable`が`true`エンティティはグローバルエンティティ（≠グローバルスコープ）となり、Workspaceを跨いで参照できる（mTLSで使う証明書やGroups、Licenseなどが該当）。
:::


ここでは以下のように作成する。
```lua:kong/plugins/myplugin/daos.lua
cat <<'EOF' > ./kong/plugins/myplugin/daos.lua
local typedefs = require "kong.db.schema.typedefs"

return {
   {
   name          = "my_plugin_table",
   primary_key   = { "id" },
   -- workspaceable = true,
   fields = {
     { id = typedefs.uuid },
     { created_at = typedefs.auto_timestamp_s },
     { mykey = { type = "string", required  = true }}
   }
   }
}
EOF
```
先ほど作成した`migrations`以下のDBの構成に合わせたフィールドを定義する。
また、作成したDAOは`kong.db.<指定したname>`でアクセスできるようになる。


### カスタムプラグインの作成
DBにアクセスするカスタムプラグインを作成する。
`schema.lua`については今回検証時にパラメータは渡さないので、空に近いものを用意する。
```lua:kong/plugins/myplugin/schema.lua
cat <<'EOF' > ./kong/plugins/myplugin/schema.lua
local PLUGIN_NAME = "myplugin"
local schema = {
  name = PLUGIN_NAME,
  fields = {
    { config = {
        type = "record",
        fields = {}
      }}
  }
}
return schema
EOF
```
`handler.lua`は以下のようにした。
```lua:kong/plugins/myplugin/handler.lua
cat <<'EOF' > ./kong/plugins/myplugin/handler.lua
local plugin = {
  PRIORITY = 1000,
  VERSION = "0.1",
}
  
function plugin:access(conf)
  local add_value = kong.request.get_header("mykey-add")
  local del_value = kong.request.get_header("mykey-del")

  if add_value then
    local ok, err = kong.db.my_plugin_table:insert({mykey = add_value})
    if not ok then
      kong.log.err("Failed to insert into my_plugin_table: ", err)
    end
  elseif del_value then
    local entity, err = kong.db.my_plugin_table:select({id = del_value})
    if not entity then
      kong.log.err("Failed to find entity with name: ", del_value, " Error: ", err)
      return kong.response.exit(404, { message = "Entity not found" })
    end

    local ok, err = kong.db.my_plugin_table:delete({ id = entity.id })
    if not ok then
      kong.log.err("Failed to delete entity: ", err)
      return kong.response.exit(500, { message = "Internal Server Error" })
    end
  end

end
return plugin
EOF
```
`:access()`を指定することで、リクエスト送信時にこの関数が動作するようにしている。
DAOの利用箇所は以下となる。
- `kong.db.my_plugin_table:insert({mykey = add_value})`：DBにエンティティをInsert
- `kong.db.my_plugin_table:select({id = del_value})`：DB内のプライマリキー（UUID）と一致するエンティティを取得
- `kong.db.my_plugin_table:delete({ id = entity.id })`：UUIDが一致するエンティティを削除

以上でカスタムプラグインの準備は終了となる。

## 動作確認
ここではPongoを使って動作確認する。
Pongoを起動し外部からアクセスできるようにする。
```sh
pongo up
pongo expose
```
shellをアタッチし、Kong Gatewayを起動する。
```sh
pongo shell
kms
```
この段階でPostgres内に定義したテーブルが作成される。

```sh
$ docker exec -it pongo-91c2117c-postgres-1 psql -U kong -d kong_tests -c "SELECT * FROM my_plugin_table;"
 id | created_at | mykey
----+------------+------
(0 rows)
```

Kong Gateway起動後、Service/Route/Pluginを作成する。
ここではhttpコマンドで作成した。
```sh
http POST :8001/services \
    name=example-service \
    url=http://httpbin.org
```
```sh
http POST :8001/services/example-service/routes \
    name=example-route \
    paths:='["/echo"]'
```
```sh
http POST :8001/services/example-service/plugins \
    name=myplugin 
```
動作確認する。
ヘッダに`mykey-add:test`したものと`mykey-add:test2`を付加したもので2回アクセスする。
```sh
http GET :8000/echo/ip \
    mykey-add:test
```
```sh
http GET :8000/echo/ip \
    mykey-add:test2
```
テーブルを確認する。
```sh
                  id                  |     created_at      | mykey
--------------------------------------+---------------------+-------
 96244c26-4778-47aa-a52e-958a43aa1f97 | 2025-01-02 06:45:06 | test
 656c428f-04ee-4851-b8bf-b29e0423f1dd | 2025-01-02 07:00:13 | test2
(2 rows)
```
正常にエンティティが追加されている。
また削除も試してみる。
```sh
http GET :8000/echo/ip \
    mykey-del:96244c26-4778-47aa-a52e-958a43aa1f97
```
テーブルの内容は以下となった。
```sh
                  id                  |     created_at      | mykey
--------------------------------------+---------------------+-------
 656c428f-04ee-4851-b8bf-b29e0423f1dd | 2025-01-02 07:00:13 | test2
(1 row)
```
問題なく操作できた。

## まとめ
無事カスタムプラグインでDAOを定義してエンティティを追加することが出来た。
参考となる文献が少ないのがややネックだが、公開されているKongのPluginのソースを参考にすれば実装で悩むことはあまりないと思われる。
データを残したいようなケースでいちいちDBを用意したくない時とかなどに使ってみると良さそうだ。