refactor worker fetch code into wrapper class

1 files changed, 9 insertions, 60 deletions

diff --git a/python/sandcrawler/grobid.py b/python/sandcrawler/grobid.py
index f329a73..d9db6c3 100644
--- a/python/sandcrawler/grobid.py
+++ b/python/sandcrawler/grobid.py
@@ -2,7 +2,7 @@
 import requests
 
 from grobid2json import teixml2json
-from .workers import SandcrawlerWorker
+from .workers import SandcrawlerWorker, SandcrawlerFetchWorker
 from .misc import gen_file_metadata
 from .ia import WaybackClient, WaybackError, PetaboxError
 
@@ -78,12 +78,11 @@ class GrobidClient(object):
                 meta[k] = tei_json[k]
         return meta
 
-class GrobidWorker(SandcrawlerWorker):
+class GrobidWorker(SandcrawlerFetchWorker):
 
     def __init__(self, grobid_client, wayback_client=None, sink=None, **kwargs):
-        super().__init__()
+        super().__init__(wayback_client=wayback_client)
         self.grobid_client = grobid_client
-        self.wayback_client = wayback_client
         self.sink = sink
         self.consolidate_mode = 2
 
@@ -98,62 +97,12 @@ class GrobidWorker(SandcrawlerWorker):
 
     def process(self, record):
         default_key = record['sha1hex']
-        if record.get('warc_path') and record.get('warc_offset'):
-            # it's a full CDX dict. fetch using WaybackClient
-            if not self.wayback_client:
-                raise Exception("wayback client not configured for this GrobidWorker")
-            try:
-                blob = self.wayback_client.fetch_petabox_body(
-                    csize=record['warc_csize'],
-                    offset=record['warc_offset'],
-                    warc_path=record['warc_path'],
-                )
-            except (WaybackError, PetaboxError) as we:
-                return dict(
-                    status="error-wayback",
-                    error_msg=str(we),
-                    source=record,
-                    key=default_key,
-                )
-        elif record.get('url') and record.get('datetime'):
-            # it's a partial CDX dict or something? fetch using WaybackClient
-            if not self.wayback_client:
-                raise Exception("wayback client not configured for this GrobidWorker")
-            try:
-                blob = self.wayback_client.fetch_replay_body(
-                    url=record['url'],
-                    datetime=record['datetime'],
-                )
-            except WaybackError as we:
-                return dict(
-                    status="error-wayback",
-                    error_msg=str(we),
-                    source=record,
-                    key=default_key,
-                )
-        elif record.get('item') and record.get('path'):
-            # it's petabox link; fetch via HTTP
-            resp = requests.get("https://archive.org/serve/{}/{}".format(
-                record['item'], record['path']))
-            try:
-                resp.raise_for_status()
-            except Exception as e:
-                return dict(
-                    status="error-petabox",
-                    error_msg=str(e),
-                    source=record,
-                    key=default_key,
-                )
-            blob = resp.content
-        else:
-            raise ValueError("not a CDX (wayback) or petabox (archive.org) dict; not sure how to proceed")
-        if not blob:
-            return dict(
-                status="error",
-                error_msg="empty blob",
-                source=record,
-                key=default_key,
-            )
+
+        fetch_result = self.fetch_blob(record)
+        if fetch_result['status'] != 'success':
+            return fetch_result
+        blob = fetch_result['blob']
+
         result = self.grobid_client.process_fulltext(blob, consolidate_mode=self.consolidate_mode)
         result['file_meta'] = gen_file_metadata(blob)
         result['source'] = record