From 6965379de10a28a44f3eb31f4a3057c8ab7a07d4 Mon Sep 17 00:00:00 2001
From: Panu Horsmalahti <phorsmalahti@mirantis.com>
Date: Tue, 18 Jan 2022 14:43:58 +0200
Subject: [PATCH] Improve entity loading for extension custom types.

Signed-off-by: Panu Horsmalahti <phorsmalahti@mirantis.com>
---
 .../extension-loader/extension-loader.ts        |  5 +++--
 src/renderer/api/catalog-entity-registry.ts     | 17 ++++++++++++++++-
 .../init-cluster-frame/init-cluster-frame.ts    |  4 ++--
 3 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/src/extensions/extension-loader/extension-loader.ts b/src/extensions/extension-loader/extension-loader.ts
index 675786dfe3..e86eafcd34 100644
--- a/src/extensions/extension-loader/extension-loader.ts
+++ b/src/extensions/extension-loader/extension-loader.ts
@@ -270,11 +270,12 @@ export class ExtensionLoader {
     });
   };
 
-  loadOnClusterRenderer = (entity: KubernetesCluster) => {
+  loadOnClusterRenderer(getCluster: () => KubernetesCluster) {
     logger.debug(`${logModule}: load on cluster renderer (dashboard)`);
 
     this.autoInitExtensions(async (extension: LensRendererExtension) => {
-      if ((await extension.isEnabledForCluster(entity)) === false) {
+      // getCluster must be a callback, as the entity might be available only after an extension has been loaded
+      if ((await extension.isEnabledForCluster(getCluster())) === false) {
         return [];
       }
 
diff --git a/src/renderer/api/catalog-entity-registry.ts b/src/renderer/api/catalog-entity-registry.ts
index e6becb6cc1..85477245a4 100644
--- a/src/renderer/api/catalog-entity-registry.ts
+++ b/src/renderer/api/catalog-entity-registry.ts
@@ -47,10 +47,25 @@ export class CatalogEntityRegistry {
     makeObservable(this);
   }
 
-  get activeEntity(): CatalogEntity | null {
+  protected getActiveEntityById() {
     return this._entities.get(this.activeEntityId) || null;
   }
 
+  get activeEntity(): CatalogEntity | null {
+    const entity = this.getActiveEntityById();
+
+    // If the entity was not found but there are rawEntities to be processed,
+    // try to process them and return the entity.
+    // This might happen if an extension registered a new Catalog category.
+    if (this.activeEntityId && !entity && this.rawEntities.length > 0) {
+      this.processRawEntities();
+
+      return this.getActiveEntityById();
+    }
+
+    return entity;
+  }
+
   set activeEntity(raw: CatalogEntity | string | null) {
     if (raw) {
       const id = typeof raw === "string"
diff --git a/src/renderer/frames/cluster-frame/init-cluster-frame/init-cluster-frame.ts b/src/renderer/frames/cluster-frame/init-cluster-frame/init-cluster-frame.ts
index 54ca0053bf..536dd2e6c0 100644
--- a/src/renderer/frames/cluster-frame/init-cluster-frame/init-cluster-frame.ts
+++ b/src/renderer/frames/cluster-frame/init-cluster-frame/init-cluster-frame.ts
@@ -19,7 +19,7 @@ import { KubeObjectStore } from "../../../../common/k8s-api/kube-object.store";
 
 interface Dependencies {
   hostedCluster: Cluster;
-  loadExtensions: (entity: CatalogEntity) => void;
+  loadExtensions: (getCluster: () => CatalogEntity) => void;
   catalogEntityRegistry: CatalogEntityRegistry;
   frameRoutingId: number;
   emitEvent: (event: AppEvent) => void;
@@ -52,7 +52,7 @@ export const initClusterFrame =
         // watch for .items, as .activeEntity might be populated only after extensions are loaded (if using custom Catalog Category)
         () => catalogEntityRegistry.items.length > 0,
         () =>
-          loadExtensions(catalogEntityRegistry.activeEntity as KubernetesCluster),
+          loadExtensions(() => catalogEntityRegistry.activeEntity as KubernetesCluster),
         {
           timeout: 15_000,
           onError: (error) => {