From 92466630b143cf95f2b8ed67d9e67ea147b3ddb1 Mon Sep 17 00:00:00 2001
From: ViacheslavKlimov <viacheslavklimov11@gmail.com>
Date: Tue, 25 Mar 2025 17:01:05 +0200
Subject: [PATCH] Refactoring for CustomerData

---
 .../org/thingsboard/server/edqs/data/CustomerData.java   | 6 +++---
 .../org/thingsboard/server/edqs/repo/TenantRepo.java     | 9 +++------
 2 files changed, 6 insertions(+), 9 deletions(-)

diff --git a/common/edqs/src/main/java/org/thingsboard/server/edqs/data/CustomerData.java b/common/edqs/src/main/java/org/thingsboard/server/edqs/data/CustomerData.java
index bf2a3f6da7..9ce300df6a 100644
--- a/common/edqs/src/main/java/org/thingsboard/server/edqs/data/CustomerData.java
+++ b/common/edqs/src/main/java/org/thingsboard/server/edqs/data/CustomerData.java
@@ -50,10 +50,10 @@ public class CustomerData extends BaseEntityData<CustomerFields> {
         entitiesById.computeIfAbsent(ed.getEntityType(), et -> new ConcurrentHashMap<>()).put(ed.getId(), ed);
     }
 
-    public boolean remove(EntityData<?> ed) {
-        var map = entitiesById.get(ed.getEntityType());
+    public boolean remove(EntityType entityType, UUID entityId) {
+        var map = entitiesById.get(entityType);
         if (map != null) {
-            return map.remove(ed.getId()) != null;
+            return map.remove(entityId) != null;
         } else {
             return false;
         }
diff --git a/common/edqs/src/main/java/org/thingsboard/server/edqs/repo/TenantRepo.java b/common/edqs/src/main/java/org/thingsboard/server/edqs/repo/TenantRepo.java
index 9ef34ce129..d8dec9fc4e 100644
--- a/common/edqs/src/main/java/org/thingsboard/server/edqs/repo/TenantRepo.java
+++ b/common/edqs/src/main/java/org/thingsboard/server/edqs/repo/TenantRepo.java
@@ -173,10 +173,7 @@ public class TenantRepo {
             if (EntityRelation.CONTAINS_TYPE.equals(entityRelation.getType()) && entityRelation.getFrom().getEntityType() == EntityType.CUSTOMER) {
                 CustomerData customerData = (CustomerData) get(entityRelation.getFrom());
                 if (customerData != null) {
-                    EntityData<?> dashboardData = get(entityRelation.getTo());
-                    if (dashboardData != null) {
-                        customerData.remove(dashboardData);
-                    }
+                    customerData.remove(EntityType.DASHBOARD, entityRelation.getTo().getId());
                 }
             }
         }
@@ -205,7 +202,7 @@ public class TenantRepo {
                 if (oldCustomerId != null) {
                     CustomerData old = (CustomerData) get(EntityType.CUSTOMER, oldCustomerId);
                     if (old != null) {
-                        old.remove(entityData);
+                        old.remove(entityType, entityId);
                     }
                 }
                 if (newCustomerId != null) {
@@ -233,7 +230,7 @@ public class TenantRepo {
                 if (customerId != null) {
                     CustomerData customerData = (CustomerData) get(EntityType.CUSTOMER, customerId);
                     if (customerData != null) {
-                        customerData.remove(removed);
+                        customerData.remove(entityType, entityId);
                     }
                 }
             }