refactoring errors on imports

cve-search · Apr 23, 2024 · d62672e · d62672e
1 parent 0c35478
commit d62672e
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 26 deletions.
diff --git a/CveXplore/database/connection/base/db_connection_base.py b/CveXplore/database/connection/base/db_connection_base.py
@@ -1,6 +1,7 @@
 import logging
 from abc import ABC, abstractmethod
 
+from CveXplore.common.config import Configuration
 from CveXplore.core.logging.logger_class import AppLogger
 
 logging.setLoggerClass(AppLogger)
@@ -9,6 +10,7 @@
 class DatabaseConnectionBase(ABC):
     def __init__(self, logger_name: str):
         self.logger = logging.getLogger(logger_name)
+        self.config = Configuration
 
     def __repr__(self):
         return f"<<{self.__class__.__name__}>>"

diff --git a/CveXplore/database/connection/database_connection.py b/CveXplore/database/connection/database_connection.py
@@ -7,6 +7,7 @@
 
 class DatabaseConnection(object):
     def __init__(self, database_type: str, database_init_parameters: dict):
+
         self.database_type = database_type
         self.database_init_parameters = database_init_parameters
 

diff --git a/CveXplore/database/connection/sqlbase/connection.py b/CveXplore/database/connection/sqlbase/connection.py
@@ -5,6 +5,7 @@
 
 config = Configuration
 
-engine = create_engine(config.SQLALCHEMY_DATABASE_URI, echo=True)
+if config.DATASOURCE_TYPE != "mongodb":
+    engine = create_engine(config.SQLALCHEMY_DATABASE_URI, echo=True)
 
-Session = sessionmaker(bind=engine)
+    Session = sessionmaker(bind=engine)
diff --git a/CveXplore/database/connection/sqlbase/sql_base.py b/CveXplore/database/connection/sqlbase/sql_base.py
@@ -1,39 +1,43 @@
 from CveXplore.database.connection.base.db_connection_base import DatabaseConnectionBase
-from CveXplore.database.connection.sqlbase.sql_client import SQLClient
 from CveXplore.database_models.models import CveXploreBase
 from CveXplore.errors import DatabaseConnectionException
 
 
 class SQLBaseConnection(DatabaseConnectionBase):
     def __init__(self, **kwargs):
         super().__init__(logger_name=__name__)
+        if self.config.DATASOURCE_TYPE != "mongodb":
+            from CveXplore.database.connection.sqlbase.sql_client import SQLClient
 
-        self._dbclient = {
-            "info": SQLClient("info"),
-            "cpe": SQLClient("cpe"),
-            "cves": SQLClient("cves"),
-            "schema": SQLClient("schema"),
-            "cwe": SQLClient("cwe"),
-            "capec": SQLClient("capec"),
-            "via4": SQLClient("via4"),
-        }
-
-        try:
-            collections = list(CveXploreBase.metadata.tables.keys())
-        except ConnectionError as err:
-            raise DatabaseConnectionException(
-                f"Connection to the database failed: {err}"
-            )
-
-        if len(collections) != 0:
-            for each in collections:
-                self.__setattr__(f"store_{each}", SQLClient(each))
+            self._dbclient = {
+                "info": SQLClient("info"),
+                "cpe": SQLClient("cpe"),
+                "cves": SQLClient("cves"),
+                "schema": SQLClient("schema"),
+                "cwe": SQLClient("cwe"),
+                "capec": SQLClient("capec"),
+                "via4": SQLClient("via4"),
+            }
+
+            try:
+                collections = list(CveXploreBase.metadata.tables.keys())
+            except ConnectionError as err:
+                raise DatabaseConnectionException(
+                    f"Connection to the database failed: {err}"
+                )
+
+            if len(collections) != 0:
+                for each in collections:
+                    self.__setattr__(f"store_{each}", SQLClient(each))
 
     @property
     def dbclient(self):
         return self._dbclient
 
     def set_handlers_for_collections(self):
-        for each in list(CveXploreBase.metadata.tables.keys()):
-            if not hasattr(self, each):
-                setattr(self, f"store_{each}", SQLClient(each))
+        if self.config.DATASOURCE_TYPE != "mongodb":
+            from CveXplore.database.connection.sqlbase.sql_client import SQLClient
+
+            for each in list(CveXploreBase.metadata.tables.keys()):
+                if not hasattr(self, each):
+                    setattr(self, f"store_{each}", SQLClient(each))