reduce queries on agent table load and bust cache on policy changes

This commit is contained in:
sadnub 2022-04-14 17:17:25 -04:00
parent 56bb206f25
commit 25e922bc4c
7 changed files with 201 additions and 78 deletions

View File

@ -337,38 +337,72 @@ class Agent(BaseAuditModel):
def get_checks_with_policies(
self, exclude_overridden: bool = False
) -> "List[Check]":
if exclude_overridden:
checks = list(self.agentchecks.filter(overridden_by_policy=False)) + self.get_checks_from_policies() # type: ignore
checks = (
list(
check
for check in self.agentchecks.all()
if not check.overridden_by_policy
)
+ self.get_checks_from_policies()
)
else:
checks = list(self.agentchecks.all()) + self.get_checks_from_policies() # type: ignore
return self.add_check_results(checks)
checks = list(self.agentchecks.all()) + self.get_checks_from_policies()
return checks
def get_tasks_with_policies(
self, exclude_synced: bool = False
) -> "List[AutomatedTask]":
from autotasks.models import TaskResult
tasks = list(self.autotasks.all()) + self.get_tasks_from_policies() # type: ignore
tasks = list(self.autotasks.all()) + self.get_tasks_from_policies()
if exclude_synced:
return [
task
for task in self.add_task_results(tasks)
for task in tasks
if not task.task_result
or isinstance(task.task_result, TaskResult)
and task.task_result.sync_status != "synced"
]
else:
return self.add_task_results(tasks)
return tasks
def add_task_results(self, tasks: "List[AutomatedTask]") -> "List[AutomatedTask]":
results = self.taskresults.all() # type: ignore
for task in tasks:
for result in results:
if result.task.id == task.pk:
task.task_result = result
break
return tasks
def add_check_results(self, checks: "List[Check]") -> "List[Check]":
results = self.checkresults.all() # type: ignore
for check in checks:
for result in results:
if result.assigned_check.id == check.pk:
check.check_result = result
break
return checks
def get_agent_policies(self) -> "Dict[str, Optional[Policy]]":
from checks.models import Check
site_policy = getattr(self.site, f"{self.monitoring_type}_policy", None)
client_policy = getattr(self.client, f"{self.monitoring_type}_policy", None)
default_policy = getattr(
get_core_settings(), f"{self.monitoring_type}_policy", None
)
# prefetch excluded objects on polices only if policy is not None
# prefetch excluded objects on polices only if policy is not Non
models.prefetch_related_objects(
[
policy
@ -378,7 +412,9 @@ class Agent(BaseAuditModel):
"excluded_agents",
"excluded_sites",
"excluded_clients",
"policychecks__script",
models.Prefetch(
"policychecks", queryset=Check.objects.select_related("script")
),
"autotasks",
)
@ -633,74 +669,52 @@ class Agent(BaseAuditModel):
self, skip_create=not self.should_create_alert(alert_template)
)
def add_task_results(self, tasks: "List[AutomatedTask]") -> "List[AutomatedTask]":
results = self.taskresults.all() # type: ignore
for task in tasks:
for result in results:
if result.task.id == task.pk:
task.task_result = result
break
if not hasattr(task, "task_result"):
task.task_result = {}
return tasks
def add_check_results(self, checks: "List[Check]") -> "List[Check]":
results = self.checkresults.all() # type: ignore
for check in checks:
for result in results:
if result.assigned_check.id == check.pk:
check.check_result = result
break
if not hasattr(check, "check_result"):
check.check_result = {}
return checks
def get_checks_from_policies(self) -> "List[Check]":
from automation.models import Policy
cache_checks = False
if not self.policy and not self.block_policy_inheritance:
cached_checks = cache.get(f"site_{self.site_id}_checks")
if cached_checks and isinstance(cached_checks, list):
return cached_checks
else:
cache_checks = True
# check if agent is blocking inheritance
if self.block_policy_inheritance or self.agentchecks.exists():
cache_key = f"agent_{self.agent_id}_checks"
# clear agent checks that have overridden_by_policy set
self.agentchecks.update(overridden_by_policy=False) # type: ignore
elif self.policy:
cache_key = f"site_{self.monitoring_type}_{self.site_id}_policy_{self.policy_id}_checks"
# get agent checks based on policies
checks = Policy.get_policy_checks(self)
else:
cache_key = f"site_{self.monitoring_type}_{self.site_id}_checks"
if cache_checks:
cache.set(f"site_{self.site_id}_checks", checks, 60)
cached_checks = cache.get(cache_key)
if cached_checks and isinstance(cached_checks, list):
return cached_checks
else:
# clear agent checks that have overridden_by_policy set
self.agentchecks.update(overridden_by_policy=False) # type: ignore
return checks
# get agent checks based on policies
checks = Policy.get_policy_checks(self)
cache.set(cache_key, checks, 600)
return checks
def get_tasks_from_policies(self) -> "List[AutomatedTask]":
from automation.models import Policy
cache_tasks = False
if not self.policy and not self.block_policy_inheritance:
cached_tasks = cache.get(f"site_{self.site_id}_tasks")
if cached_tasks and isinstance(cached_tasks, list):
return cached_tasks
else:
cache_tasks = True
# get agent tasks based on policies
tasks = Policy.get_policy_tasks(self)
# check if agent is blocking inheritance
if self.block_policy_inheritance:
cache_key = f"agent_{self.agent_id}_tasks"
if cache_tasks:
cache.set(f"site_{self.site_id}_tasks", tasks, 60)
return tasks
elif self.policy:
cache_key = f"site_{self.monitoring_type}_{self.site_id}_policy_{self.policy_id}_tasks"
else:
cache_key = f"site_{self.monitoring_type}_{self.site_id}_tasks"
cached_tasks = cache.get(cache_key)
if cached_tasks and isinstance(cached_tasks, list):
return cached_tasks
else:
# get agent tasks based on policies
tasks = Policy.get_policy_tasks(self)
cache.set(f"site_{self.site_id}_tasks", tasks, 600)
return tasks
def _do_nats_debug(self, agent: "Agent", message: str) -> None:
DebugLog.error(agent=agent, log_type="agent_issues", message=message)

View File

@ -14,7 +14,7 @@ from core.utils import (
get_core_settings,
)
from django.conf import settings
from django.db.models import Q
from django.db.models import Q, Prefetch, F
from django.http import HttpResponse
from django.shortcuts import get_object_or_404
from django.utils import timezone as djangotime
@ -70,6 +70,8 @@ class GetAgents(APIView):
permission_classes = [IsAuthenticated, AgentPerms]
def get(self, request):
from checks.models import Check, CheckResult
if "site" in request.query_params.keys():
filter = Q(site_id=request.query_params["site"])
elif "client" in request.query_params.keys():
@ -83,7 +85,6 @@ class GetAgents(APIView):
or "detail" in request.query_params.keys()
and request.query_params["detail"] == "true"
):
agents = (
Agent.objects.filter_by_role(request.user) # type: ignore
.filter(filter)
@ -97,8 +98,14 @@ class GetAgents(APIView):
"alert_template",
)
.prefetch_related(
"agentchecks__script",
"checkresults__assigned_check__script",
Prefetch(
"agentchecks",
queryset=Check.objects.select_related("script"),
),
Prefetch(
"checkresults",
queryset=CheckResult.objects.select_related("assigned_check"),
),
)
)
ctx = {"default_tz": get_default_timezone()}

View File

@ -1,6 +1,7 @@
from agents.models import Agent
from clients.models import Client, Site
from django.db import models
from django.core.cache import cache
from logs.models import BaseAuditModel
from typing import Optional, Dict, Any, List, TYPE_CHECKING
@ -45,9 +46,25 @@ class Policy(BaseAuditModel):
if old_policy:
if old_policy.alert_template != self.alert_template:
cache_agents_alert_template.delay()
elif old_policy.active != self.active and self.alert_template:
elif self.alert_template and old_policy.active != self.active:
cache_agents_alert_template.delay()
if old_policy.active != self.active or old_policy.enforced != self.enforced:
cache.delete_many_pattern(f"site_workstation_*")
cache.delete_many_pattern(f"site_server_*")
cache.delete_many_pattern("agent_*")
def delete(self, *args, **kwargs):
cache.delete_many_pattern(f"site_workstation_*")
cache.delete_many_pattern(f"site_server_*")
cache.delete_many_pattern("agent_*")
super(Policy, self).delete(
*args,
**kwargs,
)
def __str__(self) -> str:
return self.name
@ -210,6 +227,7 @@ class Policy(BaseAuditModel):
@staticmethod
def get_policy_checks(agent: "Agent") -> "List[Check]":
# Get checks added to agent directly
agent_checks = list(agent.agentchecks.all())

View File

@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, List, Dict, Any, Optional, Union
from alerts.models import SEVERITY_CHOICES
from django.core.validators import MaxValueValidator, MinValueValidator
from django.core.cache import cache
from django.utils import timezone as djangotime
from django.db import models
from django.db.models.fields import DateTimeField
@ -150,6 +151,12 @@ class AutomatedTask(BaseAuditModel):
return self.name
def save(self, *args, **kwargs) -> None:
# if task is a policy task clear cache on everything
if self.policy:
cache.delete_many_pattern("site_*_tasks")
cache.delete_many_pattern("agent_*_tasks")
# get old task if exists
old_task = AutomatedTask.objects.get(pk=self.pk) if self.pk else None
super(AutomatedTask, self).save(old_model=old_task, *args, **kwargs)
@ -167,6 +174,18 @@ class AutomatedTask(BaseAuditModel):
sync_status="notsynced"
)
def delete(self, *args, **kwargs):
# if check is a policy check clear cache on everything
if self.policy:
cache.delete_many_pattern("site_*_tasks")
cache.delete_many_pattern("agent_*_tasks")
super(AutomatedTask, self).delete(
*args,
**kwargs,
)
@property
def schedule(self) -> Optional[str]:
if self.task_type == "manual":

View File

@ -1,6 +1,7 @@
from statistics import mean
from typing import TYPE_CHECKING, Any, Dict, Union, Optional
from typing import TYPE_CHECKING, Any, Union, Dict, Optional
from django.core.cache import cache
from alerts.models import SEVERITY_CHOICES
from django.contrib.postgres.fields import ArrayField
from django.core.validators import MaxValueValidator, MinValueValidator
@ -166,7 +167,8 @@ class Check(BaseAuditModel):
# deprecated
managed_by_policy = models.BooleanField(default=False)
check_result: "Union[CheckResult, Dict]" = {}
# non-database property
check_result: "Union[CheckResult, Dict[None, None]]" = {}
def __str__(self):
if self.agent:
@ -174,6 +176,38 @@ class Check(BaseAuditModel):
else:
return f"{self.policy.name} - {self.readable_desc}"
def save(self, *args, **kwargs):
# if check is a policy check clear cache on everything
if self.policy:
cache.delete_many_pattern("site_*_checks")
cache.delete_many_pattern("agent_*_checks")
# if check is an agent check
elif self.agent:
cache.delete(f"agent_{self.agent.agent_id}_checks")
super(Check, self).save(
*args,
**kwargs,
)
def delete(self, *args, **kwargs):
# if check is a policy check clear cache on everything
if self.policy:
cache.delete_many_pattern("site_*_checks")
cache.delete_many_pattern("agent_*_checks")
# if check is an agent check
elif self.agent:
cache.delete(f"agent_{self.agent.agent_id}_checks")
super(Check, self).delete(
*args,
**kwargs,
)
@property
def readable_desc(self):
display = self.get_check_type_display() # type: ignore

View File

@ -1,6 +1,7 @@
import uuid
from agents.models import Agent
from django.core.cache import cache
from django.contrib.postgres.fields import ArrayField
from django.db import models
from logs.models import BaseAuditModel
@ -55,13 +56,25 @@ class Client(BaseAuditModel):
)
# check if polcies have changed and initiate task to reapply policies if so
if old_client:
if (
old_client.alert_template != self.alert_template
or old_client.workstation_policy != self.workstation_policy
or old_client.server_policy != self.server_policy
):
cache_agents_alert_template.delay()
if old_client and (
old_client.alert_template != self.alert_template
or old_client.workstation_policy != self.workstation_policy
or old_client.server_policy != self.server_policy
):
cache_agents_alert_template.delay()
if old_client and (
old_client.workstation_policy != self.workstation_policy
or old_client.server_policy != self.server_policy
):
sites = self.sites.all()
if old_client.workstation_policy != self.workstation_policy:
for site in sites:
cache.delete_many_pattern(f"site_workstation_{site.pk}_*")
if old_client.server_policy != self.server_policy:
for site in sites:
cache.delete_many_pattern(f"site_server_{site.pk}_*")
class Meta:
ordering = ("name",)
@ -140,6 +153,12 @@ class Site(BaseAuditModel):
):
cache_agents_alert_template.delay()
if old_site.workstation_policy != self.workstation_policy:
cache.delete_many_pattern(f"site_workstation_{self.pk}_*")
if old_site.server_policy != self.server_policy:
cache.delete_many_pattern(f"site_server_{self.pk}_*")
class Meta:
ordering = ("name",)
unique_together = (("client", "name"),)

View File

@ -121,6 +121,18 @@ class CoreSettings(BaseAuditModel):
):
cache_agents_alert_template.delay()
if old_settings.workstation_policy != self.workstation_policy:
cache.delete_many_pattern(f"site_workstation_*")
if old_settings.server_policy != self.server_policy:
cache.delete_many_pattern(f"site_server_*")
if (
old_settings.server_policy != self.server_policy
or old_settings.workstation_policy != self.workstation_policy
):
cache.delete_many_pattern("agent_*")
def __str__(self) -> str:
return "Global Site Settings"