瀏覽代碼

Add variable for filtering return values in 'query_data_and_generate_dataframe'

ogert 5 年之前
父節點
當前提交
9b55fd11a9
共有 3 個文件被更改,包括 10 次插入8 次删除
  1. 8 5
      cdplib/db_handlers/MongodbHandler.py
  2. 1 2
      cdplib/db_handlers/SQLHandler.py
  3. 1 1
      cdplib/unit_tests/TestMongodbHandler.py

+ 8 - 5
cdplib/db_handlers/MongodbHandler.py

@@ -259,15 +259,18 @@ class MongodbHandler:
 
 
     def query_data_and_generate_dataframe(self, collection_name: str, attribute: str = None,
     def query_data_and_generate_dataframe(self, collection_name: str, attribute: str = None,
                                           attribute_value: str = None, comparison_operator: str = '$eq',
                                           attribute_value: str = None, comparison_operator: str = '$eq',
-                                          index = None, return_as_dataframe: bool = True, return_id: bool = False):
+                                          index = None, return_as_dataframe: bool = True, return_id: bool = False, return_values: dict = None):
         '''
         '''
 
 
         '''
         '''
+        if return_values is None:
+            return_values = {'_id': return_id}
+
         try:
         try:
             if attribute == None or attribute_value == None:
             if attribute == None or attribute_value == None:
-                data = self._database[collection_name].find({},{'_id': return_id})
+                data = self._database[collection_name].find({},return_value)
             else:
             else:
-                data = self._database[collection_name].find({attribute: {comparison_operator: attribute_value}}, {'_id': return_id})
+                data = self._database[collection_name].find({attribute: {comparison_operator: attribute_value}}, return_value)
 
 
         except Exception as error:
         except Exception as error:
             self._log.log_and_raise_error(('An error occured trying to query data from {}, with query {}: {}:{}. \nError:{}').format(collection_name, attribute, comparison_operator, attribute_value, error))
             self._log.log_and_raise_error(('An error occured trying to query data from {}, with query {}: {}:{}. \nError:{}').format(collection_name, attribute, comparison_operator, attribute_value, error))
@@ -303,8 +306,8 @@ class MongodbHandler:
         except Exception as error:
         except Exception as error:
             self._log.log_and_raise_error(('An error occured trying to convert mongo data into pd.Dataframe. \nError: {} ').format(error))
             self._log.log_and_raise_error(('An error occured trying to convert mongo data into pd.Dataframe. \nError: {} ').format(error))
 
 
-    def update_data_in_collection(self, query_label: str, query_value: str, update_label:str, update_value: str, collection_name:str):
-        self._database[collection_name].update_one({query_label:query_value}, {"$set": {update_label: update_value}})
+    #def update_data_in_collection(self, query_label: str, query_value: str, update_label:str, update_value: str, collection_name:str):
+    #    self._database[collection_name].update_one({query_label:query_value}, {"$set": {update_label: update_value}})
 
 
 
 
     def push_data_into_collection(self, data: (dict, list, np.ndarray, pd.DataFrame, pd.Series),
     def push_data_into_collection(self, data: (dict, list, np.ndarray, pd.DataFrame, pd.Series),

+ 1 - 2
cdplib/db_handlers/SQLHandler.py

@@ -106,8 +106,7 @@ class SQLHandler:
 
 
 
 
     def __del__(self):
     def __del__(self):
-        print('Engine shouldve been disposed')
-        #self.dispose_engine()
+        self.dispose_engine()
 
 
     @property
     @property
     def _connection_params(self) -> dict:
     def _connection_params(self) -> dict:

+ 1 - 1
cdplib/unit_tests/TestMongodbHandler.py

@@ -81,7 +81,7 @@ class TestMongodbHandler(unittest.TestCase):
         Fetch data and confirms thats it is the same as was entered into the database
         Fetch data and confirms thats it is the same as was entered into the database
         Do the same with more specific query
         Do the same with more specific query
         '''
         '''
-        self.assertEqual(self.mongodb_handler.query_data_and_generate_dataframe(self.first_collection_name).to_dict()['test_value_double'][0], self.valid_input['test_value_double'])git 
+        self.assertEqual(self.mongodb_handler.query_data_and_generate_dataframe(self.first_collection_name).to_dict()['test_value_double'][0], self.valid_input['test_value_double'])
         self.assertEqual(self.mongodb_handler.query_data_and_generate_dataframe(self.first_collection_name, 'test_value_string', 'test_value').to_dict()['test_value_double'][0], self.valid_input['test_value_double'])
         self.assertEqual(self.mongodb_handler.query_data_and_generate_dataframe(self.first_collection_name, 'test_value_string', 'test_value').to_dict()['test_value_double'][0], self.valid_input['test_value_double'])
     
     
     def test_F_aggregate_data_and_generate_dataframe(self):
     def test_F_aggregate_data_and_generate_dataframe(self):