diff --git a/src/primaite/simulator/system/applications/database_client.py b/src/primaite/simulator/system/applications/database_client.py index 06d22126..e6cfa343 100644 --- a/src/primaite/simulator/system/applications/database_client.py +++ b/src/primaite/simulator/system/applications/database_client.py @@ -2,7 +2,7 @@ from __future__ import annotations from ipaddress import IPv4Address -from typing import Any, Dict, Optional +from typing import Any, Dict, Optional, Union from uuid import uuid4 from prettytable import MARKDOWN, PrettyTable @@ -54,6 +54,12 @@ class DatabaseClientConnection(BaseModel): if self.client and self.is_active: self.client._disconnect(self.connection_id) # noqa + def __str__(self) -> str: + return f"{self.__class__.__name__}(connection_id='{self.connection_id}', is_active={self.is_active})" + + def __repr__(self) -> str: + return str(self) + class DatabaseClient(Application, identifier="DatabaseClient"): """ @@ -76,7 +82,7 @@ class DatabaseClient(Application, identifier="DatabaseClient"): """Connection ID to the Database Server.""" client_connections: Dict[str, DatabaseClientConnection] = {} """Keep track of active connections to Database Server.""" - _client_connection_requests: Dict[str, Optional[str]] = {} + _client_connection_requests: Dict[str, Optional[Union[str, DatabaseClientConnection]]] = {} """Dictionary of connection requests to Database Server.""" connected: bool = False """Boolean Value for whether connected to DB Server.""" @@ -187,7 +193,7 @@ class DatabaseClient(Application, identifier="DatabaseClient"): return False return self._query("SELECT * FROM pg_stat_activity", connection_id=connection_id) - def _check_client_connection(self, connection_id: str) -> bool: + def _validate_client_connection_request(self, connection_id: str) -> bool: """Check that client_connection_id is valid.""" return True if connection_id in self._client_connection_requests else False @@ -211,23 +217,30 @@ class DatabaseClient(Application, identifier="DatabaseClient"): :type: is_reattempt: Optional[bool] """ if is_reattempt: - valid_connection = self._check_client_connection(connection_id=connection_request_id) - if valid_connection: + valid_connection_request = self._validate_client_connection_request(connection_id=connection_request_id) + if valid_connection_request: database_client_connection = self._client_connection_requests.pop(connection_request_id) - self.sys_log.info( - f"{self.name}: DatabaseClient connection to {server_ip_address} authorised." - f"Connection Request ID was {connection_request_id}." - ) - self.connected = True - self._last_connection_successful = True - return database_client_connection + if isinstance(database_client_connection, DatabaseClientConnection): + self.sys_log.info( + f"{self.name}: Connection request ({connection_request_id}) to {server_ip_address} authorised. " + f"Using connection id {database_client_connection}" + ) + self.connected = True + self._last_connection_successful = True + return database_client_connection + else: + self.sys_log.info( + f"{self.name}: Connection request ({connection_request_id}) to {server_ip_address} declined" + ) + self._last_connection_successful = False + return None else: - self.sys_log.warning( - f"{self.name}: DatabaseClient connection to {server_ip_address} declined." - f"Connection Request ID was {connection_request_id}." + self.sys_log.info( + f"{self.name}: Connection request ({connection_request_id}) to {server_ip_address} declined " + f"due to unknown client-side connection request id" ) - self._last_connection_successful = False return None + payload = {"type": "connect_request", "password": password, "connection_request_id": connection_request_id} software_manager: SoftwareManager = self.software_manager software_manager.send_payload_to_session_manager( @@ -300,9 +313,14 @@ class DatabaseClient(Application, identifier="DatabaseClient"): """ if not self._can_perform_action(): return None + connection_request_id = str(uuid4()) self._client_connection_requests[connection_request_id] = None + self.sys_log.info( + f"{self.name}: Sending new connection request ({connection_request_id}) to {self.server_ip_address}" + ) + return self._connect( server_ip_address=self.server_ip_address, password=self.server_password, diff --git a/src/primaite/simulator/system/services/database/database_service.py b/src/primaite/simulator/system/services/database/database_service.py index 22ae0ff3..74ef51ee 100644 --- a/src/primaite/simulator/system/services/database/database_service.py +++ b/src/primaite/simulator/system/services/database/database_service.py @@ -191,12 +191,16 @@ class DatabaseService(Service): :return: Response to connection request containing success info. :rtype: Dict[str, Union[int, Dict[str, bool]]] """ + self.sys_log.info(f"{self.name}: Processing new connection request ({connection_request_id}) from {src_ip}") status_code = 500 # Default internal server error connection_id = None if self.operating_state == ServiceOperatingState.RUNNING: status_code = 503 # service unavailable if self.health_state_actual == SoftwareHealthState.OVERWHELMED: - self.sys_log.error(f"{self.name}: Connect request for {src_ip=} declined. Service is at capacity.") + self.sys_log.info( + f"{self.name}: Connection request ({connection_request_id}) from {src_ip} declined, service is at " + f"capacity." + ) if self.health_state_actual in [ SoftwareHealthState.GOOD, SoftwareHealthState.FIXING, @@ -208,12 +212,16 @@ class DatabaseService(Service): # try to create connection if not self.add_connection(connection_id=connection_id, session_id=session_id): status_code = 500 - self.sys_log.warning(f"{self.name}: Connect request for {connection_id=} declined") - else: - self.sys_log.info(f"{self.name}: Connect request for {connection_id=} authorised") + self.sys_log.info( + f"{self.name}: Connection request ({connection_request_id}) from {src_ip} declined, " + f"returning status code 500" + ) else: status_code = 401 # Unauthorised - self.sys_log.warning(f"{self.name}: Connect request for {connection_id=} declined") + self.sys_log.info( + f"{self.name}: Connection request ({connection_request_id}) from {src_ip} unauthorised " + f"(incorrect password), returning status code 401" + ) else: status_code = 404 # service not found return { diff --git a/src/primaite/simulator/system/software.py b/src/primaite/simulator/system/software.py index 7c27534a..efa8c9b1 100644 --- a/src/primaite/simulator/system/software.py +++ b/src/primaite/simulator/system/software.py @@ -313,7 +313,7 @@ class IOSoftware(Software): # if over or at capacity, set to overwhelmed if len(self._connections) >= self.max_sessions: self.set_health_state(SoftwareHealthState.OVERWHELMED) - self.sys_log.warning(f"{self.name}: Connect request for {connection_id=} declined. Service is at capacity.") + self.sys_log.warning(f"{self.name}: Connection request ({connection_id}) declined. Service is at capacity.") return False else: # if service was previously overwhelmed, set to good because there is enough space for connections @@ -330,11 +330,11 @@ class IOSoftware(Software): "ip_address": session_details.with_ip_address if session_details else None, "time": datetime.now(), } - self.sys_log.info(f"{self.name}: Connect request for {connection_id=} authorised") + self.sys_log.info(f"{self.name}: Connection request ({connection_id}) authorised") return True # connection with given id already exists self.sys_log.warning( - f"{self.name}: Connect request for {connection_id=} declined. Connection already exists." + f"{self.name}: Connection request ({connection_id}) declined. Connection already exists." ) return False