Feature: custom outputs

pupa/__init__.py

@@ -1 +1 @@
-__version__ = '0.9.0'   # pragma: no cover
+__version__ = '0.9.0.dev2'   # pragma: no cover

pupa/scrape/base.py

@@ -72,37 +72,17 @@ def __init__(self, jurisdiction, datadir, *, strict_validation=True, fastmode=Fa
         self.error = self.logger.error
         self.critical = self.logger.critical
 
-    def save_object(self, obj):
-        """
-            Save object to disk as JSON.
-
-            Generally shouldn't be called directly.
-        """
-        obj.pre_save(self.jurisdiction.jurisdiction_id)
-
-        filename = '{0}_{1}.json'.format(obj._type, obj._id).replace('/', '-')
-
-        self.info('save %s %s as %s', obj._type, obj, filename)
-        self.debug(json.dumps(OrderedDict(sorted(obj.as_dict().items())),
-                              cls=utils.JSONEncoderPlus, indent=4, separators=(',', ': ')))
-
-        self.output_names[obj._type].add(filename)
-
-        with open(os.path.join(self.datadir, filename), 'w') as f:
-            json.dump(obj.as_dict(), f, cls=utils.JSONEncoderPlus)
-
-        # validate after writing, allows for inspection on failure
-        try:
-            obj.validate()
-        except ValueError as ve:
-            if self.strict_validation:
-                raise ve
-            else:
-                self.warning(ve)
-
-        # after saving and validating, save subordinate objects
-        for obj in obj._related:
-            self.save_object(obj)
+        self.output_target = self.get_output_target(os.environ.get('OUTPUT_TARGET'))
+
+    def get_output_target(self, output_target_name):
+        if output_target_name == 'GOOGLE_CLOUD_PUBSUB':
+            from pupa.scrape.outputs.google_cloud_pubsub import GoogleCloudPubSub
+            return GoogleCloudPubSub(self)
+        if output_target_name == 'AMAZON_SQS':
+            from pupa.scrape.outputs.amazon_sqs import AmazonSQS
+            return AmazonSQS(self)
+        from pupa.scrape.outputs.local_file import LocalFile
+        return LocalFile(self)
 
     def do_scrape(self, **kwargs):
         record = {'objects': defaultdict(int)}
@@ -111,9 +91,9 @@ def do_scrape(self, **kwargs):
         for obj in self.scrape(**kwargs) or []:
             if hasattr(obj, '__iter__'):
                 for iterobj in obj:
-                    self.save_object(iterobj)
+                    self.output_target.save_object(iterobj)
             else:
-                self.save_object(obj)
+                self.output_target.save_object(obj)
         record['end'] = utils.utcnow()
         record['skipped'] = getattr(self, 'skipped', 0)
         if not self.output_names:

pupa/scrape/outputs/amazon_sqs.py

@@ -0,0 +1,43 @@
+import boto3
+import os
+import json
+import uuid
+from collections import OrderedDict
+
+from pupa import utils
+from pupa.scrape.outputs.output import Output
+
+MAX_BYTE_LENGTH = 230000
+
+
+class AmazonSQS(Output):
+
+    def __init__(self, scraper):
+        super().__init__(scraper)
+
+        self.sqs = boto3.resource('sqs')
+        self.queue_name = os.environ.get('AMAZON_SQS_QUEUE')
+        self.queue = self.sqs.get_queue_by_name(QueueName=self.queue_name)
+
+        self.s3 = boto3.resource('s3')
+        self.bucket_name = os.environ.get('AMAZON_S3_BUCKET')
+
+    def handle_output(self, obj):
+        self.scraper.info('send %s %s to queue %s', obj._type, obj,
+                          self.queue_name)
+        self.debug_obj(obj)
+
+        self.add_output_name(obj, self.queue_name)
+        obj_str = self.stringify_obj(obj, True, True)
+        encoded_obj_str = obj_str.encode('utf-8')
+
+        if len(encoded_obj_str) > MAX_BYTE_LENGTH:
+            key = 'S3:{}'.format(str(uuid.uuid4()))
+
+            self.scraper.info('put %s %s to bucket %s/%s', obj._type, obj,
+                              self.bucket_name, key)
+
+            self.s3.Object(self.bucket_name, key).put(Body=encoded_obj_str)
+            self.queue.send_message(MessageBody=key)
+        else:
+            self.queue.send_message(MessageBody=obj_str)

pupa/scrape/outputs/google_cloud_pubsub.py

@@ -0,0 +1,26 @@
+import os
+
+from pupa.scrape.outputs.output import Output
+
+from google.cloud import pubsub
+
+
+class GoogleCloudPubSub(Output):
+
+    def __init__(self, scraper):
+        super().__init__(scraper)
+
+        project = os.environ.get('GOOGLE_CLOUD_PROJECT')
+        topic_name = os.environ.get('GOOGLE_CLOUD_PUBSUB_TOPIC')
+        self.publisher = pubsub.PublisherClient()
+        self.topic_path = self.publisher.topic_path(project, topic_name)
+
+    def handle_output(self, obj):
+        self.scraper.info('publish %s %s to topic %s', obj._type, obj,
+                          self.topic_path)
+        self.debug_obj(obj)
+
+        self.add_output_name(obj, self.topic_path)
+        obj_str = self.stringify_obj(obj, True, True)
+
+        self.publisher.publish(self.topic_path, obj_str.encode('utf-8'))

pupa/scrape/outputs/local_file.py

@@ -0,0 +1,19 @@
+import json
+import os
+
+from pupa import utils
+from pupa.scrape.outputs.output import Output
+
+
+class LocalFile(Output):
+
+    def handle_output(self, obj):
+        filename = '{0}_{1}.json'.format(obj._type, obj._id).replace('/', '-')
+
+        self.scraper.info('save %s %s as %s', obj._type, obj, filename)
+        self.debug_obj(obj)
+
+        self.add_output_name(obj, filename)
+
+        with open(os.path.join(self.scraper.datadir, filename), 'w') as f:
+            json.dump(obj.as_dict(), f, cls=utils.JSONEncoderPlus)

pupa/scrape/outputs/output.py

@@ -0,0 +1,60 @@
+import json
+
+from abc import ABCMeta, abstractmethod
+from collections import OrderedDict
+from datetime import datetime, timezone
+
+from pupa import utils
+
+
+class Output(metaclass=ABCMeta):
+
+    def __init__(self, scraper):
+        self.scraper = scraper
+
+    def add_output_name(self, obj, output_name):
+        self.scraper.output_names[obj._type].add(output_name)
+
+    def debug_obj(self, obj):
+        self.scraper.debug(json.dumps(OrderedDict(sorted(obj.as_dict().items())),
+                           cls=utils.JSONEncoderPlus,
+                           indent=4, separators=(',', ': ')))
+
+    def get_obj_as_dict(self, obj, add_jurisdiction=False, add_type=False):
+        obj_dict = obj.as_dict()
+        if add_jurisdiction and self.scraper.jurisdiction:
+            obj_dict['jurisdiction'] = self.scraper.jurisdiction.jurisdiction_id
+        if add_type:
+            obj_dict['type'] = obj._type
+        return obj_dict
+
+    @abstractmethod
+    def handle_output(self, obj):
+        pass
+
+    def save_object(self, obj):
+        obj.pre_save(self.scraper.jurisdiction.jurisdiction_id)
+
+        # actual output handling, to be handled by subclass
+        self.handle_output(obj)
+
+        # validate after writing, allows for inspection on failure
+        try:
+            obj.validate()
+        except ValueError as ve:
+            if self.scraper.strict_validation:
+                raise ve
+            else:
+                self.scraper.warning(ve)
+
+        # after saving and validating, save subordinate objects
+        for obj in obj._related:
+            self.save_object(obj)
+
+    def stringify_obj(self, obj, add_jurisdiction=False, add_type=False):
+        obj_dict = self.get_obj_as_dict(obj, add_jurisdiction, add_type)
+        return self.stringify_obj_dict(obj_dict)
+
+    def stringify_obj_dict(self, obj_dict):
+        return json.dumps(obj_dict, cls=utils.JSONEncoderPlus,
+                          separators=(',', ':'))

pupa/scrape/vote_event.py

@@ -1,3 +1,5 @@
+import os
+
 from ..utils import _make_pseudo_id
 from .base import BaseModel, cleanup_list, SourceMixin
 from .bill import Bill
@@ -47,7 +49,11 @@ def set_bill(self, bill_or_identifier, *, chamber=None):
         elif isinstance(bill_or_identifier, Bill):
             if chamber:
                 raise ScrapeValueError("set_bill takes no arguments when using a `Bill` object")
-            self.bill = bill_or_identifier._id
+            if os.environ.get('VOTE_EVENT_NO_BILL_UUID') == 'true':
+                kwargs = {'identifier': bill_or_identifier.identifier}
+                self.bill = _make_pseudo_id(**kwargs)
+            else:
+                self.bill = bill_or_identifier._id
         else:
             if chamber is None:
                 chamber = 'legislature'

setup.py

@@ -19,12 +19,14 @@
 pupa = pupa.cli.__main__:main''',
       install_requires=[
           'Django>=1.11',
-          'opencivicdata>=2.1.0',
+          'opencivicdata>=2.1.2',
           'dj_database_url>=0.3.0',
           'scrapelib>=1.0',
           'jsonschema>=2.6.0',
           'psycopg2',
           'pytz',
+          'google-cloud-pubsub==0.30.1',
+          'boto3==1.5.18',
       ],
       extras_require={
           'dev': [