fix(storage): Revert #478 and run migrations when item is defined and properly wait for migrations before allowing read/writes (#487)

docs/guide/storage.md

@@ -263,27 +263,4 @@ export const ignoredWebsites = storage.defineItem<IgnoredWebsiteV2[]>( // [!code
 
 ### Running Migrations
 
-To run migrations, you have two options:
-
-1. Import all versioned storage items into your background script and they will run automatically whenever your extension updates
-2. Manually call `item.migrate()`
-
-The first approach is recommended. To make importing all your storage items easy, you can define all of them in a single file, `utils/storage.ts`, and import that file into your background entrypoint:
-
-```ts
-// utils/storage.ts
-export countStorage = storage.defineItem(...);
-export themeStorage = storage.defineItem(...);
-export someOtherStorage = storage.defineItem(...);
-```
-
-```ts
-// entrypoints/background.ts
-import '@/utils/storage'; // This import runs migrations on updates
-
-export default defineBackground({
-  // ...
-});
-```
-
-> When you call `storage.defineItem`, a `browser.runtime.onInstalled` listener is added. `onInstalled` listeners are only triggered in the background, which is why you must import them into the background.
+As soon as `storage.defineItem` is called, WXT checks if migrations need to be ran, and if so, runs them. Calls to get or update the storage item's value or metadata (`getValue`, `setValue`, `removeValue`, `getMeta`, etc) will automatically wait for the migration process to finish before actually reading or writing values.

src/__tests__/storage.test.ts

@@ -3,11 +3,12 @@ import { describe, it, expect, beforeEach, vi, expectTypeOf } from 'vitest';
 import { browser } from '~/browser';
 import { WxtStorageItem, storage } from '~/storage';
 
-async function triggerUpdate() {
-  await fakeBrowser.runtime.onInstalled.trigger({
-    reason: 'update',
-    temporary: false,
-  });
+/**
+ * This works because fakeBrowser is synchronous, and is will finish any number of chained
+ * calls within a single tick of the event loop, ie: a timeout of 0.
+ */
+async function waitForMigrations() {
+  return new Promise((res) => setTimeout(res));
 }
 
 describe('Storage Utils', () => {
@@ -409,7 +410,7 @@ describe('Storage Utils', () => {
             3: migrateToV3,
           },
         });
-        await triggerUpdate();
+        await waitForMigrations();
 
         const actualValue = await item.getValue();
         const actualMeta = await item.getMeta();
@@ -436,7 +437,7 @@ describe('Storage Utils', () => {
             3: migrateToV3,
           },
         });
-        await triggerUpdate();
+        await waitForMigrations();
 
         const actualValue = await item.getValue();
         const actualMeta = await item.getMeta();
@@ -461,7 +462,7 @@ describe('Storage Utils', () => {
             2: migrateToV2,
           },
         });
-        await triggerUpdate();
+        await waitForMigrations();
 
         const actualValue = await item.getValue();
         const actualMeta = await item.getMeta();
@@ -489,7 +490,7 @@ describe('Storage Utils', () => {
             3: migrateToV3,
           },
         });
-        await triggerUpdate();
+        await waitForMigrations();
 
         expect(migrateToV2).not.toBeCalled();
         expect(migrateToV3).not.toBeCalled();
@@ -511,7 +512,7 @@ describe('Storage Utils', () => {
             3: migrateToV3,
           },
         });
-        await triggerUpdate();
+        await waitForMigrations();
 
         const actualValue = await item.getValue();
         const actualMeta = await item.getMeta();
@@ -538,7 +539,7 @@ describe('Storage Utils', () => {
           defaultValue: 0,
           version: nextVersion,
         });
-        await triggerUpdate();
+        await waitForMigrations();
 
         await expect(item.migrate()).rejects.toThrow(
           'Version downgrade detected (v2 -> v1) for "local:count"',

src/storage.ts

@@ -258,14 +258,14 @@ function createStorage(): WxtStorage {
         );
       }
       const migrate = async () => {
-        const [value, meta] = await Promise.all([
-          // TODO: Optimize with getItems
-          getItem(driver, driverKey, undefined),
-          getMeta(driver, driverKey),
+        const driverMetaKey = getMetaKey(driverKey);
+        const [{ value }, { value: meta }] = await driver.getItems([
+          driverKey,
+          driverMetaKey,
         ]);
         if (value == null) return;
 
-        const currentVersion = meta.v ?? 1;
+        const currentVersion = meta?.v ?? 1;
         if (currentVersion > targetVersion) {
           throw Error(
             `Version downgrade detected (v${currentVersion} -> v${targetVersion}) for "${key}"`,
@@ -285,24 +285,21 @@ function createStorage(): WxtStorage {
             (await migrations?.[migrateToVersion]?.(migratedValue)) ??
             migratedValue;
         }
-        await Promise.all([
-          // TODO: Optimize with `setItem`
-          setItem(driver, driverKey, migratedValue),
-          setMeta(driver, driverKey, { v: targetVersion }),
+        await driver.setItems([
+          { key: driverKey, value: migratedValue },
+          { key: driverMetaKey, value: { ...meta, v: targetVersion } },
         ]);
         logger.debug(
           `Storage migration completed for ${key} v${targetVersion}`,
           { migratedValue },
         );
       };
-      browser.runtime.onInstalled?.addListener(async ({ reason }) => {
-        if (reason !== 'update') return;
-        try {
-          await migrate();
-        } catch (err) {
-          logger.error(`Migration failed for ${key}`, err);
-        }
-      });
+      const migrationsDone =
+        opts?.migrations == null
+          ? Promise.resolve()
+          : migrate().catch((err) => {
+              logger.error(`Migration failed for ${key}`, err);
+            });
 
       const getDefaultValue = () => opts?.defaultValue ?? null;
 
@@ -310,12 +307,30 @@ function createStorage(): WxtStorage {
         get defaultValue() {
           return getDefaultValue();
         },
-        getValue: () => getItem(driver, driverKey, opts),
-        getMeta: () => getMeta(driver, driverKey),
-        setValue: (value) => setItem(driver, driverKey, value),
-        setMeta: (properties) => setMeta(driver, driverKey, properties),
-        removeValue: (opts) => removeItem(driver, driverKey, opts),
-        removeMeta: (properties) => removeMeta(driver, driverKey, properties),
+        getValue: async () => {
+          await migrationsDone;
+          return await getItem(driver, driverKey, opts);
+        },
+        getMeta: async () => {
+          await migrationsDone;
+          return await getMeta(driver, driverKey);
+        },
+        setValue: async (value) => {
+          await migrationsDone;
+          return await setItem(driver, driverKey, value);
+        },
+        setMeta: async (properties) => {
+          await migrationsDone;
+          return await setMeta(driver, driverKey, properties);
+        },
+        removeValue: async (opts) => {
+          await migrationsDone;
+          return await removeItem(driver, driverKey, opts);
+        },
+        removeMeta: async (properties) => {
+          await migrationsDone;
+          return await removeMeta(driver, driverKey, properties);
+        },
         watch: (cb) =>
           watch(driver, driverKey, (newValue, oldValue) =>
             cb(newValue ?? getDefaultValue(), oldValue ?? getDefaultValue()),
@@ -346,7 +361,10 @@ function createDriver(
       );
     }
 
-    return browser.storage[storageArea];
+    const area = browser.storage[storageArea];
+    if (area == null)
+      throw Error(`"browser.storage.${storageArea}" is undefined`);
+    return area;
   };
   const watchListeners = new Set<
     (changes: Storage.StorageAreaOnChangedChangesType) => void