From 602d32daf33cc4bf424ff462265aee26bafd4d29 Mon Sep 17 00:00:00 2001
From: Oleksandr Kuzhel <okuzhel@splunk.com>
Date: Tue, 12 Oct 2021 13:40:53 +0200
Subject: [PATCH] [ADDON-43324] mongo search improvement

---
 splunk_connect_for_snmp_mib_server/mongo.py   | 46 ++++++++++++++++---
 .../snmp_mib_server.py                        |  1 +
 tests/test_translator.py                      |  2 +
 3 files changed, 42 insertions(+), 7 deletions(-)

diff --git a/splunk_connect_for_snmp_mib_server/mongo.py b/splunk_connect_for_snmp_mib_server/mongo.py
index 4232fe14..ab0c16cd 100644
--- a/splunk_connect_for_snmp_mib_server/mongo.py
+++ b/splunk_connect_for_snmp_mib_server/mongo.py
@@ -15,6 +15,7 @@
 #   ########################################################################
 import logging
 import os
+import time
 
 import pymongo
 
@@ -22,7 +23,7 @@
 
 
 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,11 +37,14 @@ 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:
@@ -48,13 +52,28 @@ def upload_files(self, mib_files_dir):
                     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):
diff --git a/splunk_connect_for_snmp_mib_server/snmp_mib_server.py b/splunk_connect_for_snmp_mib_server/snmp_mib_server.py
index 2d690c1b..cc16cadd 100644
--- a/splunk_connect_for_snmp_mib_server/snmp_mib_server.py
+++ b/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!")
 
 
diff --git a/tests/test_translator.py b/tests/test_translator.py
index 226da864..9a84f63d 100644
--- a/tests/test_translator.py
+++ b/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):