tzinfo got lost when storing in mongodb

Author: Johannes Otepka

ipyparallel/controller/mongodb.py

@@ -15,6 +15,7 @@
 from traitlets import Dict, Instance, List, Unicode
 
 from .dictdb import BaseDB
+from datetime import datetime, timezone
 
 # we need to determine the pymongo version because of API changes. see
 # https://pymongo.readthedocs.io/en/stable/migrate-to-pymongo4.html
@@ -25,6 +26,16 @@
 # MongoDB class
 # -----------------------------------------------------------------------------
 
+def _ensure_utc(obj):
+    if isinstance(obj, datetime):
+        obj = obj.replace(tzinfo=timezone.utc)
+    return obj
+
+def _ensure_utc_for_record(rec):
+    for key in ('submitted', 'started', 'completed', 'received'):
+        if key in rec:
+            rec[key] = _ensure_utc(rec[key])
+
 
 class MongoDB(BaseDB):
     """MongoDB TaskRecord backend."""
@@ -91,6 +102,7 @@ def get_record(self, msg_id):
         if not r:
             # r will be '' if nothing is found
             raise KeyError(msg_id)
+        _ensure_utc_for_record(r)
         return r
 
     def update_record(self, msg_id, rec):
@@ -125,6 +137,7 @@ def find_records(self, check, keys=None):
         matches = list(self._records.find(check, keys))
         for rec in matches:
             rec.pop('_id')
+            _ensure_utc_for_record(rec)
         return matches
 
     def get_history(self):