wip

Author: ravenac95

lib/oso-core/oso_core/cache/__init__.py

@@ -1,2 +1,3 @@
 # ruff: noqa: F403
-from .generic import *
+from .file import *
+from .types import *

lib/oso-core/oso_core/cache/file.py

@@ -0,0 +1,96 @@
+import hashlib
+import logging
+import typing as t
+from pathlib import Path
+
+import arrow
+from pydantic import BaseModel
+
+from .types import CacheBackend, CacheInvalidError, CacheMetadata, CacheOptions
+
+T = t.TypeVar("T", bound=BaseModel)
+K = t.TypeVar("K", bound=BaseModel)
+
+logger = logging.getLogger(__name__)
+
+
+class FileCacheBackend(CacheBackend):
+    """A generic cache for pydantic models"""
+
+    def __init__(self, cache_dir: str, default_options: CacheOptions):
+        self.cache_dir = cache_dir
+        self.default_options = default_options
+
+    def store_object(
+        self, key: str, value: BaseModel, override_options: CacheOptions | None = None
+    ) -> None:
+        """Store a single object in the cache
+
+        The file cache stores everything as a jsonl file. The first object in
+        the file is the metadata, which contains the creation time and
+        expiration time of the cache entry.
+        """
+
+        # Ensure the cache directory exists
+        self._ensure_cache_dir()
+        # Create a file path based on the key
+        file_path = self._cache_key_path(key)
+
+        metadata = CacheMetadata(
+            created_at=arrow.now().isoformat(),
+            valid_until=(
+                arrow.now().shift(seconds=self.default_options.ttl).isoformat()
+                if self.default_options.ttl > 0
+                else None
+            ),
+        )
+
+        # Write the value to the file
+        with open(file_path, "w") as f:
+            f.write(metadata.model_dump_json() + "\n")
+            f.write(value.model_dump_json())
+
+    def retrieve_object(
+        self,
+        key: str,
+        model_type: type[T],
+        override_options: CacheOptions | None = None,
+    ) -> T:
+        """Retrieve a single object from the cache"""
+        self._ensure_cache_dir()
+        file_path = self._cache_key_path(key)
+
+        if not file_path.exists():
+            logger.debug(
+                f"Cache file not found: {file_path}", extra={"file_path": file_path}
+            )
+            raise CacheInvalidError(f"Cache file not found: {file_path}")
+
+        with open(file_path, "r") as f:
+            # Read the metadata and check if it is valid
+            metadata = CacheMetadata.model_validate_json(f.readline().strip())
+
+            if not metadata.is_valid(override_options):
+                logger.debug(
+                    f"Cache entry is invalid: {metadata}", extra={"metadata": metadata}
+                )
+                raise CacheInvalidError(f"Cache entry is invalid: {metadata}")
+
+            return model_type.model_validate_json(f.read())
+
+    def _cache_dir_path(self):
+        """Get the cache directory path"""
+        return Path(self.cache_dir)
+
+    def _ensure_cache_dir(self):
+        """Ensure the cache directory exists"""
+        self._cache_dir_path().mkdir(parents=True, exist_ok=True)
+
+    def _cache_key(self, key: str) -> str:
+        """Generate a cache key from the pydantic model"""
+        key_str = hashlib.sha256(key.encode()).hexdigest()
+        return f"{key_str}.json"
+
+    def _cache_key_path(self, key: str) -> Path:
+        """Get the cache file path for a given key"""
+        return self._cache_dir_path() / self._cache_key(key)

lib/oso-core/oso_core/cache/generic.py

@@ -0,0 +1,40 @@
+import tempfile
+
+import pytest
+from pydantic import BaseModel
+
+from .file import FileCacheBackend
+from .types import CacheOptions
+
+
+class FakeModel(BaseModel):
+    """A simple Pydantic model for testing purposes."""
+
+    name: str
+    value: int
+
+
+class FakeNestedModel(BaseModel):
+    """A nested Pydantic model for testing purposes."""
+
+    name: str
+    nested: FakeModel
+
+
+@pytest.fixture
+def temp_dir():
+    """Create a temporary directory for file system tests."""
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        yield tmpdirname
+
+
+def test_write_file_to_temp_dir(temp_dir):
+    """Test writing a file to the temporary directory."""
+    cache = FileCacheBackend(cache_dir=temp_dir, default_options=CacheOptions(ttl=60))
+
+    test0 = FakeModel(name="test", value=42)
+
+    cache.store_object(key="test_key0", value=test0)
+    stored_test0 = cache.retrieve_object(key="test_key0", model_type=FakeModel)
+
+    assert test0 == stored_test0, "Stored and retrieved objects should match"

lib/oso-core/oso_core/cache/test_file.py

@@ -0,0 +1,9 @@
+from .types import CacheMetadata
+
+
+def test_metadata_is_valid():
+    """Test that metadata is valid when created and not expired."""
+    metadata = CacheMetadata(
+        created_at="2024-01-01T00:00:00Z", valid_until="2024-01-01T00:01:00Z"
+    )
+    assert metadata.is_valid() is True

lib/oso-core/oso_core/cache/test_types.py

@@ -0,0 +1,90 @@
+"""Caching utilities for pydantic models
+
+Types used to describe a set of caching utilities that can be used with any
+pydantic model. This allows us to cache many python objects.
+"""
+
+import abc
+import typing as t
+
+import arrow
+from pydantic import BaseModel
+
+T = t.TypeVar("T", bound=BaseModel)
+K = t.TypeVar("K", bound=BaseModel)
+
+
+class CacheOptions(BaseModel):
+    ttl: int = 0  # Time to live for the cache in seconds. 0 means no expiration.
+
+
+class NotFoundError(Exception):
+    pass
+
+
+class CacheInvalidError(Exception):
+    pass
+
+
+class CacheMetadata(BaseModel):
+    """Metadata for the cache entry
+
+    This metadata is used to store information about the cache entry such as
+    when it was created, when it expires, etc.
+
+    Attributes:
+        created_at: The time when the cache entry was created in ISO format.
+        valid_until: The time when the cache entry expires in ISO format. If
+                    None, the entry only expires if passed in options ttl is
+                    exceeded. If it is set, it is always checked against the
+                    current time.
+    """
+
+    created_at: str  # ISO format timestamp
+    valid_until: str | None = None  # ISO format timestamp
+
+    def is_valid(self, options: CacheOptions | None = None) -> bool:
+        """Check if the cache entry is valid based on the ttl"""
+        created_at = arrow.get(self.created_at)
+        now = arrow.now()
+        age = created_at - now
+
+        if self.valid_until:
+            valid_until = arrow.get(self.valid_until)
+            if valid_until < now:
+                return False
+
+        if options:
+            if options.ttl == 0:
+                return True
+
+            if age.total_seconds() > options.ttl:
+                return False
+
+        return False  # Placeholder for actual expiration logic
+
+
+class CacheBackend(abc.ABC):
+    """A generic cache backend interface
+
+    This is a generic interface for a cache backend that can be used to store
+    and retrieve pydantic models. The actual implementation of the cache backend
+    should inherit from this class and implement the methods.
+    """
+
+    @abc.abstractmethod
+    def store_object(
+        self, key: str, value: BaseModel, override_options: CacheOptions | None = None
+    ) -> None:
+        """Store a single object in the cache"""
+        ...
+
+    @abc.abstractmethod
+    def retrieve_object(
+        self,
+        key: str,
+        model_type: type[T],
+        override_options: CacheOptions | None = None,
+    ) -> T:
+        """Retrieve a single object from the cache"""
+        ...

lib/oso-core/oso_core/cache/types.py

@@ -6,6 +6,8 @@ readme = "README.md"
 requires-python = ">=3.10"
 authors = [{ name = "OSO Team", email = "[email protected]" }]
 dependencies = [
+    "arrow<2.0.0",
+    "orjson>=3.10.18",
     "sqlmesh[trino]<1.0.0,>=0.141.1",
     "structlog>=25.4.0",
 ]