Merge pull request #4356 from ynput/feature/not_use_mongo_query_unreal_loader

Unreal: Don't use mongo queries in 'ExistingLayoutLoader'

openpype/hosts/unreal/plugins/load/load_layout_existing.py

@@ -4,9 +4,7 @@ from pathlib import Path
 import unreal
 from unreal import EditorLevelLibrary
 
-from bson.objectid import ObjectId
-
-from openpype import pipeline
+from openpype.client import get_representations
 from openpype.pipeline import (
     discover_loader_plugins,
     loaders_from_representation,
@@ -15,7 +13,6 @@ from openpype.pipeline import (
     AVALON_CONTAINER_ID,
     legacy_io,
 )
-from openpype.api import get_current_project_settings
 from openpype.hosts.unreal.api import plugin
 from openpype.hosts.unreal.api import pipeline as upipeline
 
@@ -33,6 +30,17 @@ class ExistingLayoutLoader(plugin.Loader):
     color = "orange"
     ASSET_ROOT = "/Game/OpenPype"
 
+    delete_unmatched_assets = True
+
+    @classmethod
+    def apply_settings(cls, project_settings, *args, **kwargs):
+        super(ExistingLayoutLoader, cls).apply_settings(
+            project_settings, *args, **kwargs
+        )
+        cls.delete_unmatched_assets = (
+            project_settings["unreal"]["delete_unmatched_assets"]
+        )
+
     @staticmethod
     def _create_container(
         asset_name, asset_dir, asset, representation, parent, family
@@ -179,14 +187,7 @@ class ExistingLayoutLoader(plugin.Loader):
 
         return None
 
-    def _load_asset(self, representation, version, instance_name, family):
-        valid_formats = ['fbx', 'abc']
-
-        repr_data = legacy_io.find_one({
-            "type": "representation",
-            "parent": ObjectId(version),
-            "name": {"$in": valid_formats}
-        })
+    def _load_asset(self, repr_data, representation, instance_name, family):
         repr_format = repr_data.get('name')
 
         all_loaders = discover_loader_plugins()
@@ -221,10 +222,21 @@ class ExistingLayoutLoader(plugin.Loader):
 
         return assets
 
-    def _process(self, lib_path):
-        data = get_current_project_settings()
-        delete_unmatched = data["unreal"]["delete_unmatched_assets"]
+    def _get_valid_repre_docs(self, project_name, version_ids):
+        valid_formats = ['fbx', 'abc']
 
+        repre_docs = list(get_representations(
+            project_name,
+            representation_names=valid_formats,
+            version_ids=version_ids
+        ))
+        repre_doc_by_version_id = {}
+        for repre_doc in repre_docs:
+            version_id = str(repre_doc["parent"])
+            repre_doc_by_version_id[version_id] = repre_doc
+        return repre_doc_by_version_id
+
+    def _process(self, lib_path, project_name):
         ar = unreal.AssetRegistryHelpers.get_asset_registry()
 
         actors = EditorLevelLibrary.get_all_level_actors()
@@ -232,30 +244,44 @@ class ExistingLayoutLoader(plugin.Loader):
         with open(lib_path, "r") as fp:
             data = json.load(fp)
 
-        layout_data = []
-
+        elements = []
+        repre_ids = set()
         # Get all the representations in the JSON from the database.
         for element in data:
-            if element.get('representation'):
-                layout_data.append((
-                    pipeline.legacy_io.find_one({
-                        "_id": ObjectId(element.get('representation'))
-                    }),
-                    element
-                ))
+            repre_id = element.get('representation')
+            if repre_id:
+                repre_ids.add(repre_id)
+                elements.append(element)
 
+        repre_docs = get_representations(
+            project_name, representation_ids=repre_ids
+        )
+        repre_docs_by_id = {
+            str(repre_doc["_id"]): repre_doc
+            for repre_doc in repre_docs
+        }
+        layout_data = []
+        version_ids = set()
+        for element in elements:
+            repre_id = element.get("representation")
+            repre_doc = repre_docs_by_id.get(repre_id)
+            if not repre_doc:
+                raise AssertionError("Representation not found")
+            if not (repre_doc.get('data') or repre_doc['data'].get('path')):
+                raise AssertionError("Representation does not have path")
+            if not repre_doc.get('context'):
+                raise AssertionError("Representation does not have context")
+
+            layout_data.append((repre_doc, element))
+            version_ids.add(repre_doc["parent"])
+
+        # Prequery valid repre documents for all elements at once
+        valid_repre_doc_by_version_id = self._get_valid_repre_docs(
+            project_name, version_ids)
         containers = []
         actors_matched = []
 
         for (repr_data, lasset) in layout_data:
-            if not repr_data:
-                raise AssertionError("Representation not found")
-            if not (repr_data.get('data') or
-                    repr_data.get('data').get('path')):
-                raise AssertionError("Representation does not have path")
-            if not repr_data.get('context'):
-                raise AssertionError("Representation does not have context")
-
             # For every actor in the scene, check if it has a representation in
             # those we got from the JSON. If so, create a container for it.
             # Otherwise, remove it from the scene.
@@ -339,8 +365,8 @@ class ExistingLayoutLoader(plugin.Loader):
                 continue
 
             assets = self._load_asset(
+                valid_repre_doc_by_version_id.get(lasset.get('version')),
                 lasset.get('representation'),
-                lasset.get('version'),
                 lasset.get('instance_name'),
                 lasset.get('family')
             )
@@ -360,7 +386,7 @@ class ExistingLayoutLoader(plugin.Loader):
                 continue
             if actor not in actors_matched:
                 self.log.warning(f"Actor {actor.get_name()} not matched.")
-                if delete_unmatched:
+                if self.delete_unmatched_assets:
                     EditorLevelLibrary.destroy_actor(actor)
 
         return containers
@@ -377,7 +403,8 @@ class ExistingLayoutLoader(plugin.Loader):
         if not curr_level:
             raise AssertionError("Current level not saved")
 
-        containers = self._process(self.fname)
+        project_name = context["project"]["name"]
+        containers = self._process(self.fname, project_name)
 
         curr_level_path = Path(
             curr_level.get_outer().get_path_name()).parent.as_posix()
@@ -407,7 +434,8 @@ class ExistingLayoutLoader(plugin.Loader):
         asset_dir = container.get('namespace')
 
         source_path = get_representation_path(representation)
-        containers = self._process(source_path)
+        project_name = legacy_io.active_project()
+        containers = self._process(source_path, project_name)
 
         data = {
             "representation": str(representation["_id"]),