diff --git a/CveXplore/database/maintenance/DatabaseSchemaChecker.py b/CveXplore/database/maintenance/DatabaseSchemaChecker.py index a1758272..a40ceca9 100644 --- a/CveXplore/database/maintenance/DatabaseSchemaChecker.py +++ b/CveXplore/database/maintenance/DatabaseSchemaChecker.py @@ -3,6 +3,7 @@ import os from CveXplore.database.connection.mongo_db import MongoDBConnection +from CveXplore.errors import DatabaseSchemaError runPath = os.path.dirname(os.path.realpath(__file__)) @@ -18,6 +19,33 @@ def __init__(self): self.logger = logging.getLogger("SchemaChecker") + def validate_schema(self): + if hasattr(self.dbh.connection, "store_schema"): + try: + if ( + not self.schema_version["version"] + == list(self.dbh.connection.store_schema.find({}))[0]["version"] + ): + if not self.schema_version["rebuild_needed"]: + raise DatabaseSchemaError( + "Database is not on the latest schema version; please update the database!" + ) + else: + raise DatabaseSchemaError( + "Database schema is not up to date; please re-populate the database!" + ) + else: + return True + except IndexError: + # something went wrong fetching the result from the database; assume re-populate is needed + raise DatabaseSchemaError( + "Database schema is not up to date; please re-populate the database!" + ) + else: + raise DatabaseSchemaError( + "Database schema is not up to date; please re-populate the database!" + ) + def create_indexes(self): # hack for db_updater.py to put this class in the posts variable and run the update method self.logger.info("Updating schema version") diff --git a/CveXplore/errors/database.py b/CveXplore/errors/database.py index aa2854f4..ebb4dbff 100644 --- a/CveXplore/errors/database.py +++ b/CveXplore/errors/database.py @@ -16,3 +16,7 @@ class DatabaseIllegalCollection(DatabaseException): class UpdateSourceNotFound(DatabaseException): pass + + +class DatabaseSchemaError(DatabaseException): + pass