Merge pull request #2178 from DanCech/build-index

build-index

bin/build-index

@@ -7,7 +7,7 @@ import django
 os.environ.setdefault("DJANGO_SETTINGS_MODULE", "graphite.settings")
 django.setup()
 
-from graphite.util import write_index
+from graphite.storage import write_index
 
 
 if __name__ == "__main__":
@@ -20,11 +20,7 @@ if __name__ == "__main__":
         " defaults."
     parser = OptionParser(description=description, prog=prog, version=version,
         epilog=epilog)
-    parser.add_option("-c", "--ceres_dir", default=settings.CERES_DIR,
-        help="default: %default")
-    parser.add_option("-w", "--whisper_dir", default=settings.WHISPER_DIR,
-        help="default: %default")
     parser.add_option("-i", "--index", default=settings.INDEX_FILE,
         help="default: %default")
     (options, args) = parser.parse_args()
-    write_index(options.whisper_dir, options.ceres_dir, options.index)
+    write_index(options.index)

webapp/graphite/storage.py

@@ -1,11 +1,14 @@
 from __future__ import absolute_import
+import os
 import time
 import random
 import sys
 import types
 
 from collections import defaultdict
 from copy import deepcopy
+from shutil import move
+from tempfile import mkstemp
 
 from django.conf import settings
 from django.core.cache import cache
@@ -529,4 +532,24 @@ def extractForwardHeaders(request):
     return headers
 
 
+def write_index(index=None):
+  if not index:
+    index = settings.INDEX_FILE
+  try:
+    fd, tmp = mkstemp()
+    try:
+      tmp_index = os.fdopen(fd, 'wt')
+      for metric in STORE.get_index():
+        tmp_index.write("{0}\n".format(metric))
+    finally:
+      tmp_index.close()
+    move(tmp, index)
+  finally:
+    try:
+      os.unlink(tmp)
+    except:
+      pass
+  return None
+
+
 STORE = Store()

webapp/tests/test_storage.py

@@ -1,6 +1,10 @@
+import os
 import random
+import shutil
 import time
+import whisper
 
+from django.conf import settings
 from django.test import override_settings
 from mock import patch, Mock
 
@@ -10,7 +14,7 @@
 from graphite.intervals import Interval, IntervalSet
 from graphite.node import LeafNode, BranchNode
 from graphite.readers.utils import BaseReader
-from graphite.storage import Store, extractForwardHeaders, get_finders, get_tagdb
+from graphite.storage import Store, extractForwardHeaders, get_finders, get_tagdb, write_index
 from graphite.tags.localdatabase import LocalDatabaseTagDB
 from graphite.worker_pool.pool import PoolTimeoutError
 from graphite.render.datalib import TimeSeries
@@ -518,6 +522,46 @@ def mockAutoCompleteValues(exprs, tag, valuePrefix=None, limit=None, requestCont
       with self.assertRaisesRegexp(Exception, 'Timed out in autocomplete values for \[\'tag1=value1\'\] tag2 test after [-.e0-9]+s'):
         store.tagdb_auto_complete_values(['tag1=value1'], 'tag2', 'test', 100, request_context)
 
+  # test write_index
+  hostcpu = os.path.join(settings.WHISPER_DIR, 'hosts/hostname/cpu.wsp')
+
+  def create_whisper_hosts(self):
+    worker1 = self.hostcpu.replace('hostname', 'worker1')
+    worker2 = self.hostcpu.replace('hostname', 'worker2')
+    bogus_file = os.path.join(settings.WHISPER_DIR, 'a/b/c/bogus_file.txt')
+
+    try:
+      os.makedirs(worker1.replace('cpu.wsp', ''))
+      os.makedirs(worker2.replace('cpu.wsp', ''))
+      os.makedirs(bogus_file.replace('bogus_file.txt', ''))
+    except OSError:
+      pass
+
+    open(bogus_file, 'a').close()
+    whisper.create(worker1, [(1, 60)])
+    whisper.create(worker2, [(1, 60)])
+
+    ts = int(time.time())
+    whisper.update(worker1, 1, ts)
+    whisper.update(worker2, 2, ts)
+
+  def wipe_whisper_hosts(self):
+    try:
+      os.remove(self.hostcpu.replace('hostname', 'worker1'))
+      os.remove(self.hostcpu.replace('hostname', 'worker2'))
+      os.remove(os.path.join(settings.WHISPER_DIR, 'a/b/c/bogus_file.txt'))
+      shutil.rmtree(self.hostcpu.replace('hostname/cpu.wsp', ''))
+      shutil.rmtree(os.path.join(settings.WHISPER_DIR, 'a'))
+    except OSError:
+      pass
+
+  def test_write_index(self):
+    self.create_whisper_hosts()
+    self.addCleanup(self.wipe_whisper_hosts)
+
+    self.assertEqual(None, write_index() )
+    self.assertEqual(None, write_index(settings.INDEX_FILE) )
+
 
 class DisabledFinder(object):
   disabled = True