Moved set partion state restore flag to init method in state

application/src/main/java/org/thingsboard/server/service/cf/DefaultCalculatedFieldProcessingService.java

@@ -106,10 +106,8 @@ public class DefaultCalculatedFieldProcessingService extends AbstractCalculatedF
         if (ctx.getCfType() != CalculatedFieldType.PROPAGATION) {
             return Optional.empty();
         }
-        PropagationArgumentEntry argumentEntry = (PropagationArgumentEntry)
+        return Optional.of((PropagationArgumentEntry)
-                resolveArgumentValue(PROPAGATION_CONFIG_ARGUMENT, fetchPropagationCalculatedFieldArgument(ctx, entityId));
+                resolveArgumentValue(PROPAGATION_CONFIG_ARGUMENT, fetchPropagationCalculatedFieldArgument(ctx, entityId)));
-        argumentEntry.setPartitionStateRestore(true);
-        return Optional.of(argumentEntry);
     }
 
     @Override

application/src/main/java/org/thingsboard/server/service/cf/ctx/state/propagation/PropagationCalculatedFieldState.java

@@ -69,6 +69,7 @@ public class PropagationCalculatedFieldState extends ScriptCalculatedFieldState
         super.init(restored);
         if (restored) {
             cfProcessingService.fetchPropagationArgumentFromDb(ctx, entityId).ifPresent(fromDb -> {
+                fromDb.setPartitionStateRestore(true);
                 var update = update(Map.of(PROPAGATION_CONFIG_ARGUMENT, fromDb), ctx);
                 if (update.isEmpty()) {
                     return;