Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions iceframe/cache.py
Original file line number Diff line number Diff line change
Expand Up @@ -152,10 +152,10 @@ class DiskCache(QueryCache):
def __init__(self, cache_dir: str = ".iceframe_cache", max_size_gb: float = 1.0):
super().__init__(max_size=10**9) # disk capacity bound is enforced by diskcache
try:
import diskcache # noqa: F401
import diskcache_community # noqa: F401
except ImportError:
raise ImportError(
"diskcache required for disk caching. Install with: pip install 'iceframe[cache]'"
"diskcache-community required for disk caching. Install with: pip install 'iceframe[cache]'"
) from None
import os

Expand All @@ -165,11 +165,12 @@ def __init__(self, cache_dir: str = ".iceframe_cache", max_size_gb: float = 1.0)
# diskcache index in <cache_dir>/index/ holds metadata pointers.
self._data_dir = os.path.join(self.cache_dir, "data")
os.makedirs(self._data_dir, exist_ok=True)
import diskcache
from diskcache_community import Cache, SafeDisk

self.cache = diskcache.Cache(
self.cache = Cache(
os.path.join(self.cache_dir, "index"),
size_limit=int(max_size_gb * 1024**3),
disk=SafeDisk.with_options(on_error="miss"),
)

def _payload_path(self, key: str) -> str:
Expand Down
3 changes: 2 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,7 @@ dev = [
"typer>=0.9.0",
"rich>=13.0.0",
"pandas>=2.0.0",
"diskcache-community>=5.8.4",
]
pdf = [
"fpdf2>=2.7.0",
Expand Down Expand Up @@ -116,7 +117,7 @@ agent = [
"rich>=13.0.0",
]
cache = [
"diskcache>=5.6.0",
"diskcache-community>=5.8.4",
]
streaming = [
"kafka-python>=2.0.0",
Expand Down
39 changes: 39 additions & 0 deletions tests/test_disk_cache.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
from pathlib import Path

import polars as pl
import pytest

diskcache_community = pytest.importorskip("diskcache_community")

from iceframe.cache import DiskCache

PARAMS = {"q": 1}


def _write_marker(path: str) -> None:
Path(path).write_text("ran")


class _PlantedEntry:
def __init__(self, marker: Path):
self.marker = str(marker)

def __reduce__(self):
return (_write_marker, (self.marker,))


def test_disk_cache_planted_pickle_is_a_miss(tmp_path):
cache_dir = tmp_path / "cache"
marker = tmp_path / "planted-code-ran"
df = pl.DataFrame({"id": [2, 3], "g": ["a", "b"]})
cache = DiskCache(cache_dir=str(cache_dir))
cache.put("default.t", PARAMS, df, ttl=300)

with diskcache_community.Cache(
str(cache_dir / "index"), disk=diskcache_community.Disk
) as writer:
(key,) = list(writer)
writer.set(key, _PlantedEntry(marker))

assert cache.get("default.t", PARAMS) is None
assert not marker.exists()