Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings

Commit 8727e91

Browse files
feat(firestore): Add Firestore Multi Database Support (#818)
* Added multi db support for firestore and firestore_async * Added unit and integration tests * fix docs strings
1 parent c044729 commit 8727e91

File tree

6 files changed

+396
-82
lines changed

6 files changed

+396
-82
lines changed

‎firebase_admin/firestore.py‎

Lines changed: 52 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -18,59 +18,75 @@
1818
Firebase apps. This requires the ``google-cloud-firestore`` Python module.
1919
"""
2020

21+
from __future__ import annotations
22+
from typing import Optional, Dict
23+
from firebase_admin import App
24+
from firebase_admin import _utils
25+
2126
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
2329
existing = globals().keys()
2430
for key, value in firestore.__dict__.items():
2531
if not key.startswith('_') and key not in existing:
2632
globals()[key] = value
27-
except ImportError:
33+
except ImportErroraserror:
2834
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
3236

3337

3438
_FIRESTORE_ATTRIBUTE = '_firestore'
3539

3640

37-
def client(app=None) -> firestore.Client:
41+
def client(app: Optional[App] =None, database_id: Optional[str] =None) -> firestore.Client:
3842
"""Returns a client that can be used to interact with Google Cloud Firestore.
3943
4044
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).
4249
4350
Returns:
44-
google.cloud.firestore.Firestore: A `Firestore Client`_.
51+
google.cloud.firestore.Firestore: A `Firestore Client`_.
4552
4653
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.
4957
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
5260
"""
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]

‎firebase_admin/firestore_async.py‎

Lines changed: 52 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -18,65 +18,75 @@
1818
associated with Firebase apps. This requires the ``google-cloud-firestore`` Python module.
1919
"""
2020

21-
from typing import Type
22-
23-
from firebase_admin import (
24-
App,
25-
_utils,
26-
)
27-
from firebase_admin.credentials import Base
21+
from __future__ import annotations
22+
from typing import Optional, Dict
23+
from firebase_admin import App
24+
from firebase_admin import _utils
2825

2926
try:
30-
from google.cloud import firestore # type: ignore # 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
3129
existing = globals().keys()
3230
for key, value in firestore.__dict__.items():
3331
if not key.startswith('_') and key not in existing:
3432
globals()[key] = value
35-
except ImportError:
33+
except ImportErroraserror:
3634
raise ImportError('Failed to import the Cloud Firestore library for Python. Make sure '
37-
'to install the "google-cloud-firestore" module.')
35+
'to install the "google-cloud-firestore" module.') from error
36+
3837

3938
_FIRESTORE_ASYNC_ATTRIBUTE: str = '_firestore_async'
4039

4140

42-
def client(app: App = None) -> firestore.AsyncClient:
41+
def client(app: Optional[App] =None, database_id: Optional[str] = None) -> firestore.AsyncClient:
4342
"""Returns an async client that can be used to interact with Google Cloud Firestore.
4443
4544
Args:
46-
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).
4749
4850
Returns:
49-
google.cloud.firestore.Firestore_Async: A `Firestore Async Client`_.
51+
google.cloud.firestore.Firestore_Async: A `Firestore Async Client`_.
5052
5153
Raises:
52-
ValueError: If a project ID is not specified either via options, credentials or
53-
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.
5457
55-
.. _Firestore Async Client: https://googleapis.dev/python/firestore/latest/client.html
58+
.. _Firestore Async Client: https://cloud.google.com/python/docs/reference/firestore/latest/\
59+
google.cloud.firestore_v1.async_client.AsyncClient
5660
"""
57-
fs_client = _utils.get_app_service(
58-
app, _FIRESTORE_ASYNC_ATTRIBUTE, _FirestoreAsyncClient.from_app)
59-
return fs_client.get()
60-
61-
62-
class _FirestoreAsyncClient:
63-
"""Holds a Google Cloud Firestore Async Client instance."""
64-
65-
def __init__(self, credentials: Type[Base], project: str) -> None:
66-
self._client = firestore.AsyncClient(credentials=credentials, project=project)
67-
68-
def get(self) -> firestore.AsyncClient:
69-
return self._client
70-
71-
@classmethod
72-
def from_app(cls, app: App) -> "_FirestoreAsyncClient":
73-
# Replace remove future reference quotes by importing annotations in Python 3.7+ b/238779406
74-
"""Creates a new _FirestoreAsyncClient for the specified app."""
75-
credentials = app.credential.get_credential()
76-
project = app.project_id
77-
if not project:
78-
raise ValueError(
79-
'Project ID is required to access Firestore. Either set the projectId option, '
80-
'or use service account credentials. Alternatively, set the GOOGLE_CLOUD_PROJECT '
81-
'environment variable.')
82-
return _FirestoreAsyncClient(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+
65+
fs_service = _utils.get_app_service(app, _FIRESTORE_ASYNC_ATTRIBUTE, _FirestoreAsyncService)
66+
return fs_service.get_client(database_id)
67+
68+
class _FirestoreAsyncService:
69+
"""Service that maintains a collection of firestore async clients."""
70+
71+
def __init__(self, app: App) -> None:
72+
self._app: App = app
73+
self._clients: Dict[str, firestore.AsyncClient] = {}
74+
75+
def get_client(self, database_id: Optional[str]) -> firestore.AsyncClient:
76+
"""Creates an async 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.AsyncClient(
89+
credentials=credentials, project=project, database=database_id)
90+
self._clients[database_id] = fs_client
91+
92+
return self._clients[database_id]

‎integration/test_firestore.py‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,20 @@
1717

1818
from firebase_admin import firestore
1919

20+
_CITY = {
21+
'name': u'Mountain View',
22+
'country': u'USA',
23+
'population': 77846,
24+
'capital': False
25+
}
26+
27+
_MOVIE = {
28+
'Name': u'Interstellar',
29+
'Year': 2014,
30+
'Runtime': u'2h 49m',
31+
'Academy Award Winner': True
32+
}
33+
2034

2135
def test_firestore():
2236
client = firestore.client()
@@ -35,6 +49,47 @@ def test_firestore():
3549
doc.delete()
3650
assert doc.get().exists is False
3751

52+
def test_firestore_explicit_database_id():
53+
client = firestore.client(database_id='testing-database')
54+
expected = _CITY
55+
doc = client.collection('cities').document()
56+
doc.set(expected)
57+
58+
data = doc.get()
59+
assert data.to_dict() == expected
60+
61+
doc.delete()
62+
data = doc.get()
63+
assert data.exists is False
64+
65+
def test_firestore_multi_db():
66+
city_client = firestore.client()
67+
movie_client = firestore.client(database_id='testing-database')
68+
69+
expected_city = _CITY
70+
expected_movie = _MOVIE
71+
72+
city_doc = city_client.collection('cities').document()
73+
movie_doc = movie_client.collection('movies').document()
74+
75+
city_doc.set(expected_city)
76+
movie_doc.set(expected_movie)
77+
78+
city_data = city_doc.get()
79+
movie_data = movie_doc.get()
80+
81+
assert city_data.to_dict() == expected_city
82+
assert movie_data.to_dict() == expected_movie
83+
84+
city_doc.delete()
85+
movie_doc.delete()
86+
87+
city_data = city_doc.get()
88+
movie_data = movie_doc.get()
89+
90+
assert city_data.exists is False
91+
assert movie_data.exists is False
92+
3893
def test_server_timestamp():
3994
client = firestore.client()
4095
expected = {

‎integration/test_firestore_async.py‎

Lines changed: 65 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,20 +13,31 @@
1313
# limitations under the License.
1414

1515
"""Integration tests for firebase_admin.firestore_async module."""
16+
import asyncio
1617
import datetime
1718
import pytest
1819

1920
from firebase_admin import firestore_async
2021

21-
@pytest.mark.asyncio
22-
async def test_firestore_async():
23-
client = firestore_async.client()
24-
expected = {
22+
_CITY = {
2523
'name': u'Mountain View',
2624
'country': u'USA',
2725
'population': 77846,
2826
'capital': False
2927
}
28+
29+
_MOVIE = {
30+
'Name': u'Interstellar',
31+
'Year': 2014,
32+
'Runtime': u'2h 49m',
33+
'Academy Award Winner': True
34+
}
35+
36+
37+
@pytest.mark.asyncio
38+
async def test_firestore_async():
39+
client = firestore_async.client()
40+
expected = _CITY
3041
doc = client.collection('cities').document()
3142
await doc.set(expected)
3243

@@ -37,6 +48,56 @@ async def test_firestore_async():
3748
data = await doc.get()
3849
assert data.exists is False
3950

51+
@pytest.mark.asyncio
52+
async def test_firestore_async_explicit_database_id():
53+
client = firestore_async.client(database_id='testing-database')
54+
expected = _CITY
55+
doc = client.collection('cities').document()
56+
await doc.set(expected)
57+
58+
data = await doc.get()
59+
assert data.to_dict() == expected
60+
61+
await doc.delete()
62+
data = await doc.get()
63+
assert data.exists is False
64+
65+
@pytest.mark.asyncio
66+
async def test_firestore_async_multi_db():
67+
city_client = firestore_async.client()
68+
movie_client = firestore_async.client(database_id='testing-database')
69+
70+
expected_city = _CITY
71+
expected_movie = _MOVIE
72+
73+
city_doc = city_client.collection('cities').document()
74+
movie_doc = movie_client.collection('movies').document()
75+
76+
await asyncio.gather(
77+
city_doc.set(expected_city),
78+
movie_doc.set(expected_movie)
79+
)
80+
81+
data = await asyncio.gather(
82+
city_doc.get(),
83+
movie_doc.get()
84+
)
85+
86+
assert data[0].to_dict() == expected_city
87+
assert data[1].to_dict() == expected_movie
88+
89+
await asyncio.gather(
90+
city_doc.delete(),
91+
movie_doc.delete()
92+
)
93+
94+
data = await asyncio.gather(
95+
city_doc.get(),
96+
movie_doc.get()
97+
)
98+
assert data[0].exists is False
99+
assert data[1].exists is False
100+
40101
@pytest.mark.asyncio
41102
async def test_server_timestamp():
42103
client = firestore_async.client()

0 commit comments

Comments
(0)

AltStyle によって変換されたページ (->オリジナル) /