Browse Source

Add variable for filtering return values in 'query_data_and_generate_dataframe'

ogert 4 years ago
parent
commit
9b55fd11a9

+ 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,
                                           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:
             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:
-                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:
             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:
             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),

+ 1 - 2
cdplib/db_handlers/SQLHandler.py

@@ -106,8 +106,7 @@ class SQLHandler:
 
 
     def __del__(self):
-        print('Engine shouldve been disposed')
-        #self.dispose_engine()
+        self.dispose_engine()
 
     @property
     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
         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'])
     
     def test_F_aggregate_data_and_generate_dataframe(self):