diff --git a/ravendb/documents/bulk_insert_operation.py b/ravendb/documents/bulk_insert_operation.py index 0efdd4ef..8d559b32 100644 --- a/ravendb/documents/bulk_insert_operation.py +++ b/ravendb/documents/bulk_insert_operation.py @@ -3,13 +3,12 @@ from datetime import datetime from abc import ABC -import _queue import concurrent import json +import zlib from concurrent.futures import Future -from copy import deepcopy -from queue import Queue -from threading import Lock, Semaphore +from queue import Empty, Full, Queue +from threading import Lock from typing import Optional, TYPE_CHECKING, List, TypeVar, Type, Generic, Callable import requests @@ -43,24 +42,34 @@ class BulkInsertOperation: + # bounding this is what makes the caller wait for a slow server instead of buffering for it + MAX_BUFFERS_TO_FLUSH = 8 + _ENQUEUE_TIMEOUT_IN_SECONDS = 0.05 + class _BufferExposer: - def __init__(self): + def __init__(self, max_buffers_to_flush: int): self._ongoing_operation = Future() # todo: is there any reason to use Futures? (look at error handling) - self._yield_buffer_semaphore = Semaphore(1) - self._buffers_to_flush_queue = Queue() + self._buffers_to_flush_queue = Queue(maxsize=max_buffers_to_flush) self.output_stream_mock = Future() - def enqueue_buffer_for_flush(self, buffer: bytearray): - self._buffers_to_flush_queue.put(bytes(buffer)) + def try_enqueue_buffer_for_flush(self, buffer: bytearray, timeout: float) -> bool: + """Hand a finished buffer over to be sent. False means the queue is still full. + + The buffer belongs to the queue once it is in, so it is never copied. + """ + try: + self._buffers_to_flush_queue.put(buffer, timeout=timeout) + return True + except Full: + return False - # todo: blocking semaphore acquired and released on enter and exit from bulk insert operation context manager def send_data(self): while True: try: buffer_to_flush = self._buffers_to_flush_queue.get(timeout=0.05) # todo: adjust this pooling time yield buffer_to_flush except Exception as e: - if not isinstance(e, _queue.Empty) or self.is_operation_finished(): + if not isinstance(e, Empty) or self.is_operation_finished(): break continue # expected Empty exception coming from queue, operation isn't finished yet @@ -92,13 +101,33 @@ def __init__( self._skip_overwrite_if_unchanged = skip_overwrite_if_unchanged def create_request(self, node: ServerNode) -> requests.Request: + buffers = self._buffer_exposer.send_data() + headers = {} + if self.use_compression: + buffers = self._gzip(buffers) + headers[constants.Headers.CONTENT_ENCODING] = constants.Headers.Encodings.GZIP + return requests.Request( "POST", f"{node.url}/databases/{node.database}/bulk_insert?id={self._key}" f"&skipOverwriteIfUnchanged={'true' if self._skip_overwrite_if_unchanged else 'false'}", - data=self._buffer_exposer.send_data(), + data=buffers, + headers=headers, ) + @staticmethod + def _gzip(buffers): + """Compress the outgoing buffers as one gzip stream, one flush per buffer.""" + compressor = zlib.compressobj(level=zlib.Z_BEST_SPEED, wbits=zlib.MAX_WBITS | 16) + for buffer in buffers: + compressed = compressor.compress(buffer) + compressor.flush(zlib.Z_SYNC_FLUSH) + if compressed: + yield compressed + + tail = compressor.flush() + if tail: + yield tail + def set_response(self, response: Optional[str], from_cache: bool) -> None: raise NotImplementedError("Not Implemented") @@ -112,8 +141,6 @@ def send(self, session: requests.Session, request: requests.Request) -> requests self._buffer_exposer.error_on_request_start(e) def __init__(self, database: str = None, store: "DocumentStore" = None, options: BulkInsertOptions = None): - self.use_compression = False - self._ongoing_bulk_insert_execute_task: Optional[Future] = None self._first = True self._in_progress_command: Optional[CommandType] = None @@ -127,19 +154,16 @@ def __init__(self, database: str = None, store: "DocumentStore" = None, options: if not database or database.isspace(): self._throw_no_database() - self._use_compression = options.use_compression if options else False self._options = options or BulkInsertOptions() + self.use_compression = bool(self._options.use_compression) self._request_executor = store.get_request_executor(database) - self._enqueue_current_buffer_async = Future() - self._enqueue_current_buffer_async.set_result(None) - self._max_size_in_buffer = 1024 * 1024 self._current_data_buffer = bytearray() self._time_series_batch_size = self._conventions.time_series_batch_size - self._buffer_exposer = BulkInsertOperation._BufferExposer() + self._buffer_exposer = BulkInsertOperation._BufferExposer(self.MAX_BUFFERS_TO_FLUSH) self._generate_entity_id_on_the_client = GenerateEntityIdOnTheClient( self._request_executor.conventions, @@ -164,9 +188,8 @@ def __exit__(self, exc_type, exc_val, exc_tb): if self._current_data_buffer: try: self._write_string_no_escape("]") - self._enqueue_current_buffer_async.result() # wait for enqueue - buffer = self._current_data_buffer - self._buffer_exposer.enqueue_buffer_for_flush(buffer) + self._enqueue_buffer_for_flush(self._current_data_buffer) + self._current_data_buffer = bytearray() except Exception as e: flush_ex = e @@ -287,18 +310,22 @@ def __return_func(): return __return_func def _flush_if_needed(self) -> None: - if len(self._current_data_buffer) > self._max_size_in_buffer or self._enqueue_current_buffer_async.done(): - self._enqueue_current_buffer_async.result() # wait - - buffer = deepcopy(self._current_data_buffer) - self._current_data_buffer.clear() + if len(self._current_data_buffer) <= self._max_size_in_buffer: + return - # todo: check if it's better to create a new bytearray of max size instead of clearing it (possible dealloc) + buffer = self._current_data_buffer + self._current_data_buffer = bytearray() + self._enqueue_buffer_for_flush(buffer) - def __enqueue_buffer_for_flush(flushed_buffer: bytearray): - self._buffer_exposer.enqueue_buffer_for_flush(flushed_buffer) + def _enqueue_buffer_for_flush(self, buffer: bytearray) -> None: + """Hand the buffer over to the thread sending the request, waiting for a free slot. - self._enqueue_current_buffer_async = self._thread_pool_executor.submit(__enqueue_buffer_for_flush, buffer) + Waiting here is the backpressure: a server that reads slower than this client writes slows + the client down rather than filling its memory. + """ + while not self._buffer_exposer.try_enqueue_buffer_for_flush(buffer, self._ENQUEUE_TIMEOUT_IN_SECONDS): + self._throw_if_bulk_insert_execute_task_failed() + self._throw_if_request_already_finished() def _end_previous_command_if_needed(self) -> None: if self._in_progress_command == CommandType.COUNTERS: @@ -331,11 +358,24 @@ def _ensure_ongoing_operation(self) -> None: self._get_bulk_insert_operation_id() self._start_executing_bulk_insert_command() - if ( - self._ongoing_bulk_insert_execute_task.done() and self._ongoing_bulk_insert_execute_task.exception() - ): # todo: check if isCompletedExceptionally returns false if task isn't finished + self._throw_if_bulk_insert_execute_task_failed() + + def _throw_if_request_already_finished(self) -> None: + """Waiting for a free slot only makes sense while the request carrying the data is still open.""" + request = self._ongoing_bulk_insert_execute_task + if request is None or not request.done(): + return + + raise BulkInsertAbortedException("The request carrying the bulk insert data finished before all of it was sent") + + def _throw_if_bulk_insert_execute_task_failed(self) -> None: + task = self._ongoing_bulk_insert_execute_task + if task is None: + return + + if task.done() and task.exception(): try: - self._ongoing_bulk_insert_execute_task.result() + task.result() except Exception as e: self._throw_bulk_insert_aborted(e, None) diff --git a/ravendb/tests/documents_tests/test_bulk_insert_backpressure.py b/ravendb/tests/documents_tests/test_bulk_insert_backpressure.py new file mode 100644 index 00000000..647c736b --- /dev/null +++ b/ravendb/tests/documents_tests/test_bulk_insert_backpressure.py @@ -0,0 +1,113 @@ +import threading +import time +import unittest +from concurrent.futures import Future, ThreadPoolExecutor + +from ravendb.documents.bulk_insert_operation import BulkInsertOperation +from ravendb.documents.conventions import DocumentConventions +from ravendb.exceptions.documents.bulkinsert import BulkInsertAbortedException + + +class Document: + def __init__(self, Id: str, payload: str): + self.Id = Id + self.payload = payload + + +class _FakeRequestExecutor: + def __init__(self, conventions: DocumentConventions): + self.conventions = conventions + + def execute_command(self, command, session_info=None) -> None: + command.result = {"Status": "Running"} + + +class _FakeStore: + """Enough of a document store to drive the bulk insert buffering, with no server behind it.""" + + def __init__(self): + self.thread_pool_executor = ThreadPoolExecutor(max_workers=2) + self.conventions = DocumentConventions() + + def get_request_executor(self, database: str) -> _FakeRequestExecutor: + return _FakeRequestExecutor(self.conventions) + + +class TestBulkInsertBackpressure(unittest.TestCase): + """A server that reads slower than the caller writes has to slow the caller down, not be buffered.""" + + DOCUMENT_COUNT = 200 + PAYLOAD = "x" * 256 * 1024 + + def _bulk_insert_with_a_stalled_stream(self) -> BulkInsertOperation: + bulk = BulkInsertOperation("db", _FakeStore(), None) + bulk._operation_id = 1 + bulk._ongoing_bulk_insert_execute_task = Future() + bulk._current_data_buffer += bytearray("[", encoding="utf-8") + return bulk + + def test_a_stalled_stream_stops_the_caller_instead_of_growing_memory(self): + bulk = self._bulk_insert_with_a_stalled_stream() + queue = bulk._buffer_exposer._buffers_to_flush_queue + self.assertGreater(queue.maxsize, 0, "the outbound queue is unbounded") + failures, stored = [], [] + + def store_documents(): + try: + for i in range(self.DOCUMENT_COUNT): + bulk.store(Document(f"documents/{i}", self.PAYLOAD)) + stored.append(i) + except BaseException as e: # noqa: BLE001 - reported through the assertions below + failures.append(e) + + producer = threading.Thread(target=store_documents, daemon=True) + producer.start() + + deadline = time.time() + 30 + while not queue.full() and time.time() < deadline: + time.sleep(0.01) + + self.assertTrue(queue.full(), "the queue never filled up, so nothing was throttled") + self.assertLess(len(stored), self.DOCUMENT_COUNT) + buffered = sum(len(chunk) for chunk in list(queue.queue)) + self.assertLessEqual(buffered, (queue.maxsize + 1) * bulk._max_size_in_buffer) + + # keep it full for several waiting periods: the caller waits, does not give up and does not fail + time.sleep(BulkInsertOperation._ENQUEUE_TIMEOUT_IN_SECONDS * 6) + self.assertEqual([], failures, "waiting for a free slot must not fail the bulk insert") + self.assertTrue(producer.is_alive(), "the caller stopped waiting while the queue was still full") + + while producer.is_alive(): + try: + queue.get(timeout=0.05) + except Exception: + pass + producer.join(30) + + self.assertEqual([], failures) + self.assertEqual(self.DOCUMENT_COUNT, len(stored)) + + def test_a_request_that_already_finished_does_not_leave_the_caller_waiting(self): + bulk = self._bulk_insert_with_a_stalled_stream() + queue = bulk._buffer_exposer._buffers_to_flush_queue + finished = Future() + finished.set_result(None) + bulk._ongoing_bulk_insert_execute_task = finished + for _ in range(queue.maxsize): + queue.put(bytearray(b"waiting to be sent")) + + def store_documents(): + bulk.store(Document("documents/1", self.PAYLOAD * 5)) + bulk.store(Document("documents/2", "small")) + + producer = threading.Thread(target=store_documents, daemon=True) + producer.start() + producer.join(15) + + self.assertFalse(producer.is_alive(), "the caller waited for a request that had already finished") + # the failure is reported through the operation, the way every write failure in a bulk insert is + self.assertIsInstance(bulk._buffer_exposer._ongoing_operation.exception(timeout=0), BulkInsertAbortedException) + + +if __name__ == "__main__": + unittest.main() diff --git a/ravendb/tests/documents_tests/test_bulk_insert_buffer_handover.py b/ravendb/tests/documents_tests/test_bulk_insert_buffer_handover.py new file mode 100644 index 00000000..6e7c8fb8 --- /dev/null +++ b/ravendb/tests/documents_tests/test_bulk_insert_buffer_handover.py @@ -0,0 +1,111 @@ +import json +import threading +import time +import unittest +from concurrent.futures import Future, ThreadPoolExecutor + +from ravendb.documents.bulk_insert_operation import BulkInsertOperation +from ravendb.documents.conventions import DocumentConventions +from ravendb.tests.test_base import TestBase + + +class Document: + def __init__(self, Id: str = None, payload: str = None): + self.Id = Id + self.payload = payload + + +class _FakeRequestExecutor: + def __init__(self, conventions: DocumentConventions): + self.conventions = conventions + + +class _FakeStore: + """Enough of a document store to drive the bulk insert buffering, with no server behind it.""" + + def __init__(self): + self.thread_pool_executor = ThreadPoolExecutor(max_workers=2) + self.conventions = DocumentConventions() + + def get_request_executor(self, database: str) -> _FakeRequestExecutor: + return _FakeRequestExecutor(self.conventions) + + +class TestBulkInsertBufferHandover(unittest.TestCase): + """A finished buffer is handed to the sender uncopied, which holds only while nobody writes into it after.""" + + DOCUMENT_COUNT = 120 + PAYLOAD = "p" * 40 * 1024 + + def test_a_buffer_never_changes_after_it_has_been_handed_over(self): + bulk = BulkInsertOperation("db", _FakeStore(), None) + bulk._operation_id = 1 + bulk._ongoing_bulk_insert_execute_task = Future() + bulk._current_data_buffer += bytearray("[", encoding="utf-8") + + queue = bulk._buffer_exposer._buffers_to_flush_queue + handed_over = [] # (the object the sender holds, a snapshot taken the moment it arrived) + stop = threading.Event() + + def sender(): + while not stop.is_set(): + try: + chunk = queue.get(timeout=0.02) + except Exception: + continue + handed_over.append((chunk, bytes(chunk))) + + reader = threading.Thread(target=sender, daemon=True) + reader.start() + + buffer_identities = set() + for i in range(self.DOCUMENT_COUNT): + bulk.store(Document(f"documents/{i}", f"{i}:{self.PAYLOAD}")) + buffer_identities.add(id(bulk._current_data_buffer)) + + deadline = time.time() + 30 + while queue.qsize() and time.time() < deadline: + time.sleep(0.01) + stop.set() + reader.join(30) + tail = bytes(bulk._current_data_buffer) + + self.assertGreaterEqual(len(handed_over), 3, "no mid-stream flush happened, the test proves nothing") + for index, (chunk, snapshot) in enumerate(handed_over): + self.assertEqual(snapshot, bytes(chunk), f"buffer {index} changed after it was handed over") + self.assertEqual( + len(handed_over), + len({id(chunk) for chunk, _ in handed_over}), + "the same buffer object was handed over more than once", + ) + self.assertGreater(len(buffer_identities), 1, "the caller kept writing into the same buffer object") + + commands = json.loads(b"".join(snapshot for _, snapshot in handed_over) + tail + b"]") + self.assertEqual(self.DOCUMENT_COUNT, len(commands)) + for i, command in enumerate(commands): + self.assertEqual(f"documents/{i}", command["Id"]) + self.assertEqual(f"{i}:{self.PAYLOAD}", command["Document"]["payload"]) + + +class TestBulkInsertBufferHandoverAgainstServer(TestBase): + """The same handover, but through the requests library and a real server.""" + + DOCUMENT_COUNT = 200 + PAYLOAD = "q" * 26 * 1024 # ~5 MiB in total, so the 1 MiB buffer is flushed several times + + def test_a_load_spanning_many_buffers_arrives_intact(self): + with self.store.bulk_insert() as bulk: + for i in range(self.DOCUMENT_COUNT): + bulk.store(Document(f"documents/{i}", f"{i}:{self.PAYLOAD}")) + + with self.store.open_session() as session: + self.assertEqual(self.DOCUMENT_COUNT, session.query(object_type=Document).count()) + + for i in (0, self.DOCUMENT_COUNT - 1): + with self.store.open_session() as session: + loaded = session.load(f"documents/{i}", Document) + self.assertEqual(f"{i}:{self.PAYLOAD}", loaded.payload, f"documents/{i} arrived corrupted") + + +if __name__ == "__main__": + unittest.main() diff --git a/ravendb/tests/documents_tests/test_bulk_insert_compression.py b/ravendb/tests/documents_tests/test_bulk_insert_compression.py new file mode 100644 index 00000000..7e23526b --- /dev/null +++ b/ravendb/tests/documents_tests/test_bulk_insert_compression.py @@ -0,0 +1,60 @@ +import gzip +import json +import unittest + +from ravendb.documents.bulk_insert_operation import BulkInsertOperation, BulkInsertOptions +from ravendb.http.server_node import ServerNode +from ravendb.primitives import constants +from ravendb.tests.test_base import TestBase + + +class User: + def __init__(self, Id: str = None, name: str = None, notes: str = None): + self.Id = Id + self.name = name + self.notes = notes + + +class TestBulkInsertCompressionUnit(unittest.TestCase): + def test_the_buffers_become_one_smaller_gzip_stream_written_as_it_goes(self): + buffers = [bytearray(b'[{"Id":"users/1"}'), bytearray(b',{"Id":"users/2"}'), bytearray(b"]")] + compressed = b"".join(BulkInsertOperation._BulkInsertCommand._gzip(iter(buffers))) + self.assertEqual(b"".join(bytes(buffer) for buffer in buffers), gzip.decompress(compressed)) + + bulky = bytearray(json.dumps([{"Id": f"users/{i}", "name": "the same name"} for i in range(500)]).encode()) + self.assertLess(len(b"".join(BulkInsertOperation._BulkInsertCommand._gzip(iter([bulky])))), len(bulky)) + + # a buffer must not sit in the compressor waiting for the next one + stream = BulkInsertOperation._BulkInsertCommand._gzip(iter([bytearray(b"x" * 4096), bytearray(b"y" * 4096)])) + self.assertTrue(next(stream)) + + def test_the_request_declares_the_encoding_it_used(self): + def request_with_compression(use_compression: bool): + command = BulkInsertOperation._BulkInsertCommand(1, BulkInsertOperation._BufferExposer(1), "A", False) + command.use_compression = use_compression + return command.create_request(ServerNode("http://localhost:8080", "db")) + + header = constants.Headers.CONTENT_ENCODING + self.assertEqual(constants.Headers.Encodings.GZIP, request_with_compression(True).headers.get(header)) + self.assertIsNone(request_with_compression(False).headers.get(header)) + + +class TestBulkInsertCompression(TestBase): + def test_documents_stored_with_compression_arrive_intact(self): + # a payload that actually compresses, and enough of it to cross a buffer boundary + notes = "the quick brown fox jumps over the lazy dog " * 2000 + + with self.store.bulk_insert(self.store.database, BulkInsertOptions(use_compression=True)) as bulk: + self.assertTrue(bulk.use_compression) + for i in range(60): + bulk.store(User(f"users/{i}", f"user {i}", notes)) + + with self.store.open_session() as session: + self.assertEqual(60, session.query(object_type=User).count()) + loaded = session.load("users/42", User) + self.assertEqual("user 42", loaded.name) + self.assertEqual(notes, loaded.notes) + + +if __name__ == "__main__": + unittest.main() diff --git a/ravendb/tests/session_tests/test_entity_to_dict_semantics.py b/ravendb/tests/session_tests/test_entity_to_dict_semantics.py new file mode 100644 index 00000000..9564b122 --- /dev/null +++ b/ravendb/tests/session_tests/test_entity_to_dict_semantics.py @@ -0,0 +1,123 @@ +import json +import math +import unittest +from datetime import datetime, timedelta +from enum import Enum, IntEnum + +from ravendb.documents.conventions import DocumentConventions +from ravendb.json.metadata_as_dictionary import MetadataAsDictionary +from ravendb.tools.utils import Utils + + +class Color(Enum): + RED = "red" + + +class Priority(IntEnum): + HIGH = 9 + + +class StringBacked(str, Enum): + ALPHA = "alpha" + + +class Address: + def __init__(self, city: str, zip_code: str): + self.city = city + self.zip_code = zip_code + + +class Person: + def __init__(self, name: str, address: Address = None): + self.name = name + self.address = address + + +class WithToJson: + def to_json(self): + return {"rendered": 42} + + +class SelfReferencing: + def __init__(self): + self.child = None + + +class TestEntityToDictSemantics(unittest.TestCase): + """Utils.entity_to_dict decides what a document looks like for both the session and bulk insert.""" + + def setUp(self): + self.default_method = DocumentConventions.json_default + + def test_conversion_matches_a_round_trip_through_a_json_string(self): + moment = datetime(2026, 8, 19, 14, 30, 15) + span = timedelta(days=2, minutes=23, seconds=59, milliseconds=254) + cases = { + "nested objects": ( + Person("Ayende", Address("Hadera", "38100")), + {"name": "Ayende", "address": {"city": "Hadera", "zip_code": "38100"}}, + ), + "to_json wins over the instance dict": ({"item": WithToJson()}, {"item": {"rendered": 42}}), + "datetime": ({"at": moment}, {"at": Utils.datetime_to_string(moment)}), + "timedelta": ({"took": span}, {"took": Utils.timedelta_to_str(span)}), + "enum becomes its value": ({"color": Color.RED}, {"color": "red"}), + "int enum becomes an int": ({"priority": Priority.HIGH}, {"priority": 9}), + "str backed enum becomes a str": ({"kind": StringBacked.ALPHA}, {"kind": "alpha"}), + "tuple becomes a list": ({"point": (1, 2, 3)}, {"point": [1, 2, 3]}), + "set becomes a list": ({"tags": {"a"}}, {"tags": ["a"]}), + "metadata as dictionary": ( + {"meta": MetadataAsDictionary({"@collection": "People"})}, + {"meta": {"@collection": "People"}}, + ), + "scalar dict keys": ({1: "int", 1.5: "float", None: "none"}, {"1": "int", "1.5": "float", "null": "none"}), + "bool dict keys": ({True: "yes", False: "no"}, {"true": "yes", "false": "no"}), + "types are kept": ( + {"yes": True, "i": 7, "f": 1.25, "nothing": None}, + {"yes": True, "i": 7, "f": 1.25, "nothing": None}, + ), + "non ascii text": ({"city": "Zurich, Kraków, 東京"}, {"city": "Zurich, Kraków, 東京"}), + "objects in a list": ( + {"people": [Person("A", Address("Hadera", "1")), Person("B")]}, + { + "people": [ + {"name": "A", "address": {"city": "Hadera", "zip_code": "1"}}, + {"name": "B", "address": None}, + ] + }, + ), + "the same object twice": ( + {"home": (shared := Address("Hadera", "38100")), "work": shared}, + {"home": {"city": "Hadera", "zip_code": "38100"}, "work": {"city": "Hadera", "zip_code": "38100"}}, + ), + "dict ordering": ({"b": 1, "a": 2}, {"b": 1, "a": 2}), + } + for name, (value, expected) in cases.items(): + with self.subTest(name): + converted = Utils.entity_to_dict(value, self.default_method) + self.assertEqual(expected, converted) + self.assertEqual(json.loads(json.dumps(value, default=self.default_method)), converted) + if isinstance(expected, dict): + self.assertEqual(list(expected.keys()), list(converted.keys())) + + def test_non_finite_floats_survive_as_floats(self): + converted = Utils.entity_to_dict({"nan": float("nan"), "inf": float("inf")}, self.default_method) + self.assertTrue(math.isnan(converted["nan"])) + self.assertEqual(float("inf"), converted["inf"]) + + def test_circular_reference_is_reported_and_does_not_recurse_forever(self): + first, second = SelfReferencing(), SelfReferencing() + first.child, second.child = second, first + with self.assertRaises(ValueError): + Utils.entity_to_dict(first, self.default_method) + + def test_values_that_cannot_be_converted_are_rejected(self): + for name, value in { + "unsupported dict key": {(1, 2): "tuple key"}, + "opaque value": {"opaque": object()}, + }.items(): + with self.subTest(name), self.assertRaises(TypeError): + Utils.entity_to_dict(value, self.default_method) + + +if __name__ == "__main__": + unittest.main() diff --git a/ravendb/tools/utils.py b/ravendb/tools/utils.py index 636457f4..9ee3aafa 100644 --- a/ravendb/tools/utils.py +++ b/ravendb/tools/utils.py @@ -942,7 +942,80 @@ def dictionarize(obj: object) -> dict: @staticmethod def entity_to_dict(entity, default_method) -> dict: - return json.loads(json.dumps(entity, default=default_method)) + """Build the JSON-able form of an entity, matching json.loads(json.dumps(entity)).""" + return Utils._to_json_value(entity, default_method, {}) + + @staticmethod + def _to_json_value(value, default_method, containers_in_progress: Dict[int, object]): + value_type = type(value) + if value is None or value_type is str or value_type is bool or value_type is int or value_type is float: + return value + if isinstance(value, dict): + return Utils._to_json_object(value, default_method, containers_in_progress) + if isinstance(value, (list, tuple)): + return Utils._to_json_array(value, default_method, containers_in_progress) + # json writes scalar subclasses by their built-in value, so a str or int enum never reaches default + if isinstance(value, str): + return str.__str__(value) + if isinstance(value, int): # bool cannot be subclassed, so this is not a bool + return int(value) + if isinstance(value, float): + return float(value) + return Utils._to_json_value(default_method(value), default_method, containers_in_progress) + + @staticmethod + def _to_json_object(value: dict, default_method, containers_in_progress: Dict[int, object]) -> dict: + marker = id(value) + if marker in containers_in_progress: + raise ValueError("Circular reference detected") + containers_in_progress[marker] = value + try: + return { + Utils._to_json_key(key): Utils._to_json_value(item, default_method, containers_in_progress) + for key, item in value.items() + } + finally: + del containers_in_progress[marker] + + @staticmethod + def _to_json_array(value, default_method, containers_in_progress: Dict[int, object]) -> list: + marker = id(value) + if marker in containers_in_progress: + raise ValueError("Circular reference detected") + containers_in_progress[marker] = value + try: + return [Utils._to_json_value(item, default_method, containers_in_progress) for item in value] + finally: + del containers_in_progress[marker] + + @staticmethod + def _to_json_key(key) -> str: + # mirrors the order in which json coerces object keys + if type(key) is str: + return key + if isinstance(key, str): + return str.__str__(key) + if isinstance(key, float): + return Utils._float_to_json(key) + if key is True: + return "true" + if key is False: + return "false" + if key is None: + return "null" + if isinstance(key, int): + return int.__repr__(key) + raise TypeError(f"keys must be str, int, float, bool or None, not {type(key).__name__}") + + @staticmethod + def _float_to_json(value: float) -> str: + if value != value: + return "NaN" + if value == float("inf"): + return "Infinity" + if value == float("-inf"): + return "-Infinity" + return float.__repr__(value) @staticmethod def add_hours(date: datetime, hours: int):