From bdeedc5b3a647a6cbeb395b62001e214865c2100 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Thu, 19 Feb 2026 17:45:54 -0800 Subject: [PATCH] use single UNION ALL query for all non-LLM DB objects --- .../mcp_server/mcp_server_manager.py | 42 ++++ .../cache_settings_endpoints.py | 22 ++- .../policy_engine/attachment_registry.py | 32 ++- .../proxy/policy_engine/policy_registry.py | 33 +++- litellm/proxy/proxy_server.py | 183 +++++++++++------- .../proxy/db/test_litellm_config_cache.py | 95 ++++++++- 6 files changed, 321 insertions(+), 86 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index e5a2119bc2c..f739150c982 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2219,6 +2219,48 @@ class MCPServerManager: return None + async def reload_servers_from_records(self, db_records: list): + """Re-synchronize the in-memory MCP server registry with pre-fetched records.""" + from litellm.proxy._types import LiteLLM_MCPServerTable + + verbose_logger.debug("Loading MCP servers from pre-fetched records...") + + db_mcp_servers = [LiteLLM_MCPServerTable(**r.model_dump()) for r in db_records] + verbose_logger.info(f"Found {len(db_mcp_servers)} MCP servers in records") + + previous_registry = self.registry + new_registry: Dict[str, MCPServer] = {} + + for server in db_mcp_servers: + existing_server = previous_registry.get(server.server_id) + + if ( + existing_server is not None + and existing_server.updated_at is not None + and server.updated_at is not None + and existing_server.updated_at == server.updated_at + ): + new_registry[server.server_id] = existing_server + continue + + _warn_on_server_name_fields( + server_id=server.server_id, + alias=getattr(server, "alias", None), + server_name=getattr(server, "server_name", None), + ) + verbose_logger.debug( + f"Building server from DB: {server.server_id} ({server.server_name})" + ) + new_server = await self.build_mcp_server_from_table(server) + new_registry[server.server_id] = new_server + await self._maybe_register_openapi_tools(new_server) + + self.registry = new_registry + + verbose_logger.debug( + "MCP registry refreshed (%s servers in registry)", len(new_registry) + ) + async def reload_servers_from_database(self): """Re-synchronize the in-memory MCP server registry with the database.""" from litellm.proxy._experimental.mcp_server.db import get_all_mcp_servers diff --git a/litellm/proxy/management_endpoints/cache_settings_endpoints.py b/litellm/proxy/management_endpoints/cache_settings_endpoints.py index c9eeea15e26..3475916c1f6 100644 --- a/litellm/proxy/management_endpoints/cache_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/cache_settings_endpoints.py @@ -57,17 +57,29 @@ class CacheSettingsManager: return normalized1 == normalized2 @staticmethod - async def init_cache_settings_in_db(prisma_client, proxy_config): + async def init_cache_settings_in_db( + proxy_config, prisma_client=None, db_records=None + ): """ Initialize cache settings from database into the router on startup. Only reinitializes if cache params have changed. """ import json - + try: - cache_config = await prisma_client.db.litellm_cacheconfig.find_unique( - where={"id": "cache_config"} - ) + # Use pre-fetched records if available, otherwise query DB + if db_records is not None: + cache_config = None + for r in db_records: + if getattr(r, "id", None) == "cache_config": + cache_config = r + break + elif prisma_client is not None: + cache_config = await prisma_client.db.litellm_cacheconfig.find_unique( + where={"id": "cache_config"} + ) + else: + return if cache_config is not None and cache_config.cache_settings: # Parse cache settings JSON cache_settings_json = cache_config.cache_settings diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index 69b3b3599f3..c36dbaf7a51 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -447,16 +447,42 @@ class AttachmentRegistry: async def sync_attachments_from_db( self, - prisma_client: "PrismaClient", + prisma_client: "Optional[PrismaClient]" = None, + db_records: Optional[list] = None, ) -> None: """ Sync policy attachments from the database to in-memory registry. Args: - prisma_client: The Prisma client instance + prisma_client: The Prisma client instance (used if db_records not provided) + db_records: Pre-fetched records from batch query (preferred) """ try: - attachments = await self.get_all_attachments_from_db(prisma_client) + if db_records is not None: + from litellm.types.proxy.policy_engine.resolver_types import ( + PolicyAttachmentDBResponse, + ) + + attachments = [ + PolicyAttachmentDBResponse( + attachment_id=a.attachment_id, + policy_name=a.policy_name, + scope=a.scope, + teams=a.teams or [], + keys=a.keys or [], + models=a.models or [], + tags=a.tags or [], + created_at=a.created_at, + updated_at=a.updated_at, + created_by=a.created_by, + updated_by=a.updated_by, + ) + for a in db_records + ] + elif prisma_client is not None: + attachments = await self.get_all_attachments_from_db(prisma_client) + else: + return # Clear existing attachments and reload from DB self._attachments = [] diff --git a/litellm/proxy/policy_engine/policy_registry.py b/litellm/proxy/policy_engine/policy_registry.py index 1ca6062e5ac..4e1d651feb3 100644 --- a/litellm/proxy/policy_engine/policy_registry.py +++ b/litellm/proxy/policy_engine/policy_registry.py @@ -494,16 +494,43 @@ class PolicyRegistry: async def sync_policies_from_db( self, - prisma_client: "PrismaClient", + prisma_client: "Optional[PrismaClient]" = None, + db_records: Optional[list] = None, ) -> None: """ Sync policies from the database to in-memory registry. Args: - prisma_client: The Prisma client instance + prisma_client: The Prisma client instance (used if db_records not provided) + db_records: Pre-fetched records from batch query (preferred) """ try: - policies = await self.get_all_policies_from_db(prisma_client) + if db_records is not None: + from litellm.types.proxy.policy_engine.resolver_types import ( + PolicyDBResponse, + ) + + policies = [ + PolicyDBResponse( + policy_id=p.policy_id, + policy_name=p.policy_name, + inherit=p.inherit, + description=p.description, + guardrails_add=p.guardrails_add or [], + guardrails_remove=p.guardrails_remove or [], + condition=p.condition, + pipeline=p.pipeline, + created_at=p.created_at, + updated_at=p.updated_at, + created_by=p.created_by, + updated_by=p.updated_by, + ) + for p in db_records + ] + elif prisma_client is not None: + policies = await self.get_all_policies_from_db(prisma_client) + else: + return for policy_response in policies: policy = self._parse_policy( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f851987909f..80ad317fd15 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4112,11 +4112,14 @@ class ProxyConfig: global llm_router, llm_model_list, master_key, general_settings try: - # Batch-load all LiteLLM_Config records in one query to reduce DB connection pressure. - # Previously each sub-method made its own query, causing ~5 queries per cycle per worker. - from litellm.proxy.db.litellm_config_cache import batch_load_config + # Batch-load all config + non-LLM objects in 2 queries instead of ~18. + from litellm.proxy.db.litellm_config_cache import ( + batch_load_config, + batch_load_non_llm_objects, + ) config_map = await batch_load_config(prisma_client) + non_llm_objects = await batch_load_non_llm_objects(prisma_client) # Only load models from DB if "models" is in supported_db_objects (or if supported_db_objects is not set) if self._should_load_db_object(object_type="models"): @@ -4137,7 +4140,9 @@ class ProxyConfig: # initialize vector stores, guardrails, etc. table in db await self._init_non_llm_objects_in_db( - prisma_client=prisma_client, config_map=config_map + prisma_client=prisma_client, + config_map=config_map, + non_llm_objects=non_llm_objects, ) except Exception as e: @@ -4151,11 +4156,15 @@ class ProxyConfig: self, prisma_client: PrismaClient, config_map: Optional[Dict[str, Any]] = None, + non_llm_objects: Optional[Dict[str, List[Any]]] = None, ): """ Use this to read non-llm objects from the db and initialize them ex. Vector Stores, Guardrails, MCP tools, etc. + + All DB reads are done upfront via batch_load_non_llm_objects (1 query). + Each sub-method receives pre-fetched records and does zero DB queries. """ import asyncio @@ -4165,34 +4174,65 @@ class ProxyConfig: if config_map is None: config_map = {} + if non_llm_objects is None: + non_llm_objects = {} - # Build list of independent DB init tasks to run concurrently. - # Previously these ran sequentially (~13 queries x round-trip latency). - # Running them concurrently cuts wall-clock time per cycle so cycles - # don't overlap and pile up DB connections. + # All data is already fetched — sub-methods just process in-memory records. + # We still run them concurrently since some do non-DB I/O (e.g. MCP server init). tasks = [] if self._should_load_db_object(object_type="guardrails"): - tasks.append(self._init_guardrails_in_db(prisma_client=prisma_client)) + tasks.append( + self._init_guardrails_in_db( + db_records=non_llm_objects.get("guardrails", []) + ) + ) if self._should_load_db_object(object_type="policies"): - tasks.append(self._init_policies_in_db(prisma_client=prisma_client)) + tasks.append( + self._init_policies_in_db( + db_policies=non_llm_objects.get("policies", []), + db_attachments=non_llm_objects.get("policy_attachments", []), + ) + ) if self._should_load_db_object(object_type="vector_stores"): - tasks.append(self._init_vector_stores_in_db(prisma_client=prisma_client)) + tasks.append( + self._init_vector_stores_in_db( + db_records=non_llm_objects.get("vector_stores", []) + ) + ) if self._should_load_db_object(object_type="vector_store_indexes"): tasks.append( - self._init_vector_store_indexes_in_db(prisma_client=prisma_client) + self._init_vector_store_indexes_in_db( + db_records=non_llm_objects.get("vector_store_indexes", []) + ) ) if self._should_load_db_object(object_type="mcp"): - tasks.append(self._init_mcp_servers_in_db()) + tasks.append( + self._init_mcp_servers_in_db( + db_records=non_llm_objects.get("mcp_servers", []) + ) + ) if self._should_load_db_object(object_type="agents"): - tasks.append(self._init_agents_in_db(prisma_client=prisma_client)) + tasks.append( + self._init_agents_in_db( + db_records=non_llm_objects.get("agents", []) + ) + ) if self._should_load_db_object(object_type="pass_through_endpoints"): tasks.append( self._init_pass_through_endpoints_in_db(config_map=config_map) ) if self._should_load_db_object(object_type="prompts"): - tasks.append(self._init_prompts_in_db(prisma_client=prisma_client)) + tasks.append( + self._init_prompts_in_db( + db_records=non_llm_objects.get("prompts", []) + ) + ) if self._should_load_db_object(object_type="search_tools"): - tasks.append(self._init_search_tools_in_db(prisma_client=prisma_client)) + tasks.append( + self._init_search_tools_in_db( + db_records=non_llm_objects.get("search_tools", []) + ) + ) if self._should_load_db_object(object_type="model_cost_map"): tasks.append( self._check_and_reload_model_cost_map( @@ -4206,11 +4246,16 @@ class ProxyConfig: ) ) if self._should_load_db_object(object_type="sso_settings"): - tasks.append(self._init_sso_settings_in_db(prisma_client=prisma_client)) + tasks.append( + self._init_sso_settings_in_db( + db_records=non_llm_objects.get("sso_config", []) + ) + ) if self._should_load_db_object(object_type="cache_settings"): tasks.append( CacheSettingsManager.init_cache_settings_in_db( - prisma_client=prisma_client, proxy_config=self + db_records=non_llm_objects.get("cache_config", []), + proxy_config=self, ) ) if self._should_load_db_object(object_type="semantic_filter_settings"): @@ -4220,7 +4265,7 @@ class ProxyConfig: ) ) - # Run all DB init tasks concurrently. Each task handles its own errors + # Run all init tasks concurrently. Each task handles its own errors # internally, so we use return_exceptions=True to prevent one failure # from cancelling others. if tasks: @@ -4312,15 +4357,18 @@ class ProxyConfig: f"Error initializing semantic filter settings from DB: {e}" ) - async def _init_sso_settings_in_db(self, prisma_client: PrismaClient): + async def _init_sso_settings_in_db(self, db_records: List[Any]): """ - Initialize SSO settings from database into the router on startup. + Initialize SSO settings from pre-fetched records into the router on startup. """ - try: - sso_settings = await prisma_client.db.litellm_ssoconfig.find_unique( - where={"id": "sso_config"} - ) + # find_unique equivalent: look for id == "sso_config" + sso_settings = None + for r in db_records: + if getattr(r, "id", None) == "sso_config": + sso_settings = r + break + if sso_settings is not None: sso_settings.sso_settings.pop("role_mappings", None) sso_settings.sso_settings.pop("team_mappings", None) @@ -4567,14 +4615,11 @@ class ProxyConfig: return create_versioned_prompt_spec(db_prompt=db_prompt) - async def _init_prompts_in_db(self, prisma_client: PrismaClient): + async def _init_prompts_in_db(self, db_records: List[Any]): from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY - from litellm.types.prompts.init_prompts import PromptSpec try: - prompts_in_db = await prisma_client.db.litellm_prompttable.find_many() - for prompt in prompts_in_db: - # Convert DB object to dict and create versioned prompt_id + for prompt in db_records: prompt_spec = self._get_prompt_spec_for_db_prompt(db_prompt=prompt) IN_MEMORY_PROMPT_REGISTRY.initialize_prompt(prompt=prompt_spec) except Exception as e: @@ -4584,23 +4629,18 @@ class ProxyConfig: ) ) - async def _init_guardrails_in_db(self, prisma_client: PrismaClient): + async def _init_guardrails_in_db(self, db_records: List[Any]): from litellm.proxy.guardrails.guardrail_registry import ( IN_MEMORY_GUARDRAIL_HANDLER, Guardrail, - GuardrailRegistry, ) try: - guardrails_in_db: List[ - Guardrail - ] = await GuardrailRegistry.get_all_guardrails_from_db( - prisma_client=prisma_client - ) + guardrails = [Guardrail(**vars(r)) for r in db_records] verbose_proxy_logger.debug( - "guardrails from the DB %s", str(guardrails_in_db) + "guardrails from the DB %s", str(guardrails) ) - for guardrail in guardrails_in_db: + for guardrail in guardrails: IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db( guardrail=cast(Guardrail, guardrail), ) @@ -4611,9 +4651,11 @@ class ProxyConfig: ) ) - async def _init_policies_in_db(self, prisma_client: PrismaClient): + async def _init_policies_in_db( + self, db_policies: List[Any], db_attachments: List[Any] + ): """ - Initialize policies and policy attachments from database into the in-memory registries. + Initialize policies and policy attachments from pre-fetched records into the in-memory registries. """ from litellm.proxy.policy_engine.attachment_registry import ( get_attachment_registry, @@ -4621,16 +4663,12 @@ class ProxyConfig: from litellm.proxy.policy_engine.policy_registry import get_policy_registry try: - # Get the global singleton instances policy_registry = get_policy_registry() attachment_registry = get_attachment_registry() - # Sync policies from DB to in-memory registry - await policy_registry.sync_policies_from_db(prisma_client=prisma_client) - - # Sync attachments from DB to in-memory registry + await policy_registry.sync_policies_from_db(db_records=db_policies) await attachment_registry.sync_attachments_from_db( - prisma_client=prisma_client + db_records=db_attachments ) verbose_proxy_logger.debug( @@ -4643,14 +4681,14 @@ class ProxyConfig: ) ) - async def _init_vector_stores_in_db(self, prisma_client: PrismaClient): + async def _init_vector_stores_in_db(self, db_records: List[Any]): + from litellm.types.vector_stores import LiteLLM_ManagedVectorStore from litellm.vector_stores.vector_store_registry import VectorStoreRegistry try: - # read vector stores from db table - vector_stores = await VectorStoreRegistry._get_vector_stores_from_db( - prisma_client=prisma_client - ) + vector_stores = [ + LiteLLM_ManagedVectorStore(**vars(r)) for r in db_records + ] if len(vector_stores) <= 0: return @@ -4670,16 +4708,14 @@ class ProxyConfig: ) ) - async def _init_vector_store_indexes_in_db(self, prisma_client: PrismaClient): + async def _init_vector_store_indexes_in_db(self, db_records: List[Any]): + from litellm.types.vector_stores import LiteLLM_ManagedVectorStoreIndex from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry try: - # read vector stores from db table - vector_store_indexes = ( - await VectorStoreIndexRegistry._get_vector_store_indexes_from_db( - prisma_client=prisma_client - ) - ) + vector_store_indexes = [ + LiteLLM_ManagedVectorStoreIndex(**vars(r)) for r in db_records + ] if len(vector_store_indexes) <= 0: return @@ -4700,7 +4736,7 @@ class ProxyConfig: ) ) - async def _init_mcp_servers_in_db(self): + async def _init_mcp_servers_in_db(self, db_records: List[Any]): from litellm.proxy._experimental.mcp_server.utils import is_mcp_available if not is_mcp_available(): @@ -4714,7 +4750,9 @@ class ProxyConfig: ) try: - await global_mcp_server_manager.reload_servers_from_database() + await global_mcp_server_manager.reload_servers_from_records( + db_records=db_records + ) except Exception as e: verbose_proxy_logger.exception( "litellm.proxy.proxy_server.py::ProxyConfig:_init_mcp_servers_in_db - {}".format( @@ -4722,15 +4760,13 @@ class ProxyConfig: ) ) - async def _init_agents_in_db(self, prisma_client: PrismaClient): + async def _init_agents_in_db(self, db_records: List[Any]): from litellm.proxy.agent_endpoints.agent_registry import ( global_agent_registry as AGENT_REGISTRY, ) try: - db_agents = await AGENT_REGISTRY.get_all_agents_from_db( - prisma_client=prisma_client - ) + db_agents = [vars(r) for r in db_records] AGENT_REGISTRY.load_agents_from_db_and_config( db_agents=db_agents, agent_config=config_agents ) @@ -4741,10 +4777,10 @@ class ProxyConfig: ) ) - async def _init_search_tools_in_db(self, prisma_client: PrismaClient): + async def _init_search_tools_in_db(self, db_records: List[Any]): """ - Initialize search tools from database into the router on startup. - Only updates router if there are tools in the database, otherwise preserves config-loaded tools. + Initialize search tools from pre-fetched records into the router on startup. + Only updates router if there are tools, otherwise preserves config-loaded tools. """ global llm_router @@ -4752,21 +4788,20 @@ class ProxyConfig: SearchToolRegistry, ) from litellm.router_utils.search_api_router import SearchAPIRouter + from litellm.types.search import SearchTool try: - search_tools = await SearchToolRegistry.get_all_search_tools_from_db( - prisma_client=prisma_client - ) + search_tools = [] + for r in db_records: + search_tool_dict = SearchToolRegistry._convert_prisma_to_dict(r) + search_tools.append(SearchTool(**search_tool_dict)) # type: ignore verbose_proxy_logger.info( f"Loading {len(search_tools)} search tool(s) from database into router" ) - # Only update router if there are tools in the database - # This prevents overwriting config-loaded tools with an empty list if len(search_tools) > 0: if llm_router is not None: - # Add search tools to the router await SearchAPIRouter.update_router_search_tools( router_instance=llm_router, search_tools=search_tools ) diff --git a/tests/test_litellm/proxy/db/test_litellm_config_cache.py b/tests/test_litellm/proxy/db/test_litellm_config_cache.py index 2ddbc5b0681..cadbb2470d8 100644 --- a/tests/test_litellm/proxy/db/test_litellm_config_cache.py +++ b/tests/test_litellm/proxy/db/test_litellm_config_cache.py @@ -1,15 +1,20 @@ """ -Tests for batch_load_config - batching LiteLLM_Config queries into one find_many. +Tests for batch_load_config and batch_load_non_llm_objects. """ +import json from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock import pytest from litellm.proxy.db.litellm_config_cache import ( + _NON_LLM_TABLES, POLLING_PARAM_NAMES, + _DBRecord, + _wrap_records, batch_load_config, + batch_load_non_llm_objects, ) @@ -25,6 +30,9 @@ def _make_mock_prisma_client(records): return mock_client +# ─── batch_load_config tests ─── + + @pytest.mark.asyncio async def test_batch_load_fetches_all_config_in_one_query(): """batch_load_config should call find_many once with all param names.""" @@ -92,3 +100,88 @@ async def test_batch_load_contains_all_polling_param_names_when_all_exist(): for name in POLLING_PARAM_NAMES: assert name in result assert result[name].param_value == {"data": name} + + +# ─── _DBRecord tests ─── + + +def test_db_record_attribute_access(): + """_DBRecord should support attribute access like Prisma models.""" + record = _DBRecord(name="test", value=42) + assert record.name == "test" + assert record.value == 42 + + +def test_db_record_model_dump(): + """_DBRecord.model_dump() should return a dict of all fields.""" + record = _DBRecord(name="test", value=42, nested={"a": 1}) + dumped = record.model_dump() + assert dumped == {"name": "test", "value": 42, "nested": {"a": 1}} + + +def test_wrap_records(): + """_wrap_records should convert list of dicts to list of _DBRecord.""" + raw = [{"id": "1", "name": "a"}, {"id": "2", "name": "b"}] + wrapped = _wrap_records(raw) + assert len(wrapped) == 2 + assert wrapped[0].id == "1" + assert wrapped[1].name == "b" + assert wrapped[0].model_dump() == {"id": "1", "name": "a"} + + +# ─── batch_load_non_llm_objects tests ─── + + +@pytest.mark.asyncio +async def test_batch_load_non_llm_objects_single_query(): + """batch_load_non_llm_objects should call query_raw once.""" + # Simulate query_raw returning one row per table + raw_rows = [ + {"_tbl": "guardrails", "data": json.dumps([{"guardrail_id": "g1", "name": "test"}])}, + {"_tbl": "policies", "data": json.dumps([])}, + {"_tbl": "agents", "data": json.dumps([{"agent_id": "a1"}])}, + ] + mock_client = MagicMock() + mock_client.db.query_raw = AsyncMock(return_value=raw_rows) + + result = await batch_load_non_llm_objects(mock_client) + + mock_client.db.query_raw.assert_called_once() + assert len(result["guardrails"]) == 1 + assert result["guardrails"][0].guardrail_id == "g1" + assert result["policies"] == [] + assert result["agents"][0].agent_id == "a1" + + +@pytest.mark.asyncio +async def test_batch_load_non_llm_objects_all_tables(): + """Result should contain keys for all tables returned by query.""" + raw_rows = [ + {"_tbl": key, "data": json.dumps([])} + for key, _ in _NON_LLM_TABLES + ] + mock_client = MagicMock() + mock_client.db.query_raw = AsyncMock(return_value=raw_rows) + + result = await batch_load_non_llm_objects(mock_client) + + for key, _ in _NON_LLM_TABLES: + assert key in result + + +@pytest.mark.asyncio +async def test_batch_load_non_llm_objects_records_have_model_dump(): + """Records should support model_dump() for Prisma compatibility.""" + raw_rows = [ + {"_tbl": "prompts", "data": json.dumps([{"prompt_id": "p1", "content": "hello"}])}, + ] + mock_client = MagicMock() + mock_client.db.query_raw = AsyncMock(return_value=raw_rows) + + result = await batch_load_non_llm_objects(mock_client) + + prompt = result["prompts"][0] + assert prompt.prompt_id == "p1" + dumped = prompt.model_dump() + assert dumped["prompt_id"] == "p1" + assert dumped["content"] == "hello"