|
18 | 18 | Firebase apps. This requires the ``google-cloud-firestore`` Python module. |
19 | 19 | """ |
20 | 20 |
|
| 21 | +from __future__ import annotations |
| 22 | +from typing import Optional, Dict |
| 23 | +from firebase_admin import App |
| 24 | +from firebase_admin import _utils |
| 25 | + |
21 | 26 | try: |
22 | | - from google.cloud import firestore # pylint: disable=import-error,no-name-in-module |
| 27 | + from google.cloud import firestore |
| 28 | + from google.cloud.firestore_v1.base_client import DEFAULT_DATABASE |
23 | 29 | existing = globals().keys() |
24 | 30 | for key, value in firestore.__dict__.items(): |
25 | 31 | if not key.startswith('_') and key not in existing: |
26 | 32 | globals()[key] = value |
27 | | -except ImportError: |
| 33 | +except ImportError as error: |
28 | 34 | raise ImportError('Failed to import the Cloud Firestore library for Python. Make sure ' |
29 | | - 'to install the "google-cloud-firestore" module.') |
30 | | - |
31 | | -from firebase_admin import _utils |
| 35 | + 'to install the "google-cloud-firestore" module.') from error |
32 | 36 |
|
33 | 37 |
|
34 | 38 | _FIRESTORE_ATTRIBUTE = '_firestore' |
35 | 39 |
|
36 | 40 |
|
37 | | -def client(app=None) -> firestore.Client: |
| 41 | +def client(app: Optional[App] = None, database_id: Optional[str] = None) -> firestore.Client: |
38 | 42 | """Returns a client that can be used to interact with Google Cloud Firestore. |
39 | 43 |
|
40 | 44 | Args: |
41 | | - app: An App instance (optional). |
| 45 | + app: An App instance (optional). |
| 46 | + database_id: The database ID of the Google Cloud Firestore database to be used. |
| 47 | + Defaults to the default Firestore database ID if not specified or an empty string |
| 48 | + (optional). |
42 | 49 |
|
43 | 50 | Returns: |
44 | | - google.cloud.firestore.Firestore: A `Firestore Client`_. |
| 51 | + google.cloud.firestore.Firestore: A `Firestore Client`_. |
45 | 52 |
|
46 | 53 | Raises: |
47 | | - ValueError: If a project ID is not specified either via options, credentials or |
48 | | - environment variables, or if the specified project ID is not a valid string. |
| 54 | + ValueError: If the specified database ID is not a valid string, or if a project ID is not |
| 55 | + specified either via options, credentials or environment variables, or if the specified |
| 56 | + project ID is not a valid string. |
49 | 57 |
|
50 | | - .. _Firestore Client: https://googlecloudplatform.github.io/google-cloud-python/latest\ |
51 | | - /firestore/client.html |
| 58 | + .. _Firestore Client: https://cloud.google.com/python/docs/reference/firestore/latest/\ |
| 59 | + google.cloud.firestore_v1.client.Client |
52 | 60 | """ |
53 | | - fs_client = _utils.get_app_service(app, _FIRESTORE_ATTRIBUTE, _FirestoreClient.from_app) |
54 | | - return fs_client.get() |
55 | | - |
56 | | - |
57 | | -class _FirestoreClient: |
58 | | - """Holds a Google Cloud Firestore client instance.""" |
59 | | - |
60 | | - def __init__(self, credentials, project): |
61 | | - self._client = firestore.Client(credentials=credentials, project=project) |
62 | | - |
63 | | - def get(self): |
64 | | - return self._client |
65 | | - |
66 | | - @classmethod |
67 | | - def from_app(cls, app): |
68 | | - """Creates a new _FirestoreClient for the specified app.""" |
69 | | - credentials = app.credential.get_credential() |
70 | | - project = app.project_id |
71 | | - if not project: |
72 | | - raise ValueError( |
73 | | - 'Project ID is required to access Firestore. Either set the projectId option, ' |
74 | | - 'or use service account credentials. Alternatively, set the GOOGLE_CLOUD_PROJECT ' |
75 | | - 'environment variable.') |
76 | | - return _FirestoreClient(credentials, project) |
| 61 | + # Validate database_id |
| 62 | + if database_id is not None and not isinstance(database_id, str): |
| 63 | + raise ValueError(f'database_id "{database_id}" must be a string or None.') |
| 64 | + fs_service = _utils.get_app_service(app, _FIRESTORE_ATTRIBUTE, _FirestoreService) |
| 65 | + return fs_service.get_client(database_id) |
| 66 | + |
| 67 | + |
| 68 | +class _FirestoreService: |
| 69 | + """Service that maintains a collection of firestore clients.""" |
| 70 | + |
| 71 | + def __init__(self, app: App) -> None: |
| 72 | + self._app: App = app |
| 73 | + self._clients: Dict[str, firestore.Client] = {} |
| 74 | + |
| 75 | + def get_client(self, database_id: Optional[str]) -> firestore.Client: |
| 76 | + """Creates a client based on the database_id. These clients are cached.""" |
| 77 | + database_id = database_id or DEFAULT_DATABASE |
| 78 | + if database_id not in self._clients: |
| 79 | + # Create a new client and cache it in _clients |
| 80 | + credentials = self._app.credential.get_credential() |
| 81 | + project = self._app.project_id |
| 82 | + if not project: |
| 83 | + raise ValueError( |
| 84 | + 'Project ID is required to access Firestore. Either set the projectId option, ' |
| 85 | + 'or use service account credentials. Alternatively, set the ' |
| 86 | + 'GOOGLE_CLOUD_PROJECT environment variable.') |
| 87 | + |
| 88 | + fs_client = firestore.Client( |
| 89 | + credentials=credentials, project=project, database=database_id) |
| 90 | + self._clients[database_id] = fs_client |
| 91 | + |
| 92 | + return self._clients[database_id] |
0 commit comments