Firestoreからたくさんのデータを取得してたくさんのデータを書き込むのに使えそうな関数

500件以上取得する可能性のあるクエリで取得したドキュメントをコピーしたいときがあったので作りました。一応テストもしたのでちゃんと動きます。

import { firestore } from "firebase-admin"
const db = firestore()

/// batchExecutorに最大500件ずつSnapShotを渡してくれる
/// 注意：
///   queryにはorderつけて渡すこと（query.startAtがうまく動かなさそうなので）
///   await batch.commit() はbatchExecutor内で呼ばないとデータ処理されない
export async function executeInBatch(query: firestore.Query, batchExecutor: (batch: firestore.WriteBatch, snapshot: firestore.QueryDocumentSnapshot[]) => Promise<void>): Promise<void> {
  const BATCH_LIMIT = 500

  let hasNextPage = true
  let lastDocument: firestore.DocumentSnapshot | undefined

  while (hasNextPage) {
    // 1件多めに取得して次がまだあるか判断する
    let q = query.limit(BATCH_LIMIT + 1)
    if (lastDocument) {
      // 前回の続きから取得
      q = q.startAt(lastDocument)
    }

    const snapshot = await q.get()

    // 取得できた件数が指定通りだったら次のデータがまだある
    hasNextPage = snapshot.size === (BATCH_LIMIT + 1)
    // snapshot.docsをpopしても要素が消えないのでコピーした配列に対して操作する
    const docs = snapshot.docs.concat()
    if (hasNextPage) {
      // 続きから取得する用に1件とっておく
      lastDocument = docs.pop()
    }

    // batch処理してもらう
    await batchExecutor(db.batch(), docs)
  }
}

// 使い方
const query = db.collection("users").order("createdAt", "desc")
await executeInBatch(query, async (batch, documents) => {
  documents.forEach(doc => {
    batch.set(db.collection("samples").doc(), { name: doc.name })    
  })
  await batch.commit()
})