From 63cce4c0cbdd9612b4f95103b3bd0d9fa3ced40d Mon Sep 17 00:00:00 2001
From: Oleg Strunevskiy <oleg.strunevskiy@gmail.com>
Date: Mon, 11 Sep 2023 21:22:06 -0700
Subject: [PATCH 1/2] Handle application/json in statement output

---
 sparkmagic/sparkmagic/livyclientlib/command.py |  3 +++
 sparkmagic/sparkmagic/tests/test_command.py    | 18 ++++++++++++++++++
 sparkmagic/sparkmagic/utils/constants.py       |  1 +
 3 files changed, 22 insertions(+)

diff --git a/sparkmagic/sparkmagic/livyclientlib/command.py b/sparkmagic/sparkmagic/livyclientlib/command.py
index 3a03266f..0bb7d369 100644
--- a/sparkmagic/sparkmagic/livyclientlib/command.py
+++ b/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":
diff --git a/sparkmagic/sparkmagic/tests/test_command.py b/sparkmagic/sparkmagic/tests/test_command.py
index b1d960ef..91b46816 100644
--- a/sparkmagic/sparkmagic/tests/test_command.py
+++ b/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()
diff --git a/sparkmagic/sparkmagic/utils/constants.py b/sparkmagic/sparkmagic/utils/constants.py
index 59b7c099..7c9a7562 100644
--- a/sparkmagic/sparkmagic/utils/constants.py
+++ b/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"

From e7a354c5a6ddef17c8ebc838caf8d70e8e03de9d Mon Sep 17 00:00:00 2001
From: Oleg Strunevskiy <oleg.strunevskiy@gmail.com>
Date: Mon, 11 Sep 2023 21:27:32 -0700
Subject: [PATCH 2/2] Update  CHANGELOG.md

---
 CHANGELOG.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 658831f9..758d6c34 100644
--- a/CHANGELOG.md
+++ b/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