From c6a8ae4539ec7129544a2dea62cff948a1811331 Mon Sep 17 00:00:00 2001 From: Frankie Fung Date: Sat, 25 Oct 2025 08:13:20 +0800 Subject: [PATCH] feat: Accept database_id parameter in Firestore composite index creation Fixes #5169 - Added optional database_id parameter to FirestoreDb.__init__() after project_id - Updated create_collection_indexes() to accept and pass database_id - Modified _create_composite_indexes() to use database_id in parent_path - Updated _get_or_create_collection() to pass database_id through - Defaults to '(default)' for backward compatibility This enables users to create composite indexes on non-default Firestore databases, allowing for multi-database Firestore configurations. Changes are fully backward compatible as database_id is optional and existing code using keyword arguments will continue to work. --- libs/agno/agno/db/firestore/firestore.py | 5 ++++- libs/agno/agno/db/firestore/utils.py | 10 ++++++---- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/libs/agno/agno/db/firestore/firestore.py b/libs/agno/agno/db/firestore/firestore.py index b3cd1a82ed1..f5417f527c1 100644 --- a/libs/agno/agno/db/firestore/firestore.py +++ b/libs/agno/agno/db/firestore/firestore.py @@ -39,6 +39,7 @@ def __init__( self, db_client: Optional[Client] = None, project_id: Optional[str] = None, + database_id: Optional[str] = None, session_collection: Optional[str] = None, memory_collection: Optional[str] = None, metrics_collection: Optional[str] = None, @@ -53,6 +54,7 @@ def __init__( Args: db_client (Optional[Client]): The Firestore client to use. project_id (Optional[str]): The GCP project ID for Firestore. + database_id (Optional[str]): The Firestore database ID (defaults to "(default)"). session_collection (Optional[str]): Name of the collection to store sessions. memory_collection (Optional[str]): Name of the collection to store memories. metrics_collection (Optional[str]): Name of the collection to store metrics. @@ -86,6 +88,7 @@ def __init__( self.project_id: Optional[str] = project_id self.db_client: Client = _client + self.database_id: str = database_id or "(default)" # -- DB methods -- @@ -186,7 +189,7 @@ def _get_or_create_collection( if not hasattr(self, f"_{collection_name}_initialized"): if not create_collection_if_not_found: return None - create_collection_indexes(self.db_client, collection_name, collection_type) + create_collection_indexes(self.db_client, collection_name, collection_type, database_id=self.database_id) setattr(self, f"_{collection_name}_initialized", True) return collection_ref diff --git a/libs/agno/agno/db/firestore/utils.py b/libs/agno/agno/db/firestore/utils.py index 31603ec1534..1cc2ff14ebc 100644 --- a/libs/agno/agno/db/firestore/utils.py +++ b/libs/agno/agno/db/firestore/utils.py @@ -22,7 +22,7 @@ # -- DB util methods -- -def create_collection_indexes(client: Client, collection_name: str, collection_type: str) -> None: +def create_collection_indexes(client: Client, collection_name: str, collection_type: str, database_id: Optional[str] = None) -> None: """Create all required indexes for a collection including composite indexes. This function automatically creates both single-field and composite indexes. @@ -38,14 +38,14 @@ def create_collection_indexes(client: Client, collection_name: str, collection_t # Create composite indexes programmatically if composite_indexes: - _create_composite_indexes(client, collection_name, composite_indexes) + _create_composite_indexes(client, collection_name, composite_indexes, database_id=database_id) log_debug(f"Collection '{collection_name}' initialized") except Exception as e: log_warning(f"Error processing indexes for {collection_type} collection: {e}") -def _create_composite_indexes(client: Client, collection_name: str, composite_indexes: List[Dict[str, Any]]) -> None: +def _create_composite_indexes(client: Client, collection_name: str, composite_indexes: List[Dict[str, Any]], database_id: Optional[str] = None) -> None: """Create composite indexes using Firestore Admin API.""" try: project_id = client.project @@ -53,6 +53,8 @@ def _create_composite_indexes(client: Client, collection_name: str, composite_in log_warning("Cannot create composite indexes: project_id not available from client") return + db_id = database_id or "(default)" + admin_client = FirestoreAdminClient() created_count = 0 @@ -83,7 +85,7 @@ def _create_composite_indexes(client: Client, collection_name: str, composite_in ) # Create the index - parent_path = f"projects/{project_id}/databases/(default)/collectionGroups/{collection_name}" + parent_path = f"projects/{project_id}/databases/{db_id}/collectionGroups/{collection_name}" admin_client.create_index(parent=parent_path, index=index) created_count += 1