From 8b1674619ce7004069bed2fa10d8af39ffac3cb6 Mon Sep 17 00:00:00 2001
From: Jakub Jezek <jakubjezek001@gmail.com>
Date: Thu, 12 Sep 2024 16:40:49 +0200
Subject: [PATCH] Add staging directory functionality and a new plugin for
 managing staging directories in the pipeline.

- Added import statement for 'os' in creator_plugins.py
- Implemented method 'apply_staging_dir' to apply staging directory with persistence to instance's transient data in creator_plugins.py
- Updated comments and added TODOs related to staging directories in various files
- Created a new plugin 'CollectManagedStagingDir' to manage staging directories in publish/lib.py
---
 .../pipeline/create/creator_plugins.py        | 55 +++++++++++++++++++
 client/ayon_core/pipeline/publish/lib.py      |  2 +-
 .../publish/collect_managed_staging_dir.py    | 43 +++++++++++++++
 3 files changed, 99 insertions(+), 1 deletion(-)
 create mode 100644 client/ayon_core/plugins/publish/collect_managed_staging_dir.py

diff --git a/client/ayon_core/pipeline/create/creator_plugins.py b/client/ayon_core/pipeline/create/creator_plugins.py
index 61c10ee736..1360a74519 100644
--- a/client/ayon_core/pipeline/create/creator_plugins.py
+++ b/client/ayon_core/pipeline/create/creator_plugins.py
@@ -1,4 +1,5 @@
 # -*- coding: utf-8 -*-
+import os
 import copy
 import collections
 from typing import TYPE_CHECKING, Optional
@@ -14,6 +15,7 @@ from ayon_core.pipeline.plugin_discover import (
     deregister_plugin,
     deregister_plugin_path
 )
+from ayon_core.pipeline import get_staging_dir
 
 from .constants import DEFAULT_VARIANT_VALUE
 from .product_name import get_product_name
@@ -782,6 +784,59 @@ class Creator(BaseCreator):
         """
         return self.pre_create_attr_defs
 
+    def apply_staging_dir(self, instance):
+        """Apply staging dir with persistence to instance's transient data.
+
+        Method is called on instance creation and on instance update.
+
+        Args:
+            instance (CreatedInstance): Instance for which should be staging
+                dir applied.
+
+        Returns:
+            str: Path to staging dir.
+        """
+        create_ctx = self.create_context
+        product_name = instance.get("productName")
+        product_type = instance.get("productType")
+        folder_path = instance.get("folderPath")
+        if not any([product_name, folder_path]):
+            return None
+
+        version = instance.get("version")
+        if version is not None:
+            formatting_data = {"version": version}
+
+        staging_dir_data = get_staging_dir(
+            create_ctx.host_name,
+            create_ctx.get_current_project_entity(),
+            create_ctx.get_current_folder_entity(),
+            create_ctx.get_current_task_entity(),
+            product_type,
+            product_name,
+            create_ctx.get_current_project_anatomy(),
+            create_ctx.get_current_project_settings(),
+            always_return_path=False,
+            log=self.log,
+            formatting_data=formatting_data,
+        )
+
+        if not staging_dir_data:
+            return None
+
+        staging_dir_path = staging_dir_data["stagingDir"]
+
+        # TODO: not sure if this is necessary
+        # path might be already created by get_staging_dir
+        if not os.path.exists(staging_dir_path):
+            os.makedirs(staging_dir_path)
+
+        instance.transient_data.update(staging_dir_data)
+
+        self.log.info(f"Applied staging dir to instance: {staging_dir_path}")
+
+        return staging_dir_path
+
 
 class HiddenCreator(BaseCreator):
     @abstractmethod
diff --git a/client/ayon_core/pipeline/publish/lib.py b/client/ayon_core/pipeline/publish/lib.py
index 9cfcd3f71a..714794e8f8 100644
--- a/client/ayon_core/pipeline/publish/lib.py
+++ b/client/ayon_core/pipeline/publish/lib.py
@@ -628,7 +628,7 @@ def get_publish_repre_path(instance, repre, only_published=False):
     return None
 
 
-# deprecated: backward compatibility only
+# deprecated: backward compatibility only (2024-09-12)
 # TODO: remove in the future
 def get_custom_staging_dir_info(
     project_name,
diff --git a/client/ayon_core/plugins/publish/collect_managed_staging_dir.py b/client/ayon_core/plugins/publish/collect_managed_staging_dir.py
new file mode 100644
index 0000000000..ca6d5161c1
--- /dev/null
+++ b/client/ayon_core/plugins/publish/collect_managed_staging_dir.py
@@ -0,0 +1,43 @@
+"""
+Requires:
+    anatomy
+
+
+Provides:
+    instance.data     -> stagingDir (folder path)
+                      -> stagingDir_persistent (bool)
+"""
+
+import pyblish.api
+
+from ayon_core.pipeline.publish import get_instance_staging_dir
+
+
+class CollectManagedStagingDir(pyblish.api.InstancePlugin):
+    """Apply matching Staging Dir profile to a instance.
+
+    Apply Staging dir via profiles could be useful in specific use cases
+    where is desirable to have temporary renders in specific,
+    persistent folders, could be on disks optimized for speed for example.
+
+    It is studio's responsibility to clean up obsolete folders with data.
+
+    Location of the folder is configured in:
+        `ayon+anatomy://_/templates/staging`.
+
+    Which family/task type/subset is applicable is configured in:
+        `ayon+settings://core/tools/publish/custom_staging_dir_profiles`
+    """
+
+    label = "Collect Managed Staging Directory"
+    order = pyblish.api.CollectorOrder + 0.4990
+
+    def process(self, instance):
+
+        staging_dir_path = get_instance_staging_dir(instance)
+        persistance = instance.data.get("stagingDir_persistent", False)
+
+        self.log.info((
+            f"Instance staging dir was set to `{staging_dir_path}` "
+            f"and persistence is set to `{persistance}`"
+        ))