5ee23dcb17
## Summary Refactor of `DatabaseClient` and `DatabaseService` to update how connection IDs are generated. These are now provided by DatabaseService when establishing a connection. Creation of `DatabaseClientConnection` class. This is used by `DatabaseClient` to hold a dictionary of active db connections. ## Test process Tests have been updated to reflect the changes and all pass ## Checklist - [X] PR is linked to a **work item** - [X] **acceptance criteria** of linked ticket are met - [X] performed **self-review** of the code - [X] written **tests** for any new functionality added with this PR - [X] updated the **documentation** if this PR changes or adds functionality - [ ] written/updated **design docs** if this PR implements new functionality - [X] updated the **change log** - [X] ran **pre-commit** checks for code style - [X] attended to any **TO-DOs** left in the code Related work items: #2462
394 lines
16 KiB
Python
394 lines
16 KiB
Python
from __future__ import annotations
|
|
|
|
from ipaddress import IPv4Address
|
|
from typing import Any, Dict, Optional
|
|
from uuid import uuid4
|
|
|
|
from prettytable import MARKDOWN, PrettyTable
|
|
from pydantic import BaseModel
|
|
|
|
from primaite.interface.request import RequestResponse
|
|
from primaite.simulator.core import RequestManager, RequestType
|
|
from primaite.simulator.network.hardware.nodes.host.host_node import HostNode
|
|
from primaite.simulator.network.transmission.network_layer import IPProtocol
|
|
from primaite.simulator.network.transmission.transport_layer import Port
|
|
from primaite.simulator.system.applications.application import Application
|
|
from primaite.simulator.system.core.software_manager import SoftwareManager
|
|
|
|
|
|
class DatabaseClientConnection(BaseModel):
|
|
"""
|
|
DatabaseClientConnection Class.
|
|
|
|
This class is used to record current DatabaseConnections within the DatabaseClient class.
|
|
"""
|
|
|
|
connection_id: str
|
|
"""Connection UUID."""
|
|
|
|
parent_node: HostNode
|
|
"""The parent Node that this connection was created on."""
|
|
|
|
is_active: bool = True
|
|
"""Flag to state whether the connection is still active or not."""
|
|
|
|
@property
|
|
def client(self) -> Optional[DatabaseClient]:
|
|
"""The DatabaseClient that holds this connection."""
|
|
return self.parent_node.software_manager.software.get("DatabaseClient")
|
|
|
|
def query(self, sql: str) -> bool:
|
|
"""
|
|
Query the databaseserver.
|
|
|
|
:return: Boolean value
|
|
"""
|
|
if self.is_active and self.client:
|
|
return self.client._query(connection_id=self.connection_id, sql=sql) # noqa
|
|
return False
|
|
|
|
def disconnect(self):
|
|
"""Disconnect the connection."""
|
|
if self.client and self.is_active:
|
|
self.client._disconnect(self.connection_id) # noqa
|
|
|
|
|
|
class DatabaseClient(Application):
|
|
"""
|
|
A DatabaseClient application.
|
|
|
|
Extends the Application class to provide functionality for connecting, querying, and disconnecting from a
|
|
Database Service. It mainly operates over TCP protocol.
|
|
|
|
:ivar server_ip_address: The IPv4 address of the Database Service server, defaults to None.
|
|
"""
|
|
|
|
server_ip_address: Optional[IPv4Address] = None
|
|
server_password: Optional[str] = None
|
|
_last_connection_successful: Optional[bool] = None
|
|
_query_success_tracker: Dict[str, bool] = {}
|
|
"""Keep track of connections that were established or verified during this step. Used for rewards."""
|
|
last_query_response: Optional[Dict] = None
|
|
"""Keep track of the latest query response. Used to determine rewards."""
|
|
_server_connection_id: Optional[str] = None
|
|
"""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]] = {}
|
|
"""Dictionary of connection requests to Database Server."""
|
|
connected: bool = False
|
|
"""Boolean Value for whether connected to DB Server."""
|
|
native_connection: Optional[DatabaseClientConnection] = None
|
|
"""Native Client Connection for using the client directly (similar to psql in a terminal)."""
|
|
|
|
def __init__(self, **kwargs):
|
|
kwargs["name"] = "DatabaseClient"
|
|
kwargs["port"] = Port.POSTGRES_SERVER
|
|
kwargs["protocol"] = IPProtocol.TCP
|
|
super().__init__(**kwargs)
|
|
|
|
def _init_request_manager(self) -> RequestManager:
|
|
"""
|
|
Initialise the request manager.
|
|
|
|
More information in user guide and docstring for SimComponent._init_request_manager.
|
|
"""
|
|
rm = super()._init_request_manager()
|
|
rm.add_request("execute", RequestType(func=lambda request, context: RequestResponse.from_bool(self.execute())))
|
|
return rm
|
|
|
|
def execute(self) -> bool:
|
|
"""Execution definition for db client: perform a select query."""
|
|
if not self._can_perform_action():
|
|
return False
|
|
|
|
self.num_executions += 1 # trying to connect counts as an execution
|
|
|
|
if not self.native_connection:
|
|
self.connect()
|
|
|
|
if self.native_connection:
|
|
return self.check_connection(connection_id=self.native_connection.connection_id)
|
|
|
|
return False
|
|
|
|
def describe_state(self) -> Dict:
|
|
"""
|
|
Describes the current state of the ACLRule.
|
|
|
|
:return: A dictionary representing the current state.
|
|
"""
|
|
state = super().describe_state()
|
|
# list of connections that were established or verified during this step.
|
|
state["last_connection_successful"] = self._last_connection_successful
|
|
return state
|
|
|
|
def show(self, markdown: bool = False):
|
|
"""
|
|
Display the client connections in tabular format.
|
|
|
|
:param markdown: Whether to display the table in Markdown format or not. Default is `False`.
|
|
"""
|
|
table = PrettyTable(["Connection ID", "Active"])
|
|
if markdown:
|
|
table.set_style(MARKDOWN)
|
|
table.align = "l"
|
|
table.title = f"{self.sys_log.hostname} {self.name} Client Connections"
|
|
if self.native_connection:
|
|
table.add_row([self.native_connection.connection_id, self.native_connection.is_active])
|
|
for connection_id, connection in self.client_connections.items():
|
|
table.add_row([connection_id, connection.is_active])
|
|
print(table.get_string(sortby="Connection ID"))
|
|
|
|
def configure(self, server_ip_address: IPv4Address, server_password: Optional[str] = None):
|
|
"""
|
|
Configure the DatabaseClient to communicate with a DatabaseService.
|
|
|
|
:param server_ip_address: The IP address of the Node the DatabaseService is on.
|
|
:param server_password: The password on the DatabaseService.
|
|
"""
|
|
self.server_ip_address = server_ip_address
|
|
self.server_password = server_password
|
|
self.sys_log.info(f"{self.name}: Configured the {self.name} with {server_ip_address=}, {server_password=}.")
|
|
|
|
def connect(self) -> bool:
|
|
"""Connect the native client connection."""
|
|
if self.native_connection:
|
|
return True
|
|
self.native_connection = self.get_new_connection()
|
|
return self.native_connection is not None
|
|
|
|
def disconnect(self):
|
|
"""Disconnect the native client connection."""
|
|
if self.native_connection:
|
|
self._disconnect(self.native_connection.connection_id)
|
|
self.native_connection = None
|
|
|
|
def check_connection(self, connection_id: str) -> bool:
|
|
"""Check whether the connection can be successfully re-established.
|
|
|
|
:param connection_id: connection ID to check
|
|
:type connection_id: str
|
|
:return: Whether the connection was successfully re-established.
|
|
:rtype: bool
|
|
"""
|
|
if not self._can_perform_action():
|
|
return False
|
|
return self._query("SELECT * FROM pg_stat_activity", connection_id=connection_id)
|
|
|
|
def _check_client_connection(self, connection_id: str) -> bool:
|
|
"""Check that client_connection_id is valid."""
|
|
return True if connection_id in self._client_connection_requests else False
|
|
|
|
def _connect(
|
|
self,
|
|
server_ip_address: IPv4Address,
|
|
connection_request_id: str,
|
|
password: Optional[str] = None,
|
|
is_reattempt: bool = False,
|
|
) -> Optional[DatabaseClientConnection]:
|
|
"""
|
|
Connects the DatabaseClient to the DatabaseServer.
|
|
|
|
:param: server_ip_address: IP address of the database server
|
|
:type: server_ip_address: IPv4Address
|
|
|
|
:param: password: Password used to connect to the database server. Optional.
|
|
:type: password: Optional[str]
|
|
|
|
:param: is_reattempt: True if the connect request has been reattempted. Default False
|
|
:type: is_reattempt: Optional[bool]
|
|
"""
|
|
if is_reattempt:
|
|
valid_connection = self._check_client_connection(connection_id=connection_request_id)
|
|
if valid_connection:
|
|
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
|
|
else:
|
|
self.sys_log.warning(
|
|
f"{self.name}: DatabaseClient connection to {server_ip_address} declined."
|
|
f"Connection Request ID was {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(
|
|
payload=payload, dest_ip_address=server_ip_address, dest_port=self.port
|
|
)
|
|
return self._connect(
|
|
server_ip_address=server_ip_address,
|
|
password=password,
|
|
is_reattempt=True,
|
|
connection_request_id=connection_request_id,
|
|
)
|
|
|
|
def _disconnect(self, connection_id: str) -> bool:
|
|
"""Disconnect from the Database Service.
|
|
|
|
If no connection_id is provided, connect from first ID in
|
|
self.client_connections.
|
|
|
|
:param: connection_id: connection ID to disconnect.
|
|
:type: connection_id: str
|
|
|
|
:return: bool
|
|
"""
|
|
if not self._can_perform_action():
|
|
return False
|
|
|
|
# if there are no connections - nothing to disconnect
|
|
if len(self.client_connections) == 0:
|
|
self.sys_log.warning(f"{self.name}: Unable to disconnect, no active connections.")
|
|
return False
|
|
if not self.client_connections.get(connection_id):
|
|
return False
|
|
software_manager: SoftwareManager = self.software_manager
|
|
software_manager.send_payload_to_session_manager(
|
|
payload={"type": "disconnect", "connection_id": connection_id},
|
|
dest_ip_address=self.server_ip_address,
|
|
dest_port=self.port,
|
|
)
|
|
connection = self.client_connections.pop(connection_id)
|
|
self.terminate_connection(connection_id=connection_id)
|
|
|
|
connection.is_active = False
|
|
|
|
self.sys_log.info(f"{self.name}: DatabaseClient disconnected {connection_id} from {self.server_ip_address}")
|
|
self.connected = False
|
|
return True
|
|
|
|
def uninstall(self) -> None:
|
|
"""
|
|
Uninstall the DatabaseClient.
|
|
|
|
Calls disconnect on all client connections to ensure that both client and server connections are killed.
|
|
"""
|
|
while self.client_connections.values():
|
|
client_connection = self.client_connections[next(iter(self.client_connections.keys()))]
|
|
client_connection.disconnect()
|
|
super().uninstall()
|
|
|
|
def get_new_connection(self) -> Optional[DatabaseClientConnection]:
|
|
"""Get a new connection to the DatabaseServer.
|
|
|
|
:return: DatabaseClientConnection object
|
|
"""
|
|
if not self._can_perform_action():
|
|
return None
|
|
connection_request_id = str(uuid4())
|
|
self._client_connection_requests[connection_request_id] = None
|
|
|
|
return self._connect(
|
|
server_ip_address=self.server_ip_address,
|
|
password=self.server_password,
|
|
connection_request_id=connection_request_id,
|
|
)
|
|
|
|
def _create_client_connection(self, connection_id: str, connection_request_id: str) -> None:
|
|
"""Create a new DatabaseClientConnection Object."""
|
|
client_connection = DatabaseClientConnection(
|
|
connection_id=connection_id, client=self, parent_node=self.software_manager.node
|
|
)
|
|
self.client_connections[connection_id] = client_connection
|
|
self._client_connection_requests[connection_request_id] = client_connection
|
|
|
|
def _query(self, sql: str, connection_id: str, query_id: Optional[str] = False, is_reattempt: bool = False) -> bool:
|
|
"""
|
|
Send a query to the connected database server.
|
|
|
|
:param: sql: SQL query to send to the database server.
|
|
:type: sql: str
|
|
|
|
:param: query_id: ID of the query, used as reference
|
|
:type: query_id: str
|
|
|
|
:param: connection_id: ID of the connection to the database server.
|
|
:type: connection_id: str
|
|
|
|
:param: is_reattempt: True if the query request has been reattempted. Default False
|
|
:type: is_reattempt: Optional[bool]
|
|
"""
|
|
if not query_id:
|
|
query_id = str(uuid4())
|
|
if is_reattempt:
|
|
success = self._query_success_tracker.get(query_id)
|
|
if success:
|
|
self.sys_log.info(f"{self.name}: Query successful {sql}")
|
|
self._last_connection_successful = True
|
|
return True
|
|
self.sys_log.error(f"{self.name}: Unable to run query {sql}")
|
|
self._last_connection_successful = False
|
|
return False
|
|
else:
|
|
software_manager: SoftwareManager = self.software_manager
|
|
software_manager.send_payload_to_session_manager(
|
|
payload={"type": "sql", "sql": sql, "uuid": query_id, "connection_id": connection_id},
|
|
dest_ip_address=self.server_ip_address,
|
|
dest_port=self.port,
|
|
)
|
|
return self._query(sql=sql, query_id=query_id, connection_id=connection_id, is_reattempt=True)
|
|
|
|
def run(self) -> None:
|
|
"""Run the DatabaseClient."""
|
|
super().run()
|
|
|
|
def query(self, sql: str) -> bool:
|
|
"""
|
|
Send a query to the Database Service.
|
|
|
|
:param: sql: The SQL query.
|
|
:type: sql: str
|
|
|
|
:return: True if the query was successful, otherwise False.
|
|
"""
|
|
if not self._can_perform_action():
|
|
return False
|
|
|
|
if not self.native_connection:
|
|
return False
|
|
|
|
# reset last query response
|
|
self.last_query_response = None
|
|
|
|
uuid = str(uuid4())
|
|
self._query_success_tracker[uuid] = False
|
|
return self.native_connection.query(sql)
|
|
|
|
def receive(self, session_id: str, payload: Any, **kwargs) -> bool:
|
|
"""
|
|
Receive a payload from the Software Manager.
|
|
|
|
:param payload: A payload to receive.
|
|
:param session_id: The session id the payload relates to.
|
|
:return: True.
|
|
"""
|
|
if not self._can_perform_action():
|
|
return False
|
|
if isinstance(payload, dict) and payload.get("type"):
|
|
if payload["type"] == "connect_response":
|
|
if payload["response"] is True:
|
|
# add connection
|
|
connection_id = payload["connection_id"]
|
|
self._create_client_connection(
|
|
connection_id=connection_id, connection_request_id=payload["connection_request_id"]
|
|
)
|
|
elif payload["type"] == "sql":
|
|
self.last_query_response = payload
|
|
query_id = payload.get("uuid")
|
|
status_code = payload.get("status_code")
|
|
self._query_success_tracker[query_id] = status_code == 200
|
|
if self._query_success_tracker[query_id]:
|
|
self.sys_log.debug(f"Received {payload=}")
|
|
elif payload["type"] == "disconnect":
|
|
connection_id = payload["connection_id"]
|
|
self.sys_log.info(f"{self.name}: Received disconnect command for {connection_id=} from the server")
|
|
self._disconnect(payload["connection_id"])
|
|
return True
|