From 125d0bbeed7b07640bc34dd877dac2e4c814895f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?F=C3=A9lix=20David?= <felixg.david@gmail.com>
Date: Mon, 31 Oct 2022 13:28:50 +0100
Subject: [PATCH] Feature: Auto download last published workfile as first
 workfile

---
 .../hooks/pre_copy_last_published_workfile.py | 124 ++++++++++++++++++
 openpype/modules/sync_server/sync_server.py   |   9 +-
 2 files changed, 132 insertions(+), 1 deletion(-)
 create mode 100644 openpype/hooks/pre_copy_last_published_workfile.py

diff --git a/openpype/hooks/pre_copy_last_published_workfile.py b/openpype/hooks/pre_copy_last_published_workfile.py
new file mode 100644
index 0000000000..004f9d25e7
--- /dev/null
+++ b/openpype/hooks/pre_copy_last_published_workfile.py
@@ -0,0 +1,124 @@
+import gc
+import os
+import shutil
+from openpype.client.entities import (
+    get_last_version_by_subset_id,
+    get_representations,
+    get_subsets,
+)
+from openpype.lib import PreLaunchHook
+from openpype.modules.base import ModulesManager
+from openpype.pipeline.load.utils import get_representation_path
+
+
+class CopyLastPublishedWorkfile(PreLaunchHook):
+    """Copy last published workfile as first workfile.
+
+    Prelaunch hook works only if last workfile leads to not existing file.
+        - That is possible only if it's first version.
+    """
+
+    # Before `AddLastWorkfileToLaunchArgs`
+    order = -1
+    app_groups = ["blender", "photoshop", "tvpaint", "aftereffects"]
+
+    def execute(self):
+        """Check if local workfile doesn't exist, else copy it.
+
+        1- Check if setting for this feature is enabled
+        2- Check if workfile in work area doesn't exist
+        3- Check if published workfile exists and is copied locally in publish
+
+        Returns:
+            None: This is a void method.
+        """
+        # TODO setting
+        self.log.info("Trying to fetch last published workfile...")
+
+        last_workfile = self.data.get("last_workfile_path")
+        if os.path.exists(last_workfile):
+            self.log.debug(
+                "Last workfile exists. Skipping {} process.".format(
+                    self.__class__.__name__
+                )
+            )
+            return
+
+        project_name = self.data["project_name"]
+        task_name = self.data["task_name"]
+
+        project_doc = self.data.get("project_doc")
+        asset_doc = self.data.get("asset_doc")
+        anatomy = self.data.get("anatomy")
+        if project_doc and asset_doc:
+            # Get subset id
+            subset_id = next(
+                (
+                    subset["_id"]
+                    for subset in get_subsets(
+                        project_name,
+                        asset_ids=[asset_doc["_id"]],
+                        fields=["_id", "data.family"],
+                    )
+                    if subset["data"]["family"] == "workfile"
+                ),
+                None,
+            )
+            if not subset_id:
+                return
+
+            # Get workfile representation
+            workfile_representation = next(
+                (
+                    representation
+                    for representation in get_representations(
+                        project_name,
+                        version_ids=[
+                            get_last_version_by_subset_id(
+                                project_name, subset_id, fields=["_id"]
+                            )["_id"]
+                        ],
+                    )
+                    if representation["context"]["task"]["name"] == task_name
+                ),
+                None,
+            )
+
+            if workfile_representation:  # TODO add setting
+                # Get sync server from Tray, which handles the asynchronous thread instance
+                sync_server = next(
+                    (
+                        t["sync_server"]
+                        for t in [
+                            obj
+                            for obj in gc.get_objects()
+                            if isinstance(obj, ModulesManager)
+                        ]
+                        if t["sync_server"].sync_server_thread
+                    ),
+                    None,
+                )
+
+                # Add site and reset timer
+                active_site = sync_server.get_active_site(project_name)
+                sync_server.add_site(
+                    project_name,
+                    workfile_representation["_id"],
+                    active_site,
+                    force=True,
+                )
+                sync_server.reset_timer()
+
+                # Wait for the download loop to end
+                sync_server.sync_server_thread.files_processed.wait()
+
+                # Get paths
+                published_workfile_path = get_representation_path(
+                    workfile_representation, root=anatomy.roots
+                )
+                local_workfile_dir = os.path.dirname(last_workfile)
+
+                # Copy file and substitute path
+                self.data["last_workfile_path"] = shutil.copy(
+                    published_workfile_path, local_workfile_dir
+                )
diff --git a/openpype/modules/sync_server/sync_server.py b/openpype/modules/sync_server/sync_server.py
index 8b11055e65..def9e6cfd8 100644
--- a/openpype/modules/sync_server/sync_server.py
+++ b/openpype/modules/sync_server/sync_server.py
@@ -236,7 +236,11 @@ class SyncServerThread(threading.Thread):
     """
     def __init__(self, module):
         self.log = Logger.get_logger(self.__class__.__name__)
-        super(SyncServerThread, self).__init__()
+
+        # Event to trigger files have been processed
+        self.files_processed = threading.Event()
+        
+        super(SyncServerThread, self).__init__(args=(self.files_processed,))
         self.module = module
         self.loop = None
         self.is_running = False
@@ -396,6 +400,8 @@ class SyncServerThread(threading.Thread):
                                               representation,
                                               site,
                                               error)
+                    # Trigger files are processed
+                    self.files_processed.set()
 
                 duration = time.time() - start_time
                 self.log.debug("One loop took {:.2f}s".format(duration))
@@ -454,6 +460,7 @@ class SyncServerThread(threading.Thread):
 
     async def run_timer(self, delay):
         """Wait for 'delay' seconds to start next loop"""
+        self.files_processed.clear()
         await asyncio.sleep(delay)
 
     def reset_timer(self):