fix: Snowflake tool leaks memory by default

[organic1337]

Snowflake tool cache is enabled by default, and the storage use for the caching the queries is simply a dict. Changed it so it'll use LRU cache instead.

This PR:

• Change to cachetools LRUCache object
• Add cachetools to README.md and pyproject.toml
• Add test coverage for caching

Test strategy:

• Validate that query is cached
• Mock execution the exceeds the cache size, and validate that the least recently used is deleted from the cache.

[joaomdmoura]

Disclaimer: This review was made by a crew of AI Agents.

Code Review Comment: PR #218 - Fix Caching Memory Leak

Overview

This pull request introduces significant improvements to the caching mechanism within the Snowflake search tool. By replacing an unbounded caching dictionary with an LRU (Least Recently Used) cache, the risk of memory leaks is mitigated. This is complemented by enhancements in test coverage and dependency management documentation.

Code Quality Insights

Positive Changes:

• Effective Memory Management: Replacing the unbounded _query_cache = {} with LRUCache significantly reduces potential memory issues.
• Maintainability Enhancements: The introduction of a constant CACHE_SIZE supports easier adjustments to cache size across the codebase.
• Proper Dependency Management: The inclusion of cachetools reflects a step towards using established libraries for better functionality.

Suggested Improvements:

1. Cache Size Configuration: It would enhance the flexibility of the caching mechanism to allow the cache size to be configurable through the constructor:

   class SnowflakeSearchTool(BaseTool):
       def __init__(self, cache_size: int = 128, *args, **kwargs):
           super().__init__(*args, **kwargs)
           self._query_cache = LRUCache(maxsize=cache_size)
   

2. Cache Monitoring: Implementing a way to monitor cache statistics would provide valuable insights into the cache's performance:

   @property
   def cache_stats(self):
       return {
           'size': len(self._query_cache),
           'maxsize': self._query_cache.maxsize,
           'hits': self._query_cache.hits,
           'misses': self._query_cache.misses
       }
   

3. Error Handling in Tests: Extend unit tests to include error handling scenarios, ensuring robustness in the caching implementation:

   @pytest.mark.asyncio
   async def test_cache_error_handling(snowflake_tool, mock_snowflake_connection, clear_cache):
       with patch.object(snowflake_tool, '_query_cache.get') as mock_cache_get:
           mock_cache_get.side_effect = Exception("Cache error")
           result = await snowflake_tool._run(query="SELECT * FROM test_table", timeout=300)
           assert result is not None
   

Historical Context and Implications

The modifications resonate with recent trends in developing more efficient memory management practices seen in earlier PRs, emphasizing test coverage and stability. It is critical to ensure that the caching logic aligns with the broader system performance, particularly in environments subjected to high query variability.

Recommendations:

1. Documentation: Update the class docstring to clearly document cache behavior, including how it manages memory.
2. Cache Warm-Up: Consider implementing a mechanism to "warm up" the cache for critical queries after initialization.
3. Monitoring Logs: Integrate logging functionality for cache hits and misses to facilitate performance monitoring:

   def _get_from_cache(self, key: str) -> Optional[Any]:
       result = self._query_cache.get(key)
       logger.debug(
           "Cache %s for key %s",
           "hit" if result is not None else "miss",
           key
       )
       return result
   

Security Considerations

The implementation thoughtfully addresses the management of sensitive data by ensuring cached data is evicted automatically, thereby guarding against information leakage.

Performance Impact

The proposed changes are expected to strengthen performance by maintaining frequently used queries in memory while limiting the overall size of the cache, with potential adjustments to the cache size based on usage patterns recommended.

Conclusion

The code changes presented in PR #218 effectively tackle the memory leak issue while adhering to good coding practices. The inclusion of comprehensive test cases and dependency management further bolsters the quality of the implementation. The PR stands ready for merging with the consideration of the proposed enhancements.