From b5d5475572aaf5afd156a94f087962e9c22fec06 Mon Sep 17 00:00:00 2001 From: KSokhal Date: Fri, 15 Sep 2023 18:53:45 +0100 Subject: [PATCH 1/2] Added multiple firestore support --- firebase_admin/_utils.py | 4 +++- firebase_admin/firestore.py | 19 ++++++++++++++----- 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/firebase_admin/_utils.py b/firebase_admin/_utils.py index dcfb520d..ab4bd9a1 100644 --- a/firebase_admin/_utils.py +++ b/firebase_admin/_utils.py @@ -93,8 +93,10 @@ def _get_initialized_app(app): -def get_app_service(app, name, initializer): +def get_app_service(app, options, name, initializer): app = _get_initialized_app(app) + if options: + app.options._options.update(options) return app._get_service(name, initializer) # pylint: disable=protected-access diff --git a/firebase_admin/firestore.py b/firebase_admin/firestore.py index 32c9897d..0a66bd1f 100644 --- a/firebase_admin/firestore.py +++ b/firebase_admin/firestore.py @@ -34,11 +34,12 @@ _FIRESTORE_ATTRIBUTE = '_firestore' -def client(app=None): +def client(app=None, database_id=None): """Returns a client that can be used to interact with Google Cloud Firestore. Args: app: An App instance (optional). + database_id: The ID of the Google Cloud Firestore database to use. If none provided, default database will be used (optional). Returns: google.cloud.firestore.Firestore: A `Firestore Client`_. @@ -50,15 +51,22 @@ def client(app=None): .. _Firestore Client: https://googlecloudplatform.github.io/google-cloud-python/latest\ /firestore/client.html """ - fs_client = _utils.get_app_service(app, _FIRESTORE_ATTRIBUTE, _FirestoreClient.from_app) + + options = {"database_id": database_id} if database_id else None + + fs_client = _utils.get_app_service(app, options, _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 __init__(self, credentials, project, database_id=None): + if database_id: + self._client = firestore.Client(credentials=credentials, project=project, database=database_id) + else: + self._client = firestore.Client(credentials=credentials, project=project) + def get(self): return self._client @@ -68,9 +76,10 @@ def from_app(cls, app): """Creates a new _FirestoreClient for the specified app.""" credentials = app.credential.get_credential() project = app.project_id + database_id = app.options.get('database_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) + return _FirestoreClient(credentials, project, database_id) From 2cc4a829ce2319ca51cdc8b6d1890315e63c3f0c Mon Sep 17 00:00:00 2001 From: KSokhal Date: Sat, 3 Feb 2024 14:03:52 +0000 Subject: [PATCH 2/2] Moved options to optional kwarg Check firestore ver is > 2.12.0 is using database id Added test for explicit database id --- firebase_admin/_utils.py | 2 +- firebase_admin/firestore.py | 9 +++++++-- tests/test_firestore.py | 7 +++++++ 3 files changed, 15 insertions(+), 3 deletions(-) diff --git a/firebase_admin/_utils.py b/firebase_admin/_utils.py index ab4bd9a1..6da57650 100644 --- a/firebase_admin/_utils.py +++ b/firebase_admin/_utils.py @@ -93,7 +93,7 @@ def _get_initialized_app(app): -def get_app_service(app, options, name, initializer): +def get_app_service(app, name, initializer, options=None): app = _get_initialized_app(app) if options: app.options._options.update(options) diff --git a/firebase_admin/firestore.py b/firebase_admin/firestore.py index 0a66bd1f..18065676 100644 --- a/firebase_admin/firestore.py +++ b/firebase_admin/firestore.py @@ -29,7 +29,7 @@ 'to install the "google-cloud-firestore" module.') from firebase_admin import _utils - +from packaging.version import Version _FIRESTORE_ATTRIBUTE = '_firestore' @@ -54,7 +54,7 @@ def client(app=None, database_id=None): options = {"database_id": database_id} if database_id else None - fs_client = _utils.get_app_service(app, options, _FIRESTORE_ATTRIBUTE, _FirestoreClient.from_app) + fs_client = _utils.get_app_service(app, _FIRESTORE_ATTRIBUTE, _FirestoreClient.from_app, options=options) return fs_client.get() @@ -63,6 +63,11 @@ class _FirestoreClient: def __init__(self, credentials, project, database_id=None): if database_id: + if Version(firestore.__version__) < Version("2.12.0"): + raise ValueError( + 'The database_id parameter is only supported for google-cloud-firestore version ' + '2.12.0 and above. Please upgrade your google-cloud-firestore package to use ' + 'this feature.') self._client = firestore.Client(credentials=credentials, project=project, database=database_id) else: self._client = firestore.Client(credentials=credentials, project=project) diff --git a/tests/test_firestore.py b/tests/test_firestore.py index 768eb637..0489a84a 100644 --- a/tests/test_firestore.py +++ b/tests/test_firestore.py @@ -65,6 +65,13 @@ def test_service_account(self): assert client is not None assert client.project == 'mock-project-id' + def test_service_account_with_explicit_database_id(self): + cred = credentials.Certificate(testutils.resource_filename('service_account.json')) + firebase_admin.initialize_app(cred) + client = firestore.client(database_id='explicit-database-id') + assert client is not None + assert client._database == 'explicit-database-id' + def test_service_account_with_explicit_app(self): cred = credentials.Certificate(testutils.resource_filename('service_account.json')) app = firebase_admin.initialize_app(cred)