JSON serializer for Elasticsearch Python client that uses the 'orjson' library for performance

elasticsearch-orjson-serializer.py

"""JSON serializer for Elasticsearch Python client that
uses the 'orjson' library for performance.

https://github.com/ijl/orjson
"""

# Implement the Serializer.loads() and .dumps() APIs w/ orjson:
import orjson
from elasticsearch import SerializationError, JSONSerializer


class OrjsonSerializer(JSONSerializer):

    # Options to pass to orjson.dumps()
    ORJSON_OPTIONS = (
        orjson.OPT_SERIALIZE_UUID
        | orjson.OPT_SERIALIZE_NUMPY
        | orjson.OPT_PASSTHROUGH_DATETIME
    )

    def dumps(self, data):
        try:
            return orjson.dumps(data, self.ORJSON_OPTIONS)
        except Exception as e:
            raise SerializationError(data, e)

    def loads(self, s):
        try:
            return orjson.loads(s)
        except Exception as e:
            raise SerializationError(s, e)


# Use the OrjsonSerializer with Elasticsearch client:
import datetime
import uuid
from elasticsearch import Elasticsearch


client = Elasticsearch(serializer=OrjsonSerializer())
print(
    client.index(
        index="test-index",
        body={
            "string": "Hello, world!",
            "dt": datetime.datetime.now(),
            "int": 1,
            "float": 1.0,
            "uuid": uuid.uuid4(),
            "object": {"key": "val"},
            "list": [1, 2, 3],
        },
    )
)

"""
{
  "_id": "oxOLA3UBlcwG7i27SSTB",
  "_index": "test-index",
  "_primary_term": 1,
  "_seq_no": 1,
  "_shards": {
    "failed": 0,
    "successful": 1,
    "total": 2
  },
  "_type": "_doc",
  "_version": 1,
  "result": "created"
}
"""

print(client.search(index="test-index"))

"""
{
  "_shards": {
    "failed": 0,
    "skipped": 0,
    "successful": 1,
    "total": 1
  },
  "hits": {
    "hits": [
      {
        "_id": "oxOLA3UBlcwG7i27SSTB",
        "_index": "test-index",
        "_score": 1.0,
        "_source": {
          "dt": "2020-10-07T09:51:01.439270",
          "float": 1.0,
          "int": 1,
          "list": [
            1,
            2,
            3
          ],
          "object": {
            "key": "val"
          },
          "string": "Hello, world!",
          "uuid": "08ce3041-eb33-478b-a0ca-0587b9be6411"
        },
        "_type": "_doc"
      }
    ],
    "max_score": 1.0,
    "total": {
      "relation": "eq",
      "value": 1
    }
  },
  "timed_out": false,
  "took": 0
}
"""