Fix compressed Store lookups and corrupt file handling

The store manager caches by filename, so a compressed store asking for
"key.zst" was told the file does not exist whenever only the uncompressed
predecessor was on disk. That short-circuited the load before the fallback
in _load_data_from_disk could run, so enabling compress on an existing
store read as empty on any real start, where the manager is initialized.
Only report "does not exist" when neither filename is on disk.

Rename the file that actually failed to parse when handling corruption.
The compressed path was renamed unconditionally, which raised
FileNotFoundError out of the executor when the uncompressed fallback was
the corrupt one, and named the wrong file in the repair issue.

Preload compressed stores too. Callers only know the plain key, so look
for both spellings before intersecting with the files on disk.

Also invalidate both filenames when writing or removing, since both drop
the uncompressed predecessor, and use a sentinel to tell "no file" apart
from a file holding an empty dict.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

homeassistant/helpers/storage.py

@@ -46,22 +46,30 @@ STORAGE_MANAGER: HassKey[_StoreManager] = HassKey("storage_manager")
 
 MANAGER_CLEANUP_DELAY = 60
 
+COMPRESSED_SUFFIX = ".zst"
 
-def _load_json_file(path: str | Path) -> json_util.JsonValueType:
+# Distinguishes "no file on disk" from a file whose content is an empty dict.
+_NOT_FOUND = object()
+
+
+def _load_json_file(
+    path: str | Path,
+    default: json_util.JsonValueType = _NOT_FOUND,  # type: ignore[assignment]
+) -> json_util.JsonValueType:
     """Load JSON from a file, transparently decompressing .zst files.
 
-    Returns ``{}`` (the same sentinel as :func:`json_util.load_json`) when
-    the file does not exist.  Raises :class:`HomeAssistantError` wrapping the
-    original exception when the file is corrupt or cannot be read.
+    Returns ``default`` when the file does not exist, mirroring
+    :func:`json_util.load_json`.  Raises :class:`HomeAssistantError` wrapping
+    the original exception when the file is corrupt or cannot be read.
     """
-    if not str(path).endswith(".zst"):
-        return json_util.load_json(path)
+    if Path(path).suffix != COMPRESSED_SUFFIX:
+        return json_util.load_json(path, default=default)
     try:
         with open(path, "rb") as fh:
             raw = zstd.decompress(fh.read())
     except FileNotFoundError:
         _LOGGER.debug("JSON file not found: %s", path)
-        return {}
+        return default
     except zstd.ZstdError as err:
         _LOGGER.exception("Could not decompress storage file: %s", path)
         raise HomeAssistantError(f"Error decompressing {path}: {err}") from err
@@ -230,7 +238,14 @@ class _StoreManager:
     async def async_preload(self, keys: Iterable[str]) -> None:
         """Cache the keys."""
         # If async_initialize has not been called yet, we can't preload
-        if self._files is not None and (existing := self._files.intersection(keys)):
+        if self._files is None:
+            return
+        # Callers pass plain keys; a compressed store is on disk as key + ".zst"
+        # and is cached under that name, so look for both spellings.
+        candidates = {
+            spelling for key in keys for spelling in (key, f"{key}{COMPRESSED_SUFFIX}")
+        }
+        if existing := self._files.intersection(candidates):
             await self._hass.async_add_executor_job(self._preload, existing)
 
     def _preload(self, keys: Iterable[str]) -> None:
@@ -241,7 +256,7 @@ class _StoreManager:
             storage_file: Path = storage_path.joinpath(key)
             try:
                 if storage_file.is_file():
-                    data_preload[key] = _load_json_file(storage_file)
+                    data_preload[key] = _load_json_file(storage_file, default={})
             except Exception as ex:  # noqa: BLE001
                 _LOGGER.debug("Error loading %s: %s", key, ex)
 
@@ -276,6 +291,17 @@ class Store[_T: Mapping[str, Any] | Sequence[Any]]:
             to version. Set higher than version to support forward compatibility,
             allowing reading data written by newer versions (e.g., after downgrade).
 
+            compress: Whether to store the data zstd compressed, under the key
+            with a ".zst" suffix appended. No migration is needed to turn this
+            on or off: an already existing uncompressed file is read as a
+            fallback and removed once the compressed file has been written, so
+            a user can decompress, edit and save a file by hand at any time.
+
+            Compression trades CPU for I/O and disk space. It is paid on every
+            write, while it is only earned back once per load, so it is meant
+            for stores which grow large rather than for ones which are written
+            frequently.
+
             serialize_in_event_loop: Whether to serialize data in the event loop.
             Set to True (default) if data passed to async_save and data produced by
             data_func passed to async_delay_save needs to be serialized in the event
@@ -316,7 +342,7 @@ class Store[_T: Mapping[str, Any] | Sequence[Any]]:
     def path(self):
         """Return the config path."""
         if self._compress:
-            return self.hass.config.path(STORAGE_DIR, self.key + ".zst")
+            return self.hass.config.path(STORAGE_DIR, self._cache_key)
         return self.hass.config.path(STORAGE_DIR, self.key)
 
     @cached_property
@@ -329,7 +355,7 @@ class Store[_T: Mapping[str, Any] | Sequence[Any]]:
         suffix here to match the real filename.
         """
         if self._compress:
-            return self.key + ".zst"
+            return f"{self.key}{COMPRESSED_SUFFIX}"
         return self.key
 
     @cached_property
@@ -394,6 +420,36 @@ class Store[_T: Mapping[str, Any] | Sequence[Any]]:
         async with self.hass.data[STORAGE_SEMAPHORE]:
             return await self._async_load_data()
 
+    @callback
+    def _async_invalidate_cache(self) -> None:
+        """Invalidate every filename this store may be cached under."""
+        self._manager.async_invalidate(self._cache_key)
+        if self._compress:
+            # Writing and removing both drop the uncompressed predecessor.
+            self._manager.async_invalidate(self.key)
+
+    @callback
+    def _async_fetch_cached(
+        self,
+    ) -> tuple[bool, json_util.JsonValueType | None] | None:
+        """Fetch preloaded data, accounting for the uncompressed fallback.
+
+        Returns None when the manager cannot answer and the file has to be
+        read from disk. A compressed store has two candidate filenames, so
+        only report "does not exist" when neither of them is on disk.
+        """
+        primary = self._manager.async_fetch(self._cache_key)
+        if not self._compress:
+            return primary
+        if primary is not None and primary[0]:
+            return primary
+        fallback = self._manager.async_fetch(self.key)
+        if fallback is not None and fallback[0]:
+            return fallback
+        if primary is not None and fallback is not None:
+            return (False, None)
+        return None
+
     async def _async_load_data(self):
         """Load the data."""
         # When load_empty is set, skip loading storage files and use empty
@@ -413,7 +469,7 @@ class Store[_T: Mapping[str, Any] | Sequence[Any]]:
             # We make a copy because code might assume it's safe to mutate loaded data
             # and we don't want that to mess with what we're trying to store.
             data = deepcopy(data)
-        elif cache := self._manager.async_fetch(self._cache_key):
+        elif cache := self._async_fetch_cached():
             exists, data = cache
             if not exists:
                 return None
@@ -577,7 +633,7 @@ class Store[_T: Mapping[str, Any] | Sequence[Any]]:
     async def _async_handle_write_data(self, *_args):
         """Handle writing the config."""
         async with self._write_lock:
-            self._manager.async_invalidate(self._cache_key)
+            self._async_invalidate_cache()
             self._async_cleanup_delay_listener()
             self._async_cleanup_final_write_listener()
 
@@ -625,10 +681,9 @@ class Store[_T: Mapping[str, Any] | Sequence[Any]]:
         neither file exists.
         """
         data = _load_json_file(self.path)
-        if data == {} and self._compress:
-            # .zst not found - fall back to the plain file.
+        if data is _NOT_FOUND and self._compress:
             data = _load_json_file(self._uncompressed_path)
-        return data
+        return {} if data is _NOT_FOUND else data
 
     def _write_prepared_data(self, mode: str, json_data: str | bytes) -> None:
         """Write the data."""
@@ -658,8 +713,18 @@ class Store[_T: Mapping[str, Any] | Sequence[Any]]:
 
         isotime = dt_util.utcnow().isoformat()
         corrupt_postfix = f".corrupt.{isotime}"
-        corrupt_path = f"{self.path}{corrupt_postfix}"
-        await self.hass.async_add_executor_job(os.rename, self.path, corrupt_path)
+
+        def _rename_corrupt_file() -> tuple[str, str]:
+            # The compressed path is read first, so when it is absent the file
+            # we failed to read is the uncompressed fallback.
+            path = self.path if os.path.isfile(self.path) else self._uncompressed_path
+            corrupt_path = f"{path}{corrupt_postfix}"
+            os.rename(path, corrupt_path)
+            return path, corrupt_path
+
+        path, corrupt_path = await self.hass.async_add_executor_job(
+            _rename_corrupt_file
+        )
         storage_key = self.key
         _LOGGER.error(
             "Unrecoverable error decoding storage %s at %s; "
@@ -668,7 +733,7 @@ class Store[_T: Mapping[str, Any] | Sequence[Any]]:
             "The corrupt file has been saved as %s; "
             "It is recommended to restore from backup: %s",
             storage_key,
-            self.path,
+            path,
             corrupt_path,
             err,
         )
@@ -690,7 +755,7 @@ class Store[_T: Mapping[str, Any] | Sequence[Any]]:
             severity=IssueSeverity.CRITICAL,
             translation_placeholders={
                 "storage_key": storage_key,
-                "original_path": self.path,
+                "original_path": path,
                 "corrupt_path": corrupt_path,
                 "error": str(err),
             },
@@ -702,7 +767,7 @@ class Store[_T: Mapping[str, Any] | Sequence[Any]]:
 
     async def async_remove(self) -> None:
         """Remove all data."""
-        self._manager.async_invalidate(self._cache_key)
+        self._async_invalidate_cache()
         self._async_cleanup_delay_listener()
         self._async_cleanup_final_write_listener()
 

tests/helpers/test_storage.py

@@ -1,6 +1,7 @@
 """Tests for the storage helper."""
 
 import asyncio
+from collections.abc import AsyncGenerator
 from compression import zstd
 from datetime import timedelta
 import json
@@ -1385,128 +1386,210 @@ async def test_load_empty_returns_none_and_read_only(
     assert hass_storage[MOCK_KEY]["version"] == 99
 
 
-async def test_compress_save_load_round_trip(tmpdir: py.path.local) -> None:
-    """Test that a compressed store saves a .zst file and loads back correctly."""
+def _storage_file(hass: HomeAssistant, name: str) -> Path:
+    """Return the path of a file in the .storage dir."""
+    return Path(hass.config.config_dir) / storage.STORAGE_DIR / name
+
+
+def _write_store_file(path: Path, data: dict[str, Any], compress: bool) -> None:
+    """Write a store envelope to disk the way a previous run would have."""
+    payload = json.dumps(
+        {
+            "version": MOCK_VERSION,
+            "minor_version": 1,
+            "key": MOCK_KEY,
+            "data": data,
+        }
+    ).encode()
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_bytes(zstd.compress(payload) if compress else payload)
+
+
+@pytest.fixture
+async def disk_hass(tmpdir: py.path.local) -> AsyncGenerator[HomeAssistant]:
+    """Yield a Home Assistant instance backed by a real config directory."""
     loop = asyncio.get_running_loop()
     config_dir = await loop.run_in_executor(None, tmpdir.mkdir, "temp_storage")
 
     async with async_test_home_assistant(config_dir=config_dir.strpath) as hass:
-        store = storage.Store(hass, MOCK_VERSION, MOCK_KEY, compress=True)
-        await store.async_save(MOCK_DATA)
-
-        storage_path = Path(config_dir.strpath) / ".storage"
-        zst_file = storage_path / (MOCK_KEY + ".zst")
-        plain_file = storage_path / MOCK_KEY
-
-        assert zst_file.is_file()
-        assert not plain_file.exists()
-
-        raw = zstd.decompress(zst_file.read_bytes())
-        on_disk = json.loads(raw)
-        assert on_disk["data"] == MOCK_DATA
-
-        loaded = await store.async_load()
-        assert loaded == MOCK_DATA
-
+        yield hass
         await hass.async_stop(force=True)
 
 
-async def test_compress_migrates_plain_to_compressed(tmpdir: py.path.local) -> None:
-    """Test that saving with compress=True removes an existing plain file."""
-    loop = asyncio.get_running_loop()
-    config_dir = await loop.run_in_executor(None, tmpdir.mkdir, "temp_storage")
-
-    async with async_test_home_assistant(config_dir=config_dir.strpath) as hass:
-        plain_store = storage.Store(hass, MOCK_VERSION, MOCK_KEY)
-        await plain_store.async_save(MOCK_DATA)
-
-        storage_path = Path(config_dir.strpath) / ".storage"
-        plain_file = storage_path / MOCK_KEY
-        assert plain_file.is_file()
-
-        compressed_store = storage.Store(hass, MOCK_VERSION, MOCK_KEY, compress=True)
-
-        # Before the first compressed write the plain file is still the fallback.
-        loaded = await compressed_store.async_load()
-        assert loaded == MOCK_DATA
-
-        # Saving with compress=True should write .zst and remove the plain file.
-        await compressed_store.async_save(MOCK_DATA2)
-
-        zst_file = storage_path / (MOCK_KEY + ".zst")
-        assert zst_file.is_file()
-        assert not plain_file.exists()
-
-        loaded = await compressed_store.async_load()
-        assert loaded == MOCK_DATA2
-
-        await hass.async_stop(force=True)
-
-
-async def test_compress_corrupt_file(
-    tmpdir: py.path.local, caplog: pytest.LogCaptureFixture
+@pytest.mark.parametrize(
+    ("atomic_writes", "serialize_in_event_loop"),
+    [
+        pytest.param(False, True, id="default"),
+        pytest.param(True, True, id="atomic_writes"),
+        pytest.param(False, False, id="serialize_in_executor"),
+    ],
+)
+async def test_compress_save_load_round_trip(
+    disk_hass: HomeAssistant, atomic_writes: bool, serialize_in_event_loop: bool
 ) -> None:
-    """Test that a corrupt .zst file is handled gracefully."""
-    loop = asyncio.get_running_loop()
-    config_dir = await loop.run_in_executor(None, tmpdir.mkdir, "temp_storage")
+    """Test that a compressed store saves a .zst file and loads back correctly."""
+    store = storage.Store(
+        disk_hass,
+        MOCK_VERSION,
+        MOCK_KEY,
+        compress=True,
+        atomic_writes=atomic_writes,
+        serialize_in_event_loop=serialize_in_event_loop,
+    )
+    await store.async_save(MOCK_DATA)
 
-    async with async_test_home_assistant(config_dir=config_dir.strpath) as hass:
-        store = storage.Store(hass, MOCK_VERSION, MOCK_KEY, compress=True)
-        await store.async_save(MOCK_DATA)
+    zst_file = _storage_file(disk_hass, MOCK_KEY + ".zst")
+    plain_file = _storage_file(disk_hass, MOCK_KEY)
 
-        storage_path = Path(config_dir.strpath) / ".storage"
-        zst_file = storage_path / (MOCK_KEY + ".zst")
+    assert zst_file.is_file()
+    assert not plain_file.exists()
 
-        def _corrupt_file() -> None:
-            zst_file.write_bytes(b"this is not valid zstd data")
+    on_disk = json.loads(zstd.decompress(zst_file.read_bytes()))
+    assert on_disk["data"] == MOCK_DATA
 
-        await hass.async_add_executor_job(_corrupt_file)
-
-        loaded = await store.async_load()
-        assert loaded is None
-        assert "Unrecoverable error decoding storage" in caplog.text
-
-        files = await hass.async_add_executor_job(os.listdir, storage_path)
-        corrupt_files = [f for f in files if ".corrupt" in f]
-        assert len(corrupt_files) == 1
-
-        await hass.async_stop(force=True)
+    assert await store.async_load() == MOCK_DATA
 
 
-async def test_compress_store_manager_cache(tmpdir: py.path.local) -> None:
-    """Test that compressed stores are cached and served by the store manager."""
-    loop = asyncio.get_running_loop()
+async def test_compress_migrates_plain_to_compressed(disk_hass: HomeAssistant) -> None:
+    """Test that saving with compress=True removes an existing plain file."""
+    plain_store = storage.Store(disk_hass, MOCK_VERSION, MOCK_KEY)
+    await plain_store.async_save(MOCK_DATA)
 
-    def _setup_mock_storage() -> py.path.local:
-        config_dir = tmpdir.mkdir("temp_config")
-        tmp_storage = config_dir.mkdir(".storage")
-        payload = json.dumps(
-            {
-                "version": MOCK_VERSION,
-                "minor_version": 1,
-                "key": MOCK_KEY,
-                "data": MOCK_DATA,
-            }
-        ).encode()
-        tmp_storage.join(MOCK_KEY + ".zst").write_binary(zstd.compress(payload))
-        return config_dir
+    plain_file = _storage_file(disk_hass, MOCK_KEY)
+    assert plain_file.is_file()
 
-    config_dir = await loop.run_in_executor(None, _setup_mock_storage)
+    compressed_store = storage.Store(disk_hass, MOCK_VERSION, MOCK_KEY, compress=True)
 
-    async with async_test_home_assistant(config_dir=config_dir.strpath) as hass:
-        store_manager = storage.get_internal_store_manager(hass)
-        await store_manager.async_initialize()
-        await store_manager.async_preload([MOCK_KEY + ".zst"])
+    # Before the first compressed write the plain file is still the fallback.
+    assert await compressed_store.async_load() == MOCK_DATA
 
-        # The cache key for a compressed store is key + ".zst".
-        result = store_manager.async_fetch(MOCK_KEY + ".zst")
-        assert result is not None
-        exists, cached_data = result
-        assert exists is True
-        assert cached_data["data"] == MOCK_DATA  # type: ignore[index]
+    # Saving with compress=True should write .zst and remove the plain file.
+    await compressed_store.async_save(MOCK_DATA2)
 
-        store = storage.Store(hass, MOCK_VERSION, MOCK_KEY, compress=True)
-        loaded = await store.async_load()
-        assert loaded == MOCK_DATA
+    assert _storage_file(disk_hass, MOCK_KEY + ".zst").is_file()
+    assert not plain_file.exists()
+    assert await compressed_store.async_load() == MOCK_DATA2
 
-        await hass.async_stop(force=True)
+
+async def test_compress_falls_back_with_initialized_manager(
+    disk_hass: HomeAssistant,
+) -> None:
+    """Test the plain fallback is used when the store manager knows the files.
+
+    The manager caches by filename, so it reports the .zst as non-existent.
+    That must not short-circuit the load, because the uncompressed file the
+    store falls back to is still on disk.
+    """
+    await disk_hass.async_add_executor_job(
+        _write_store_file, _storage_file(disk_hass, MOCK_KEY), MOCK_DATA, False
+    )
+    await storage.get_internal_store_manager(disk_hass).async_initialize()
+
+    store = storage.Store(disk_hass, MOCK_VERSION, MOCK_KEY, compress=True)
+    assert await store.async_load() == MOCK_DATA
+
+
+async def test_compress_no_file_at_all(disk_hass: HomeAssistant) -> None:
+    """Test a compressed store with neither file on disk loads as empty."""
+    await storage.get_internal_store_manager(disk_hass).async_initialize()
+
+    store = storage.Store(disk_hass, MOCK_VERSION, MOCK_KEY, compress=True)
+    assert await store.async_load() is None
+
+
+@pytest.mark.parametrize(
+    ("corrupt_name", "absent_name"),
+    [
+        pytest.param(MOCK_KEY + ".zst", MOCK_KEY, id="compressed_file"),
+        pytest.param(MOCK_KEY, MOCK_KEY + ".zst", id="uncompressed_fallback"),
+    ],
+)
+async def test_compress_corrupt_file(
+    disk_hass: HomeAssistant,
+    caplog: pytest.LogCaptureFixture,
+    corrupt_name: str,
+    absent_name: str,
+) -> None:
+    """Test that the corrupt file is the one renamed, not the compressed path."""
+    corrupt_file = _storage_file(disk_hass, corrupt_name)
+
+    def _write_corrupt() -> None:
+        corrupt_file.parent.mkdir(parents=True, exist_ok=True)
+        corrupt_file.write_bytes(b"this is not valid zstd data")
+
+    await disk_hass.async_add_executor_job(_write_corrupt)
+    assert not _storage_file(disk_hass, absent_name).exists()
+
+    store = storage.Store(disk_hass, MOCK_VERSION, MOCK_KEY, compress=True)
+    assert await store.async_load() is None
+    assert "Unrecoverable error decoding storage" in caplog.text
+    assert not corrupt_file.exists()
+
+    storage_path = corrupt_file.parent
+    files = await disk_hass.async_add_executor_job(os.listdir, storage_path)
+    assert [f for f in files if ".corrupt" in f] == [
+        f for f in files if f.startswith(corrupt_name + ".corrupt")
+    ]
+
+    issue_registry = ir.async_get(disk_hass)
+    issues = [
+        entry
+        for entry in issue_registry.issues.values()
+        if entry.translation_key == "storage_corruption"
+    ]
+    assert len(issues) == 1
+    assert issues[0].translation_placeholders["original_path"] == str(corrupt_file)
+
+
+async def test_compress_store_manager_preload(disk_hass: HomeAssistant) -> None:
+    """Test that a compressed store is preloaded when its plain key is asked for."""
+    await disk_hass.async_add_executor_job(
+        _write_store_file,
+        _storage_file(disk_hass, MOCK_KEY + ".zst"),
+        MOCK_DATA,
+        True,
+    )
+
+    store_manager = storage.get_internal_store_manager(disk_hass)
+    await store_manager.async_initialize()
+    # Callers such as bootstrap only know the plain key.
+    await store_manager.async_preload([MOCK_KEY])
+
+    result = store_manager.async_fetch(MOCK_KEY + ".zst")
+    assert result is not None
+    exists, cached_data = result
+    assert exists is True
+    assert cached_data["data"] == MOCK_DATA  # type: ignore[index]
+
+
+async def test_compress_store_manager_cache(disk_hass: HomeAssistant) -> None:
+    """Test that compressed stores are served from the store manager cache."""
+    await disk_hass.async_add_executor_job(
+        _write_store_file,
+        _storage_file(disk_hass, MOCK_KEY + ".zst"),
+        MOCK_DATA,
+        True,
+    )
+
+    store_manager = storage.get_internal_store_manager(disk_hass)
+    await store_manager.async_initialize()
+    await store_manager.async_preload([MOCK_KEY + ".zst"])
+
+    store = storage.Store(disk_hass, MOCK_VERSION, MOCK_KEY, compress=True)
+    assert await store.async_load() == MOCK_DATA
+
+
+async def test_compress_async_remove(disk_hass: HomeAssistant) -> None:
+    """Test that removing a compressed store removes both files."""
+    zst_file = _storage_file(disk_hass, MOCK_KEY + ".zst")
+    plain_file = _storage_file(disk_hass, MOCK_KEY)
+    await disk_hass.async_add_executor_job(_write_store_file, zst_file, MOCK_DATA, True)
+    await disk_hass.async_add_executor_job(
+        _write_store_file, plain_file, MOCK_DATA, False
+    )
+
+    store = storage.Store(disk_hass, MOCK_VERSION, MOCK_KEY, compress=True)
+    await store.async_remove()
+
+    assert not zst_file.exists()
+    assert not plain_file.exists()