1 files changed, 48 insertions, 22 deletions

diff --git a/lib/storage.rb b/lib/storage.rb
index 6ba4026..a22450d 100644
--- a/lib/storage.rb
+++ b/lib/storage.rb
@@ -5,6 +5,10 @@ require 'pp'
 module Ag::Storage
   module_function
 
+  def index_names(list)
+    return 'ml-' + list
+  end
+
   def index_status(indexname)
     $es.cluster.health(level: :indices)['indices'][indexname]['status'] || 'FAIL'
   end
@@ -18,13 +22,13 @@ module Ag::Storage
   # throws Elasticsearch::Transport::Transport::Errors::NotFound
   # if the list does not exist
   def delete_index(list)
-      $es.indices.delete(index: 'ml-' + list)
+      $es.indices.delete(index: index_name(list))
   end
 
   # Create an index for list
   # sleep in 5ms intervals until it is ready
   def create_index(list, sleeptime: 0.005)
-    indexname = 'ml-' + list
+    indexname = index_name(list)
     $es.indices.create(
       index: indexname,
       body: {
@@ -94,7 +98,14 @@ module Ag::Storage
               },
               to: {
                 type: 'string'
-              }
+              },
+              hidden: {
+                type: 'boolean',
+              },
+              comment: {
+                type: 'string',
+                index: 'not_analyzed',
+              },
             }
           }
         }
@@ -124,7 +135,7 @@ module Ag::Storage
     return nil if value == nil
 
     result = $es.search(
-      index: 'ml-' + list,
+      index: index_name(list),
       body: {
         query: {
           filtered: {
@@ -196,7 +207,7 @@ module Ag::Storage
     end
 
     $es.index(
-      index: 'ml-' + list,
+      index: index_name(list),
       type: 'message',
       id: identifier,
       body: {
@@ -212,14 +223,14 @@ module Ag::Storage
         attachments: attachments,
         received: received,
         raw_parent: raw_parent,
-        raw_filename: filename
+        raw_filename: filename,
       }
     )
   end
 
   def fix_threading(list, pass)
     result = $es.search(
-      index: 'ml-' + list,
+      index: index_name(list),
       size: 5000,
       body: {
         size: 5000,
@@ -250,31 +261,46 @@ module Ag::Storage
     opts[:progress] = "Calculating Threading (Pass #{pass})" if $options.progress
     Parallel.each(result['hits']['hits'], opts) do |hit|
       msg = resolve_message_id(list, hit['_source']['raw_parent'])
-
-      unless msg == nil
-        $es.update(
-          index: 'ml-' + list,
-          type: 'message',
-          id: hit['_id'],
-          body: {
-            doc: {
-              parent: msg
-            }
-          }
-        )
-      end
+      update(list, hit['_id'], { doc: { parent: msg } }) unless msg.nil?
     end
 
     result['hits']['total']
   end
 
   def delete(list, id)
-    $es.delete(index: 'ml-' + list, type: 'message', id: id)
+    $es.delete(index: index_name(list), type: 'message', id: id)
+  end
+
+  def _hide_unhide(list, id, hide_status, comment)
+    doc = { hidden: hide_status }
+    doc[:comment] = comment if comment # Only modify it if it exists
+    update(list, id, {
+        doc: doc,
+        detect_noop: true,
+    })
+  end
+
+  def unhide(list, id, comment)
+      hide_unhide(list, id, true, comment)
+  end
+
+  def unhide(list, id, comment)
+      _hide_unhide(list, id, false, comment)
+  end
+
+  def update(list, id, doc_changes)
+    raise "Invalid update for #{list}/#{id} #{doc_changes.to_s}" unless doc_changes.is_a?(Hash)
+    $es.update(
+      index: index_name(list),
+      type: 'message',
+      id: id,
+      body: doc_changes,
+    )
   end
 
   def get(list, id)
     result = $es.search(
-      index: 'ml-' + list,
+      index: index_name(list),
       size: 1,
       body: {
         query: {