From 047f57897941e7b69c495fd13dc8f28a0586f1b2 Mon Sep 17 00:00:00 2001
From: Kirill Roskolii <kirill.roskolii@catalyst.net.nz>
Date: Thu, 5 May 2016 10:38:42 +1200
Subject: [PATCH] Fixed local IDs loss on message type revert.

---
 includes/message.admin.inc        | 27 +++++++++++++++++++++++++++
 includes/message.message_type.inc | 23 +++++++++++++++++++++++
 message.module                    |  1 +
 3 files changed, 51 insertions(+)

diff --git a/includes/message.admin.inc b/includes/message.admin.inc
index c7ae1d0..be5ea39 100644
--- a/includes/message.admin.inc
+++ b/includes/message.admin.inc
@@ -35,6 +35,33 @@ class MessageTypeUIController extends EntityDefaultUIController {
     field_attach_submit('message_type', $entity, $form['message_type_fields'], $form_state);
     return $entity;
   }
+
+  /**
+   * @inheritdoc
+   *
+   * Overridden in order to prevent message types from loosing local IDs.
+   */
+  public function applyOperation($op, $message_type) {
+    if ($op == 'revert' && !empty($message_type->module)) {
+      $function = $message_type->module . '_default_message_type';
+      if (function_exists($function)) {
+        $defaults = $function();
+        $default = $defaults[$message_type->name];
+        unset($default->is_new);
+        $default->id = $message_type->id;
+        $default->status = ENTITY_IN_CODE;
+        $default->is_rebuild = TRUE;
+        entity_save('message_type', $default);
+        $label = entity_label($this->entityType, $message_type);
+        $vars = array('%entity' => $this->entityInfo['label'], '%label' => $label);
+        return t('Reverted %entity %label to the defaults.', $vars);
+      }
+    }
+    else {
+      return parent::applyOperation($op, $message_type);
+    }
+  }
+
 }
 
 /**
diff --git a/includes/message.message_type.inc b/includes/message.message_type.inc
index 5e31927..1fdf0a9 100644
--- a/includes/message.message_type.inc
+++ b/includes/message.message_type.inc
@@ -236,3 +236,26 @@ class MessageType extends Entity {
     return $property;
   }
 }
+
+/**
+ * Class MessageTypeFeaturesController.
+ */
+class MessageTypeFeaturesController extends EntityDefaultFeaturesController {
+
+  /**
+   * Overridden to update message type instead of deleting.
+   *
+   * Deleting leads to new local IDs which leads to data loss when message type
+   * is reference with entity reference module.
+   */
+  public function revert($module = NULL) {
+    if ($defaults = features_get_default($this->type, $module)) {
+      foreach ($defaults as $name => $message_type) {
+        unset($message_type->is_new);
+        $modified = entity_load_single($this->type, $name);
+        $message_type->id = $modified->id;
+        entity_save($this->type, $message_type);
+      }
+    }
+  }
+}
diff --git a/message.module b/message.module
index 2edfef2..3ce9bc0 100644
--- a/message.module
+++ b/message.module
@@ -177,6 +177,7 @@ function message_entity_info() {
   $items['message_type'] = array(
     'label' => t('Message type'),
     'controller class' => 'EntityAPIControllerExportable',
+    'features controller class' => 'MessageTypeFeaturesController',
     'entity class' => 'MessageType',
     'base table' => 'message_type',
     'fieldable' => TRUE,