From 6a9efabe1560d5137986df70f1b4f79731deac02 Mon Sep 17 00:00:00 2001
From: HemangChothani <50404902+HemangChothani@users.noreply.github.com>
Date: Fri, 28 Feb 2020 03:30:43 +0530
Subject: [PATCH] feat(datastore): add return query object in add filter method
 (#12)

---
 google/cloud/datastore/query.py | 12 ++++++++----
 tests/unit/test_query.py        |  8 ++++++++
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/google/cloud/datastore/query.py b/google/cloud/datastore/query.py
index f7979d12..4b3daa66 100644
--- a/google/cloud/datastore/query.py
+++ b/google/cloud/datastore/query.py
@@ -220,8 +220,8 @@ def add_filter(self, property_name, operator, value):
           >>> from google.cloud import datastore
           >>> client = datastore.Client()
           >>> query = client.query(kind='Person')
-          >>> query.add_filter('name', '=', 'James')
-          >>> query.add_filter('age', '>', 50)
+          >>> query = query.add_filter('name', '=', 'James')
+          >>> query = query.add_filter('age', '>', 50)
 
         :type property_name: str
         :param property_name: A property name.
@@ -235,6 +235,9 @@ def add_filter(self, property_name, operator, value):
                      :class:`google.cloud.datastore.key.Key`
         :param value: The value to filter on.
 
+        :rtype: :class:`~google.cloud.datastore.query.Query`
+        :returns: A query object.
+
         :raises: :class:`ValueError` if ``operation`` is not one of the
                  specified values, or if a filter names ``'__key__'`` but
                  passes an invalid value (a key is required).
@@ -248,6 +251,7 @@ def add_filter(self, property_name, operator, value):
             raise ValueError('Invalid key: "%s"' % value)
 
         self._filters.append((property_name, operator, value))
+        return self
 
     @property
     def projection(self):
@@ -348,8 +352,8 @@ def fetch(
           >>> from google.cloud import datastore
           >>> client = datastore.Client()
           >>> query = client.query(kind='Person')
-          >>> query.add_filter('name', '=', 'Sally')
-          >>> list(query.fetch())
+          >>> result = query.add_filter('name', '=', 'Sally').fetch()
+          >>> list(result)
           [<Entity object>, <Entity object>, ...]
           >>> list(query.fetch(1))
           [<Entity object>]
diff --git a/tests/unit/test_query.py b/tests/unit/test_query.py
index 01a005f4..ed6cbc9d 100644
--- a/tests/unit/test_query.py
+++ b/tests/unit/test_query.py
@@ -214,6 +214,14 @@ def test_add_filter___key__valid_key(self):
         query.add_filter("__key__", "=", key)
         self.assertEqual(query.filters, [("__key__", "=", key)])
 
+    def test_add_filter_return_query_obj(self):
+        from google.cloud.datastore.query import Query
+
+        query = self._make_one(self._make_client())
+        query_obj = query.add_filter("firstname", "=", u"John")
+        self.assertIsInstance(query_obj, Query)
+        self.assertEqual(query_obj.filters, [("firstname", "=", u"John")])
+
     def test_filter___key__not_equal_operator(self):
         from google.cloud.datastore.key import Key