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

Author: wendigo

setup.py

@@ -84,6 +84,7 @@
     python_requires=">=3.9",
     install_requires=[
         "lz4",
+        "orjson",
         "python-dateutil",
         "pytz",
         # requests CVE https://github.com/advisories/GHSA-j8r2-6x86-q33q

tests/unit/oauth_test_utils.py

@@ -9,7 +9,7 @@
 # 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 orjson
 import re
 import uuid
 from collections import namedtuple
@@ -51,7 +51,7 @@ def __init__(self, redirect_server, token_server, tokens, sample_post_response_d
     def __call__(self, request, uri, response_headers):
         authorization = request.headers.get("Authorization")
         if authorization and authorization.replace("Bearer ", "") in self.tokens:
-            return [200, response_headers, json.dumps(self.sample_post_response_data)]
+            return [200, response_headers, orjson.dumps(self.sample_post_response_data)]
         elif self.redirect_server is None and self.token_server is not None:
             return [401,
                     {
@@ -127,7 +127,7 @@ def post_statement_callback(self, request, uri, response_headers):
         authorization = request.headers.get("Authorization")
 
         if authorization and authorization.replace("Bearer ", "") in self.tokens:
-            return [200, response_headers, json.dumps(self.sample_post_response_data)]
+            return [200, response_headers, orjson.dumps(self.sample_post_response_data)]
 
         challenge_id = str(uuid.uuid4())
         token = str(uuid.uuid4())

tests/unit/test_client.py

@@ -9,7 +9,7 @@
 # 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 orjson
 import threading
 import time
 import urllib
@@ -61,8 +61,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 = orjson.dumps(sample_post_response_data)
     req = TrinoRequest(
         host="coordinator",
         port=8080,
@@ -558,7 +557,7 @@ def test_oauth2_header_parsing(header, sample_post_response_data):
     def post_statement(request, uri, response_headers):
         authorization = request.headers.get("Authorization")
         if authorization and authorization.replace("Bearer ", "") in token:
-            return [200, response_headers, json.dumps(sample_post_response_data)]
+            return [200, response_headers, orjson.dumps(sample_post_response_data)]
         return [401, {'Www-Authenticate': header.format(redirect_server=redirect_server, token_server=token_server),
                       'Basic realm': '"Trino"'}, ""]
 
@@ -692,7 +691,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 = orjson.dumps(sample_get_response_data)
 
     req = TrinoRequest(
         host="coordinator",
@@ -718,7 +717,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 = orjson.dumps(sample_get_response_data_none)
 
     req = TrinoRequest(
         host="coordinator",
@@ -744,7 +743,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 = orjson.dumps(sample_get_error_response_data)
 
     req = TrinoRequest(
         host="coordinator",
@@ -1154,8 +1153,9 @@ def headers(self):
                 'X-Trino-Fake-2': 'two',
             }
 
-        def json(self):
-            return sample_get_response_data
+        @property
+        def text(self):
+            return orjson.dumps(sample_get_response_data)
 
     req = TrinoRequest(
         host="coordinator",

trino/auth.py

@@ -11,7 +11,7 @@
 # limitations under the License.
 import abc
 import importlib
-import json
+import orjson
 import os
 import re
 import threading
@@ -368,7 +368,7 @@ def get_token_from_cache(self, key: Optional[str]) -> Optional[str]:
         password = self._keyring.get_password(key, "token")
 
         try:
-            password_as_dict = json.loads(str(password))
+            password_as_dict = orjson.loads(str(password))
             if password_as_dict.get("sharded_password"):
                 # if password was stored shared, reconstruct it
                 shard_count = int(password_as_dict.get("shard_count"))
@@ -404,7 +404,7 @@ def store_token_to_cache(self, key: Optional[str], token: str) -> None:
                 }
 
                 # store the "shard info" as the "base" password
-                self._keyring.set_password(key, "token", json.dumps(shard_info))
+                self._keyring.set_password(key, "token", orjson.dumps(shard_info))
                 # then store all shards with the shard number as postfix
                 for i, s in enumerate(password_shards):
                     self._keyring.set_password(key, f"token__{i}", s)
@@ -521,7 +521,7 @@ def _get_token(self, token_server: str, response: Response, **kwargs: Any) -> st
             with response.connection.send(Request(
                     method='GET', url=token_server).prepare(), **kwargs) as response:
                 if response.status_code == 200:
-                    token_response = json.loads(response.text)
+                    token_response = orjson.loads(response.text)
                     token = token_response.get('token')
                     if token:
                         return token

trino/client.py

@@ -39,7 +39,7 @@
 import base64
 import copy
 import functools
-import json
+import orjson
 import os
 import random
 import re
@@ -687,7 +687,7 @@ def process(self, http_response: Response) -> TrinoStatus:
             self.raise_response_error(http_response)
 
         http_response.encoding = "utf-8"
-        response = http_response.json()
+        response = orjson.loads(http_response.text)
         if "error" in response and response["error"]:
             raise self._process_error(response["error"], response.get("id"))
 
@@ -1285,7 +1285,7 @@ def __init__(self, mapper: RowMapper) -> None:
         self._mapper = mapper
 
     def decode(self, data: bytes, metadata: Dict[str, Any]) -> List[List[Any]]:
-        return self._mapper.map(json.loads(data.decode("utf8")))
+        return self._mapper.map(orjson.loads(data.decode("utf8")))
 
 
 class CompressedQueryDataDecoder(QueryDataDecoder):