feat(firestore): Add Firestore Multi Database Support (#818)

jonathanedey · web-flow · commit 8727e91739f2 · 2024-10-24T10:18:05.000-04:00
* Added multi db support for firestore and firestore_async

* Added unit and integration tests

* fix docs strings
diff --git a/firebase_admin/firestore.py b/firebase_admin/firestore.py
@@ -18,59 +18,75 @@
 Firebase apps. This requires the ``google-cloud-firestore`` Python module.
 """
 
+from __future__ import annotations
+from typing import Optional, Dict
+from firebase_admin import App
+from firebase_admin import _utils
+
 try:
-    from google.cloud import firestore # pylint: disable=import-error,no-name-in-module
+    from google.cloud import firestore
+    from google.cloud.firestore_v1.base_client import DEFAULT_DATABASE
     existing = globals().keys()
     for key, value in firestore.__dict__.items():
         if not key.startswith('_') and key not in existing:
             globals()[key] = value
-except ImportError:
+except ImportError as error:
     raise ImportError('Failed to import the Cloud Firestore library for Python. Make sure '
-                      'to install the "google-cloud-firestore" module.')
-
-from firebase_admin import _utils
+                      'to install the "google-cloud-firestore" module.') from error
 
 
 _FIRESTORE_ATTRIBUTE = '_firestore'
 
 
-def client(app=None) -> firestore.Client:
+def client(app: Optional[App] = None, database_id: Optional[str] = None) -> firestore.Client:
     """Returns a client that can be used to interact with Google Cloud Firestore.
 
     Args:
-      app: An App instance (optional).
+        app: An App instance (optional).
+        database_id: The database ID of the Google Cloud Firestore database to be used.
+            Defaults to the default Firestore database ID if not specified or an empty string
+            (optional).
 
     Returns:
-      google.cloud.firestore.Firestore: A `Firestore Client`_.
+        google.cloud.firestore.Firestore: A `Firestore Client`_.
 
     Raises:
-      ValueError: If a project ID is not specified either via options, credentials or
-          environment variables, or if the specified project ID is not a valid string.
+        ValueError: If the specified database ID is not a valid string, or if a project ID is not
+            specified either via options, credentials or environment variables, or if the specified
+            project ID is not a valid string.
 
-    .. _Firestore Client: https://googlecloudplatform.github.io/google-cloud-python/latest\
-          /firestore/client.html
+    .. _Firestore Client: https://cloud.google.com/python/docs/reference/firestore/latest/\
+        google.cloud.firestore_v1.client.Client
     """
-    fs_client = _utils.get_app_service(app, _FIRESTORE_ATTRIBUTE, _FirestoreClient.from_app)
-    return fs_client.get()
-
-
-class _FirestoreClient:
-    """Holds a Google Cloud Firestore client instance."""
-
-    def __init__(self, credentials, project):
-        self._client = firestore.Client(credentials=credentials, project=project)
-
-    def get(self):
-        return self._client
-
-    @classmethod
-    def from_app(cls, app):
-        """Creates a new _FirestoreClient for the specified app."""
-        credentials = app.credential.get_credential()
-        project = app.project_id
-        if not project:
-            raise ValueError(
-                'Project ID is required to access Firestore. Either set the projectId option, '
-                'or use service account credentials. Alternatively, set the GOOGLE_CLOUD_PROJECT '
-                'environment variable.')
-        return _FirestoreClient(credentials, project)
+    # Validate database_id
+    if database_id is not None and not isinstance(database_id, str):
+        raise ValueError(f'database_id "{database_id}" must be a string or None.')
+    fs_service = _utils.get_app_service(app, _FIRESTORE_ATTRIBUTE, _FirestoreService)
+    return fs_service.get_client(database_id)
+
+
+class _FirestoreService:
+    """Service that maintains a collection of firestore clients."""
+
+    def __init__(self, app: App) -> None:
+        self._app: App = app
+        self._clients: Dict[str, firestore.Client] = {}
+
+    def get_client(self, database_id: Optional[str]) -> firestore.Client:
+        """Creates a client based on the database_id. These clients are cached."""
+        database_id = database_id or DEFAULT_DATABASE
+        if database_id not in self._clients:
+            # Create a new client and cache it in _clients
+            credentials = self._app.credential.get_credential()
+            project = self._app.project_id
+            if not project:
+                raise ValueError(
+                    'Project ID is required to access Firestore. Either set the projectId option, '
+                    'or use service account credentials. Alternatively, set the '
+                    'GOOGLE_CLOUD_PROJECT environment variable.')
+
+            fs_client = firestore.Client(
+                credentials=credentials, project=project, database=database_id)
+            self._clients[database_id] = fs_client
+
+        return self._clients[database_id]
diff --git a/firebase_admin/firestore_async.py b/firebase_admin/firestore_async.py
@@ -18,65 +18,75 @@
 associated with Firebase apps. This requires the ``google-cloud-firestore`` Python module.
 """
 
-from typing import Type
-
-from firebase_admin import (
-    App,
-    _utils,
-)
-from firebase_admin.credentials import Base
+from __future__ import annotations
+from typing import Optional, Dict
+from firebase_admin import App
+from firebase_admin import _utils
 
 try:
-    from google.cloud import firestore # type: ignore # pylint: disable=import-error,no-name-in-module
+    from google.cloud import firestore
+    from google.cloud.firestore_v1.base_client import DEFAULT_DATABASE
     existing = globals().keys()
     for key, value in firestore.__dict__.items():
         if not key.startswith('_') and key not in existing:
             globals()[key] = value
-except ImportError:
+except ImportError as error:
     raise ImportError('Failed to import the Cloud Firestore library for Python. Make sure '
-                      'to install the "google-cloud-firestore" module.')
+                      'to install the "google-cloud-firestore" module.') from error
+
 
 _FIRESTORE_ASYNC_ATTRIBUTE: str = '_firestore_async'
 
 
-def client(app: App = None) -> firestore.AsyncClient:
+def client(app: Optional[App] = None, database_id: Optional[str] = None) -> firestore.AsyncClient:
     """Returns an async client that can be used to interact with Google Cloud Firestore.
 
     Args:
-      app: An App instance (optional).
+        app: An App instance (optional).
+        database_id: The database ID of the Google Cloud Firestore database to be used.
+            Defaults to the default Firestore database ID if not specified or an empty string
+            (optional).
 
     Returns:
-      google.cloud.firestore.Firestore_Async: A `Firestore Async Client`_.
+        google.cloud.firestore.Firestore_Async: A `Firestore Async Client`_.
 
     Raises:
-      ValueError: If a project ID is not specified either via options, credentials or
-          environment variables, or if the specified project ID is not a valid string.
+        ValueError: If the specified database ID is not a valid string, or if a project ID is not
+            specified either via options, credentials or environment variables, or if the specified
+            project ID is not a valid string.
 
-    .. _Firestore Async Client: https://googleapis.dev/python/firestore/latest/client.html
+    .. _Firestore Async Client: https://cloud.google.com/python/docs/reference/firestore/latest/\
+        google.cloud.firestore_v1.async_client.AsyncClient
     """
-    fs_client = _utils.get_app_service(
-        app, _FIRESTORE_ASYNC_ATTRIBUTE, _FirestoreAsyncClient.from_app)
-    return fs_client.get()
-
-
-class _FirestoreAsyncClient:
-    """Holds a Google Cloud Firestore Async Client instance."""
-
-    def __init__(self, credentials: Type[Base], project: str) -> None:
-        self._client = firestore.AsyncClient(credentials=credentials, project=project)
-
-    def get(self) -> firestore.AsyncClient:
-        return self._client
-
-    @classmethod
-    def from_app(cls, app: App) -> "_FirestoreAsyncClient":
-        # Replace remove future reference quotes by importing annotations in Python 3.7+ b/238779406
-        """Creates a new _FirestoreAsyncClient for the specified app."""
-        credentials = app.credential.get_credential()
-        project = app.project_id
-        if not project:
-            raise ValueError(
-                'Project ID is required to access Firestore. Either set the projectId option, '
-                'or use service account credentials. Alternatively, set the GOOGLE_CLOUD_PROJECT '
-                'environment variable.')
-        return _FirestoreAsyncClient(credentials, project)
+    # Validate database_id
+    if database_id is not None and not isinstance(database_id, str):
+        raise ValueError(f'database_id "{database_id}" must be a string or None.')
+
+    fs_service = _utils.get_app_service(app, _FIRESTORE_ASYNC_ATTRIBUTE, _FirestoreAsyncService)
+    return fs_service.get_client(database_id)
+
+class _FirestoreAsyncService:
+    """Service that maintains a collection of firestore async clients."""
+
+    def __init__(self, app: App) -> None:
+        self._app: App = app
+        self._clients: Dict[str, firestore.AsyncClient] = {}
+
+    def get_client(self, database_id: Optional[str]) -> firestore.AsyncClient:
+        """Creates an async client based on the database_id. These clients are cached."""
+        database_id = database_id or DEFAULT_DATABASE
+        if database_id not in self._clients:
+            # Create a new client and cache it in _clients
+            credentials = self._app.credential.get_credential()
+            project = self._app.project_id
+            if not project:
+                raise ValueError(
+                    'Project ID is required to access Firestore. Either set the projectId option, '
+                    'or use service account credentials. Alternatively, set the '
+                    'GOOGLE_CLOUD_PROJECT environment variable.')
+
+            fs_client = firestore.AsyncClient(
+                credentials=credentials, project=project, database=database_id)
+            self._clients[database_id] = fs_client
+
+        return self._clients[database_id]
diff --git a/integration/test_firestore.py b/integration/test_firestore.py
@@ -17,6 +17,20 @@
 
 from firebase_admin import firestore
 
+_CITY = {
+        'name': u'Mountain View',
+        'country': u'USA',
+        'population': 77846,
+        'capital': False
+    }
+
+_MOVIE = {
+        'Name': u'Interstellar',
+        'Year': 2014,
+        'Runtime': u'2h 49m',
+        'Academy Award Winner': True
+    }
+
 
 def test_firestore():
     client = firestore.client()
@@ -35,6 +49,47 @@ def test_firestore():
     doc.delete()
     assert doc.get().exists is False
 
+def test_firestore_explicit_database_id():
+    client = firestore.client(database_id='testing-database')
+    expected = _CITY
+    doc = client.collection('cities').document()
+    doc.set(expected)
+
+    data = doc.get()
+    assert data.to_dict() == expected
+
+    doc.delete()
+    data = doc.get()
+    assert data.exists is False
+
+def test_firestore_multi_db():
+    city_client = firestore.client()
+    movie_client = firestore.client(database_id='testing-database')
+
+    expected_city = _CITY
+    expected_movie = _MOVIE
+
+    city_doc = city_client.collection('cities').document()
+    movie_doc = movie_client.collection('movies').document()
+
+    city_doc.set(expected_city)
+    movie_doc.set(expected_movie)
+
+    city_data = city_doc.get()
+    movie_data = movie_doc.get()
+
+    assert city_data.to_dict() == expected_city
+    assert movie_data.to_dict() == expected_movie
+
+    city_doc.delete()
+    movie_doc.delete()
+
+    city_data = city_doc.get()
+    movie_data = movie_doc.get()
+
+    assert city_data.exists is False
+    assert movie_data.exists is False
+
 def test_server_timestamp():
     client = firestore.client()
     expected = {
diff --git a/integration/test_firestore_async.py b/integration/test_firestore_async.py
@@ -13,20 +13,31 @@
 # limitations under the License.
 
 """Integration tests for firebase_admin.firestore_async module."""
+import asyncio
 import datetime
 import pytest
 
 from firebase_admin import firestore_async
 
-@pytest.mark.asyncio
-async def test_firestore_async():
-    client = firestore_async.client()
-    expected = {
+_CITY = {
         'name': u'Mountain View',
         'country': u'USA',
         'population': 77846,
         'capital': False
     }
+
+_MOVIE = {
+        'Name': u'Interstellar',
+        'Year': 2014,
+        'Runtime': u'2h 49m',
+        'Academy Award Winner': True
+    }
+
+
+@pytest.mark.asyncio
+async def test_firestore_async():
+    client = firestore_async.client()
+    expected = _CITY
     doc = client.collection('cities').document()
     await doc.set(expected)
 
@@ -37,6 +48,56 @@ async def test_firestore_async():
     data = await doc.get()
     assert data.exists is False
 
+@pytest.mark.asyncio
+async def test_firestore_async_explicit_database_id():
+    client = firestore_async.client(database_id='testing-database')
+    expected = _CITY
+    doc = client.collection('cities').document()
+    await doc.set(expected)
+
+    data = await doc.get()
+    assert data.to_dict() == expected
+
+    await doc.delete()
+    data = await doc.get()
+    assert data.exists is False
+
+@pytest.mark.asyncio
+async def test_firestore_async_multi_db():
+    city_client = firestore_async.client()
+    movie_client = firestore_async.client(database_id='testing-database')
+
+    expected_city = _CITY
+    expected_movie = _MOVIE
+
+    city_doc = city_client.collection('cities').document()
+    movie_doc = movie_client.collection('movies').document()
+
+    await asyncio.gather(
+        city_doc.set(expected_city),
+        movie_doc.set(expected_movie)
+    )
+
+    data = await asyncio.gather(
+        city_doc.get(),
+        movie_doc.get()
+    )
+
+    assert data[0].to_dict() == expected_city
+    assert data[1].to_dict() == expected_movie
+
+    await asyncio.gather(
+        city_doc.delete(),
+        movie_doc.delete()
+    )
+
+    data = await asyncio.gather(
+        city_doc.get(),
+        movie_doc.get()
+    )
+    assert data[0].exists is False
+    assert data[1].exists is False
+
 @pytest.mark.asyncio
 async def test_server_timestamp():
     client = firestore_async.client()
diff --git a/tests/test_firestore.py b/tests/test_firestore.py
diff --git a/tests/test_firestore_async.py b/tests/test_firestore_async.py
