From 7351593cbb03c1501b514f114b56b94f77986f08 Mon Sep 17 00:00:00 2001 From: Gao Date: Mon, 12 Jun 2023 12:52:07 +0200 Subject: [PATCH] read and write to MongoDB --- DataContainer/MongoDB.py | 76 ++++++++++++++++++++++++++++- testMongoDB.ipynb | 103 ++++++++++++++++++++++++++++++++++++--- 2 files changed, 171 insertions(+), 8 deletions(-) diff --git a/DataContainer/MongoDB.py b/DataContainer/MongoDB.py index 08c484b..e6f8ba2 100644 --- a/DataContainer/MongoDB.py +++ b/DataContainer/MongoDB.py @@ -3,6 +3,8 @@ import xarray_mongodb import bson import builtins +import xarray as xr + from ToolFunction.ToolFunction import get_date @@ -123,6 +125,28 @@ class MongoDB: } self.mongoDB['global'].update_one(filter, {"$set": data_label}, upsert=False) + + def _add_data_additional(self, shotNum, data, name): + filter = { + 'year': self.year, + 'month': self.month, + 'day': self.day, + 'shotNum': shotNum, + } + + mongoID = self.mongoDB.additional.insert_one(data).inserted_id + + data_label = { + name: + { + 'name': name, + 'mongoID': mongoID, + 'engine': 'additional', + 'dtype': 'dict', + } + } + + self.mongoDB['global'].update_one(filter, {"$set": data_label}, upsert=False) def add_data(self, shotNum, data, date=None, name=None, engine='normal'): if not date is None: @@ -131,10 +155,58 @@ class MongoDB: if engine == 'normal': self._add_data_normal(shotNum=shotNum, data=data) elif engine == 'xarray': - if isinstance(dataSet, type(xr.Dataset())): + if isinstance(data, type(xr.Dataset())): self._add_data_xarray_dataSet(shotNum=shotNum, data=data, name=name) else: self._add_data_xarray_dataArray(shotNum=shotNum, data=data, name=name) + elif engine == 'additional': + self._add_data_additional(shotNum=shotNum, data=data, name=name) + + def read_global(self, shotNum, date=None): + + if not date is None: + self.set_date(date) + + filter = { + 'year': self.year, + 'month': self.month, + 'day': self.day, + 'shotNum': shotNum, + } + + return self.mongoDB['global'].find_one(filter) + + def _load_data_single(self, mongoID, engine): + if engine == 'xarray': + return self.xdb.get(mongoID) + if engine == 'additional': + return self.mongoDB.additional.find_one({'_id': mongoID}) - def read_data(self, shotNum, data=None): + def load_data_single(self, shotNum=None, globalDict=None, date=None, field=None): + + if not date is None: + self.set_date(date) + + if globalDict is None: + globalDict = self.read_global(shotNum) + + if field is None: + field = globalDict + + res = {} + + for key in field: + if isinstance(globalDict[key], dict) and ('mongoID' in globalDict[key]): + mongoID = globalDict[key]['mongoID'] + engine = globalDict[key]['engine'] + + res.update( + { + key: self._load_data_single(mongoID=mongoID, engine=engine) + } + ) + + return res + + def load_data(self, shotNum=None, globalDict=None, date=None, field=None): pass \ No newline at end of file diff --git a/testMongoDB.ipynb b/testMongoDB.ipynb index 2264b97..8a5baf8 100644 --- a/testMongoDB.ipynb +++ b/testMongoDB.ipynb @@ -1957,7 +1957,7 @@ }, { "cell_type": "code", - "execution_count": 146, + "execution_count": 166, "metadata": {}, "outputs": [], "source": [ @@ -2086,6 +2086,28 @@ " }\n", " \n", " self.mongoDB['global'].update_one(filter, {\"$set\": data_label}, upsert=False)\n", + " \n", + " def _add_data_additional(self, shotNum, data, name):\n", + " filter = { \n", + " 'year': self.year,\n", + " 'month': self.month,\n", + " 'day': self.day,\n", + " 'shotNum': shotNum,\n", + " }\n", + " \n", + " mongoID = self.mongoDB.additional.insert_one(data).inserted_id\n", + " \n", + " data_label = {\n", + " name: \n", + " {\n", + " 'name': name,\n", + " 'mongoID': mongoID,\n", + " 'engine': 'additional',\n", + " 'dtype': 'dict',\n", + " }\n", + " }\n", + " \n", + " self.mongoDB['global'].update_one(filter, {\"$set\": data_label}, upsert=False)\n", " \n", " def add_data(self, shotNum, data, date=None, name=None, engine='normal'):\n", " if not date is None:\n", @@ -2094,12 +2116,60 @@ " if engine == 'normal':\n", " self._add_data_normal(shotNum=shotNum, data=data)\n", " elif engine == 'xarray':\n", - " if isinstance(dataSet, type(xr.Dataset())):\n", + " if isinstance(data, type(xr.Dataset())):\n", " self._add_data_xarray_dataSet(shotNum=shotNum, data=data, name=name)\n", " else:\n", " self._add_data_xarray_dataArray(shotNum=shotNum, data=data, name=name)\n", + " elif engine == 'additional':\n", + " self._add_data_additional(shotNum=shotNum, data=data, name=name)\n", + " \n", + " def read_global(self, shotNum, date=None):\n", + " \n", + " if not date is None:\n", + " self.set_date(date)\n", + " \n", + " filter = { \n", + " 'year': self.year,\n", + " 'month': self.month,\n", + " 'day': self.day,\n", + " 'shotNum': shotNum,\n", + " }\n", + " \n", + " return self.mongoDB['global'].find_one(filter)\n", + " \n", + " def _load_data_single(self, mongoID, engine):\n", + " if engine == 'xarray':\n", + " return self.xdb.get(mongoID)\n", + " if engine == 'additional':\n", + " return self.mongoDB.additional.find_one({'_id': mongoID})\n", " \n", - " def read_data(self, shotNum, data=None):\n", + " def load_data_single(self, shotNum=None, globalDict=None, date=None, field=None):\n", + " \n", + " if not date is None:\n", + " self.set_date(date)\n", + " \n", + " if globalDict is None:\n", + " globalDict = self.read_global(shotNum)\n", + " \n", + " if field is None:\n", + " field = globalDict\n", + " \n", + " res = {}\n", + " \n", + " for key in field:\n", + " if isinstance(globalDict[key], dict) and ('mongoID' in globalDict[key]):\n", + " mongoID = globalDict[key]['mongoID']\n", + " engine = globalDict[key]['engine']\n", + " \n", + " res.update(\n", + " {\n", + " key: self._load_data_single(mongoID=mongoID, engine=engine)\n", + " }\n", + " )\n", + " \n", + " return res\n", + " \n", + " def load_data(self, shotNum=None, globalDict=None, date=None, field=None):\n", " pass\n", " \n", "DB = MongoDB(mongoClient, mongoDB, date='2023/05/23')" @@ -2107,14 +2177,35 @@ }, { "cell_type": "code", - "execution_count": 105, + "execution_count": 169, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'test': {'_id': ObjectId('6486e8ed129ec5a1a24a737a'), 'test': 10}}" + ] + }, + "execution_count": 169, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "res = DB.load_data_single(shotNum='0069', field=['test'])\n", + "res" + ] + }, + { + "cell_type": "code", + "execution_count": 170, "metadata": {}, "outputs": [], "source": [ "data = {\n", - " 'test':10\n", + " 'test':100\n", "}\n", - "DB._add_data_normal(shotNum, data)" + "DB._add_data_additional(shotNum, data, name='test')" ] }, {