fix: keep the market-hours database when a refresh returns malformed JSON

lean/components/cloud/data_downloader.py

@@ -11,6 +11,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+import shutil
+import tempfile
 from pathlib import Path
 from datetime import datetime, timedelta
 from typing import Any, List, Callable
@@ -23,8 +26,30 @@
 
 def _store_local_file(file_content: bytes, file_path: Path):
     file_path.parent.mkdir(parents=True, exist_ok=True)
-    with file_path.open("wb+") as f:
-        f.write(file_content)
+    # Write to a temporary file first so a failed write never truncates the existing file
+    fd, temp_name = tempfile.mkstemp(dir=file_path.parent, prefix=f".{file_path.name}.", suffix=".tmp")
+    try:
+        with os.fdopen(fd, "wb") as f:
+            f.write(file_content)
+        if file_path.exists():
+            shutil.copymode(file_path, temp_name)
+        os.replace(temp_name, file_path)
+    except BaseException:
+        if os.path.exists(temp_name):
+            os.remove(temp_name)
+        raise
+
+
+def _validate_market_hours_database(file_content: bytes) -> None:
+    """Raises if the content is not a market hours database that MarketHoursDatabase can read."""
+    from re import sub, DOTALL
+    from json import loads
+
+    # Same parsing as MarketHoursDatabase._get_all_entries()
+    text = sub(r"(/\*.*\*/)", "", file_content.decode("utf-8"), flags=DOTALL)
+    database = loads(text)
+    if not isinstance(database, dict) or not isinstance(database.get("entries"), dict):
+        raise ValueError("missing 'entries' object")
 
 
 def parse_timedelta(database_update_frequency: str):
@@ -87,9 +112,15 @@ def update_database_files(self):
                 _store_local_file(self._api_client.data.download_public_file(
                     "https://raw.githubusercontent.com/QuantConnect/Lean/master/Data/symbol-properties/symbol-properties-database.csv"),
                     data_dir / "symbol-properties" / "symbol-properties-database.csv")
-                _store_local_file(self._api_client.data.download_public_file(
-                    "https://raw.githubusercontent.com/QuantConnect/Lean/master/Data/market-hours/market-hours-database.json"),
-                    data_dir / "market-hours" / "market-hours-database.json")
+                market_hours_database = self._api_client.data.download_public_file(
+                    "https://raw.githubusercontent.com/QuantConnect/Lean/master/Data/market-hours/market-hours-database.json")
+                try:
+                    _validate_market_hours_database(market_hours_database)
+                except ValueError as e:
+                    # Keep the existing local file; not a ValueError for the config option handler below
+                    raise RuntimeError(f"Downloaded market-hours database is invalid ({e}), "
+                                       f"keeping the existing local file") from e
+                _store_local_file(market_hours_database, data_dir / "market-hours" / "market-hours-database.json")
         except MoreInfoError as e:
             if "not found" in str(e):
                 pass

tests/components/cloud/test_data_downloader.py

@@ -0,0 +1,68 @@
+from datetime import datetime, timedelta
+from pathlib import Path
+from unittest import mock
+
+import pytest
+
+from lean.components.cloud.data_downloader import DataDownloader, _store_local_file
+
+VALID_DATABASE = '{"entries": {}}'
+
+
+@pytest.fixture
+def tmp_path() -> Path:
+    # The test suite runs on a fake filesystem (see conftest.py), so use a fixed path instead of pytest's tmp_path
+    path = Path("/data")
+    path.mkdir(parents=True)
+    return path
+
+
+def _create_downloader(data_dir: Path, market_hours: bytes) -> DataDownloader:
+    lean_config_manager = mock.Mock()
+    lean_config_manager.get_lean_config.return_value = {}
+    lean_config_manager.get_data_directory.return_value = data_dir
+    api_client = mock.Mock()
+    api_client.data.download_public_file.side_effect = \
+        lambda url: market_hours if url.endswith(".json") else b"symbol,properties\n"
+    return DataDownloader(mock.Mock(), api_client, lean_config_manager, "0.00:00:01")
+
+
+def _market_hours_file(data_dir: Path) -> Path:
+    path = data_dir / "market-hours" / "market-hours-database.json"
+    path.parent.mkdir(parents=True)
+    path.write_text(VALID_DATABASE, encoding="utf-8")
+    return path
+
+
+def test_update_database_files_keeps_existing_file_when_response_is_malformed(tmp_path: Path) -> None:
+    path = _market_hours_file(tmp_path)
+
+    _create_downloader(tmp_path, b"not json").update_database_files()
+
+    assert path.read_text(encoding="utf-8") == VALID_DATABASE
+
+
+def test_update_database_files_keeps_existing_file_when_entries_are_missing(tmp_path: Path) -> None:
+    path = _market_hours_file(tmp_path)
+
+    _create_downloader(tmp_path, b'{"other": 1}').update_database_files()
+
+    assert path.read_text(encoding="utf-8") == VALID_DATABASE
+
+
+def test_update_database_files_replaces_file_with_valid_commented_response(tmp_path: Path) -> None:
+    path = _market_hours_file(tmp_path)
+    new_content = b'/* comment */ {"entries": {"Equity-usa-[*]": {}}}'
+
+    _create_downloader(tmp_path, new_content).update_database_files()
+
+    assert path.read_bytes() == new_content
+
+
+def test_store_local_file_leaves_no_temporary_files(tmp_path: Path) -> None:
+    path = tmp_path / "dir" / "file.json"
+
+    _store_local_file(b"content", path)
+
+    assert path.read_bytes() == b"content"
+    assert [p.name for p in path.parent.iterdir()] == ["file.json"]