Added cassandra revisiting backend and tests for it

Author: voith

frontera/contrib/backends/__init__.py

@@ -6,6 +6,7 @@
 from frontera import Backend
 from frontera.core.components import States, Queue as BaseQueue, DistributedBackend
 from frontera.core.models import Request, Response
+from frontera.utils.misc import utcnow_timestamp
 
 from w3lib.util import to_native_str
 
@@ -182,3 +183,26 @@ def _modify_page(self, obj):
             db_page.cookies = obj.request.cookies
             db_page.status_code = obj.status_code
         return db_page
+
+
+class CommonRevisitingStorageBackendMixin(object):
+
+    def _schedule(self, requests):
+        batch = []
+        for request in requests:
+            if request.meta[b'state'] in [States.NOT_CRAWLED]:
+                request.meta[b'crawl_at'] = utcnow_timestamp()
+            elif request.meta[b'state'] in [States.CRAWLED, States.ERROR]:
+                request.meta[b'crawl_at'] = utcnow_timestamp() + self.interval
+            else:
+                continue    # QUEUED
+            batch.append((request.meta[b'fingerprint'], self._get_score(request), request, True))
+        self.queue.schedule(batch)
+        self.metadata.update_score(batch)
+        self.queue_size += len(batch)
+
+    def page_crawled(self, response):
+        super(CommonRevisitingStorageBackendMixin, self).page_crawled(response)
+        self.states.set_states(response.request)
+        self._schedule([response.request])
+        self.states.update_cache(response.request)

frontera/contrib/backends/cassandra/__init__.py

@@ -1,14 +1,14 @@
 from __future__ import absolute_import
 
 import six
-from cassandra.cluster import Cluster
 from cassandra.cqlengine import connection
-from cassandra.cqlengine.management import drop_table, sync_table
+from cassandra.cqlengine.management import drop_table
 
 from frontera.contrib.backends import (CommonDistributedStorageBackend,
                                        CommonStorageBackend)
-from frontera.contrib.backends.cassandra.components import (Metadata, Queue,
-                                                            States)
+from frontera.contrib.backends.cassandra.components import (Metadata,
+                                                            BroadCrawlingQueue,
+                                                            Queue, States)
 from frontera.utils.misc import load_object
 
 
@@ -99,17 +99,22 @@ def __init__(self, manager):
         settings = manager.settings
         cluster_hosts = settings.get('CASSANDRABACKEND_CLUSTER_HOSTS')
         cluster_port = settings.get('CASSANDRABACKEND_CLUSTER_PORT')
-        keyspace = settings.get('CASSANDRABACKEND_KEYSPACE')
+        drop_all_tables = settings.get('CASSANDRABACKEND_DROP_ALL_TABLES')
         models = settings.get('CASSANDRABACKEND_MODELS')
+        keyspace = settings.get('CASSANDRABACKEND_KEYSPACE')
+
+        self.models = dict([(name, load_object(cls)) for name, cls in six.iteritems(models)])
         cluster_kwargs = {
             'port': cluster_port,
-            'compression': True
+            'compression': True,
         }
-        self.cluster = Cluster(cluster_hosts, **cluster_kwargs)
-        self.models = dict([(name, load_object(cls)) for name, cls in six.iteritems(models)])
+        if not connection.cluster:
+            connection.setup(cluster_hosts, keyspace, **cluster_kwargs)
+            connection.session.default_timeout = settings.get('CASSANDRABACKEND_REQUEST_TIMEOUT')
 
-        self.session.set_keyspace(keyspace)
-        connection.set_session(self.session)
+        if drop_all_tables:
+            for name, table in six.iteritems(self.models):
+                drop_table(table)
 
         self._metadata = None
         self._queue = None
@@ -119,32 +124,17 @@ def __init__(self, manager):
     def strategy_worker(cls, manager):
         b = cls(manager)
         settings = manager.settings
-        drop_all_tables = settings.get('CASSANDRABACKEND_DROP_ALL_TABLES')
-        model = b.models['StateModel']
-
-        if drop_all_tables:
-            drop_table(model)
-
-        sync_table(model)
-
-        b._states = States(b.session, model, settings.get('STATE_CACHE_SIZE_LIMIT'))
+        b._states = States(b.models['StateModel'], settings.get('STATE_CACHE_SIZE_LIMIT'))
         return b
 
     @classmethod
     def db_worker(cls, manager):
         b = cls(manager)
         settings = manager.settings
-        drop = settings.get('CASSANDRABACKEND_DROP_ALL_TABLES')
-        metadata_m = b.models['MetadataModel']
-        queue_m = b.models['QueueModel']
-
-        if drop:
-            drop_table(metadata_m)
-            drop_table(queue_m)
-
-        sync_table(metadata_m)
-        sync_table(queue_m)
-
-        b._metadata = Metadata(metadata_m)
-        b._queue = Queue(queue_m, settings.get('SPIDER_FEED_PARTITIONS'))
+        b._metadata = Metadata(b.models['MetadataModel'], settings.get('CASSANDRABACKEND_CACHE_SIZE'))
+        b._queue = BroadCrawlingQueue(b.models['QueueModel'], settings.get('SPIDER_FEED_PARTITIONS'))
         return b
+
+    def frontier_stop(self):
+        super(Distributed, self).frontier_stop()
+        connection.unregister_connection('default')

frontera/contrib/backends/cassandra/components.py

@@ -132,7 +132,8 @@ def get_next_requests(self, max_n_requests, partition_id, **kwargs):
         """
         results = []
         try:
-            for item in self._order_by(self.queue_model.filter(partition_id=partition_id).allow_filtering()).limit(max_n_requests):
+            for item in self._order_by(self.queue_model.filter(partition_id=partition_id).
+                                               allow_filtering()).limit(max_n_requests):
                 method = item.method or b'GET'
                 r = Request(item.url, method=method, meta=item.meta, headers=item.headers, cookies=item.cookies)
                 r.meta[b'fingerprint'] = to_bytes(item.fingerprint)
@@ -207,7 +208,8 @@ def get_next_requests(self, max_n_requests, partition_id, **kwargs):
                               tries, limit, count, len(queue.keys()))
             queue.clear()
             count = 0
-            for item in self._order_by(self.queue_model.filter(partition_id=partition_id).allow_filtering()).limit(max_n_requests):
+            for item in self._order_by(self.queue_model.filter(partition_id=partition_id).
+                                               allow_filtering()).limit(max_n_requests):
                 if item.host_crc32 not in queue:
                     queue[item.host_crc32] = []
                 if max_requests_per_host is not None and len(queue[item.host_crc32]) > max_requests_per_host:
@@ -227,8 +229,11 @@ def get_next_requests(self, max_n_requests, partition_id, **kwargs):
         for items in six.itervalues(queue):
             for item in items:
                 method = item.method or b'GET'
-                results.append(Request(item.url, method=method,
-                                       meta=item.meta, headers=item.headers, cookies=item.cookies))
+                results.append(Request(item.url,
+                                       method=method,
+                                       meta=item.meta,
+                                       headers=item.headers,
+                                       cookies=item.cookies))
                 item.batch(self.batch).delete()
         self.batch.execute()
         return results

frontera/contrib/backends/cassandra/models.py

@@ -70,13 +70,9 @@ def __repr__(self):
         return '<State:%s=%s>' % (self.fingerprint, self.state)
 
 
-class QueueModel(Model):
-    __table_name__ = 'queue'
+class BaseQueueModel(Model):
+    __abstract__ = True
 
-    partition_id = Integer(primary_key=True)
-    score = Float(primary_key=True)
-    created_at = BigInt(primary_key=True)
-    id = UUID(primary_key=True)
     url = Text(required=True)
     fingerprint = Text(required=True)
     host_crc32 = Integer(required=True)
@@ -90,25 +86,32 @@ def __repr__(self):
         return '<Queue:%s (%s)>' % (self.url, self.id)
 
 
-class FifoOrLIfoQueueModel(Model):
-    # Separate models are needed as
-    # order_by is supported on columns
-    # only in the order, the clustering
-    # keys were created
+class QueueModel(BaseQueueModel):
+    __abstract__ = False
+    __table_name__ = 'queue'
 
-    # Also Inheriting model has some runtime issues
-    # mostly a bug in the driver
-    # Hence the duplicate code
+    partition_id = Integer(primary_key=True)
+    score = Float(primary_key=True)
+    created_at = BigInt(primary_key=True)
+    id = UUID(primary_key=True)
+
+
+class FifoOrLIfoQueueModel(BaseQueueModel):
+    __abstract__ = False
+    __table_name__ = 'fifo_lifo_queue'
 
     partition_id = Integer(primary_key=True)
     score = Float(required=True)
     created_at = BigInt(primary_key=True)
     id = UUID(primary_key=True)
-    url = Text(required=True)
-    fingerprint = Text(required=True)
-    host_crc32 = Integer(required=True)
-    meta = PickleDict()
-    headers = PickleDict()
-    cookies = PickleDict()
-    method = Text()
-    depth = SmallInt()
+
+
+class RevisitingQueueModel(BaseQueueModel):
+    __abstract__ = False
+    __table_name__ = 'revisiting_queue'
+
+    partition_id = Integer(primary_key=True)
+    crawl_at = BigInt(primary_key=True)
+    id = UUID(primary_key=True)
+    score = Float(required=True)
+    created_at = BigInt(required=True)

frontera/contrib/backends/cassandra/revisiting.py

@@ -1,54 +1,53 @@
 # -*- coding: utf-8 -*-
-import json
 import logging
-from datetime import datetime, timedelta
+import uuid
+from datetime import timedelta
 from time import time
 
-from cassandra.cqlengine import columns
-from cassandra.cqlengine.models import Model
+from cassandra.cqlengine.management import sync_table
+from cassandra.cqlengine.query import BatchQuery
+from w3lib.util import to_native_str
 
 from frontera import Request
+from frontera.contrib.backends import CommonRevisitingStorageBackendMixin
 from frontera.contrib.backends.cassandra import CassandraBackend
+from frontera.contrib.backends.cassandra.models import RevisitingQueueModel
 from frontera.contrib.backends.partitioners import Crc32NamePartitioner
 from frontera.core.components import Queue as BaseQueue
 from frontera.core.components import States
-from frontera.utils.misc import get_crc32
+from frontera.utils.misc import get_crc32, utcnow_timestamp
 from frontera.utils.url import parse_domain_from_url_fast
 
 
-class RevisitingQueueModel(Model):
-    __table_name__ = 'revisiting_queue'
-
-    crawl_at = columns.DateTime(required=True, default=datetime.now(), index=True)
-
-
 class RevisitingQueue(BaseQueue):
-    def __init__(self, session, queue_cls, partitions):
-        self.session = session()
+    def __init__(self, queue_cls, partitions):
         self.queue_model = queue_cls
-        self.logger = logging.getLogger("frontera.contrib.backends.sqlalchemy.revisiting.RevisitingQueue")
+        self.logger = logging.getLogger("frontera.contrib.backends.cassandra.revisiting.RevisitingQueue")
         self.partitions = [i for i in range(0, partitions)]
         self.partitioner = Crc32NamePartitioner(self.partitions)
+        self.batch = BatchQuery()
+        sync_table(queue_cls)
 
     def frontier_stop(self):
         pass
 
     def get_next_requests(self, max_n_requests, partition_id, **kwargs):
         results = []
         try:
-            for item in self.queue_model.objects.filter(crawl_at=datetime.utcnow(), partition_id=partition_id).\
-                    limit(max_n_requests):
+            for item in self.queue_model.objects.filter(partition_id=partition_id,
+                                                        crawl_at__lte=utcnow_timestamp()).limit(max_n_requests):
                 method = 'GET' if not item.method else item.method
                 results.append(Request(item.url, method=method, meta=item.meta, headers=item.headers,
                                        cookies=item.cookies))
-                item.delete()
+                item.batch(self.batch).delete()
+                self.batch.execute()
         except Exception as exc:
             self.logger.exception(exc)
         return results
 
     def schedule(self, batch):
-        for fprint, score, request, schedule_at in batch:
-            if schedule_at:
+        for fprint, score, request, schedule in batch:
+            if schedule:
                 _, hostname, _, _, _, _ = parse_domain_from_url_fast(request.url)
                 if not hostname:
                     self.logger.error("Can't get hostname for URL %s, fingerprint %s" % (request.url, fprint))
@@ -57,65 +56,46 @@ def schedule(self, batch):
                 else:
                     partition_id = self.partitioner.partition(hostname, self.partitions)
                     host_crc32 = get_crc32(hostname)
-                created_at = time()*1E+6
-                q = self._create_queue(request, fprint, score, partition_id, host_crc32, created_at)
-
-                q.save()
-                request.meta['state'] = States.QUEUED
-
-    def _create_queue(self, obj, fingerprint, score, partition_id, host_crc32, created_at):
-        db_queue = self.queue_model()
-        db_queue.fingerprint = fingerprint
-        db_queue.score = score
-        db_queue.partition_id = partition_id
-        db_queue.host_crc32 = host_crc32
-        db_queue.url = obj.url
-        db_queue.created_at = created_at
-
-        new_dict = {}
-        for kmeta, vmeta in obj.meta.iteritems():
-            if type(vmeta) is dict:
-                new_dict[kmeta] = json.dumps(vmeta)
-            else:
-                new_dict[kmeta] = str(vmeta)
-
-        db_queue.meta = new_dict
-        db_queue.depth = 0
-
-        db_queue.headers = obj.headers
-        db_queue.method = obj.method
-        db_queue.cookies = obj.cookies
-
-        return db_queue
+                schedule_at = request.meta[b'crawl_at'] if b'crawl_at' in request.meta else utcnow_timestamp()
+                q = self.queue_model(id=uuid.uuid4(),
+                                     fingerprint=to_native_str(fprint),
+                                     score=score,
+                                     url=request.url,
+                                     meta=request.meta,
+                                     headers=request.headers,
+                                     cookies=request.cookies,
+                                     method=to_native_str(request.method),
+                                     partition_id=partition_id,
+                                     host_crc32=host_crc32,
+                                     created_at=time() * 1E+6,
+                                     crawl_at=schedule_at)
+                q.batch(self.batch).save()
+                request.meta[b'state'] = States.QUEUED
+        self.batch.execute()
+
+    def _create_queue_obj(self, fprint, score, request, partition_id, host_crc32, schedule_at):
+        q = self.queue_model(id=uuid.uuid4(),
+                             fingerprint=to_native_str(fprint),
+                             score=score,
+                             url=request.url,
+                             meta=request.meta,
+                             headers=request.headers,
+                             cookies=request.cookies,
+                             method=to_native_str(request.method),
+                             partition_id=partition_id,
+                             host_crc32=host_crc32,
+                             created_at=time() * 1E+6,
+                             crawl_at=schedule_at)
+        return q
 
     def count(self):
-        return self.session.query(self.queue_model).count()
+        return self.queue_model.all().count()
 
 
-class Backend(CassandraBackend):
+class Backend(CommonRevisitingStorageBackendMixin, CassandraBackend):
 
     def _create_queue(self, settings):
-        self.interval = settings.get("SQLALCHEMYBACKEND_REVISIT_INTERVAL")
+        self.interval = settings.get("CASSANDRABACKEND_REVISIT_INTERVAL")
         assert isinstance(self.interval, timedelta)
-        return RevisitingQueue(self.session, RevisitingQueueModel, settings.get('SPIDER_FEED_PARTITIONS'))
-
-    def _schedule(self, requests):
-        batch = []
-        queue_incr = 0
-        for request in requests:
-            if request.meta['state'] in [States.NOT_CRAWLED, None]:
-                schedule_at = datetime.utcnow()
-            elif request.meta['state'] in [States.CRAWLED, States.ERROR]:
-                schedule_at = datetime.utcnow() + self.interval
-            else:  # QUEUED
-                schedule_at = None
-            batch.append((request.meta['fingerprint'], self._get_score(request), request, schedule_at))
-            if schedule_at:
-                queue_incr += 1
-        self.queue.schedule(batch)
-        self.metadata.update_score(batch)
-        self.queue_size += queue_incr
-
-    def page_crawled(self, response, links):
-        super(Backend, self).page_crawled(response, links)
-        self._schedule([response.request])
+        self.interval = self.interval.total_seconds()
+        return RevisitingQueue(RevisitingQueueModel, settings.get('SPIDER_FEED_PARTITIONS'))