Allow user to specify the violation payload format (csv / json) (#1301)

* Added test for GCS upload with json/csv respectively

* pylint fixes

* implement csv/json uploads, fix tests/test data

* test for invalid data format

* minor fix

* Make _get_output_filename() observe the data format

* Make _get_output_filename() simpler

* Move _get_output_filename() into the base class

This will facilitate reuse of the same method in the EmailViolations
class.

* Use the base class _get_output_filename() method

Also, add a method that creates the attachment in json format.

* implement csv/json format; fix tests & test data

* pylint fix

* Better test coverage

* better naming

* review comments, batch #1

* review comments, batch #2

* fix custom exception class

* review comments, batch #3

* minor fix

* Remove redundant error log

* Update email_violations.py

google/cloud/forseti/notifier/notifiers/base_notification.py

@@ -16,15 +16,32 @@
 
 import abc
 
+from google.cloud.forseti.common.util import date_time
 from google.cloud.forseti.common.util import logger
+from google.cloud.forseti.common.util import string_formats
 
 LOGGER = logger.get_logger(__name__)
 
 
+class InvalidDataFormatError(Exception):
+    """Raised in case of an invalid notifier data format."""
+
+    def __init__(self, notifier, invalid_data_format):
+        """Constructor for the base notifier.
+
+        Args:
+            notifier (str): the notifier module/name
+            invalid_data_format (str): the invalid data format in question.
+        """
+        super(InvalidDataFormatError, self).__init__(
+            '%s: invalid data format: %s' % (notifier, invalid_data_format))
+
+
 class BaseNotification(object):
     """Base notifier to perform notifications"""
 
     __metaclass__ = abc.ABCMeta
+    supported_data_formats = ['csv', 'json']
 
     def __init__(self, resource, cycle_timestamp,
                  violations, global_configs, notifier_config,
@@ -55,3 +72,20 @@ def __init__(self, resource, cycle_timestamp,
     def run(self):
         """Runs the notifier."""
         pass
+
+    def _get_output_filename(self, filename_template):
+        """Create the output filename.
+
+        Args:
+            filename_template (string): template to use for the output filename
+
+        Returns:
+            str: The output filename for the violations CSV file.
+        """
+        utc_now_datetime = date_time.get_utc_now_datetime()
+        output_timestamp = utc_now_datetime.strftime(
+            string_formats.TIMESTAMP_TIMEZONE_FILES)
+
+        output_filename = filename_template.format(
+            self.resource, self.cycle_timestamp, output_timestamp)
+        return output_filename

google/cloud/forseti/notifier/notifiers/email_violations.py

@@ -14,11 +14,14 @@
 
 """Email notifier to perform notifications"""
 
+import tempfile
+
 from google.cloud.forseti.common.data_access import csv_writer
-from google.cloud.forseti.common.util import email
 from google.cloud.forseti.common.util import date_time
+from google.cloud.forseti.common.util import email
 from google.cloud.forseti.common.util import errors as util_errors
 from google.cloud.forseti.common.util import logger
+from google.cloud.forseti.common.util import parser
 from google.cloud.forseti.common.util import string_formats
 from google.cloud.forseti.notifier.notifiers import base_notification
 
@@ -52,36 +55,42 @@ def __init__(self, resource, cycle_timestamp,
         self.mail_util = email.EmailUtil(
             self.notification_config['sendgrid_api_key'])
 
-    def _get_output_filename(self):
-        """Create the output filename.
-
-        Returns:
-            str: The output filename for the violations json.
-        """
-        now_utc = date_time.get_utc_now_datetime()
-        output_timestamp = now_utc.strftime(
-            string_formats.TIMESTAMP_TIMEZONE_FILES)
-        output_filename = string_formats.VIOLATION_CSV_FMT.format(
-            self.resource,
-            self.cycle_timestamp,
-            output_timestamp)
-        return output_filename
-
-    def _make_attachment(self):
-        """Create the attachment object.
+    def _make_attachment_csv(self):
+        """Create the attachment object in csv format.
 
         Returns:
             attachment: SendGrid attachment object.
         """
-        output_file_name = self._get_output_filename()
+        output_filename = self._get_output_filename(
+            string_formats.VIOLATION_CSV_FMT)
         with csv_writer.write_csv(resource_name='violations',
                                   data=self.violations,
                                   write_header=True) as csv_file:
             output_csv_name = csv_file.name
             LOGGER.info('CSV filename: %s', output_csv_name)
             attachment = self.mail_util.create_attachment(
                 file_location=csv_file.name,
-                content_type='text/csv', filename=output_file_name,
+                content_type='text/csv', filename=output_filename,
+                content_id='Violations')
+
+        return attachment
+
+    def _make_attachment_json(self):
+        """Create the attachment object json format.
+
+        Returns:
+            attachment: SendGrid attachment object.
+        """
+        output_filename = self._get_output_filename(
+            string_formats.VIOLATION_JSON_FMT)
+        with tempfile.NamedTemporaryFile() as tmp_violations:
+            tmp_violations.write(parser.json_stringify(self.violations))
+            tmp_violations.flush()
+            LOGGER.info('JSON filename: %s', tmp_violations.name)
+            attachment = self.mail_util.create_attachment(
+                file_location=tmp_violations.name,
+                content_type='application/json',
+                filename=output_filename,
                 content_id='Violations')
 
         return attachment
@@ -121,11 +130,21 @@ def _compose(self, **kwargs):
 
         email_map = {}
 
-        attachment = self._make_attachment()
+        data_format = self.notification_config.get('data_format', 'csv')
+        if data_format not in self.supported_data_formats:
+            raise base_notification.InvalidDataFormatError(
+                'Email notifier', data_format)
+
+        attachment = None
+        if data_format == 'csv':
+            attachment = self._make_attachment_csv()
+        else:
+            attachment = self._make_attachment_json()
         subject, content = self._make_content()
         email_map['subject'] = subject
         email_map['content'] = content
         email_map['attachment'] = attachment
+
         return email_map
 
     def _send(self, **kwargs):
@@ -156,4 +175,5 @@ def _send(self, **kwargs):
     def run(self):
         """Run the email notifier"""
         email_notification = self._compose()
-        self._send(notification=email_notification)
+        if email_notification:
+            self._send(notification=email_notification)

google/cloud/forseti/notifier/notifiers/gcs_violations.py

@@ -14,11 +14,12 @@
 
 """Upload violations to GCS."""
 
+import tempfile
 
 from google.cloud.forseti.common.data_access import csv_writer
 from google.cloud.forseti.common.gcp_api import storage
-from google.cloud.forseti.common.util import date_time
 from google.cloud.forseti.common.util import logger
+from google.cloud.forseti.common.util import parser
 from google.cloud.forseti.common.util import string_formats
 from google.cloud.forseti.notifier.notifiers import base_notification
 
@@ -29,30 +30,50 @@
 class GcsViolations(base_notification.BaseNotification):
     """Upload violations to GCS."""
 
-    def _get_output_filename(self):
-        """Create the output filename.
+    def _upload_json(self, gcs_upload_path):
+        """Upload violations in json format.
 
-        Returns:
-            str: The output filename for the violations CSV file.
+        Args:
+            gcs_upload_path (string): the GCS upload path.
         """
-        now_utc = date_time.get_utc_now_datetime()
-        output_timestamp = now_utc.strftime(
-            string_formats.TIMESTAMP_TIMEZONE_FILES)
-        output_filename = string_formats.VIOLATION_CSV_FMT.format(
-            self.resource, self.cycle_timestamp, output_timestamp)
-        return output_filename
+        with tempfile.NamedTemporaryFile() as tmp_violations:
+            tmp_violations.write(parser.json_stringify(self.violations))
+            tmp_violations.flush()
+            storage_client = storage.StorageClient()
+            storage_client.put_text_file(tmp_violations.name, gcs_upload_path)
 
-    def run(self):
-        """Generate the temporary CSV file and upload to GCS."""
+    def _upload_csv(self, gcs_upload_path):
+        """Upload violations in csv format.
+
+        Args:
+            gcs_upload_path (string): the GCS upload path.
+        """
         with csv_writer.write_csv(resource_name='violations',
                                   data=self.violations,
                                   write_header=True) as csv_file:
             LOGGER.info('CSV filename: %s', csv_file.name)
+            storage_client = storage.StorageClient()
+            storage_client.put_text_file(csv_file.name, gcs_upload_path)
+
+    def run(self):
+        """Generate the temporary (CSV xor JSON) file and upload to GCS."""
+        if not self.notification_config['gcs_path'].startswith('gs://'):
+            return
 
+        data_format = self.notification_config.get('data_format', 'csv')
+        if data_format not in self.supported_data_formats:
+            raise base_notification.InvalidDataFormatError(
+                'GCS uploader', data_format)
+
+        if data_format == 'csv':
             gcs_upload_path = '{}/{}'.format(
                 self.notification_config['gcs_path'],
-                self._get_output_filename())
-
-            if gcs_upload_path.startswith('gs://'):
-                storage_client = storage.StorageClient()
-                storage_client.put_text_file(csv_file.name, gcs_upload_path)
+                self._get_output_filename(
+                    string_formats.VIOLATION_CSV_FMT))
+            self._upload_csv(gcs_upload_path)
+        else:
+            gcs_upload_path = '{}/{}'.format(
+                self.notification_config['gcs_path'],
+                self._get_output_filename(
+                    string_formats.VIOLATION_JSON_FMT))
+            self._upload_json(gcs_upload_path)