Merge pull request #180 from center-for-threat-informed-defense/175-f…

…ix-json-export

Fix JSON export for reports (#175)

README.md

@@ -154,7 +154,7 @@ docker.errors.DockerException: Error while fetching server API version: ('Connec
 [97438] Failed to execute script docker-compose
 ```
 
-Then most likely Docker is not running and you need to start Docker.
+Then most likely reason is that Docker is not running and you need to start it.
 
 ## Report Troubleshooting
 

src/tram/renderers.py

@@ -0,0 +1,29 @@
+import io
+
+from rest_framework import renderers
+
+import tram.report.docx
+
+
+class DocxReportRenderer(renderers.BaseRenderer):
+    """This custom renderer exports mappings into Word .docx format."""
+
+    media_type = (
+        "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+    )
+    format = "docx"
+
+    def render(self, data, accepted_media_type=None, renderer_context=None):
+        """
+        Export report mappings into Word .docx format.
+
+        :param data: the report mappings dict
+        :param accepted_media_type: the content type negotiated by DRF
+        :param renderer_context: optional additional data
+        :returns bytes: .docx binary data
+        """
+        document = tram.report.docx.build(data)
+        buffer = io.BytesIO()
+        document.save(buffer)
+        buffer.seek(0)
+        return buffer.read()

src/tram/templates/index.html

@@ -71,10 +71,8 @@ <h1>Reports</h1>
               Export
             </button>
             <div class="dropdown-menu" aria-labelledby="export-dropdown">
-              <a class="dropdown-item btn btn-sm btn-outline-secondary"
-                href="/api/report-export/{{ report.id }}/?type=json">JSON</a>
-              <a class="dropdown-item btn btn-sm btn-outline-secondary"
-                href="/api/report-export/{{ report.id }}/?type=docx">DOCX</a>
+              <a class="dropdown-item btn btn-sm btn-outline-secondary" href="{% url 'report-mapping-detail' report.id %}?format=json">JSON</a>
+              <a class="dropdown-item btn btn-sm btn-outline-secondary" href="{% url 'report-mapping-detail' report.id %}?format=docx">DOCX</a>
             </div>
           </div>
           {% if report.document_id is not None %}

src/tram/urls.py

@@ -29,7 +29,9 @@
 router.register(r"jobs", views.DocumentProcessingJobViewSet)
 router.register(r"mappings", views.MappingViewSet)
 router.register(r"reports", views.ReportViewSet)
-router.register(r"report-export", views.ReportExportViewSet)
+router.register(
+    r"report-mappings", views.ReportMappingViewSet, basename="report-mapping"
+)
 router.register(r"sentences", views.SentenceViewSet)
 
 

src/tram/views.py

@@ -1,25 +1,17 @@
-import io
 import json
 import logging
 import time
 from urllib.parse import quote
 
 from constance import config
 from django.contrib.auth.decorators import login_required
-from django.http import (
-    Http404,
-    HttpResponse,
-    HttpResponseBadRequest,
-    JsonResponse,
-    StreamingHttpResponse,
-)
+from django.http import Http404, HttpResponse, HttpResponseBadRequest, JsonResponse
 from django.shortcuts import render
 from django.views.decorators.http import require_POST
-from rest_framework import viewsets
+from rest_framework import renderers, viewsets
 from rest_framework.decorators import api_view
 from rest_framework.response import Response
 
-import tram.report.docx
 from tram import serializers
 from tram.ml import base
 from tram.models import (
@@ -30,6 +22,7 @@
     Report,
     Sentence,
 )
+from tram.renderers import DocxReportRenderer
 
 logger = logging.getLogger(__name__)
 
@@ -62,56 +55,41 @@ class ReportViewSet(viewsets.ModelViewSet):
     serializer_class = serializers.ReportSerializer
 
 
-class ReportExportViewSet(viewsets.ModelViewSet):
-    queryset = Report.objects.all()
+class ReportMappingViewSet(viewsets.ModelViewSet):
+    """
+    This viewset provides access to report mappings.
+    """
+
     serializer_class = serializers.ReportExportSerializer
+    renderer_classes = [renderers.JSONRenderer, DocxReportRenderer]
 
     def get_queryset(self):
-        queryset = ReportViewSet.queryset
+        """
+        Override parent implementation to support lookup by document ID.
+        """
+        queryset = Report.objects.all()
         document_id = self.request.query_params.get("doc-id", None)
         if document_id:
             queryset = queryset.filter(document__id=document_id)
 
         return queryset
 
-    def retrieve(self, request, *args, **kwargs):
-
-        report_format = request.GET.get("type", "")
-
-        # If an invalid report_format is given, just default to json
-        if report_format not in ["json", "docx"]:
-            report_format = "json"
-            logger.warning("Invalid File Type. Defaulting to JSON.")
-
-        # Retrieve report data as json
-        response = super().retrieve(request, *args, **kwargs)
-        basename = quote(self.get_object().name, safe="")
+    def retrieve(self, request, pk=None):
+        """
+        Get the mappings for a report.
 
-        if report_format == "json":
-            response["Content-Disposition"] = f'attachment; filename="{basename}.json"'
-
-        elif report_format == "docx":
-            # Uses json dictionary to create formatted document
-            document = tram.report.docx.build(response.data)
-
-            # save document info
-            buffer = io.BytesIO()
-            document.save(buffer)  # save your memory stream
-            buffer.seek(0)  # rewind the stream
-
-            # put them to streaming content response within docx content_type
-            content_type = (
-                "application/"
-                "vnd.openxmlformats-officedocument.wordprocessingml.document"
-            )
-            response = StreamingHttpResponse(
-                streaming_content=buffer,  # use the stream's content
-                content_type=content_type,
-            )
-
-            response["Content-Disposition"] = f'attachment; filename="{basename}.docx"'
-            response["Content-Encoding"] = "UTF-8"
+        Overrides the parent implementation to add a Content-Disposition header
+        so that the browser will download instead of displaying inline.
 
+        :param request: HTTP request
+        :param pk: primary key of a report
+        """
+        response = super().retrieve(request, request, pk)
+        report = self.get_object()
+        filename = "{}.{}".format(
+            quote(report.name, safe=""), request.accepted_renderer.format
+        )
+        response["Content-Disposition"] = f'attachment; filename="{filename}"'
         return response
 
 

tests/tram/test_views.py

@@ -254,28 +254,28 @@ def test_get_sentences_by_technique(self, logged_in_client):
 
 
 @pytest.mark.django_db
-class TestReportExport:
-    def test_get_report_export_succeeds(self, logged_in_client, mapping):
+class TestReportMappings:
+    def test_get_json(self, logged_in_client, mapping):
         # Act
-        response = logged_in_client.get("/api/report-export/1/")
+        response = logged_in_client.get("/api/report-mappings/1/?format=json")
         json_response = json.loads(response.content)
 
         # Assert
-        assert "sentences" in json_response
+        assert json_response["id"] == 1
         assert len(json_response["sentences"][0]["mappings"]) == 1
 
-    def test_export_docx_report(self, logged_in_client, mapping):
+    def test_get_docx(self, logged_in_client, mapping):
         """
         Check that something that looks like a Word doc was returned.
 
         There are separate unit tests for the doc's content.
         """
         # Act
-        response = logged_in_client.get("/api/report-export/1/?type=docx")
-        data = list(response.streaming_content)
+        response = logged_in_client.get("/api/report-mappings/1/?format=docx")
+        data = response.content
 
         # Assert
-        assert data[0].startswith(b"PK\x03\x04")
+        assert data.startswith(b"PK\x03\x04")
 
     def test_bootstrap_training_data_can_be_posted_as_json_report(
         self, logged_in_client
@@ -286,7 +286,7 @@ def test_bootstrap_training_data_can_be_posted_as_json_report(
 
         # Act
         response = logged_in_client.post(
-            "/api/report-export/", json_string, content_type="application/json"
+            "/api/report-mappings/", json_string, content_type="application/json"
         )
 
         # Assert
@@ -295,7 +295,7 @@ def test_bootstrap_training_data_can_be_posted_as_json_report(
     def test_report_export_update_not_implemented(self, logged_in_client):
         # Act
         response = logged_in_client.post(
-            "/api/report-export/1/", "{}", content_type="application/json"
+            "/api/report-mappings/1/", "{}", content_type="application/json"
         )
 
         # Assert
@@ -311,7 +311,7 @@ def test_download_original_report(self, logged_in_client, document):
     def test_get_reports_by_doc_id(self, logged_in_client, report_with_document):
         # Act
         doc_id = report_with_document.document.id
-        response = logged_in_client.get(f"/api/report-export/?doc-id={doc_id}")
+        response = logged_in_client.get(f"/api/report-mappings/?doc-id={doc_id}")
         json_response = json.loads(response.content)
 
         # Assert