fix(migrate): buffer exports through file

sanity-io · Jan 30, 2024 · dd47e2b · dd47e2b
1 parent ed7c904
commit dd47e2b
Show file tree

Hide file tree

Showing 6 changed files with 41 additions and 47 deletions.
diff --git a/packages/@sanity/migrate/src/fetch-utils/fetchStream.ts b/packages/@sanity/migrate/src/fetch-utils/fetchStream.ts
@@ -24,9 +24,13 @@ export async function assert2xx(res: Response): Promise<void> {
   }
 }
 
-export async function fetchAsyncIterator({url, init}: FetchOptions) {
+export async function fetchStream({url, init}: FetchOptions) {
   const response = await fetch(url, init)
   await assert2xx(response)
   if (response.body === null) throw new Error('No response received')
-  return streamAsyncIterator(response.body)
+  return response.body
+}
+
+export async function fetchAsyncIterator(options: FetchOptions) {
+  return streamAsyncIterator(await fetchStream(options))
 }
diff --git a/packages/@sanity/migrate/src/runner/dryRun.ts b/packages/@sanity/migrate/src/runner/dryRun.ts
@@ -3,25 +3,31 @@ import {SanityDocument} from '@sanity/types'
 import {APIConfig, Migration} from '../types'
 import {ndjson} from '../it-utils/ndjson'
 import {fromExportEndpoint, safeJsonParser} from '../sources/fromExportEndpoint'
+import {streamAsyncIterator} from '../utils/streamToAsyncIterator'
+import {bufferThroughFile} from '../fs-webstream/bufferThroughFile'
 import {collectMigrationMutations} from './collectMigrationMutations'
+import {getBufferFilePath} from './utils/getBufferFile'
 
 interface MigrationRunnerOptions {
   api: APIConfig
 }
 
 export async function* dryRun(config: MigrationRunnerOptions, migration: Migration) {
+  const exportStream = bufferThroughFile(
+    await fromExportEndpoint({...config.api, documentTypes: migration.documentTypes}),
+    getBufferFilePath(),
+  )
+
   const mutations = collectMigrationMutations(
     migration,
-    ndjson<SanityDocument>(
-      await fromExportEndpoint({...config.api, documentTypes: migration.documentTypes}),
-      {
-        parse: safeJsonParser,
-      },
-    ),
+    ndjson<SanityDocument>(streamAsyncIterator(exportStream), {
+      parse: safeJsonParser,
+    }),
   )
 
   for await (const mutation of mutations) {
     if (!mutation) continue
     yield CompactFormatter.format(Array.isArray(mutation) ? mutation : [mutation])
   }
+  await exportStream.cancel()
 }
diff --git a/packages/@sanity/migrate/src/runner/run.ts b/packages/@sanity/migrate/src/runner/run.ts
@@ -13,6 +13,8 @@ import {lastValueFrom} from '../it-utils/lastValueFrom'
 import {decodeText, parseJSON} from '../it-utils'
 import {concatStr} from '../it-utils/concatStr'
 import {fetchAsyncIterator, FetchOptions} from '../fetch-utils/fetchStream'
+import {bufferThroughFile} from '../fs-webstream/bufferThroughFile'
+import {streamAsyncIterator} from '../utils/streamToAsyncIterator'
 import {toSanityMutations} from './utils/toSanityMutations'
 import {
   DEFAULT_MUTATION_CONCURRENCY,
@@ -21,6 +23,7 @@ import {
 } from './constants'
 import {batchMutations} from './utils/batchMutations'
 import {collectMigrationMutations} from './collectMigrationMutations'
+import {getBufferFilePath} from './utils/getBufferFile'
 
 type MigrationProgress = {
   documents: number
@@ -61,11 +64,14 @@ export async function run(config: MigrationRunnerOptions, migration: Migration)
     completedTransactions: [],
     currentMutations: [],
   }
+
+  const exportStream = bufferThroughFile(
+    await fromExportEndpoint({...config.api, documentTypes: migration.documentTypes}),
+    getBufferFilePath(),
+  )
+
   const documents = tap(
-    ndjson<SanityDocument>(
-      await fromExportEndpoint({...config.api, documentTypes: migration.documentTypes}),
-      {parse: safeJsonParser},
-    ),
+    ndjson<SanityDocument>(streamAsyncIterator(exportStream), {parse: safeJsonParser}),
     () => {
       config.onProgress?.({...stats, documents: ++stats.documents})
     },
@@ -109,6 +115,8 @@ export async function run(config: MigrationRunnerOptions, migration: Migration)
       ...stats,
     })
   }
+  // Cancel export/buffer stream, it's not needed anymore
+  await exportStream.cancel()
   config.onProgress?.({
     ...stats,
     done: true,

diff --git a/packages/@sanity/migrate/src/runner/utils/getBufferFile.ts b/packages/@sanity/migrate/src/runner/utils/getBufferFile.ts
@@ -0,0 +1,6 @@
+import {tmpdir} from 'node:os'
+import path from 'node:path'
+
+export function getBufferFilePath() {
+  return path.join(tmpdir(), `/export-buffer-${Date.now()}.tmp`)
+}
diff --git a/packages/@sanity/migrate/src/sources/fromExportEndpoint.ts b/packages/@sanity/migrate/src/sources/fromExportEndpoint.ts
@@ -1,14 +1,12 @@
 import {createSafeJsonParser} from '@sanity/util/createSafeJsonParser'
 import {SanityDocument} from '@sanity/types'
-import {fetchAsyncIterator} from '../fetch-utils/fetchStream'
+import {fetchStream} from '../fetch-utils/fetchStream'
 import {toFetchOptions} from '../fetch-utils/sanityRequestOptions'
 import {endpoints} from '../fetch-utils/endpoints'
 import {ExportAPIConfig} from '../types'
 
-export function fromExportEndpoint(
-  options: ExportAPIConfig,
-): Promise<AsyncGenerator<Uint8Array, void, unknown>> {
-  return fetchAsyncIterator(
+export function fromExportEndpoint(options: ExportAPIConfig) {
+  return fetchStream(
     toFetchOptions({
       projectId: options.projectId,
       apiVersion: options.apiVersion,

diff --git a/packages/sanity/src/_internal/cli/commands/migration/runMigrationCommand.ts b/packages/sanity/src/_internal/cli/commands/migration/runMigrationCommand.ts
@@ -2,13 +2,11 @@ import path from 'path'
 import type {CliCommandContext, CliCommandDefinition} from '@sanity/cli'
 import {register} from 'esbuild-register/dist/node'
 import {
-  APIConfig,
   collectMigrationMutations,
+  DEFAULT_MUTATION_CONCURRENCY,
+  dryRun,
   fromExportArchive,
-  fromExportEndpoint,
-  safeJsonParser,
   MAX_MUTATION_CONCURRENCY,
-  DEFAULT_MUTATION_CONCURRENCY,
   Migration,
   ndjson,
   run,
@@ -227,30 +225,4 @@ async function runFromArchive(
 
   output.print('Done!')
 }
-
-interface MigrationRunnerOptions {
-  api: APIConfig
-}
-
-async function* dryRun(
-  config: MigrationRunnerOptions,
-  migration: Migration,
-  {chalk}: CliCommandContext,
-) {
-  const mutations = collectMigrationMutations(
-    migration,
-    ndjson<SanityDocument>(
-      await fromExportEndpoint({...config.api, documentTypes: migration.documentTypes}),
-      {
-        parse: safeJsonParser,
-      },
-    ),
-  )
-
-  for await (const mutation of mutations) {
-    if (!mutation) continue
-    yield format(chalk, Array.isArray(mutation) ? mutation : ([mutation] as any))
-  }
-}
-
 export default createMigrationCommand