refactor nats-api / optimize queries

api/tacticalrmm/agents/models.py

@@ -322,7 +322,7 @@ class Agent(BaseAuditModel):
                 online = [
                     agent
                     for agent in Agent.objects.only(
-                        "pk", "last_seen", "overdue_time", "offline_time"
+                        "pk", "agent_id", "last_seen", "overdue_time", "offline_time"
                     )
                     if agent.status == "online"
                 ]

api/tacticalrmm/agents/tasks.py

@@ -1,9 +1,6 @@
 import asyncio
 import datetime as dt
-import json
 import random
-import subprocess
-import tempfile
 from time import sleep
 from typing import Union
 
@@ -17,6 +14,7 @@ from core.models import CoreSettings
 from logs.models import PendingAction
 from scripts.models import Script
 from tacticalrmm.celery import app
+from tacticalrmm.utils import run_nats_api_cmd
 
 logger.configure(**settings.LOG_CONFIG)
 
@@ -257,30 +255,13 @@ def run_script_email_results_task(
         logger.error(e)
 
 
-def _get_nats_config() -> dict:
-    return {
-        "key": settings.SECRET_KEY,
-        "natsurl": f"tls://{settings.ALLOWED_HOSTS[0]}:4222",
-    }
-
-
 @app.task
 def monitor_agents_task() -> None:
     agents = Agent.objects.only(
         "pk", "agent_id", "last_seen", "overdue_time", "offline_time"
     )
-    ret = [i.agent_id for i in agents if i.status != "online"]
-    config = _get_nats_config()
-    config["agents"] = ret
-    with tempfile.NamedTemporaryFile() as fp:
-        with open(fp.name, "w") as f:
-            json.dump(config, f)
-
-        cmd = ["/usr/local/bin/nats-api", "-c", fp.name, "-m", "monitor"]
-        try:
-            subprocess.run(cmd, capture_output=True, timeout=30)
-        except Exception as e:
-            logger.error(e)
+    ids = [i.agent_id for i in agents if i.status != "online"]
+    run_nats_api_cmd("monitor", ids)
 
 
 @app.task
@@ -288,15 +269,5 @@ def get_wmi_task() -> None:
     agents = Agent.objects.only(
         "pk", "agent_id", "last_seen", "overdue_time", "offline_time"
     )
-    ret = [i.agent_id for i in agents if i.status == "online"]
-    config = _get_nats_config()
-    config["agents"] = ret
-    with tempfile.NamedTemporaryFile() as fp:
-        with open(fp.name, "w") as f:
-            json.dump(config, f)
-
-        cmd = ["/usr/local/bin/nats-api", "-c", fp.name, "-m", "wmi"]
-        try:
-            subprocess.run(cmd, capture_output=True, timeout=30)
-        except Exception as e:
-            logger.error(e)
+    ids = [i.agent_id for i in agents if i.status == "online"]
+    run_nats_api_cmd("wmi", ids)

api/tacticalrmm/agents/views.py

@@ -40,7 +40,7 @@ logger.configure(**settings.LOG_CONFIG)
 
 @api_view()
 def get_agent_versions(request):
-    agents = Agent.objects.only("pk")
+    agents = Agent.objects.prefetch_related("site").only("pk", "hostname")
     return Response(
         {
             "versions": [settings.LATEST_AGENT_VER],

api/tacticalrmm/tacticalrmm/utils.py

@@ -228,3 +228,20 @@ class KnoxAuthMiddlewareInstance:
 KnoxAuthMiddlewareStack = lambda inner: KnoxAuthMiddlewareInstance(
     AuthMiddlewareStack(inner)
 )
+
+
+def run_nats_api_cmd(mode: str, ids: list[str], timeout: int = 30) -> None:
+    config = {
+        "key": settings.SECRET_KEY,
+        "natsurl": f"tls://{settings.ALLOWED_HOSTS[0]}:4222",
+    }
+    config["agents"] = ids
+    with tempfile.NamedTemporaryFile() as fp:
+        with open(fp.name, "w") as f:
+            json.dump(config, f)
+
+        cmd = ["/usr/local/bin/nats-api", "-c", fp.name, "-m", mode]
+        try:
+            subprocess.run(cmd, capture_output=True, timeout=timeout)
+        except Exception as e:
+            logger.error(e)