[feature]Implement session cache for Doris connections (#44)
* [feature]Implement session cache for Doris connections This PR introduces a `DorisSessionCache` to cache and reuse `DorisConnection` objects in memory. This helps to reduce the overhead of creating new connections, especially for frequently used system sessions like "query" and "system", and avoid not calling release_connection leads to `Connection acquisition timed out` when the number of connection pools reaches the maximum value. The PR #34 fixed the issue when calling the tool `exec_query`, but in the codebase, a large number of other tools directly using get_connection ("query") to get connection object but without calling the release_connection method will cause the connection to fail to be obtained after a certain number of times. Key changes: - Added `DorisSessionCache` class to manage the lifecycle of cached sessions. - The cache is configurable to store system sessions, user sessions, or both. By default, only system sessions are cached. - Integrated the session cache into `DorisConnectionManager`. - `get_connection` now checks the cache before creating a new connection. - `release_connection` removes the connection from the cache. * Add tests
This commit is contained in:
78
test/utils/test_db.py
Normal file
78
test/utils/test_db.py
Normal file
@@ -0,0 +1,78 @@
|
||||
from unittest.mock import MagicMock
|
||||
import pytest
|
||||
|
||||
from doris_mcp_server.utils.db import DorisConnection, DorisSessionCache
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def session_cache():
|
||||
"""Provides a DorisSessionCache instance with a mock connection manager."""
|
||||
connection_manager = MagicMock()
|
||||
cache = DorisSessionCache(connection_manager=connection_manager)
|
||||
yield cache, connection_manager
|
||||
|
||||
|
||||
class TestDorisSessionCache:
|
||||
|
||||
def test_initialization(self, session_cache):
|
||||
cache, _ = session_cache
|
||||
assert cache.cache_system_session is True
|
||||
assert cache.cache_user_session is False
|
||||
assert not cache.cached
|
||||
|
||||
def test_should_cache(self, session_cache):
|
||||
cache, _ = session_cache
|
||||
assert cache._should_cache("query") is True
|
||||
assert cache._should_cache("system") is True
|
||||
assert cache._should_cache("user-test-session-id") is False
|
||||
|
||||
cache.cache_user_session = True
|
||||
assert cache._should_cache("user-test-session-id") is True
|
||||
|
||||
def test_save_and_get_session(self, session_cache):
|
||||
cache, _ = session_cache
|
||||
mock_connection = MagicMock(spec=DorisConnection)
|
||||
mock_connection.session_id = "query"
|
||||
|
||||
cache.save(mock_connection)
|
||||
retrieved_conn = cache.get("query")
|
||||
assert retrieved_conn is mock_connection
|
||||
|
||||
mock_user_connection = MagicMock(spec=DorisConnection)
|
||||
mock_user_connection.session_id = "user-test-session-id"
|
||||
cache.save(mock_user_connection)
|
||||
assert cache.get("user-test-session-id") is None
|
||||
|
||||
cache.cache_user_session = True
|
||||
cache.save(mock_user_connection)
|
||||
retrieved_user_conn = cache.get("user-test-session-id")
|
||||
assert retrieved_user_conn is mock_user_connection
|
||||
|
||||
def test_remove_session(self, session_cache):
|
||||
cache, _ = session_cache
|
||||
mock_connection = MagicMock(spec=DorisConnection)
|
||||
mock_connection.session_id = "system"
|
||||
|
||||
cache.save(mock_connection)
|
||||
assert cache.get("system") is not None
|
||||
|
||||
cache.remove("system")
|
||||
assert cache.get("system") is None
|
||||
|
||||
def test_clear_cache(self, session_cache):
|
||||
cache, connection_manager = session_cache
|
||||
mock_conn1 = MagicMock(spec=DorisConnection)
|
||||
mock_conn1.session_id = "query"
|
||||
mock_conn2 = MagicMock(spec=DorisConnection)
|
||||
mock_conn2.session_id = "system"
|
||||
|
||||
cache.save(mock_conn1)
|
||||
cache.save(mock_conn2)
|
||||
assert len(cache.cached) == 2
|
||||
|
||||
cache.clear()
|
||||
|
||||
assert not cache.cached
|
||||
connection_manager.release_connection.assert_any_call("query", mock_conn1)
|
||||
connection_manager.release_connection.assert_any_call("system", mock_conn2)
|
||||
assert connection_manager.release_connection.call_count == 2
|
||||
Reference in New Issue
Block a user