From eb81f5a4d208598979c168e79b2043b7e7a6220d Mon Sep 17 00:00:00 2001 From: Steve Nyemba Date: Tue, 2 Apr 2024 12:31:41 -0500 Subject: [PATCH] bug fix: mongodb inserts of structured objects with lists as elements --- transport/common.py | 3 ++- transport/nosql/mongodb.py | 20 +++++++++++++++++--- 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/transport/common.py b/transport/common.py index e17c615..f439ea7 100644 --- a/transport/common.py +++ b/transport/common.py @@ -1,5 +1,6 @@ import json - +import numpy as np +from datetime import datetime class IEncoder (json.JSONEncoder): def default (self,object): diff --git a/transport/nosql/mongodb.py b/transport/nosql/mongodb.py index 2784cd2..2b94311 100644 --- a/transport/nosql/mongodb.py +++ b/transport/nosql/mongodb.py @@ -213,12 +213,26 @@ class Writer(Mongo): _uid = self.collection if 'doc' not in _args else _args['doc'] if self._lock : Mongo.lock.acquire() + if type(info) == list or type(info) == pd.DataFrame : - info if type(info) == list else info.to_dict(orient='records') - info = json.loads(json.dumps(info,cls=IEncoder)) + if type(info) == pd.DataFrame : + info = info.to_dict(orient='records') + # info if type(info) == list else info.to_dict(orient='records') + info = json.loads(json.dumps(info)) self.db[_uid].insert_many(info) else: - self.db[_uid].insert_one(json.loads(json.dumps(info,cls=IEncoder))) + # + # sometimes a dictionary can have keys with arrays (odd shaped) + # + _keycount = len(info.keys()) + _arraycount = [len(info[key]) for key in info if type(info[key]) in (list,np.array,np.ndarray)] + if _arraycount and len(_arraycount) == _keycount and np.max(_arraycount) == np.min(_arraycount) : + # + # In case an object with consistent structure is passed, we store it accordingly + # + self.write(pd.DataFrame(info),**_args) + else: + self.db[_uid].insert_one(json.loads(json.dumps(info,cls=IEncoder))) finally: if self._lock : Mongo.lock.release()