Shallow reads support in database API (firebase#109)

hiranya911 · web-flow · commit cd4e7d694ba3 · 2018-01-10T11:47:23.000-08:00
* Shallow reads support in database API

* Updated docstring
diff --git a/firebase_admin/db.py b/firebase_admin/db.py
@@ -125,26 +125,33 @@ def child(self, path):
         full_path = self._pathurl + '/' + path
         return Reference(client=self._client, path=full_path)
 
-    def get(self, etag=False):
+    def get(self, etag=False, shallow=False):
         """Returns the value, and optionally the ETag, at the current location of the database.
 
         Args:
           etag: A boolean indicating whether the Etag value should be returned or not (optional).
+          shallow: A boolean indicating whether to execute a shallow read (optional). Shallow
+              reads do not retrieve the child nodes of the current database location. Cannot be
+              set to True if ``etag`` is also set to True.
 
         Returns:
           object: If etag is False returns the decoded JSON value of the current database location.
           If etag is True, returns a 2-tuple consisting of the decoded JSON value and the Etag
           associated with the current database location.
 
         Raises:
+          ValueError: If both ``etag`` and ``shallow`` are set to True.
           ApiCallError: If an error occurs while communicating with the remote database server.
         """
         if etag:
+            if shallow:
+                raise ValueError('etag and shallow cannot both be set to True.')
             headers, data = self._client.headers_and_body(
                 'get', self._add_suffix(), headers={'X-Firebase-ETag' : 'true'})
             return data, headers.get('ETag')
         else:
-            return self._client.body('get', self._add_suffix())
+            params = 'shallow=true' if shallow else None
+            return self._client.body('get', self._add_suffix(), params=params)
 
     def get_if_changed(self, etag):
         """Gets data in this location only if the specified ETag does not match.
diff --git a/integration/test_db.py b/integration/test_db.py
@@ -89,6 +89,11 @@ def test_get_value_and_etag(self, testref, testdata):
         assert testdata == value
         assert isinstance(etag, six.string_types)
 
+    def test_get_shallow(self, testref):
+        value = testref.get(shallow=True)
+        assert isinstance(value, dict)
+        assert value == {'dinosaurs': True, 'scores': True}
+
     def test_get_if_changed(self, testref, testdata):
         success, data, etag = testref.get_if_changed('wrong_etag')
         assert success is True
diff --git a/tests/test_db.py b/tests/test_db.py
@@ -155,6 +155,22 @@ def test_get_with_etag(self, data):
         assert recorder[0].headers['User-Agent'] == db._USER_AGENT
         assert recorder[0].headers['X-Firebase-ETag'] == 'true'
 
+    @pytest.mark.parametrize('data', valid_values)
+    def test_get_shallow(self, data):
+        ref = db.reference('/test')
+        recorder = self.instrument(ref, json.dumps(data))
+        assert ref.get(shallow=True) == data
+        assert len(recorder) == 1
+        assert recorder[0].method == 'GET'
+        assert recorder[0].url == 'https://test.firebaseio.com/test.json?shallow=true'
+        assert recorder[0].headers['Authorization'] == 'Bearer mock-token'
+        assert recorder[0].headers['User-Agent'] == db._USER_AGENT
+
+    def test_get_with_etag_and_shallow(self):
+        ref = db.reference('/test')
+        with pytest.raises(ValueError):
+            ref.get(etag=True, shallow=True)
+
     @pytest.mark.parametrize('data', valid_values)
     def test_get_if_changed(self, data):
         ref = db.reference('/test')
