DynamoDB copy tool

main.ts

import {
  DynamoDBClient,
  DescribeTableCommand,
  CreateTableCommand,
  ScanCommand,
  PutItemCommand,
  ListTablesCommand,
  waitUntilTableExists,
  ScanCommandOutput,
  UpdateTableCommand,
} from "@aws-sdk/client-dynamodb";
import { WaiterConfiguration } from "@smithy/util-waiter";

// コピー元のAWSアカウント
const sourceClient = new DynamoDBClient({
  profile: process.env.SRC_AWS_PROFILE,
  region: "ap-northeast-1",
});

// コピー先のAWSアカウント
const destinationClient = new DynamoDBClient({
  profile: process.env.DEST_AWS_PROFILE,
  region: "ap-northeast-1",
});

async function copyTable(
  sourceTableName: string,
  destinationTableName: string
) {
  // テーブルのスキーマを取得
  const sourceTable = await sourceClient.send(
    new DescribeTableCommand({ TableName: sourceTableName })
  );
  const tableSchema = sourceTable.Table;
  if (!tableSchema) return;

  // 新しいテーブルを作成
  await destinationClient.send(
    new CreateTableCommand({
      TableName: destinationTableName,
      KeySchema: tableSchema.KeySchema,
      AttributeDefinitions:
        tableSchema.KeySchema?.length ===
        tableSchema.AttributeDefinitions?.length
          ? tableSchema.AttributeDefinitions
          : [{ AttributeName: "id", AttributeType: "S" }],
      BillingMode: "PAY_PER_REQUEST",
      OnDemandThroughput: {
        MaxReadRequestUnits: 1000,
        MaxWriteRequestUnits: 1000,
      },
    })
  );

  // テーブルがアクティブになるまで待機
  const waiterConfig: WaiterConfiguration<DynamoDBClient> = {
    client: destinationClient,
    maxWaitTime: 120,
  };
  await waitUntilTableExists(waiterConfig, { TableName: destinationTableName });

  // データをコピー
  let lastEvaluatedKey: Record<string, any> | undefined = undefined;

  do {
    const response: ScanCommandOutput = await sourceClient.send(
      new ScanCommand({
        TableName: sourceTableName,
        ExclusiveStartKey: lastEvaluatedKey,
      })
    );
    const items = response.Items || [];
    lastEvaluatedKey = response.LastEvaluatedKey;

    for (const item of items) {
      await destinationClient.send(
        new PutItemCommand({
          TableName: destinationTableName,
          Item: item,
        })
      );
    }
  } while (lastEvaluatedKey);
}

async function modifyTable(
  sourceTableName: string,
  destinationTableName: string
) {
  // テーブルのスキーマを取得
  const sourceTable = await sourceClient.send(
    new DescribeTableCommand({ TableName: sourceTableName })
  );
  const tableSchema = sourceTable.Table;
  if (!tableSchema) return;

  await destinationClient.send(
    new UpdateTableCommand({
      TableName: destinationTableName,
      BillingMode: tableSchema.BillingModeSummary?.BillingMode || "PROVISIONED",
      OnDemandThroughput:
        tableSchema.OnDemandThroughput &&
        tableSchema.OnDemandThroughput.MaxWriteRequestUnits &&
        tableSchema.OnDemandThroughput.MaxWriteRequestUnits > 0
          ? {
              MaxReadRequestUnits:
                tableSchema.OnDemandThroughput.MaxReadRequestUnits,
              MaxWriteRequestUnits:
                tableSchema.OnDemandThroughput.MaxWriteRequestUnits,
            }
          : undefined,
      ProvisionedThroughput:
        tableSchema.ProvisionedThroughput &&
        tableSchema.ProvisionedThroughput.WriteCapacityUnits &&
        tableSchema.ProvisionedThroughput.WriteCapacityUnits > 0
          ? {
              ReadCapacityUnits:
                tableSchema.ProvisionedThroughput.ReadCapacityUnits,
              WriteCapacityUnits:
                tableSchema.ProvisionedThroughput.WriteCapacityUnits,
            }
          : undefined,
    })
  );
}

async function verifyTableData(
  sourceTableName: string,
  destinationTableName: string
): Promise<boolean> {
  const sourceCount = await getTableItemCount(sourceClient, sourceTableName);
  const destinationCount = await getTableItemCount(
    destinationClient,
    destinationTableName
  );

  console.log(`Source table (${sourceTableName}) item count: ${sourceCount}`);
  console.log(
    `Destination table (${destinationTableName}) item count: ${destinationCount}`
  );

  return sourceCount === destinationCount;
}

async function getTableItemCount(
  client: DynamoDBClient,
  tableName: string
): Promise<number> {
  let itemCount = 0;
  let lastEvaluatedKey: Record<string, any> | undefined = undefined;

  do {
    const response: ScanCommandOutput = await client.send(
      new ScanCommand({
        TableName: tableName,
        Select: "COUNT",
        ExclusiveStartKey: lastEvaluatedKey,
      })
    );
    itemCount += response.Count || 0;
    lastEvaluatedKey = response.LastEvaluatedKey;
  } while (lastEvaluatedKey);

  return itemCount;
}

async function main() {
  try {
    // ソースアカウントのテーブルリストを取得
    const response = await sourceClient.send(new ListTablesCommand({}));
    const tableNames = response.TableNames || [];

    for (const tableName of tableNames.filter((name) => name !== "Customer")) {
      console.log(`Copying table ${tableName}...`);
      await copyTable(tableName, tableName); // 同じ名前でコピー

      // データサイズが小さいテーブルのみの場合はコメントアウトを解除
      // const isDataVerified = await verifyTableData(tableName, tableName);
      // if (isDataVerified) {
      //   console.log(
      //     `Table ${tableName} copied successfully and data verified.`
      //   );
      // } else {
      //   console.error(`Data verification failed for table ${tableName}.`);
      // }

      await modifyTable(tableName, tableName); // キャパシティモードを更新
    }

    console.log("All tables copied and verified successfully.");
  } catch (e) {
    console.error(e);
  }
}

main();