Multi DB support (#178)

Author: hiranya911

CHANGELOG.md

@@ -1,6 +1,8 @@
 # Unreleased
 
--
+- [added] The `db.reference()` method now optionally takes a `url`
+  parameter. This can be used to access multiple Firebase Databases
+  in the same project more easily. 
 
 # v2.12.0
 

firebase_admin/db.py

@@ -41,22 +41,28 @@
 _TRANSACTION_MAX_RETRIES = 25
 
 
-def reference(path='/', app=None):
+def reference(path='/', app=None, url=None):
     """Returns a database Reference representing the node at the specified path.
 
     If no path is specified, this function returns a Reference that represents the database root.
+    By default, the returned References provide access to the Firebase Database specified at
+    app initialization. To connect to a different Database instance in the same Firebase project,
+    specify the ``url`` parameter.
 
     Args:
       path: Path to a node in the Firebase realtime database (optional).
       app: An App instance (optional).
+      url: Base URL of the Firebase Database instance (optional). When specified, takes
+          precedence over the the ``databaseURL`` option set at app initialization.
 
     Returns:
       Reference: A newly initialized Reference.
 
     Raises:
       ValueError: If the specified path or app is invalid.
     """
-    client = _utils.get_app_service(app, _DB_ATTRIBUTE, _Client.from_app)
+    service = _utils.get_app_service(app, _DB_ATTRIBUTE, _DatabaseService)
+    client = service.get_client(url)
     return Reference(client=client, path=path)
 
 def _parse_path(path):
@@ -662,65 +668,50 @@ def __eq__(self, other):
         return self._compare(other) is 0
 
 
-class _Client(_http_client.JsonHttpClient):
-    """HTTP client used to make REST calls.
-
-    _Client maintains an HTTP session, and handles authenticating HTTP requests along with
-    marshalling and unmarshalling of JSON data.
-    """
+class _DatabaseService(object):
+    """Service that maintains a collection of database clients."""
 
     _DEFAULT_AUTH_OVERRIDE = '_admin_'
 
-    def __init__(self, credential, base_url, auth_override=_DEFAULT_AUTH_OVERRIDE, timeout=None):
-        """Creates a new _Client from the given parameters.
-
-        This exists primarily to enable testing. For regular use, obtain _Client instances by
-        calling the from_app() class method.
-
-        Args:
-          credential: A Google credential that can be used to authenticate requests.
-          base_url: A URL prefix to be added to all outgoing requests. This is typically the
-              Firebase Realtime Database URL.
-          auth_override: A dictionary representing auth variable overrides or None (optional).
-              Default value provides admin privileges. A None value here provides un-authenticated
-              guest privileges.
-          timeout: HTTP request timeout in seconds (optional). If not set connections will never
-              timeout, which is the default behavior of the underlying requests library.
-        """
-        _http_client.JsonHttpClient.__init__(
-            self, credential=credential, base_url=base_url, headers={'User-Agent': _USER_AGENT})
+    def __init__(self, app):
+        self._credential = app.credential.get_credential()
+        db_url = app.options.get('databaseURL')
+        if db_url:
+            self._db_url = _DatabaseService._validate_url(db_url)
+        else:
+            self._db_url = None
+        auth_override = _DatabaseService._get_auth_override(app)
         if auth_override != self._DEFAULT_AUTH_OVERRIDE and auth_override != {}:
             encoded = json.dumps(auth_override, separators=(',', ':'))
             self._auth_override = 'auth_variable_override={0}'.format(encoded)
         else:
             self._auth_override = None
-        self._timeout = timeout
+        self._timeout = app.options.get('httpTimeout')
+        self._clients = {}
+
+    def get_client(self, base_url=None):
+        if base_url is None:
+            base_url = self._db_url
+        base_url = _DatabaseService._validate_url(base_url)
+        if base_url not in self._clients:
+            client = _Client(self._credential, base_url, self._auth_override, self._timeout)
+            self._clients[base_url] = client
+        return self._clients[base_url]
 
     @classmethod
-    def from_app(cls, app):
-        """Creates a new _Client for a given App"""
-        credential = app.credential.get_credential()
-        db_url = cls._get_db_url(app)
-        auth_override = cls._get_auth_override(app)
-        timeout = app.options.get('httpTimeout')
-        return _Client(credential, db_url, auth_override, timeout)
-
-    @classmethod
-    def _get_db_url(cls, app):
-        """Retrieves and parses the database URL option."""
-        url = app.options.get('databaseURL')
+    def _validate_url(cls, url):
+        """Parses and validates a given database URL."""
         if not url or not isinstance(url, six.string_types):
             raise ValueError(
-                'Invalid databaseURL option: "{0}". databaseURL must be a non-empty URL '
-                'string.'.format(url))
-
+                'Invalid database URL: "{0}". Database URL must be a non-empty '
+                'URL string.'.format(url))
         parsed = urllib.parse.urlparse(url)
         if parsed.scheme != 'https':
             raise ValueError(
-                'Invalid databaseURL option: "{0}". databaseURL must be an HTTPS URL.'.format(url))
+                'Invalid database URL: "{0}". Database URL must be an HTTPS URL.'.format(url))
         elif not parsed.netloc.endswith('.firebaseio.com'):
             raise ValueError(
-                'Invalid databaseURL option: "{0}". databaseURL must be a valid URL to a '
+                'Invalid database URL: "{0}". Database URL must be a valid URL to a '
                 'Firebase Realtime Database instance.'.format(url))
         return 'https://{0}'.format(parsed.netloc)
 
@@ -735,6 +726,39 @@ def _get_auth_override(cls, app):
         else:
             return auth_override
 
+    def close(self):
+        for value in self._clients.values():
+            value.close()
+        self._clients = {}
+
+
+class _Client(_http_client.JsonHttpClient):
+    """HTTP client used to make REST calls.
+
+    _Client maintains an HTTP session, and handles authenticating HTTP requests along with
+    marshalling and unmarshalling of JSON data.
+    """
+
+    def __init__(self, credential, base_url, auth_override, timeout):
+        """Creates a new _Client from the given parameters.
+
+        This exists primarily to enable testing. For regular use, obtain _Client instances by
+        calling the from_app() class method.
+
+        Args:
+          credential: A Google credential that can be used to authenticate requests.
+          base_url: A URL prefix to be added to all outgoing requests. This is typically the
+              Firebase Realtime Database URL.
+          auth_override: The encoded auth_variable_override query parameter to be included in
+              outgoing requests.
+          timeout: HTTP request timeout in seconds. If not set connections will never
+              timeout, which is the default behavior of the underlying requests library.
+        """
+        _http_client.JsonHttpClient.__init__(
+            self, credential=credential, base_url=base_url, headers={'User-Agent': _USER_AGENT})
+        self._auth_override = auth_override
+        self._timeout = timeout
+
     @property
     def auth_override(self):
         return self._auth_override

tests/test_db.py

@@ -562,21 +562,45 @@ def test_invalid_db_url(self, url):
         firebase_admin.initialize_app(testutils.MockCredential(), {'databaseURL' : url})
         with pytest.raises(ValueError):
             db.reference()
+        other_app = firebase_admin.initialize_app(testutils.MockCredential(), name='otherApp')
+        with pytest.raises(ValueError):
+            db.reference(app=other_app, url=url)
+
+    def test_multi_db_support(self):
+        default_url = 'https://test.firebaseio.com'
+        firebase_admin.initialize_app(testutils.MockCredential(), {
+            'databaseURL' : default_url,
+        })
+        ref = db.reference()
+        assert ref._client.base_url == default_url
+        assert ref._client.auth_override is None
+        assert ref._client.timeout is None
+        assert ref._client is db.reference()._client
+        assert ref._client is db.reference(url=default_url)._client
+
+        other_url = 'https://other.firebaseio.com'
+        other_ref = db.reference(url=other_url)
+        assert other_ref._client.base_url == other_url
+        assert other_ref._client.auth_override is None
+        assert other_ref._client.timeout is None
+        assert other_ref._client is db.reference(url=other_url)._client
+        assert other_ref._client is db.reference(url=other_url + '/')._client
 
     @pytest.mark.parametrize('override', [{}, {'uid':'user1'}, None])
     def test_valid_auth_override(self, override):
         firebase_admin.initialize_app(testutils.MockCredential(), {
             'databaseURL' : 'https://test.firebaseio.com',
             'databaseAuthVariableOverride': override
         })
-        ref = db.reference()
-        assert ref._client.base_url == 'https://test.firebaseio.com'
-        assert ref._client.timeout is None
-        if override == {}:
-            assert ref._client.auth_override is None
-        else:
-            encoded = json.dumps(override, separators=(',', ':'))
-            assert ref._client.auth_override == 'auth_variable_override={0}'.format(encoded)
+        default_ref = db.reference()
+        other_ref = db.reference(url='https://other.firebaseio.com')
+        for ref in [default_ref, other_ref]:
+            assert ref._client.timeout is None
+            if override == {}:
+                assert ref._client.auth_override is None
+            else:
+                encoded = json.dumps(override, separators=(',', ':'))
+                assert ref._client.auth_override == 'auth_variable_override={0}'.format(encoded)
 
     @pytest.mark.parametrize('override', [
         '', 'foo', 0, 1, True, False, list(), tuple(), _Object()])
@@ -587,33 +611,43 @@ def test_invalid_auth_override(self, override):
         })
         with pytest.raises(ValueError):
             db.reference()
+        other_app = firebase_admin.initialize_app(testutils.MockCredential(), {
+            'databaseAuthVariableOverride': override
+        }, name='otherApp')
+        with pytest.raises(ValueError):
+            db.reference(app=other_app, url='https://other.firebaseio.com')
 
     def test_http_timeout(self):
         test_url = 'https://test.firebaseio.com'
         firebase_admin.initialize_app(testutils.MockCredential(), {
             'databaseURL' : test_url,
             'httpTimeout': 60
         })
-        ref = db.reference()
-        recorder = []
-        adapter = MockAdapter('{}', 200, recorder)
-        ref._client.session.mount(test_url, adapter)
-        assert ref._client.base_url == test_url
-        assert ref._client.timeout == 60
-        assert ref.get() == {}
-        assert len(recorder) == 1
-        assert recorder[0]._extra_kwargs['timeout'] == 60
+        default_ref = db.reference()
+        other_ref = db.reference(url='https://other.firebaseio.com')
+        for ref in [default_ref, other_ref]:
+            recorder = []
+            adapter = MockAdapter('{}', 200, recorder)
+            ref._client.session.mount(ref._client.base_url, adapter)
+            assert ref._client.timeout == 60
+            assert ref.get() == {}
+            assert len(recorder) == 1
+            assert recorder[0]._extra_kwargs['timeout'] == 60
 
     def test_app_delete(self):
         app = firebase_admin.initialize_app(
             testutils.MockCredential(), {'databaseURL' : 'https://test.firebaseio.com'})
         ref = db.reference()
-        assert ref is not None
+        other_ref = db.reference(url='https://other.firebaseio.com')
         assert ref._client.session is not None
+        assert other_ref._client.session is not None
         firebase_admin.delete_app(app)
         with pytest.raises(ValueError):
             db.reference()
+        with pytest.raises(ValueError):
+            db.reference(url='https://other.firebaseio.com')
         assert ref._client.session is None
+        assert other_ref._client.session is None
 
     def test_user_agent_format(self):
         expected = 'Firebase/HTTP/{0}/{1}.{2}/AdminPython'.format(