Include 'received_command_ids' and 'executed_command_ids' into heartbeat request body and backend logic

Author: artursarlo

scripts/setup/postgres/gprofiler_recreate.sql

@@ -256,6 +256,8 @@ CREATE TABLE HostHeartbeats (
     ip_address inet NOT NULL,
     service_name text NOT NULL,
     last_command_id uuid NULL,
+    received_command_ids uuid[] NULL,
+    executed_command_ids uuid[] NULL,
     status HostStatus NOT NULL DEFAULT 'active',
     heartbeat_timestamp timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP,
     created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP,

src/gprofiler-dev/gprofiler_dev/postgres/db_manager.py

@@ -855,6 +855,8 @@ def upsert_host_heartbeat(
         ip_address: str,
         service_name: str,
         last_command_id: Optional[str] = None,
+        received_command_ids: Optional[List[str]] = None,
+        executed_command_ids: Optional[List[str]] = None,
         status: str = "active",
         heartbeat_timestamp: Optional[datetime] = None,
     ) -> bool:
@@ -865,16 +867,21 @@ def upsert_host_heartbeat(
         query = """
         INSERT INTO HostHeartbeats (
             hostname, ip_address, service_name, last_command_id,
+            received_command_ids, executed_command_ids,
             status, heartbeat_timestamp, created_at, updated_at
         ) VALUES (
             %(hostname)s, %(ip_address)s::inet, %(service_name)s,
-            %(last_command_id)s::uuid, %(status)s::HostStatus,
+            %(last_command_id)s::uuid, %(received_command_ids)s::uuid[],
+            %(executed_command_ids)s::uuid[],
+            %(status)s::HostStatus,
             %(heartbeat_timestamp)s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP
         )
         ON CONFLICT (hostname, service_name)
         DO UPDATE SET
             ip_address = EXCLUDED.ip_address,
             last_command_id = EXCLUDED.last_command_id,
+            received_command_ids = EXCLUDED.received_command_ids,
+            executed_command_ids = EXCLUDED.executed_command_ids,
             status = EXCLUDED.status,
             heartbeat_timestamp = EXCLUDED.heartbeat_timestamp,
             updated_at = CURRENT_TIMESTAMP
@@ -885,6 +892,8 @@ def upsert_host_heartbeat(
             "ip_address": ip_address,
             "service_name": service_name,
             "last_command_id": last_command_id,
+            "received_command_ids": received_command_ids,
+            "executed_command_ids": executed_command_ids,
             "status": status,
             "heartbeat_timestamp": heartbeat_timestamp,
         }
@@ -897,6 +906,7 @@ def get_host_heartbeat(self, hostname: str) -> Optional[Dict]:
         query = """
         SELECT
             hostname, ip_address, service_name, last_command_id,
+            received_command_ids, executed_command_ids,
             status, heartbeat_timestamp, created_at, updated_at
         FROM HostHeartbeats
         WHERE hostname = %(hostname)s
@@ -911,6 +921,7 @@ def get_active_hosts(self, service_name: Optional[str] = None) -> List[Dict]:
         query = """
         SELECT
             hostname, ip_address, service_name, last_command_id,
+            received_command_ids, executed_command_ids,
             status, heartbeat_timestamp
         FROM HostHeartbeats
         WHERE status = 'active'
@@ -995,6 +1006,7 @@ def get_all_host_heartbeats(self, limit: Optional[int] = None, offset: Optional[
         query = """
         SELECT
             ID, hostname, ip_address, service_name, last_command_id,
+            received_command_ids, executed_command_ids,
             status, heartbeat_timestamp, created_at, updated_at
         FROM HostHeartbeats
         ORDER BY heartbeat_timestamp DESC
@@ -1024,6 +1036,7 @@ def get_host_heartbeats_by_service(self, service_name: str, limit: Optional[int]
         query = f"""
         SELECT
             ID, hostname, ip_address, service_name, last_command_id,
+            received_command_ids, executed_command_ids,
             status, heartbeat_timestamp, created_at, updated_at
         FROM HostHeartbeats
         {where_clause}
@@ -1042,6 +1055,7 @@ def get_host_heartbeats_by_status(self, status: str, limit: Optional[int] = None
         query = """
         SELECT
             ID, hostname, ip_address, service_name, last_command_id,
+            received_command_ids, executed_command_ids,
             status, heartbeat_timestamp, created_at, updated_at
         FROM HostHeartbeats
         WHERE status = %(status)s

src/gprofiler/backend/models/metrics_models.py

@@ -225,6 +225,8 @@ class HeartbeatRequest(BaseModel):
     hostname: str
     service_name: str
     last_command_id: Optional[str] = None
+    received_command_ids: Optional[List[str]] = None
+    executed_command_ids: Optional[List[str]] = None
     status: str = "active"  # active, idle, error
     timestamp: Optional[datetime] = None
 

src/gprofiler/backend/routers/metrics_routes.py

@@ -702,6 +702,8 @@ def receive_heartbeat(heartbeat: HeartbeatRequest):
                 ip_address=heartbeat.ip_address,
                 service_name=heartbeat.service_name,
                 last_command_id=heartbeat.last_command_id,
+                received_command_ids=heartbeat.received_command_ids,
+                executed_command_ids=heartbeat.executed_command_ids,
                 status=heartbeat.status,
                 heartbeat_timestamp=heartbeat.timestamp,
             )