@tldraw/store 4.3.0-next.f4772c19540d → 4.4.0-canary.6f91153ede5e

package/src/lib/StoreSchema.ts

@@ -1,13 +1,14 @@
 import {
-	Result,
 	assert,
 	exhaustiveSwitchError,
 	getOwnProperty,
+	isEqual,
+	objectMapEntries,
+	Result,
 	structuredClone,
 } from '@tldraw/utils'
 import { UnknownRecord } from './BaseRecord'
-import { RecordType } from './RecordType'
-import { SerializedStore, Store, StoreSnapshot } from './Store'
+import { devFreeze } from './devFreeze'
 import {
 	Migration,
 	MigrationFailureReason,
@@ -16,8 +17,11 @@ import {
 	MigrationSequence,
 	parseMigrationId,
 	sortMigrations,
+	SynchronousStorage,
 	validateMigrations,
 } from './migrate'
+import { RecordType } from './RecordType'
+import { SerializedStore, Store, StoreSnapshot } from './Store'
 
 /**
  * Version 1 format for serialized store schema information.
@@ -530,7 +534,7 @@ export class StoreSchema<R extends UnknownRecord, P = unknown> {
 			return { type: 'success', value: record }
 		}
 
-		if (migrationsToApply.some((m) => m.scope === 'store')) {
+		if (!migrationsToApply.every((m) => m.scope === 'record')) {
 			return {
 				type: 'error',
 				reason:
@@ -541,7 +545,7 @@ export class StoreSchema<R extends UnknownRecord, P = unknown> {
 		}
 
 		if (direction === 'down') {
-			if (!migrationsToApply.every((m) => m.down)) {
+			if (!migrationsToApply.every((m) => m.scope === 'record' && m.down)) {
 				return {
 					type: 'error',
 					reason: MigrationFailureReason.TargetVersionTooOld,
@@ -554,6 +558,7 @@ export class StoreSchema<R extends UnknownRecord, P = unknown> {
 		try {
 			for (const migration of migrationsToApply) {
 				if (migration.scope === 'store') throw new Error(/* won't happen, just for TS */)
+				if (migration.scope === 'storage') throw new Error(/* won't happen, just for TS */)
 				const shouldApply = migration.filter ? migration.filter(record) : true
 				if (!shouldApply) continue
 				const result = migration[direction]!(record)
@@ -569,6 +574,71 @@ export class StoreSchema<R extends UnknownRecord, P = unknown> {
 		return { type: 'success', value: record }
 	}
 
+	migrateStorage(storage: SynchronousStorage<R>) {
+		const schema = storage.getSchema()
+		assert(schema, 'Schema is missing.')
+
+		const migrations = this.getMigrationsSince(schema)
+		if (!migrations.ok) {
+			console.error('Error migrating store', migrations.error)
+			throw new Error(migrations.error)
+		}
+		const migrationsToApply = migrations.value
+		if (migrationsToApply.length === 0) {
+			return
+		}
+
+		storage.setSchema(this.serialize())
+
+		for (const migration of migrationsToApply) {
+			if (migration.scope === 'record') {
+				// Collect updates during iteration, then apply them after.
+				// This avoids issues with live iterators (e.g., SQLite) where updating
+				// records during iteration can cause them to be visited multiple times.
+				const updates: [string, R][] = []
+				for (const [id, state] of storage.entries()) {
+					const shouldApply = migration.filter ? migration.filter(state) : true
+					if (!shouldApply) continue
+					const record = structuredClone(state)
+					const result = migration.up!(record as any) ?? record
+					if (!isEqual(result, state)) {
+						updates.push([id, result as R])
+					}
+				}
+				for (const [id, record] of updates) {
+					storage.set(id, record)
+				}
+			} else if (migration.scope === 'store') {
+				// legacy
+				const prevStore = Object.fromEntries(storage.entries())
+				let nextStore = structuredClone(prevStore)
+				nextStore = (migration.up!(nextStore) as any) ?? nextStore
+				for (const [id, state] of Object.entries(nextStore)) {
+					if (!state) continue // these will be deleted in the next loop
+					if (!isEqual(state, prevStore[id])) {
+						storage.set(id, state)
+					}
+				}
+				for (const id of Object.keys(prevStore)) {
+					if (!nextStore[id]) {
+						storage.delete(id)
+					}
+				}
+			} else if (migration.scope === 'storage') {
+				migration.up!(storage)
+			} else {
+				exhaustiveSwitchError(migration)
+			}
+		}
+		// Clean up by filtering out any non-document records.
+		// This is mainly legacy support for extremely early days tldraw.
+		for (const [id, state] of storage.entries()) {
+			if (this.getType(state.typeName).scope !== 'document') {
+				storage.delete(id)
+			}
+		}
+	}
+
 	/**
 	 * Migrates an entire store snapshot to match the current schema version.
 	 *
@@ -604,7 +674,6 @@ export class StoreSchema<R extends UnknownRecord, P = unknown> {
 		snapshot: StoreSnapshot<R>,
 		opts?: { mutateInputStore?: boolean }
 	): MigrationResult<SerializedStore<R>> {
-		let { store } = snapshot
 		const migrations = this.getMigrationsSince(snapshot.schema)
 		if (!migrations.ok) {
 			// TODO: better error
@@ -613,39 +682,37 @@ export class StoreSchema<R extends UnknownRecord, P = unknown> {
 		}
 		const migrationsToApply = migrations.value
 		if (migrationsToApply.length === 0) {
-			return { type: 'success', value: store }
-		}
-
-		if (!opts?.mutateInputStore) {
-			store = structuredClone(store)
+			return { type: 'success', value: snapshot.store }
 		}
-
+		const store = Object.assign(
+			new Map<string, R>(objectMapEntries(snapshot.store).map(devFreeze)),
+			{
+				getSchema: () => snapshot.schema,
+				setSchema: (_: SerializedSchema) => {},
+			}
+		)
 		try {
-			for (const migration of migrationsToApply) {
-				if (migration.scope === 'record') {
-					for (const [id, record] of Object.entries(store)) {
-						const shouldApply = migration.filter ? migration.filter(record as UnknownRecord) : true
-						if (!shouldApply) continue
-						const result = migration.up!(record as any)
-						if (result) {
-							store[id as keyof typeof store] = result as any
-						}
-					}
-				} else if (migration.scope === 'store') {
-					const result = migration.up!(store)
-					if (result) {
-						store = result as any
+			this.migrateStorage(store)
+			if (opts?.mutateInputStore) {
+				for (const [id, record] of store.entries()) {
+					snapshot.store[id as keyof typeof snapshot.store] = record
+				}
+				for (const id of Object.keys(snapshot.store)) {
+					if (!store.has(id)) {
+						delete snapshot.store[id as keyof typeof snapshot.store]
 					}
-				} else {
-					exhaustiveSwitchError(migration)
+				}
+				return { type: 'success', value: snapshot.store }
+			} else {
+				return {
+					type: 'success',
+					value: Object.fromEntries(store.entries()) as SerializedStore<R>,
 				}
 			}
 		} catch (e) {
 			console.error('Error migrating store', e)
 			return { type: 'error', reason: MigrationFailureReason.MigrationError }
 		}
-
-		return { type: 'success', value: store }
 	}
 
 	/**

package/src/lib/StoreSideEffects.ts

@@ -493,7 +493,7 @@ export class StoreSideEffects<R extends UnknownRecord> {
 	 * ```ts
 	 * editor.sideEffects.registerAfterCreateHandler('page', (page, source) => {
 	 *     // Automatically create a shape when a page is created
-	 *     editor.createShape<TLTextShape>({
+	 *     editor.createShape({
 	 *         id: createShapeId(),
 	 *         type: 'text',
 	 *         props: { richText: toRichText(page.name) },

package/src/lib/devFreeze.test.ts

@@ -1,8 +1,12 @@
-import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
+import { afterAll, beforeEach, describe, expect, it, MockedFunction, vi } from 'vitest'
 import { devFreeze } from './devFreeze'
+import { isDev } from './isDev'
 
 // Mock process.env for testing
 const originalEnv = process.env.NODE_ENV
+vi.mock('./isDev', () => ({
+	isDev: vi.fn(() => true),
+}))
 
 describe('devFreeze', () => {
 	beforeEach(() => {
@@ -13,7 +17,7 @@ describe('devFreeze', () => {
 	describe('production mode behavior', () => {
 		beforeEach(() => {
 			// Mock production environment
-			vi.stubGlobal('process', { env: { NODE_ENV: 'production' } })
+			;(isDev as MockedFunction<typeof isDev>).mockReturnValue(false)
 		})
 
 		it('should return objects unchanged in production mode', () => {

package/src/lib/devFreeze.ts

@@ -1,4 +1,5 @@
 import { STRUCTURED_CLONE_OBJECT_PROTOTYPE } from '@tldraw/utils'
+import { isDev } from './isDev'
 
 /**
  * Freeze an object when in development mode. Copied from
@@ -15,9 +16,8 @@ import { STRUCTURED_CLONE_OBJECT_PROTOTYPE } from '@tldraw/utils'
  * @public
  */
 export function devFreeze<T>(object: T): T {
-	if (process.env.NODE_ENV === 'production') {
-		return object
-	}
+	if (!isDev()) return object
+
 	const proto = Object.getPrototypeOf(object)
 	if (
 		proto &&
@@ -32,6 +32,10 @@ export function devFreeze<T>(object: T): T {
 		throw new Error('cannot include non-js data in a record')
 	}
 
+	if (Object.isFrozen(object)) {
+		return object
+	}
+
 	// Retrieve the property names defined on object
 	const propNames = Object.getOwnPropertyNames(object)
 

package/src/lib/isDev.ts

@@ -0,0 +1,20 @@
+let _isDev = false
+try {
+	_isDev = process.env.NODE_ENV === 'development' || process.env.NODE_ENV === 'test'
+} catch (_e) {
+	/* noop */
+}
+try {
+	_isDev =
+		_isDev ||
+		(import.meta as any).env.DEV ||
+		(import.meta as any).env.TEST ||
+		(import.meta as any).env.MODE === 'development' ||
+		(import.meta as any).env.MODE === 'test'
+} catch (_e) {
+	/* noop */
+}
+
+export function isDev() {
+	return _isDev
+}

package/src/lib/migrate.ts

@@ -1,6 +1,7 @@
 import { assert, objectMapEntries } from '@tldraw/utils'
 import { UnknownRecord } from './BaseRecord'
 import { SerializedStore } from './Store'
+import { SerializedSchema } from './StoreSchema'
 
 function squashDependsOn(sequence: Array<Migration | StandaloneDependsOn>): Migration[] {
 	const result: Migration[] = []
@@ -219,8 +220,36 @@ export type Migration = {
 				newState: SerializedStore<UnknownRecord>
 			) => void | SerializedStore<UnknownRecord>
 	  }
+	| {
+			readonly scope: 'storage'
+			// eslint-disable-next-line @typescript-eslint/method-signature-style
+			readonly up: (storage: SynchronousRecordStorage<UnknownRecord>) => void
+			readonly down?: never
+	  }
 )
 
+/**
+ * Abstraction over the store that can be used to perform migrations.
+ * @public
+ */
+export interface SynchronousRecordStorage<R extends UnknownRecord> {
+	get(id: string): R | undefined
+	set(id: string, record: R): void
+	delete(id: string): void
+	keys(): Iterable<string>
+	values(): Iterable<R>
+	entries(): Iterable<[string, R]>
+}
+
+/**
+ * Abstraction over the storage that can be used to perform migrations.
+ * @public
+ */
+export interface SynchronousStorage<R extends UnknownRecord> extends SynchronousRecordStorage<R> {
+	getSchema(): SerializedSchema
+	setSchema(schema: SerializedSchema): void
+}
+
 /**
  * Base interface for legacy migration information.
  *

package/src/lib/test/recordStore.test.ts

@@ -883,6 +883,188 @@ describe('snapshots', () => {
 		expect(up).toHaveBeenCalledTimes(1)
 		expect(store2.get(Book.createId('lotr'))!.numPages).toBe(42)
 	})
+
+	it('migrates the snapshot with storage scope', () => {
+		const snapshot1 = store.getStoreSnapshot()
+		const up = vi.fn((storage: any) => {
+			const book = storage.get('book:lotr')
+			storage.set('book:lotr', { ...book, numPages: 42 })
+		})
+
+		expect((snapshot1.store as any)['book:lotr'].numPages).toBe(1000)
+
+		const store2 = new Store({
+			props: {},
+			schema: StoreSchema.create<Book | Author>(
+				{
+					book: Book,
+					author: Author,
+				},
+				{
+					migrations: [
+						createMigrationSequence({
+							sequenceId: 'com.tldraw',
+							retroactive: true,
+							sequence: [
+								{
+									id: `com.tldraw/1`,
+									scope: 'storage',
+									up,
+								},
+							],
+						}),
+					],
+				}
+			),
+		})
+
+		expect(() => {
+			store2.loadStoreSnapshot(snapshot1)
+		}).not.toThrow()
+
+		expect(up).toHaveBeenCalledTimes(1)
+		expect(store2.get(Book.createId('lotr'))!.numPages).toBe(42)
+	})
+
+	it('storage scope migration can delete records', () => {
+		const snapshot1 = store.getStoreSnapshot()
+		const up = vi.fn((storage: any) => {
+			storage.delete('author:mcavoy')
+		})
+
+		expect((snapshot1.store as any)['author:mcavoy']).toBeDefined()
+
+		const store2 = new Store({
+			props: {},
+			schema: StoreSchema.create<Book | Author>(
+				{
+					book: Book,
+					author: Author,
+				},
+				{
+					migrations: [
+						createMigrationSequence({
+							sequenceId: 'com.tldraw',
+							retroactive: true,
+							sequence: [
+								{
+									id: `com.tldraw/1`,
+									scope: 'storage',
+									up,
+								},
+							],
+						}),
+					],
+				}
+			),
+		})
+
+		expect(() => {
+			store2.loadStoreSnapshot(snapshot1)
+		}).not.toThrow()
+
+		expect(up).toHaveBeenCalledTimes(1)
+		expect(store2.get(Author.createId('mcavoy'))).toBeUndefined()
+	})
+
+	it('storage scope migration can iterate records', () => {
+		const snapshot1 = store.getStoreSnapshot()
+		const up = vi.fn((storage: any) => {
+			for (const [id, record] of storage.entries()) {
+				if (record.typeName === 'book') {
+					storage.set(id, { ...record, numPages: record.numPages + 100 })
+				}
+			}
+		})
+
+		expect((snapshot1.store as any)['book:lotr'].numPages).toBe(1000)
+		expect((snapshot1.store as any)['book:hobbit'].numPages).toBe(300)
+
+		const store2 = new Store({
+			props: {},
+			schema: StoreSchema.create<Book | Author>(
+				{
+					book: Book,
+					author: Author,
+				},
+				{
+					migrations: [
+						createMigrationSequence({
+							sequenceId: 'com.tldraw',
+							retroactive: true,
+							sequence: [
+								{
+									id: `com.tldraw/1`,
+									scope: 'storage',
+									up,
+								},
+							],
+						}),
+					],
+				}
+			),
+		})
+
+		expect(() => {
+			store2.loadStoreSnapshot(snapshot1)
+		}).not.toThrow()
+
+		expect(up).toHaveBeenCalledTimes(1)
+		expect(store2.get(Book.createId('lotr'))!.numPages).toBe(1100)
+		expect(store2.get(Book.createId('hobbit'))!.numPages).toBe(400)
+	})
+
+	it('storage scope migration can use values() and keys()', () => {
+		const snapshot1 = store.getStoreSnapshot()
+		const keysCollected: string[] = []
+		const valuesCollected: any[] = []
+
+		const up = vi.fn((storage: any) => {
+			for (const key of storage.keys()) {
+				keysCollected.push(key)
+			}
+			for (const value of storage.values()) {
+				valuesCollected.push(value)
+			}
+		})
+
+		const store2 = new Store({
+			props: {},
+			schema: StoreSchema.create<Book | Author>(
+				{
+					book: Book,
+					author: Author,
+				},
+				{
+					migrations: [
+						createMigrationSequence({
+							sequenceId: 'com.tldraw',
+							retroactive: true,
+							sequence: [
+								{
+									id: `com.tldraw/1`,
+									scope: 'storage',
+									up,
+								},
+							],
+						}),
+					],
+				}
+			),
+		})
+
+		expect(() => {
+			store2.loadStoreSnapshot(snapshot1)
+		}).not.toThrow()
+
+		expect(up).toHaveBeenCalledTimes(1)
+		expect(keysCollected).toContain('book:lotr')
+		expect(keysCollected).toContain('book:hobbit')
+		expect(keysCollected).toContain('author:tolkein')
+		expect(keysCollected).toContain('author:mcavoy')
+		expect(keysCollected).toContain('author:cassidy')
+		expect(valuesCollected.length).toBe(5)
+	})
 })
 
 describe('diffs', () => {