Skip to content

Commit

Permalink
fix(migrate): limit client methods during migration
Browse files Browse the repository at this point in the history
  • Loading branch information
bjoerge committed Jan 30, 2024
1 parent 5b690a8 commit da7e0ad
Show file tree
Hide file tree
Showing 5 changed files with 78 additions and 12 deletions.
5 changes: 2 additions & 3 deletions packages/@sanity/migrate/src/runner/dryRun.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
import {SanityDocument} from '@sanity/types'
import {createClient} from '@sanity/client'
import {APIConfig, Migration, MigrationProgress} from '../types'
import {fromExportEndpoint, safeJsonParser} from '../sources/fromExportEndpoint'
import {streamToAsyncIterator} from '../utils/streamToAsyncIterator'
Expand All @@ -11,7 +10,7 @@ import {collectMigrationMutations} from './collectMigrationMutations'
import {getBufferFilePath} from './utils/getBufferFile'
import {createFilteredDocumentsClient} from './utils/createFilteredDocumentsClient'
import {applyFilters} from './utils/applyFilters'
import {limitClientConcurrency} from './utils/limitClientConcurrency'
import {createContextClient} from './utils/createContextClient'

interface MigrationRunnerOptions {
api: APIConfig
Expand Down Expand Up @@ -40,7 +39,7 @@ export async function* dryRun(config: MigrationRunnerOptions, migration: Migrati
)

// Create a client exposed to the migration script. This will have a max concurrency of 10
const client = limitClientConcurrency(createClient({...config.api, useCdn: false}))
const client = createContextClient({...config.api, useCdn: false})

const filteredDocumentsClient = createFilteredDocumentsClient(createReader)
const context = {
Expand Down
12 changes: 7 additions & 5 deletions packages/@sanity/migrate/src/runner/run.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import {SanityDocument} from '@sanity/types'
import {createClient, MultipleMutationResult} from '@sanity/client'
import {MultipleMutationResult} from '@sanity/client'
import arrify from 'arrify'
import {APIConfig, Migration, MigrationProgress} from '../types'
import {parse, stringify} from '../it-utils/ndjson'
Expand All @@ -26,7 +26,7 @@ import {collectMigrationMutations} from './collectMigrationMutations'
import {getBufferFilePath} from './utils/getBufferFile'
import {createFilteredDocumentsClient} from './utils/createFilteredDocumentsClient'
import {applyFilters} from './utils/applyFilters'
import {limitClientConcurrency} from './utils/limitClientConcurrency'
import {createContextClient} from './utils/createContextClient'

export interface MigrationRunnerConfig {
api: APIConfig
Expand Down Expand Up @@ -80,9 +80,11 @@ export async function run(config: MigrationRunnerConfig, migration: Migration) {
{signal: abortController.signal},
)

const client = limitClientConcurrency(
createClient({...config.api, useCdn: false, requestTagPrefix: 'sanity.migration'}),
)
const client = createContextClient({
...config.api,
useCdn: false,
requestTagPrefix: 'sanity.migration',
})

const filteredDocumentsClient = createFilteredDocumentsClient(createReader)
const context = {
Expand Down
5 changes: 2 additions & 3 deletions packages/@sanity/migrate/src/runner/runFromArchive.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
import {SanityDocument} from '@sanity/types'
import arrify from 'arrify'
import {createClient} from '@sanity/client'
import {Migration, MigrationProgress} from '../types'
import {decodeText} from '../it-utils'
import {fromExportArchive} from '../sources/fromExportArchive'
Expand All @@ -23,7 +22,7 @@ import {getBufferFilePath} from './utils/getBufferFile'
import {collectMigrationMutations} from './collectMigrationMutations'
import {MigrationRunnerConfig, toFetchOptionsIterable} from './run'
import {applyFilters} from './utils/applyFilters'
import {limitClientConcurrency} from './utils/limitClientConcurrency'
import {createContextClient} from './utils/createContextClient'

export async function runFromArchive(
migration: Migration,
Expand Down Expand Up @@ -62,7 +61,7 @@ export async function runFromArchive(
},
)

const client = limitClientConcurrency(createClient({...config.api, useCdn: false}))
const client = createContextClient({...config.api, useCdn: false})

const filteredDocumentsClient = createFilteredDocumentsClient(createReader)
const context = {
Expand Down
65 changes: 65 additions & 0 deletions packages/@sanity/migrate/src/runner/utils/createContextClient.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
import {createClient, type SanityClient} from '@sanity/client'
import {limitClientConcurrency} from './limitClientConcurrency'

export function createContextClient(config: Parameters<typeof createClient>[0]): RestrictedClient {
return restrictClient(
limitClientConcurrency(
createClient({...config, useCdn: false, requestTagPrefix: 'sanity.migration'}),
),
)
}

const ALLOWED_PROPERTIES = [
'fetch',
'clone',
'config',
'withConfig',
'getDocument',
'getDocuments',
'users',
'projects',
] as const

type AllowedMethods = (typeof ALLOWED_PROPERTIES)[number]

export type RestrictedClient = Pick<SanityClient, AllowedMethods>

function restrictClient(client: SanityClient): RestrictedClient {
return new Proxy(client, {
get: (target, property) => {
switch (property) {
case 'clone': {
return (...args: Parameters<SanityClient['clone']>) => {
return restrictClient(target.clone(...args))
}
}
case 'config': {
return (...args: Parameters<SanityClient['config']>) => {
const result = target.config(...args)

// if there is a config, it returns a client so we need to wrap again
if (args[0]) return restrictClient(result)
return result
}
}
case 'withConfig': {
return (...args: Parameters<SanityClient['withConfig']>) => {
return restrictClient(target.withConfig(...args))
}
}
default: {
if (ALLOWED_PROPERTIES.includes(property as any)) {
return target[property as keyof SanityClient]
}
throw new Error(
`Client method "${String(
property,
)}" can not be called during a migration. Only ${ALLOWED_PROPERTIES.join(
', ',
)} are allowed.`,
)
}
}
},
})
}
3 changes: 2 additions & 1 deletion packages/@sanity/migrate/src/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ import type {Path, SanityDocument} from '@sanity/types'
import {MultipleMutationResult, Mutation as RawMutation, SanityClient} from '@sanity/client'
import {JsonArray, JsonObject, JsonValue} from './json'
import {Mutation, NodePatch, Operation, Transaction} from './mutations'
import {RestrictedClient} from './runner/utils/createContextClient'

export type {Path}
export type * from './json'
Expand Down Expand Up @@ -47,7 +48,7 @@ export type MigrationProgress = {
}

export interface MigrationContext {
client: SanityClient
client: RestrictedClient
filtered: {
getDocument<T extends SanityDocument>(id: string): Promise<T | undefined>
getDocuments<T extends SanityDocument>(ids: string[]): Promise<T[]>
Expand Down

0 comments on commit da7e0ad

Please sign in to comment.