Add method for creating indexes in collections

cdplib/db_handlers/MongodbHandler.py

@@ -14,7 +14,7 @@ import simplejson
 import sys
 import os
 
-
+import pymongo
 from pymongo import MongoClient
 import pandas as pd
 import numpy as np
@@ -97,7 +97,24 @@ class MongodbHandler:
         '''
         self._database[collection_name].drop()
 
-   
+    def create_index(self, collection_name: str, key: str, direction: str='text'):
+        '''
+        :param str collection_name: name on the collection for which the schema will be set.
+        :param str key: Which value should be used as the index.
+        :param str direction: see https://api.mongodb.com/python/current/api/pymongo/collection.html for reference.
+        '''
+
+        allowed_directions = [1, -1, '2d', 'geoHaystack', '2dsphere', 'hashed', 'text']
+
+        assert(isinstance(collection_name, str)),\
+            "Parameter 'collection_name' must be a string type"
+        assert(isinstance(key, str) and key in allowed_directions),\
+            "Parameter 'key' must be a string type and must be one of these values: 1, -1, '2d', 'geoHaystack', '2dsphere', 'hashed', 'text' "
+        assert(isinstance(validation_level, str)),\
+            "Parameter 'validdirectionation_lever' must be a string type"
+            
+        self._database[collection_name].create_index([key, direction])
+
     def set_collection_schema(self, collection_name: str, schema_path: str,
                               validation_level: str = 'moderate',validation_action: str = 'error'):
         '''
@@ -111,7 +128,7 @@ class MongodbHandler:
         assert(isinstance(schema_path, str)),\
             "Parameter 'schema_path' must be a string type"
         assert(isinstance(validation_level, str)),\
-            "Parameter 'validation_lever' must be a string type"
+            "Parameter 'validation_level' must be a string type"
         assert(isinstance(validation_action, str)),\
             "Parameter 'validation_action' must be a string type"
             
@@ -257,7 +274,14 @@ if __name__ == "__main__":
 
     log.info('Script started')
 
-    db_handler = MongodbHandler()
+    database_url = "mongodb://{0}:{1}@{2}:{3}"\
+                           .format('root',
+                                   'oebb',
+                                   'localhost',
+                                   27017)
+    database_name = 'test_database'
+
+    db_handler = MongodbHandler(database_name=database_name, database_url=database_url)
 
     # Create a colleciton for the wheelsets and give it its schema.
     for schema_path in [
@@ -270,6 +294,8 @@ if __name__ == "__main__":
             collection_name = os.path.basename(schema_path).lstrip("_schema").split(".")[0]
 
             db_handler.create_collection_and_set_schema(collection_name, schema_path)
+        else:
+            log.log_and_raise_warning(('No file exists at path: {}').format(schema_path))
 
     log.info(("Existing databases: {}, Collection in OEBB database {}")\
              .format(db_handler._client.list_database_names(), db_handler._database.list_collection_names()))