Add possibility to return in mongodb shape instead of dataframe to query_data_and_return_dataframe

cdplib/db_handlers/MongodbHandler.py

@@ -246,7 +246,8 @@ class MongodbHandler:
         self.set_collection_schema(collection_name=collection_name, schema_path=schema_path)
 
     def query_data_and_generate_dataframe(self, collection_name: str, attribute: str = None,
-                                          attribute_value: str = None, comparison_operator: str = '$eq', index = None):
+                                          attribute_value: str = None, comparison_operator: str = '$eq',
+                                          index = None, return_as_dataframe: bool = True):
         '''
 
         '''
@@ -258,8 +259,10 @@ class MongodbHandler:
 
         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))
-
-        return self.convert_mongo_data_into_dataframe(data, index, collection_name)
+        if return_as_dataframe:
+            return self.convert_mongo_data_into_dataframe(data, index, collection_name)
+        else:
+            return data
 
     def aggregate_data_and_generate_dataframe(self, collection_name: str, aggregation_pipeline: list, index: str = None):