From 4f0bf9ca67d251a5b6595e5626c8b3b8c137fa60 Mon Sep 17 00:00:00 2001
From: "Mateusz \"Serafin\" Gajewski" <github@wendigo.pl>
Date: Fri, 1 Aug 2025 15:34:49 +0200
Subject: [PATCH] Migrate from built-in json to faster orjson

It is way faster than the current json parser: https://github.com/ijl/orjson?tab=readme-ov-file#performance
---
 setup.py                  |  7 ++++++-
 tests/unit/test_client.py | 19 +++++++++++--------
 trino/client.py           |  8 ++++++--
 3 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/setup.py b/setup.py
index e497ab36..cfe8787f 100755
--- a/setup.py
+++ b/setup.py
@@ -11,6 +11,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
+import platform
 from codecs import open
 from typing import Any
 from typing import Dict
@@ -26,6 +27,10 @@
 with open(os.path.join(here, "README.md"), "r", "utf-8") as f:
     readme = f.read()
 
+orjson_require = []
+if platform.python_implementation() != 'PyPy':
+    orjson_require = ['orjson >= 3.11.0']
+
 kerberos_require = ["requests_kerberos"]
 gssapi_require = [""
                   "requests_gssapi",
@@ -90,7 +95,7 @@
         "requests>=2.31.0",
         "tzlocal",
         "zstandard",
-    ],
+    ] + orjson_require,
     extras_require={
         "all": all_require,
         "kerberos": kerberos_require,
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
index 85ac25f7..ba5f7f28 100644
--- a/tests/unit/test_client.py
+++ b/tests/unit/test_client.py
@@ -9,7 +9,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import json
 import threading
 import time
 import urllib
@@ -26,6 +25,10 @@
 import gssapi
 import httpretty
 import keyring
+try:
+    import orjson as json
+except ImportError:
+    import json
 import pytest
 import requests
 from httpretty import httprettified
@@ -61,8 +64,7 @@
 
 @mock.patch("trino.client.TrinoRequest.http")
 def test_trino_initial_request(mock_requests, sample_post_response_data):
-    mock_requests.Response.return_value.json.return_value = sample_post_response_data
-
+    mock_requests.Response.return_value.text = json.dumps(sample_post_response_data)
     req = TrinoRequest(
         host="coordinator",
         port=8080,
@@ -692,7 +694,7 @@ def run(self) -> None:
 
 @mock.patch("trino.client.TrinoRequest.http")
 def test_trino_fetch_request(mock_requests, sample_get_response_data):
-    mock_requests.Response.return_value.json.return_value = sample_get_response_data
+    mock_requests.Response.return_value.text = json.dumps(sample_get_response_data)
 
     req = TrinoRequest(
         host="coordinator",
@@ -718,7 +720,7 @@ def test_trino_fetch_request(mock_requests, sample_get_response_data):
 
 @mock.patch("trino.client.TrinoRequest.http")
 def test_trino_fetch_request_data_none(mock_requests, sample_get_response_data_none):
-    mock_requests.Response.return_value.json.return_value = sample_get_response_data_none
+    mock_requests.Response.return_value.text = json.dumps(sample_get_response_data_none)
 
     req = TrinoRequest(
         host="coordinator",
@@ -744,7 +746,7 @@ def test_trino_fetch_request_data_none(mock_requests, sample_get_response_data_n
 
 @mock.patch("trino.client.TrinoRequest.http")
 def test_trino_fetch_error(mock_requests, sample_get_error_response_data):
-    mock_requests.Response.return_value.json.return_value = sample_get_error_response_data
+    mock_requests.Response.return_value.text = json.dumps(sample_get_error_response_data)
 
     req = TrinoRequest(
         host="coordinator",
@@ -1154,8 +1156,9 @@ def headers(self):
                 'X-Trino-Fake-2': 'two',
             }
 
-        def json(self):
-            return sample_get_response_data
+        @property
+        def text(self):
+            return json.dumps(sample_get_response_data)
 
     req = TrinoRequest(
         host="coordinator",
diff --git a/trino/client.py b/trino/client.py
index cf8ec9d1..0271b79d 100644
--- a/trino/client.py
+++ b/trino/client.py
@@ -39,7 +39,6 @@
 import base64
 import copy
 import functools
-import json
 import os
 import random
 import re
@@ -66,6 +65,11 @@
 from zoneinfo import ZoneInfo
 
 import lz4.block
+try:
+    import orjson as json
+except ImportError:
+    import json
+
 import requests
 import zstandard
 from requests import Response
@@ -687,7 +691,7 @@ def process(self, http_response: Response) -> TrinoStatus:
             self.raise_response_error(http_response)
 
         http_response.encoding = "utf-8"
-        response = http_response.json()
+        response = json.loads(http_response.text)
         if "error" in response and response["error"]:
             raise self._process_error(response["error"], response.get("id"))