Fix formatting script for fuzzy stats

1 files changed, 107 insertions, 26 deletions

diff --git a/src/rspamadm/fuzzy_stat.lua b/src/rspamadm/fuzzy_stat.lua
index f4c4455a9..32435ab96 100644
--- a/src/rspamadm/fuzzy_stat.lua
+++ b/src/rspamadm/fuzzy_stat.lua
@@ -70,11 +70,50 @@ local function sort_ips(tbl, opts)
   return res
 end
 
+local function add_result(dst, src)
+  if type(src) == 'table' then
+    if type(dst) == 'number' then
+      -- Convert dst to table
+      dst = {dst}
+    elseif type(dst) == 'nil' then
+      dst = {}
+    end
+
+    for i,v in ipairs(src) do
+      if dst[i] then
+        dst[i] = dst[i] + v
+      else
+        dst[i] = v
+      end
+    end
+  else
+    if type(dst) == 'table' then
+      dst[1] = dst[1] + src
+    else
+      if dst then
+        dst = dst + src
+      else
+        dst = src
+      end
+    end
+  end
+
+  return dst
+end
+
+local function print_result(r)
+  if type(r) == 'table' then
+    return table.concat(r, ', ')
+  end
+
+  return tostring(r)
+end
+
 --.USE "getopt"
 
 return function(args, res)
-  local res_keys = {}
   local res_ips = {}
+  local res_databases = {}
   local wrk = res['workers']
   local opts = getopt(args, '')
 
@@ -82,19 +121,47 @@ return function(args, res)
     for i,pr in pairs(wrk) do
       -- processes cycle
       if pr['data'] then
-        for k,elts in pairs(pr['data']) do
-          -- keys cycle
-          if not res_keys[k] then
-            res_keys[k] = {}
+        local id = pr['id']
+
+        if id then
+          local res_db = res_databases[id]
+          if not res_db then
+            res_db = {
+              keys = {}
+            }
+            res_databases[id] = res_db
+          end
+
+          -- General stats
+          for k,v in pairs(pr['data']) do
+            if k ~= 'keys' then
+              res_db[k] = add_result(res_databases[k], v)
+            end
           end
-          add_data(res_keys[k], elts)
 
-          if elts['ips'] then
-            for ip,v in pairs(elts['ips']) do
-              if not res_ips[ip] then
-                res_ips[ip] = {}
+          if pr['data']['keys'] then
+            local res_keys = res_db['keys']
+            if not res_keys then
+              res_keys = {}
+              res_db['keys'] = res_keys
+            end
+            -- Go through keys in input
+            for k,elts in pairs(pr['data']['keys']) do
+              -- keys cycle
+              if not res_keys[k] then
+                res_keys[k] = {}
+              end
+
+              add_data(res_keys[k], elts)
+
+              if elts['ips'] then
+                for ip,v in pairs(elts['ips']) do
+                  if not res_ips[ip] then
+                    res_ips[ip] = {}
+                  end
+                  add_data(res_ips[ip], v)
+                end
               end
-              add_data(res_ips[ip], v)
             end
           end
         end
@@ -102,29 +169,43 @@ return function(args, res)
     end
   end
 
-  if not opts['no-keys'] then
-    print('Keys statistics:')
-    for k,st in pairs(res_keys) do
-      print(string.format('Key id: %s', k))
-      print_stat(st, '\t')
+  -- General stats
+  for db,st in pairs(res_databases) do
+    print(string.format('Statistics for storage %s', db))
 
-      if st['ips'] then
-        print('')
-        print('\tIPs stat:')
-        local sorted_ips = sort_ips(st['ips'], opts)
+    for k,v in pairs(st) do
+      if k ~= 'keys' then
+        print(string.format('%s: %s', k, print_result(v)))
+      end
+    end
+    print('')
+
+    local res_keys = st['keys']
+    if res_keys and not opts['no-keys'] and not opts['short'] then
+      print('Keys statistics:')
+      for k,st in pairs(res_keys) do
+        print(string.format('Key id: %s', k))
+        print_stat(st, '\t')
 
-        for i,v in ipairs(sorted_ips) do
-          print(string.format('\t%s', v['ip']))
-          print_stat(v['data'], '\t\t')
+        if st['ips'] and not opts['no-ips'] then
           print('')
+          print('\tIPs stat:')
+          local sorted_ips = sort_ips(st['ips'], opts)
+
+          for i,v in ipairs(sorted_ips) do
+            print(string.format('\t%s', v['ip']))
+            print_stat(v['data'], '\t\t')
+            print('')
+          end
         end
-      end
 
-      print('')
+        print('')
+      end
     end
+
   end
 
-  if not opts['no-ips'] then
+  if not opts['no-ips'] and not opts['short'] then
     print('')
     print('IPs statistics:')