Update stats processing

sitescripts/stats/bin/datamerger.py

 # coding: utf-8
 
 # This file is part of the Adblock Plus web scripts,
 # Copyright (C) 2006-2013 Eyeo GmbH
 #
 # Adblock Plus is free software: you can redistribute it and/or modify
 # it under the terms of the GNU General Public License version 3 as
 # published by the Free Software Foundation.
 #
 # Adblock Plus is distributed in the hope that it will be useful,
 # but WITHOUT ANY WARRANTY; without even the implied warranty of
 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 # GNU General Public License for more details.
 #
 # You should have received a copy of the GNU General Public License
 # along with Adblock Plus.  If not, see <http://www.gnu.org/licenses/>.
 
-import os, re, subprocess, urllib
+import os, sys, re, codecs, subprocess, urllib, simplejson, traceback

[Sebastian Noack, 2013/08/26 16:05:22]

Since Python 2.6 you can use the built-in "json" module, and don't need to rely
on the third-party "simplejson" module anymore.

[Wladimir Palant, 2013/08/27 07:34:28]

I am aware of that of course. However, simplejson is around 40% faster than the
built-in json module.

[Wladimir Palant, 2013/08/27 11:59:47]

Reply by Sebastian:

That is strange since the original simplejson module was just merged into
upstream Python and became the built-in json module. However, it might be that
they have added some performance improvements to simplejson afterwards, that
aren't merged into upstream Python yet. You might also want to look at cjson,
which is known to be even faster, if performance is an issue.

[Wladimir Palant, 2013/08/27 12:42:01]

I did. However, looking at
http://www.tablix.org/~avian/blog/archives/2011/07/a_case_for_cjson/ we probably
don't want to use it - an unmaintained project with a known security
vulnerability isn't a good solution for performance issues.

+import sitescripts.stats.common as common
 from sitescripts.utils import get_config, setupStderr
-from ConfigParser import SafeConfigParser, NoOptionError
-from StringIO import StringIO
 
-def readStatsFile(path):
-  result = SafeConfigParser()
-  match = re.search(r'^ssh://(\w+)@([^/:]+)(?::(\d+))?', path)
+def read_stats_file(path):
+  match = re.search(r"^ssh://(\w+)@([^/:]+)(?::(\d+))?", path)
   if match:
-    command = ['ssh', '-q', '-o' 'NumberOfPasswordPrompts 0', '-T', '-k', '-l', match.group(1), match.group(2)]
+    command = ["ssh", "-q", "-o", "NumberOfPasswordPrompts 0", "-T", "-k", "-l", match.group(1), match.group(2)]
     if match.group(3):
-      command[1:1] = ['-P', match.group(3)]
+      command[1:1] = ["-P", match.group(3)]
     data = subprocess.check_output(command)
-    result.readfp(StringIO(data))
-  elif path.startswith('http://') or path.startswith('https://'):
-    result.readfp(urllib.urlopen(path))
+    return simplejson.loads(data.decode("utf-8"))
+  elif path.startswith("http://") or path.startswith("https://"):
+    return simplejson.load(urllib.urlopen(path).read().decode("utf-8"))
   elif os.path.exists(path):
-    result.read(path)
-  return result
+    with codecs.open(path, "rb", encoding="utf-8") as file:
+      return simplejson.load(file)
 
-def getStatsFiles():
+  raise IOError("Path '%s' not recognized" % path)
+
+def get_stats_files(mirrors):
   config = get_config()
 
-  for option in config.options('subscriptionStats'):
-    match = re.search(r'^mirror_(.*)', option, re.I)
-    if match:
-      yield (match.group(1), config.get('subscriptionStats', option))
+  if len(mirrors) > 0:
+    options = map(lambda m: "mirror_" + m, mirrors)
+  else:
+    options = filter(lambda o: o.startswith("mirror_"), config.options("stats"))
+  for option in options:
+    if config.has_option("stats", option):
+      value = config.get("stats", option)
+      if " " in value:
+        yield re.split(r"\s+", value, 1)

[Sebastian Noack, 2013/08/26 16:05:22]

Did you know that value.split(None, 1) would do the same, just more efficient
because no regular expression is involved.

[Wladimir Palant, 2013/08/27 07:34:28]

No, I didn't. Fixed.

+      else:
+        print >>sys.stderr, "Option '%s' has invalid value: '%s'" % (option, value)
+    else:
+      print >>sys.stderr, "Option '%s' not found in the configuration" % option
 
-def mergeStatsFile(mirrorName, config1, config2):
-  def increaseOption(section, option, increase):
-    if config1.has_option(section, option):
-      oldval = config1.getint(section, option)
-      config1.set(section, option, str(oldval + increase))
+def merge_objects(object1, object2):
+  for key, value in object2.iteritems():
+    if key in object1:
+      if isinstance(value, int):
+        object1[key] += value
+      else:
+        merge_objects(object1[key], object2[key])
     else:
-      config1.set(section, option, str(increase))
+      object1[key] = value
 
-  for section in config2.sections():
-    if not config1.has_section(section):
-      config1.add_section(section)
-    for option in config2.options(section):
-      increase = config2.getint(section, option)
-      increaseOption(section, option, increase)
+def merge_stats_file(server_type, data):
+  base_dir = os.path.join(get_config().get("stats", "dataDirectory"), common.filename_encode(server_type))
+  for month, month_data in data.iteritems():
+    for name, file_data in month_data.iteritems():
+      path = os.path.join(base_dir, common.filename_encode(month), common.filename_encode(name + ".json"))
+      if os.path.exists(path):
+        with codecs.open(path, "rb", encoding="utf-8") as file:
+          existing = simplejson.load(file)
+      else:
+        existing = {}
 
-      match = re.search(r'^(\S+) (hits|bandwidth)$', option, re.I)
-      if match:
-        increaseOption(section, '%s %s mirror %s' % (match.group(1), match.group(2), mirrorName), increase)
+      merge_objects(existing, file_data)
 
-if __name__ == '__main__':
+      dir = os.path.dirname(path)
+      if not os.path.exists(dir):
+        os.makedirs(dir)
+
+      with codecs.open(path, "wb", encoding="utf-8") as file:
+        simplejson.dump(existing, file, indent=2, sort_keys=True)
+
+def merge_mirror_stats(mirrors):
+  for server_type, path in get_stats_files(mirrors):
+    try:
+      merge_stats_file(server_type, read_stats_file(path))
+    except:
+      print >>sys.stderr, "Unable to merge stats for '%s'" % path
+      traceback.print_exc()
+
+if __name__ == "__main__":
   setupStderr()
-
-  result = readStatsFile(get_config().get('subscriptionStats', 'mainFile'))
-  for (mirrorName, statsFile) in getStatsFiles():
-    mergeStatsFile(mirrorName, result, readStatsFile(statsFile))
-  file = open(get_config().get('subscriptionStats', 'mainFile'), 'wb')
-  result.write(file)
+  merge_mirror_stats(sys.argv[1:])