Add support for application/json mime type in statement output

CHANGELOG.md

@@ -6,6 +6,7 @@
 - Added support for http requests session adapter configuration
 - Added support default spark session config settings. Any settings set in as defaults
   are preserved unless explicitly overridden by the user
+- Added support for application/json mime type in statement output
 
 ## 0.20.5
 

sparkmagic/sparkmagic/livyclientlib/command.py

@@ -19,6 +19,7 @@
     MIMETYPE_TEXT_PLAIN,
     COMMAND_INTERRUPTED_MSG,
     COMMAND_CANCELLATION_FAILED_MSG,
+    MIMETYPE_APPLICATION_JSON
 )
 from .exceptions import (
     LivyUnexpectedStatusException,
@@ -144,6 +145,8 @@ def _get_statement_output(self, session, statement_id):
                         return (True, image, MIMETYPE_IMAGE_PNG)
                     elif MIMETYPE_TEXT_HTML in data:
                         return (True, data[MIMETYPE_TEXT_HTML], MIMETYPE_TEXT_HTML)
+                    elif MIMETYPE_APPLICATION_JSON in data:
+                        return (True, data[MIMETYPE_APPLICATION_JSON], MIMETYPE_APPLICATION_JSON)
                     else:
                         return (True, data[MIMETYPE_TEXT_PLAIN], MIMETYPE_TEXT_PLAIN)
                 elif statement_output["status"] == "error":

sparkmagic/sparkmagic/tests/test_command.py

@@ -14,6 +14,7 @@
     MIMETYPE_TEXT_HTML,
     MIMETYPE_TEXT_PLAIN,
     COMMAND_INTERRUPTED_MSG,
+    MIMETYPE_APPLICATION_JSON
 )
 from sparkmagic.livyclientlib.command import Command
 from sparkmagic.livyclientlib.livysession import LivySession
@@ -117,6 +118,23 @@ def test_execute():
     assert "<p>out</p>" == result[1]
     assert MIMETYPE_TEXT_HTML == result[2]
 
+    # Now try with JSON result:
+    http_client.get_statement.return_value = {
+        "id": 0,
+        "state": "available",
+        "output": {
+            "status": "ok",
+            "execution_count": 0,
+            "data": {"application/json": {
+                "key_test": "value_test"
+            }},
+        },
+    }
+    result = command.execute(session)
+    assert result[0]
+    assert {"key_test": "value_test"} == result[1]
+    assert MIMETYPE_APPLICATION_JSON == result[2]
+
 
 def test_execute_waiting():
     spark_events = MagicMock()

sparkmagic/sparkmagic/utils/constants.py

@@ -138,3 +138,4 @@
 MIMETYPE_IMAGE_PNG = "image/png"
 MIMETYPE_TEXT_HTML = "text/html"
 MIMETYPE_TEXT_PLAIN = "text/plain"
+MIMETYPE_APPLICATION_JSON = "application/json"