diff --git a/utils/database/database.py b/utils/database/database.py index 4458de4..a8994ea 100644 --- a/utils/database/database.py +++ b/utils/database/database.py @@ -1,6 +1,7 @@ from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker from utils.models.models import Client, VMImage +from utils.exceptions.DatabaseException import DatabaseException import logging @@ -71,3 +72,33 @@ class Database: f"Error getting list of clients with VM installed: {ex}") result = [] return result + + def add_client(self, client: Client): + try: + with self.session.begin(): + self.session.add(client) + self.session.flush() + self.session.commit() + except Exception as ex: + self.logger.error(f"Error adding entity to database: {ex}") + raise DatabaseException("Error adding entity to database") + + def modify_client(self, client: Client) -> Client: + try: + old_object = self.get_client_by_mac_address(client.mac_address) + with self.session.begin(): + old_object = client + self.session.merge(old_object) + self.session.flush() + self.session.commit() + return old_object + except Exception as ex: + self.logger.error(f"Error modifying object in the database: {ex}") + raise DatabaseException("Error modifying entity in database") + + def delete_client(self, client: Client): + try: + with self.session.begin(): + self.session.delete(client) + except Exception as ex: + self.logger.error(f"Error deleting client from database: {ex}") diff --git a/utils/exceptions/DatabaseException.py b/utils/exceptions/DatabaseException.py index c613867..a05b64b 100644 --- a/utils/exceptions/DatabaseException.py +++ b/utils/exceptions/DatabaseException.py @@ -1,2 +1,2 @@ class DatabaseException(Exception): - \ No newline at end of file + pass diff --git a/utils/exceptions/__init__.py b/utils/exceptions/__init__.py new file mode 100644 index 0000000..4857659 --- /dev/null +++ b/utils/exceptions/__init__.py @@ -0,0 +1 @@ +from DatabaseException import DatabaseException