Merge pull request #66 from splunk/bug/addon-43531-changed-mongo

[ADDON-43531] new mongo way

Author: okuzhel

splunk_connect_for_snmp_mib_server/mib_server.py

@@ -59,9 +59,11 @@ def translator():
             logger.debug(request.json)
             var_binds = request.json.get("var_binds")
             data_format = request.args.get("data_format")
-            logger.debug(f"type of var_binds: {type(var_binds)}")
-            logger.debug(f"var_binds: {var_binds}")
-            logger.debug(f"requested data format: {str(data_format)} ")
+            logger.debug(
+                f"type of var_binds: {type(var_binds)}\n"
+                f"var_binds: {var_binds}\n"
+                f"requested data format: {str(data_format)}"
+            )
             if data_format in self.data_format:
                 method = self.data_format.get(data_format)
                 result = method(var_binds)

splunk_connect_for_snmp_mib_server/mongo.py

@@ -15,6 +15,7 @@
 #   ########################################################################
 import logging
 import os
+import re
 import time
 
 import pymongo
@@ -37,7 +38,7 @@ def __init__(self):
 
 
 class MibsRepository(MongoRepository):
-    is_text_index_created = False
+    oid_regex = re.compile(r"\(((?:1, 3, 6, 1, )(?:\d+,?\s?)+)\)")
 
     def __init__(self, mongo_config):
         super().__init__()
@@ -49,31 +50,41 @@ def upload_files(self, mib_files_dir):
             file_path = mib_files_dir + "/" + filename
             with open(file_path, "r") as mib_file:
                 try:
+                    read_file = mib_file.read()
                     self._mibs.insert_one(
-                        dict(content=mib_file.read(), filename=filename, _id=filename)
+                        dict(
+                            content=read_file,
+                            filename=filename,
+                            _id=filename,
+                            mibs=self.parse_iods(read_file),
+                        )
                     )
+                    self.parse_iods(read_file)
                 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):
+    @staticmethod
+    def parse_iods(read_file):
+        return re.findall(MibsRepository.oid_regex, read_file)
+
+    def create_index_on_oids(self):
         tic = time.perf_counter()
         try:
-            self._mibs.create_index([("content", pymongo.TEXT)], name="oid_index")
+            self._mibs.create_index("mibs", name="oid_index")
         except Exception:
             logger.exception(
-                "Failed to create the index, searches will be performed with regex"
+                "Failed to create the index, searches will be performed without it"
             )
             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.perform_correct_search(oid)
+        data = self.perform_search(oid)
         if data:
             mib_list = []
             for item in data:
@@ -82,31 +93,23 @@ 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"
-        )
+    def perform_search(self, oid):
         tic = time.perf_counter()
-        queried_cursor = list(data)
+        data = list(self._mibs.find({"mibs": oid}))
         toc = time.perf_counter()
         logger.debug(
-            f"data that we found {len(queried_cursor)}"
-            f"and the query of the cursor take - {toc - tic:0.7f} seconds"
+            f"We searched for oid - {oid} and it took - {toc - tic:0.7f} seconds"
         )
-        return queried_cursor
+        return data
 
     def delete_mib(self, filename):
         self._mibs.delete_many({"filename": {"$regex": filename}})
 
     def clear(self):
-        self._mibs.delete_many({})
+        tic = time.perf_counter()
+        self._mibs.drop()
+        toc = time.perf_counter()
+        logger.info(f"Cleaning took - {toc - tic:0.4f} seconds")
 
 
 class OidsRepository(MongoRepository):

splunk_connect_for_snmp_mib_server/snmp_mib_server.py

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

splunk_connect_for_snmp_mib_server/translator.py

@@ -129,27 +129,28 @@ def find_mib_file(self, oid):
         mib_list_without_index = None
 
         try:
-            mib_list = self._mongo_mibs_coll.search_oid(value_tuple)
             # Find mib module for OID without index (remove the last part of OID)
             # Handle the scenario that tries to translate an OID which has index append at the end.
             # e.g 1.3.6.1.2.1.25.1.6.0, where 0 is index and it's not part of the OID object
             # So we cannot find mapping MIBs for it
             # Instead, 1.3.6.1.2.1.25.1.6 is actually the OID that needed to be used for searching MIBs
             # Therefore, we should remove index (0) and search the real oid (1.3.6.1.2.1.25.1.6) to detect the MIBs
-            logger.debug(f"[-] oid_without_index: {oid_without_index}")
+
+            mib_list = self._mongo_mibs_coll.search_oid(value_tuple)
             mib_list_without_index = self._mongo_mibs_coll.search_oid(
                 changed_oid_without_index
             )
         except Exception:
             logger.exception(
-                f"Error happened during search for the oid - {oid} in mongo mibs collection"
+                f"Error happened during search for the oids - {oid},"
+                f" {changed_oid_without_index} in mongo mibs collection"
             )
         self.add_oid_to_db(mib_list, oid)
-        self.add_oid_to_db(mib_list_without_index, oid)
+        self.add_oid_to_db(mib_list_without_index, changed_oid_without_index)
         if mib_list:
             self.load_extra_mibs(mib_list, oid)
         if mib_list_without_index:
-            self.load_extra_mibs(mib_list_without_index, oid)
+            self.load_extra_mibs(mib_list_without_index, changed_oid_without_index)
 
     def add_oid_to_db(self, mib_list, oid):
         if not mib_list:

tests/test_translator.py

@@ -27,7 +27,6 @@
 )
 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
 
@@ -101,7 +100,6 @@ 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):