read and write to MongoDB

DataContainer/MongoDB.py

@@ -3,6 +3,8 @@ import xarray_mongodb
 import bson
 import builtins
 
+import xarray as xr
+
 from ToolFunction.ToolFunction import get_date
 
 
@@ -124,6 +126,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:
             self.set_date(date)
@@ -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_data(self, shotNum, data=None):
+    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 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

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_data(self, shotNum, data=None):\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 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')"
    ]
   },
   {