Support multiple index sets in Graylog

graylog_archiver/cli.py

@@ -1,7 +1,7 @@
 import argparse
 from elasticsearch import Elasticsearch
-from graylog_archiver.graylog_archiver import GraylogArchiver
-from graylog_archiver.config import Config
+from .graylog_archiver import GraylogArchiver
+from .config import Config
 
 
 def parse():

graylog_archiver/graylog_archiver.py

@@ -1,7 +1,7 @@
 import shutil
-from graylog_archiver import logs
-from graylog_archiver.graylog_elasticsearch import GraylogElasticsearch
-from graylog_archiver.utils import compress_directory
+from . import logs
+from .graylog_elasticsearch import GraylogElasticsearch
+from .utils import compress_directory
 
 
 class GraylogArchiver:

graylog_archiver/graylog_elasticsearch.py

@@ -2,8 +2,11 @@
 import time
 import re
 import shutil
+from itertools import groupby, chain
 
-from graylog_archiver import utils
+from . import utils
+
+INDEX_NAME_PATTERN = re.compile(r"(\S+)_\d+")
 
 
 def extract_number(index):
@@ -15,6 +18,13 @@ def sort_indices(indices):
     return sorted(indices, key=extract_number, reverse=True)
 
 
+def group_and_sort_indices(indices):
+    index_sets = [(re.match(INDEX_NAME_PATTERN, x).group(1), x) for x in indices if re.match(INDEX_NAME_PATTERN, x)]
+    return dict(
+        (x[0], sort_indices([y[1] for y in x[1]])) for x in groupby(sorted(index_sets), lambda x: x[0])
+    )
+
+
 class GraylogElasticsearch:
     def __init__(self, es, max_indices, backup_dir):
         self.es = es
@@ -25,8 +35,9 @@ def indices(self):
         return list(self.es.indices.get_mapping().keys())
 
     def indices_to_archive(self):
-        indices_sorted = sort_indices(self.indices())
-        return indices_sorted[self.max_indices:]  # keeps the max indices
+        indices_sorted = group_and_sort_indices(self.indices())
+        to_be_archieved = [idx[self.max_indices:] for idx in indices_sorted.values()]
+        return list(chain(*to_be_archieved)) # keeps the max indices from each index set
 
     def create_backup_repository(self, repository, location):
         return self.es.snapshot.create_repository(

graylog_archiver/utils.py

@@ -1,6 +1,5 @@
 import uuid
 from shutil import make_archive
-import os
 
 
 def random_string():

setup.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python
+# coding=utf-8
 
 from setuptools import setup, find_packages
 

test/tests.py

@@ -1,5 +1,5 @@
 import unittest
-from graylog_archiver.graylog_elasticsearch import sort_indices
+from graylog_archiver.graylog_elasticsearch import sort_indices, group_and_sort_indices
 
 
 class TestStringMethods(unittest.TestCase):
@@ -9,6 +9,17 @@ def test_sort_from_oldest(self):
         indices_sorted = sort_indices(indices)
         self.assertEqual(['graylog_36', 'graylog_9'], indices_sorted)
 
+    def test_group_and_sort_indexsets(self):
+        indices = [
+            'graylog_9',
+            'graylog_internal_3',
+            'graylog_internal_5',
+            'graylog_36',
+            'graylog_internal_6'
+        ]
+        indices_grouped = group_and_sort_indices(indices)
+        self.assertEqual(['graylog_36', 'graylog_9'], indices_grouped['graylog'])
+        self.assertEqual(['graylog_internal_6', 'graylog_internal_5', 'graylog_internal_3'], indices_grouped['graylog_internal'])
 
 if __name__ == '__main__':
     unittest.main()