Merge pull request #61 from splunk/feature/addon-43324-mongo-search-improv

Feature/[ADDON-43324] mongo search improvement

Author: okuzhel

splunk_connect_for_snmp_mib_server/mongo.py

@@ -15,14 +15,15 @@
 #   ########################################################################
 import logging
 import os
+import time
 
 import pymongo
 
 logger = logging.getLogger(__name__)
 
 
 class MongoRepository:
-    def __init__(self, mongo_config):
+    def __init__(self):
         self._client = pymongo.MongoClient(
             os.environ["MONGO_SERVICE_SERVICE_HOST"],
             int(os.environ["MONGO_SERVICE_SERVICE_PORT"]),
@@ -36,25 +37,43 @@ def __init__(self, mongo_config):
 
 
 class MibsRepository(MongoRepository):
+    is_text_index_created = False
+
     def __init__(self, mongo_config):
-        super().__init__(mongo_config)
+        super().__init__()
         self._mibs = self._client[mongo_config["database"]][mongo_config["collection"]]
 
     def upload_files(self, mib_files_dir):
+        tic = time.perf_counter()
         for filename in os.listdir(mib_files_dir):
             file_path = mib_files_dir + "/" + filename
             with open(file_path, "r") as mib_file:
                 try:
                     self._mibs.insert_one(
                         dict(content=mib_file.read(), filename=filename, _id=filename)
                     )
-                except Exception as e:
-                    logger.error(
-                        f"Error happened during insert mib files {filename} into mongo: {e}"
+                except Exception:
+                    logger.exception(
+                        "Error happened during insert mib files %s into mongo", filename
                     )
+        toc = time.perf_counter()
+        logger.info(f"Uploading files took - {toc - tic:0.4f} seconds")
+
+    def create_text_index(self):
+        tic = time.perf_counter()
+        try:
+            self._mibs.create_index([("content", pymongo.TEXT)], name="oid_index")
+        except Exception:
+            logger.exception(
+                "Failed to create the index, searches will be performed with regex"
+            )
+            return
+        toc = time.perf_counter()
+        MibsRepository.is_text_index_created = True
+        logger.info(f"Creating index took - {toc - tic:0.4f} seconds")
 
     def search_oid(self, oid):
-        data = self._mibs.find({"content": {"$regex": oid}})
+        data = self.perform_correct_search(oid)
         if data:
             mib_list = []
             for item in data:
@@ -63,6 +82,19 @@ def search_oid(self, oid):
         else:
             return None
 
+    def perform_correct_search(self, oid):
+        tic = time.perf_counter()
+        if MibsRepository.is_text_index_created:
+            data = self._mibs.find({"$text": {"$search": f'"{oid}"'}})
+        else:
+            data = self._mibs.find({"content": {"$regex": oid}})
+        toc = time.perf_counter()
+        logger.debug(
+            f"We searched with {'Index' if MibsRepository.is_text_index_created else 'Regex'} "
+            f"and the search took - {toc - tic:0.7f} seconds"
+        )
+        return data
+
     def delete_mib(self, filename):
         self._mibs.delete_many({"filename": {"$regex": filename}})
 
@@ -72,7 +104,7 @@ def clear(self):
 
 class OidsRepository(MongoRepository):
     def __init__(self, mongo_config):
-        super().__init__(mongo_config)
+        super().__init__()
         self._oids = self._client[mongo_config["database"]][mongo_config["collection"]]
 
     def contains_oid(self, oid):

splunk_connect_for_snmp_mib_server/snmp_mib_server.py

@@ -33,6 +33,7 @@ def upload_mibs(server_config):
     mibs_collection.clear()
     # Upload all mib files in specific dir into mongo
     mibs_collection.upload_files(mib_files_dir)
+    mibs_collection.create_text_index()
     logger.debug("Uploaded all mib files into mongo!")
 
 

tests/test_translator.py

@@ -27,6 +27,7 @@
 )
 from pysnmp.smi.rfc1902 import ObjectIdentity, ObjectType
 
+from splunk_connect_for_snmp_mib_server.mongo import MibsRepository
 from splunk_connect_for_snmp_mib_server.snmp_mib_server import upload_mibs
 from splunk_connect_for_snmp_mib_server.translator import Translator
 
@@ -100,6 +101,7 @@ def setUp(self):
         # server_config["snmp"]["mibs"]["dir"] = "../mibs/pysnmp"
         self.my_translator = Translator(server_config)
         upload_mibs(server_config)
+        MibsRepository.is_text_index_created = False
 
     @mongomock.patch()
     def test_format_trap(self):