Amplify（Gen2）でマルチテナントを実装する

今までのAmplify（Gen1）では、マルチテナントを実装しようとすると普段触りなれないVTLを使ってカスタムリゾルバーを実装する必要がありとても面倒でした。しかしながらAmplify（Gen2）ではカスタムリゾルバーを直接実装することなくこれを実現することができたので、備忘録として記事にします。

ゴール状態

共通するtenantIdを持つユーザー間であれば、データを共有できる。
（逆に、異なるtenantIdのユーザーが作ったデータを見ることができない。）

全体の流れ

1. Amplify（Gen2）のプロジェクトをセットアップする
2. CognitoのCustom AttributeにtenantIdを追加する
3. tenantIdを使ったauth ruleをデータモデルに追加する
4. カスタムヘッダーを実装したGraphQLクライアントでCRUD処理を行う

1. Amplify（Gen2）のプロジェクトをセットアップする

今回は、Vite + ReactでAmplify（Gen2）のプロジェクトを作成します。

npm create vite@latest
✔ Project name: … <プロジェクト名>
✔ Select a framework: › React
✔ Select a variant: › TypeScript
cd <プロジェクト名>
npm install
npm create amplify@latest
npm add --save-dev @aws-amplify/backend@latest @aws-amplify/backend-cli@latest typescript

2. CognitoのCustom AttributeにtenantIdを追加する

テナントを識別するtenantIdをCustom Attributeに追加します。
※tenantIdを設定するロジックは別で実装が必要です。（ここでは触れません）

amplify/backend.ts

import { defineBackend } from '@aws-amplify/backend';
import { auth } from './auth/resource';
import { data } from './data/resource';

- defineBackend({
+ const backend = defineBackend({
    auth,
    data,
});

+ const { cfnUserPool } = backend.auth.resources.cfnResources;
+ if (Array.isArray(cfnUserPool.schema)) {
+    cfnUserPool.schema.push({
+        name: 'tenantId',
+        attributeDataType: 'String',
+        required: false,
+        mutable: true,
+    });
+}

※参考
https://docs.amplify.aws/react/build-a-backend/auth/modify-resources-with-cdk/#custom-attributes
https://docs.aws.amazon.com/cdk/api/v1/docs/@aws-cdk_aws-cognito.CfnUserPool.html

3. tenantIdを使ったauth ruleをデータモデルに追加する

amplify/data/resource.ts

import { type ClientSchema, a, defineData } from '@aws-amplify/backend';
const schema = a.schema({
    Todo: a
        .model({
+           tenantId: a.string(),
            content: a.string()
        })
-        .authorization(allow => [allow.guest()])
+        .authorization(allow => [allow.ownerDefinedIn('tenantId').identityClaim('custom:tenantId')]),
});

export type Schema = ClientSchema<typeof schema>;

export const data = defineData({
    schema,
    authorizationModes: {
-        defaultAuthorizationMode: 'iam',
+        defaultAuthorizationMode: 'userPool',
    },
});

※参考
https://docs.amplify.aws/react/build-a-backend/data/customize-authz/configure-custom-identity-and-group-claim/
https://medium.com/@dantasfiles/multi-tenant-aws-amplify-method-1-cognito-custom-attributes-6719d27d84cc

4. カスタムヘッダーを実装したGraphQLクライアントでCRUD処理を行う

GraphQLのクライアントはデフォルトではAuthenticationヘッダーにaccessTokenを使いますが、accessTokenにはUserAttributeの情報が含まれないので、それが含まれるidTokenを使うように修正します。

import { generateClient } from "aws-amplify/data";
+import { fetchAuthSession } from "aws-amplify/auth";
import type { Schema } from "amplify/data/resource";

const client = generateClient<Schema>({
+    headers: async (requestOptions) => {
+        const session = await fetchAuthSession();
+        return {
+            'Authorization': session.tokens?.idToken?.toString() | '',
+        };
+    }
});

export const listTodos = async () => {
    const response = await client.models.Todo.list({authMode: 'userPool'});
    return response.data;
    //tenantIdが一致するTodoレコードのみが返される
};

※参考
https://docs.amplify.aws/react/build-a-backend/data/connect-to-API/#set-custom-request-headers

感想

カスタムリゾルバーが必要だったGen1と比べてかなり簡単にマルチテナントを実装できるようになったと思います。
Gen2ではGen1から引き続き、データモデルやフィールド別に複数のauth strategyを設定できるので、例えば「テナント内でreadはできるが、editは作った本人だけ」といったことができないか実験してみたいと思います。