Add count query samples with limit

samples/snippets/requirements.txt

@@ -1 +1 @@
-google-cloud-datastore==2.10.0
+google-cloud-datastore==2.11.0

samples/snippets/snippets.py

@@ -116,7 +116,7 @@ def count_query_on_kind(client):
     task2 = datastore.Entity(client.key("Task", "task2"))
 
     tasks = [task1, task2]
-    client.put_multi([task1, task2])
+    client.put_multi(tasks)
     all_tasks_query = client.query(kind="Task")
     all_tasks_count_query = client.aggregation_query(all_tasks_query).count()
     query_result = all_tasks_count_query.fetch()
@@ -127,6 +127,24 @@ def count_query_on_kind(client):
     return tasks
 
 
+def count_query_with_limit(client):
+    # [START datastore_count_with_limit]
+    task1 = datastore.Entity(client.key("Task", "task1"))
+    task2 = datastore.Entity(client.key("Task", "task2"))
+    task3 = datastore.Entity(client.key("Task", "task3"))
+
+    tasks = [task1, task2, task3]
+    client.put_multi(tasks)
+    all_tasks_query = client.query(kind="Task")
+    all_tasks_count_query = client.aggregation_query(all_tasks_query).count()
+    query_result = all_tasks_count_query.fetch(limit=2)
+    for aggregation_results in query_result:
+        for aggregation in aggregation_results:
+            print(f"We have at least {aggregation.value} tasks")
+    # [END datastore_count_with_limit]
+    return tasks
+
+
 def count_query_property_filter(client):
     # [START datastore_count_with_property_filter]
     task1 = datastore.Entity(client.key("Task", "task1"))

samples/snippets/snippets_test.py

@@ -88,6 +88,15 @@ def test_count_query_on_kind(self, capsys, client):
 
         client.entities_to_delete.extend(tasks)
 
+    @backoff.on_exception(backoff.expo, AssertionError, max_time=240)
+    def test_count_query_with_limit(self, capsys, client):
+        tasks = snippets.count_query_with_limit(client)
+        captured = capsys.readouterr()
+        assert captured.out.strip() == "We have at least 2 tasks"
+        assert captured.err == ""
+
+        client.entities_to_delete.extend(tasks)
+
     @backoff.on_exception(backoff.expo, AssertionError, max_time=240)
     def test_count_query_property_filter(self, capsys, client):
         tasks = snippets.count_query_property_filter(client)