From 4ea0b6a291c63805b62bfbf3e8d51bf719974a47 Mon Sep 17 00:00:00 2001
From: Johan Marcusson <johan.marcusson@sunet.se>
Date: Mon, 28 Oct 2024 10:57:24 +0100
Subject: [PATCH] when looking for existing worktrees, try to do git pull on
 each of them and mark devices as unsync if there are updates

---
 src/cnaas_nms/db/git.py           |  4 ++--
 src/cnaas_nms/db/git_worktrees.py | 40 ++++++++++++++++++++++++++++---
 src/cnaas_nms/db/settings.py      |  9 +++++++
 3 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/src/cnaas_nms/db/git.py b/src/cnaas_nms/db/git.py
index ba4b8971..be070ea0 100644
--- a/src/cnaas_nms/db/git.py
+++ b/src/cnaas_nms/db/git.py
@@ -11,7 +11,7 @@
 from cnaas_nms.app_settings import app_settings
 from cnaas_nms.db.device import Device, DeviceType
 from cnaas_nms.db.exceptions import ConfigException, RepoStructureException
-from cnaas_nms.db.git_worktrees import WorktreeError, clean_templates_worktree
+from cnaas_nms.db.git_worktrees import WorktreeError, refresh_existing_templates_worktrees
 from cnaas_nms.db.job import Job, JobStatus
 from cnaas_nms.db.joblock import Joblock, JoblockError
 from cnaas_nms.db.session import redis_session, sqla_session
@@ -302,7 +302,7 @@ def _refresh_repo_task(repo_type: RepoType = RepoType.TEMPLATES, job_id: Optiona
             devtype: DeviceType
             for devtype, platform in updated_devtypes:
                 Device.set_devtype_syncstatus(session, devtype, ret, "templates", platform, job_id)
-        clean_templates_worktree()
+        refresh_existing_templates_worktrees(ret, job_id)
 
     return ret
 
diff --git a/src/cnaas_nms/db/git_worktrees.py b/src/cnaas_nms/db/git_worktrees.py
index 33f3952e..0832b284 100644
--- a/src/cnaas_nms/db/git_worktrees.py
+++ b/src/cnaas_nms/db/git_worktrees.py
@@ -1,9 +1,13 @@
 import os
 import shutil
-from typing import Optional
+from typing import List, Optional
 
 import git.exc
 from cnaas_nms.app_settings import app_settings
+from cnaas_nms.db.device import Device
+from cnaas_nms.db.session import sqla_session
+from cnaas_nms.db.settings import get_device_primary_groups, get_groups_using_branch
+from cnaas_nms.devicehandler.sync_history import add_sync_event
 from cnaas_nms.tools.log import get_logger
 from git import Repo
 
@@ -12,10 +16,40 @@ class WorktreeError(Exception):
     pass
 
 
-def clean_templates_worktree():
+def refresh_existing_templates_worktrees(by: str, job_id: int):
+    """Look for existing worktrees and refresh them"""
+    logger = get_logger()
+    updated_groups: List[str] = []
     if os.path.isdir("/tmp/worktrees"):
         for subdir in os.listdir("/tmp/worktrees"):
-            shutil.rmtree("/tmp/worktrees/" + subdir, ignore_errors=True)
+            try:
+                logger.info("Pulling worktree for branch {}".format(subdir))
+                wt_repo = Repo("/tmp/worktrees/" + subdir)
+                diff = wt_repo.remotes.origin.pull()
+                if not diff:
+                    continue
+            except Exception as e:
+                logger.exception(e)
+                shutil.rmtree("/tmp/worktrees/" + subdir, ignore_errors=True)
+            updated_groups.append(get_groups_using_branch(subdir))
+
+    # find all devices that are using these branches and mark them as unsynchronized
+    updated_hostnames: List[str] = []
+    with sqla_session() as session:
+        for hostname, primary_group in get_device_primary_groups():
+            if hostname in updated_hostnames:
+                continue
+            if primary_group in updated_groups:
+                dev: Device = session.query(Device).filter_by(hostname=hostname).one_or_none()
+                if dev:
+                    dev.synchronized = False
+                    add_sync_event(hostname, "refresh_templates", by, job_id)
+                    updated_hostnames.append(hostname)
+    logger.debug(
+        "Devices marked as unsynchronized because git worktree branches were refreshed: {}".format(
+            ", ".join(updated_hostnames)
+        )
+    )
 
     local_repo = Repo(app_settings.TEMPLATES_LOCAL)
     local_repo.git.worktree("prune")
diff --git a/src/cnaas_nms/db/settings.py b/src/cnaas_nms/db/settings.py
index c6354f85..26a249f8 100644
--- a/src/cnaas_nms/db/settings.py
+++ b/src/cnaas_nms/db/settings.py
@@ -767,6 +767,15 @@ def get_group_templates_branch(group_name: str) -> Optional[str]:
     return get_group_settings_asdict().get(group_name, {}).get("templates_branch")
 
 
+def get_groups_using_branch(branch_name: str) -> List[str]:
+    """Returns a list of group names that use the specified branch name"""
+    groups = []
+    for group_name, group_data in get_group_settings_asdict().items():
+        if group_data.get("templates_branch") == branch_name:
+            groups.append(group_name)
+    return groups
+
+
 @redis_lru_cache
 def get_group_settings_asdict() -> Dict[str, Dict[str, Any]]:
     """Returns a dict with group name as key and other parameters as values"""