Merge pull request #51 from scrapinghub/kumo1719-newcount

Frontier newcount counter per slot

Author: vshlapakov

README.rst

@@ -454,6 +454,15 @@ Add a request to the slot::
 
     >>> slot.queue.add([{'fp': '/some/path.html'}])
     >>> slot.flush()
+    >>> slot.newcount
+    1
+
+``newcount`` is defined per slot, but also available per frontier and globally::
+
+    >>> frontier.newcount
+    1
+    >>> frontiers.newcount
+    3
 
 Add a fingerprint only to the slot::
 

scrapinghub/client.py

@@ -1,5 +1,7 @@
 import json
 import collections
+from functools import partial
+from collections import defaultdict
 
 from six import string_types
 from requests.compat import urljoin
@@ -9,6 +11,7 @@
 
 from .hubstorage.resourcetype import DownloadableResource
 from .hubstorage.resourcetype import ItemsResourceType
+from .hubstorage.utils import urlpathjoin
 
 # scrapinghub.hubstorage classes to use as-is
 from .hubstorage.job import JobMeta
@@ -227,7 +230,7 @@ def __init__(self, client, projectid):
         # proxied sub-resources
         self.activity = Activity(_Activity, client, projectid)
         self.collections = Collections(_Collections, client, projectid)
-        self.frontiers = Frontiers(_Frontier, client, projectid)
+        self.frontiers = Frontiers(_HSFrontier, client, projectid)
         self.settings = Settings(client._hsclient, projectid)
 
 
@@ -1051,6 +1054,34 @@ def post(self, _value, **kwargs):
         self._origin.post(_value, **kwargs)
 
 
+class _HSFrontier(_Frontier):
+    """Modified hubstorage Frontier with newcount per slot."""
+
+    def __init__(self, *args, **kwargs):
+        super(_HSFrontier, self).__init__(*args, **kwargs)
+        self.newcount = defaultdict(int)
+
+    def _get_writer(self, frontier, slot):
+        key = (frontier, slot)
+        writer = self._writers.get(key)
+        if not writer:
+            writer = self.client.batchuploader.create_writer(
+                url=urlpathjoin(self.url, frontier, 's', slot),
+                auth=self.auth,
+                size=self.batch_size,
+                start=self.batch_start,
+                interval=self.batch_interval,
+                qsize=self.batch_qsize,
+                content_encoding=self.batch_content_encoding,
+                callback=partial(self._writer_callback, key),
+            )
+            self._writers[key] = writer
+        return writer
+
+    def _writer_callback(self, key, response):
+        self.newcount[key] += response.json()["newcount"]
+
+
 class Frontiers(_Proxy):
     """Frontiers collection for a project.
 
@@ -1074,6 +1105,10 @@ class Frontiers(_Proxy):
     - flush data of all frontiers of a project
         >>> project.frontiers.flush()
 
+    - show amount of new requests added for all frontiers
+        >>> project.frontiers.newcount
+        3
+
     - close batch writers of all frontiers of a project
         >>> project.frontiers.close()
     """
@@ -1095,7 +1130,7 @@ def list(self):
 
     @property
     def newcount(self):
-        return self._origin.newcount
+        return sum(self._origin.newcount.values())
 
 
 class Frontier(object):
@@ -1120,6 +1155,10 @@ class Frontier(object):
 
     - flush frontier data
         >>> frontier.flush()
+
+    - show amount of new requests added to frontier
+        >>> frontier.newcount
+        3
     """
     def __init__(self, client, frontiers, name):
         self.key = name
@@ -1145,6 +1184,12 @@ def flush(self):
             if fname == self.key:
                 writer.flush()
 
+    @property
+    def newcount(self):
+        newcount_values = self._frontiers._origin.newcount
+        return sum(v for (frontier, _), v in newcount_values.items()
+                   if frontier == self.key)
+
 
 class FrontierSlot(object):
     """Representation of a frontier slot object.
@@ -1164,6 +1209,10 @@ class FrontierSlot(object):
     - flush data for a slot
         >>> slot.flush()
 
+    - show amount of new requests added to a slot
+        >>> slot.newcount
+        2
+
     - read requests from a slot
         >>> slot.q.iter()
         <generator object jldecode at 0x1049aa9e8>
@@ -1202,7 +1251,8 @@ def q(self):
     def delete(self):
         """Delete the slot."""
         origin = self._frontier._frontiers._origin
-        return origin.delete_slot(self._frontier.key, self.key)
+        origin.delete_slot(self._frontier.key, self.key)
+        origin.newcount.pop((self._frontier.key, self.key), None)
 
     def flush(self):
         """Flush data for the slot."""
@@ -1211,6 +1261,11 @@ def flush(self):
         if writer:
             writer.flush()
 
+    @property
+    def newcount(self):
+        newcount_values = self._frontier._frontiers._origin.newcount
+        return newcount_values.get((self._frontier.key, self.key), 0)
+
 
 class FrontierSlotFingerprints(object):
 

tests/client/conftest.py

@@ -177,4 +177,4 @@ def clean_collection(collection):
 # Frontier helpers section
 
 def clean_frontier_slot(frontier):
-    frontier.delete(TEST_FRONTIER_SLOT)
+    frontier.get(TEST_FRONTIER_SLOT).delete()

tests/client/test_frontier.py

@@ -1,22 +1,22 @@
-
-from six import string_types
+import time
 from types import GeneratorType
 from collections import Iterable
 
-from scrapinghub.client import Frontiers, Frontier, FrontierSlot
+from six import string_types
 
+from scrapinghub.client import Frontiers, Frontier, FrontierSlot
 from .conftest import TEST_FRONTIER_NAME, TEST_FRONTIER_SLOT
 
 
 def _add_test_requests_to_frontier(frontier):
     slot = frontier.get(TEST_FRONTIER_SLOT)
-    slot.add([{'fp': '/some/path.html'}, {'fp': '/other/path.html'}])
+    slot.q.add([{'fp': '/some/path.html'}, {'fp': '/other/path.html'}])
     slot.flush()
 
 
 def test_frontiers(project, frontier):
     # reset a test slot and add some requests to init it
-    frontier.delete(TEST_FRONTIER_SLOT)
+    frontier.get(TEST_FRONTIER_SLOT).delete()
     _add_test_requests_to_frontier(frontier)
 
     assert isinstance(project.frontiers, Frontiers)
@@ -42,7 +42,7 @@ def test_frontiers(project, frontier):
 
 def test_frontier(project, frontier):
     # add some requests to test frontier to init a test slot
-    frontier.delete(TEST_FRONTIER_SLOT)
+    frontier.get(TEST_FRONTIER_SLOT).delete()
     _add_test_requests_to_frontier(frontier)
 
     slots = frontier.iter()
@@ -60,13 +60,13 @@ def test_frontier(project, frontier):
 
 def test_frontier_slot(project, frontier):
     # add some requests to test frontier to init a test slot
-    frontier.delete(TEST_FRONTIER_SLOT)
+    frontier.get(TEST_FRONTIER_SLOT).delete()
     _add_test_requests_to_frontier(frontier)
 
     slot = frontier.get(TEST_FRONTIER_SLOT)
 
     # get all batches from slot and validate its content
-    batches_iter = slot.iter()
+    batches_iter = slot.q.iter()
     assert isinstance(batches_iter, GeneratorType)
     batches = list(batches_iter)
     assert len(batches) == 1
@@ -78,21 +78,53 @@ def test_frontier_slot(project, frontier):
     assert requests == [['/some/path.html', None],
                         ['/other/path.html', None]]
 
-    # validate that slot.list() returns same data as slot.iter()
-    batches_list = slot.list()
+    # validate that slot.list() returns same data as slot.q.iter()
+    batches_list = slot.q.list()
     assert isinstance(batches, list)
     assert batches_list == batches
 
     # add a requests with additional parameters
-    slot.add([{'fp': 'page1.html', 'p': 1, 'qdata': {'depth': 1}}])
+    slot.q.add([{'fp': 'page1.html', 'p': 1, 'qdata': {'depth': 1}}])
     slot.flush()
-    batches = slot.list()
+    batches = slot.q.list()
     assert len(batches) == 2
     assert batches[1]['requests'] == [['page1.html', {'depth': 1}]]
 
     # drop all batches and validate that slot is empty
-    slot.delete([batch['id'] for batch in batches])
-    assert slot.list() == []
+    slot.q.delete([batch['id'] for batch in batches])
+    assert slot.q.list() == []
 
     slot.delete()
     assert TEST_FRONTIER_SLOT not in frontier.list()
+
+
+def test_frontier_newcount(project):
+    # add some requests to test frontier to init a test slot
+    frontier = project.frontiers.get(TEST_FRONTIER_NAME)
+    first_slot = frontier.get(TEST_FRONTIER_SLOT)
+    first_slot.delete()
+
+    assert frontier._frontiers.newcount == 0
+    assert frontier.newcount == 0
+    assert first_slot.newcount == 0
+
+    # shorter batch interval for faster tests
+    frontier._frontiers._origin.batch_interval = 0.1
+    _add_test_requests_to_frontier(frontier)
+    time.sleep(0.5)
+
+    assert frontier._frontiers.newcount == 2
+    assert frontier.newcount == 2
+    assert first_slot.newcount == 2
+
+    second_slot = frontier.get('test2.com')
+    second_slot.delete()
+    second_slot.q.add([{'fp': '/different_path.html'}])
+    second_slot.flush()
+
+    assert frontier._frontiers.newcount == 3
+    assert frontier.newcount == 3
+    assert second_slot.newcount == 1
+    assert first_slot.newcount == 2
+
+    frontier._frontiers.close()