Repository: pmndrs/zustand
Pull request #292 (2021-01-30 09:45:12)

Fix persist error forwarding
In current implementation, errors from persist middleware could not be received by the caller. There are chances that errors will happen when reading or writing to storage, so error handling is important. In this PR I changed the way persist works with error forwarding so the caller can catch errors and handle it accordingly.

There are multiple changes I made:

### Tests

Before changing anything, it would be better if we have tests for persist. There are currently another PR with tests for devtools middleware, #233 . I think it would be better if we separate the test file, so I put persist tests in `persist.test.tsx` rather than in `middleware.test.tsx`. Also, to prevent merge conflict with that PR.

### ~~`storage.setItem`~~

~~When the `setItem` function is an async function, such as in AsyncStorage for React Native, the caller would not receive the error because persist middleware call `setItem` without await-ing. Discussion here #287 .~~

~~To solve that, I change all functions that call `storage.setItem` to await or return the promise .~~

### Rehydration

Rehydration is performed asynchronously in store initialization ([middleware.ts#L229-L251](https://github.com/pmndrs/zustand/blob/cf4d1365b66895e6fc776e27ef1f1d25476aed01/src/middleware.ts#L229-L251). Here, the error is thrown, but since it is asynchronous, caller will not receive this exception.

I initially want to change the way `onRehydrateStorage` to be able to send error with the way a callback commonly does: `callback(error, data)`. But it may cause problem (backward incompatibility) because in current implementation the first parameter is the data. If there are already users that utilize this callback, and they expect the first parameter to be the data, it will break their code.

Then I decided to introduce another callback, `onRehydrateError`. This callback will be called if there is any error while performing hydration. The error can be from reading the storage, deserialization, or failed migration (this will be explained below).

### Migration

When studying persist code, I think the way we handle version mismatch is very dangerous. "if versions mismatch, clear storage by storing the new initial state" is basically clear everything and start from new. If I have a store abour user login data, and then I change the storage to a new version, the login data would be cleared and users have to re-login.

Other state persistence libraries (in other state management libraries) commonly use migration to handle this. Migration interface that I introduce here is similar to redux persist. In persist options I introduce `migrate(persistedState, version)`, a function that can be defined by developers to put their migration logic.

codesandbox[bot]: This pull request is automatically built and testable in [CodeSandbox](https://codesandbox.io).

  To see build info of the built libraries, click [here](https://ci.codesandbox.io/status/pmndrs/zustand/pr/292/builds/97008) or the icon next to each commit SHA.

Latest deployment of this branch, based on commit b410f0ecd6128fc7375c595309c18fb884133827:

|Sandbox| Source |
|--|--|
|[React](https://codesandbox.io/s/react-12skb)| Configuration |
|[React Typescript](https://codesandbox.io/s/react-typescript-4ixli)| Configuration |

sumartoyo: >Is there a reason not to do something like onRehydrateStorage(data, error) instead of a new callback?

No reasons other than to follow convention of `cb(error, data)` 😄. Maybe users will be confused because the parameters are reversed.
What do you think? Should I implement it with `onRehydrateStorage(data, error)`? If we document it well we can minimize the confusion.

>It seems like you also need to alter the return type of set and api.setState so typescript doesn't complain about a "useless" await

Yes, I didn't catch that. Thank you for pointing that out. I will fix it.
AnatoleLucet: > What do you think? Should I implement it with onRehydrateStorage(data, error)? If we document it well we can minimize the confusion.

Well, it's true that it breaks conventions, but I think it should be ok now that jsdoc, lsp completion and all that as become mainstream in most editors.
sumartoyo: Hmm I changed return type of `SetState` to `void | Promise<void>` but I'm afraid it may cause problems.

For example, I have to change `act(() => setState({}))` to `act(() => { setState({}) })` because act is confused. It can accept void or Promise but not `void | Promise<void>`.

Is it okay or is there a better way to do it? How do people commonly type a function that can return a promise or not a promise?
AnatoleLucet: It doesn't work because `void | Promise<void>` doesn't match any of `act()`'s overloads.
One way to fix this would be to define two overloads for `SetState`. One for `void` and another one for `Promise<void>`.
E.g.:
```ts
export type SetState<T extends State> = {
  (partial: PartialState<T>, replace?: boolean): void
  (partial: PartialState<T>, replace?: boolean): Promise<void>
}
```
sumartoyo: > One way to fix this would be to define two overloads for `SetState`. One for `void` and another one for `Promise<void>`.

I have tried that and it still has problems. When I try  
`const mySetState: SetState<TState> = (partial: PartialState<TState>) => {}`  
it says "Type 'void' is not assignable to type 'Promise<void>'". It has no error if I set the function as async.
`const mySetState: SetState<TState> = async (partial: PartialState<TState>) => {}`

But then, when I try calling it with  
`await mySetState({})`  
it says "'await' has no effect on the type of this expression".

I think we can't have a function that can returns two different types when they have identical parameters. At least that's how understand it in different langsuage (Java, C#), unless we use generics to the function.

It seems that this is more complex than what I thought. How about we take out things related to `storage.setItem` from this PR, so that the scope of this PR is only about rehydration and migration? I will revert the code related to the `setItem` and we continue our discussion in another discussion?
AnatoleLucet: > It seems that this is more complex than what I thought. How about we take out things related to storage.setItem from this PR, so that the scope of this PR is only about rehydration and migration? I will revert the code related to the setItem and we continue our discussion in another discussion?

Sure. Maybe a new drafted pr will be easier to propose change and discuss the code instead of a discussion.

Patch:
diff --git a/.size-snapshot.json b/.size-snapshot.json
index 96658c4..3faabfd 100644
--- a/.size-snapshot.json
+++ b/.size-snapshot.json
@@ -68,9 +68,9 @@
     "gzipped": 251
   },
   "middleware.js": {
-    "bundled": 7993,
-    "minified": 3568,
-    "gzipped": 1397
+    "bundled": 8192,
+    "minified": 3730,
+    "gzipped": 1464
   },
   "vanilla.js": {
     "bundled": 1825,
diff --git a/readme.md b/readme.md
index dc9e091..8f8e712 100644
--- a/readme.md
+++ b/readme.md
@@ -6,7 +6,7 @@
 [![Build Size](https://img.shields.io/bundlephobia/min/zustand?label=bundle%20size&style=flat&colorA=000000&colorB=000000)](https://bundlephobia.com/result?p=zustand)
 [![Version](https://img.shields.io/npm/v/zustand?style=flat&colorA=000000&colorB=000000)](https://www.npmjs.com/package/zustand)
 [![Downloads](https://img.shields.io/npm/dt/zustand.svg?style=flat&colorA=000000&colorB=000000)](https://www.npmjs.com/package/zustand)
-[![Discord Shield](https://img.shields.io/discord/740090768164651008?style=flat&colorA=000000&colorB=000000&label=discord&logo=discord&logoColor=ffffff)](https://discord.gg/ZZjjNvJ)
+[![Discord Shield](https://img.shields.io/discord/740090768164651008?style=flat&colorA=000000&colorB=000000&label=discord&logo=discord&logoColor=ffffff)](https://discord.gg/poimandres)
 
 A small, fast and scaleable bearbones state-management solution. Has a comfy api based on hooks, isn't boilerplatey or opinionated, but still just enough to be explicit and flux-like.
 
@@ -272,10 +272,8 @@ const useStore = create(
 )
 ```
 
-
-<overview>
-<summary>How to pipe middlewares</summary>
 <details>
+<summary>How to pipe middlewares</summary>
 
 ```js
 import create from "zustand"
@@ -295,11 +293,9 @@ export default useStore
 ```
 For a TS example see the following [discussion](https://github.com/pmndrs/zustand/discussions/224#discussioncomment-118208)
 </details>
-</overview>
 
-<overview>
-<summary>How to type immer middleware in TypeScript</summary>
 <details>
+<summary>How to type immer middleware in TypeScript</summary>
 
 ```ts
 import { State, StateCreator } from 'zustand'
@@ -312,7 +308,6 @@ const immer = <T extends State>(
 ```
 
 </details>
-</overview>
 
 ## Persist middleware
 
@@ -363,6 +358,28 @@ import { redux } from 'zustand/middleware'
 const useStore = create(redux(reducer, initialState))
 ```
 
+## Calling actions outside a React event handler
+
+Because React handles `setState` synchronously if it's called outside an event handler. Updating the state outside an event handler will force react to update the components synchronously, therefore adding the risk of encountering the zombie-child effect.
+In order to fix this, the action needs to be wrapped in `unstable_batchedUpdates`
+
+```jsx
+import { unstable_batchedUpdates } from 'react-dom' // or 'react-native'
+
+const useStore = create((set) => ({
+  fishes: 0,
+  increaseFishes: () => set((prev) => ({ fishes: prev.fishes + 1 }))
+}))
+
+const nonReactCallback = () => {
+  unstable_batchedUpdates(() => {
+    useStore.getState().increaseFishes()
+  })
+}
+```
+
+More details: https://github.com/pmndrs/zustand/issues/302
+
 ## Redux devtools
 
 ```jsx
diff --git a/src/middleware.ts b/src/middleware.ts
index f7df6f5..938b97b 100644
--- a/src/middleware.ts
+++ b/src/middleware.ts
@@ -158,14 +158,19 @@ type PersistOptions<S> = {
   /**
    * A function returning another (optional) function.
    * The main function will be called before the state rehydration.
-   * The returned function will be called after the state rehydration.
+   * The returned function will be called after the state rehydration or when an error occured.
    */
-  onRehydrateStorage?: (state: S) => ((state: S) => void) | void
+  onRehydrateStorage?: (state: S) => ((state?: S, error?: Error) => void) | void
   /**
    * If the stored state's version mismatch the one specified here, the storage will not be used.
    * This is useful when adding a breaking change to your store.
    */
   version?: number
+  /**
+   * A function to perform persisted state migration.
+   * This function will be called when persisted state versions mismatch with the one specified here.
+   */
+  migrate?: (persistedState: any, version: number) => S | Promise<S>
 }
 
 export const persist = <S extends State>(
@@ -181,6 +186,7 @@ export const persist = <S extends State>(
     whitelist,
     onRehydrateStorage,
     version = 0,
+    migrate,
   } = options || {}
 
   let storage: StateStorage | undefined
@@ -216,7 +222,7 @@ export const persist = <S extends State>(
       blacklist.forEach((key) => delete state[key])
     }
 
-    storage?.setItem(name, await serialize({ state, version }))
+    return storage?.setItem(name, await serialize({ state, version }))
   }
 
   const savedSetState = api.setState
@@ -236,18 +242,26 @@ export const persist = <S extends State>(
       if (storageValue) {
         const deserializedStorageValue = await deserialize(storageValue)
 
-        // if versions mismatch, clear storage by storing the new initial state
+        // if versions mismatch, run migration
         if (deserializedStorageValue.version !== version) {
-          setItem()
+          const migratedState = await migrate?.(
+            deserializedStorageValue.state,
+            deserializedStorageValue.version
+          )
+          if (migratedState) {
+            set(migratedState)
+            await setItem()
+          }
         } else {
           set(deserializedStorageValue.state)
         }
       }
     } catch (e) {
-      throw new Error(`Unable to get to stored state in "${name}"`)
-    } finally {
-      postRehydrationCallback?.(get())
+      postRehydrationCallback?.(undefined, e)
+      return
     }
+
+    postRehydrationCallback?.(get(), undefined)
   })()
 
   return config(


Test patch:
diff --git a/tests/persist.test.tsx b/tests/persist.test.tsx
new file mode 100644
index 0000000..96b5002
--- /dev/null
+++ b/tests/persist.test.tsx
@@ -0,0 +1,212 @@
+import React from 'react'
+import { act, cleanup, render } from '@testing-library/react'
+import create from '../src/index'
+import { persist } from '../src/middleware'
+
+const consoleError = console.error
+afterEach(() => {
+  cleanup()
+  console.error = consoleError
+})
+
+it('can rehydrate state', async () => {
+  let postRehydrationCallbackCallCount = 0
+
+  const storage = {
+    getItem: async (name: string) =>
+      JSON.stringify({
+        state: { count: 42, name },
+        version: 0,
+      }),
+    setItem: () => {},
+  }
+
+  const useStore = create(
+    persist(
+      () => ({
+        count: 0,
+        name: 'empty',
+      }),
+      {
+        name: 'test-storage',
+        getStorage: () => storage,
+        onRehydrateStorage: () => (state) => {
+          postRehydrationCallbackCallCount++
+          expect(state?.count).toBe(42)
+          expect(state?.name).toBe('test-storage')
+        },
+      }
+    )
+  )
+
+  function Counter() {
+    const { count, name } = useStore()
+    return (
+      <div>
+        count: {count}, name: {name}
+      </div>
+    )
+  }
+
+  const { findByText } = render(<Counter />)
+
+  await findByText('count: 0, name: empty')
+  await findByText('count: 42, name: test-storage')
+  expect(postRehydrationCallbackCallCount).toBe(1)
+})
+
+it('can throw rehydrate error', async () => {
+  let postRehydrationCallbackCallCount = 0
+
+  const storage = {
+    getItem: async () => {
+      throw new Error('getItem error')
+    },
+    setItem: () => {},
+  }
+
+  const useStore = create(
+    persist(() => ({ count: 0 }), {
+      name: 'test-storage',
+      getStorage: () => storage,
+      onRehydrateStorage: () => (_, e) => {
+        postRehydrationCallbackCallCount++
+        expect(e?.message).toBe('getItem error')
+      },
+    })
+  )
+
+  function Counter() {
+    const { count } = useStore()
+    return <div>count: {count}</div>
+  }
+
+  const { findByText } = render(<Counter />)
+
+  await findByText('count: 0')
+  expect(postRehydrationCallbackCallCount).toBe(1)
+})
+
+it('can persist state', async () => {
+  let setItemCallCount = 0
+
+  const storage = {
+    getItem: () => null,
+    setItem: (name: string, value: string) => {
+      setItemCallCount++
+      expect(name).toBe('test-storage')
+      expect(value).toBe(
+        JSON.stringify({
+          state: { count: 42 },
+          version: 0,
+        })
+      )
+    },
+  }
+
+  const useStore = create<any>(
+    persist(() => ({ count: 0 }), {
+      name: 'test-storage',
+      getStorage: () => storage,
+    })
+  )
+
+  function Counter() {
+    const { count } = useStore()
+    return <div>count: {count}</div>
+  }
+
+  const { findByText } = render(<Counter />)
+
+  await findByText('count: 0')
+  act(() => useStore.setState({ count: 42 }))
+
+  await findByText('count: 42')
+  expect(setItemCallCount).toBe(1)
+})
+
+it('can migrate persisted state', async () => {
+  let migrateCallCount = 0
+  let setItemCallCount = 0
+
+  const storage = {
+    getItem: async () =>
+      JSON.stringify({
+        state: { count: 42 },
+        version: 12,
+      }),
+    setItem: (_: string, value: string) => {
+      setItemCallCount++
+      expect(value).toBe(
+        JSON.stringify({
+          state: { count: 99 },
+          version: 13,
+        })
+      )
+    },
+  }
+
+  const useStore = create(
+    persist(() => ({ count: 0 }), {
+      name: 'test-storage',
+      version: 13,
+      getStorage: () => storage,
+      migrate: (state, version) => {
+        migrateCallCount++
+        expect(state.count).toBe(42)
+        expect(version).toBe(12)
+        return { count: 99 }
+      },
+    })
+  )
+
+  function Counter() {
+    const { count } = useStore()
+    return <div>count: {count}</div>
+  }
+
+  const { findByText } = render(<Counter />)
+
+  await findByText('count: 0')
+  await findByText('count: 99')
+  expect(migrateCallCount).toBe(1)
+  expect(setItemCallCount).toBe(1)
+})
+
+it('can throw migrate error', async () => {
+  let postRehydrationCallbackCallCount = 0
+
+  const storage = {
+    getItem: async () =>
+      JSON.stringify({
+        state: {},
+        version: 12,
+      }),
+    setItem: () => {},
+  }
+
+  const useStore = create(
+    persist(() => ({ count: 0 }), {
+      name: 'test-storage',
+      version: 13,
+      getStorage: () => storage,
+      migrate: () => {
+        throw new Error('migrate error')
+      },
+      onRehydrateStorage: () => (_, e) => {
+        postRehydrationCallbackCallCount++
+        expect(e?.message).toBe('migrate error')
+      },
+    })
+  )
+
+  function Counter() {
+    const { count } = useStore()
+    return <div>count: {count}</div>
+  }
+
+  const { findByText } = render(<Counter />)
+
+  await findByText('count: 0')
+  expect(postRehydrationCallbackCallCount).toBe(1)
+})