Centralize stats processing, have the stats server pull in logs

sitescripts/stats/bin/logprocessor.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 argparse
 import codecs
 from collections import OrderedDict
 from datetime import datetime, timedelta
 import errno
 import gzip
 import json
 import math
+import multiprocessing
+import numbers
 import os
 import re
 import pygeoip
 import socket
-from StringIO import StringIO
 import subprocess
 import sys
 import traceback
 import urllib
 import urlparse
 
 import sitescripts.stats.common as common
 from sitescripts.utils import get_config, setupStderr
 
 log_regexp = None
 gecko_apps = None
 
 def open_stats_file(path):
-  match = re.search(r"^ssh://(\w+)@([^/:]+)(?::(\d+))?/([^/]+)", path)
-  if match:
-    user, host, port, filename = match.groups()
-    command = ["ssh", "-q", "-o", "NumberOfPasswordPrompts 0", "-T", "-k", "-l", user, host, filename]
-    if port:
-      command[1:1] = ["-P", port]
-
-    # Not using StringIO here would be better but gzip module needs seeking
-    result = StringIO(subprocess.check_output(command))
-  elif path.startswith("http://") or path.startswith("https://"):
-    result = StringIO(urllib.urlopen(path).read())
+  parseresult = urlparse.urlparse(path)
+  if parseresult.scheme == "ssh" and parseresult.username and parseresult.hostname and parseresult.path:
+    command = [
+      "ssh", "-q", "-o", "NumberOfPasswordPrompts 0", "-T", "-k",
+      "-l", parseresult.username,
+      parseresult.hostname,
+      parseresult.path.lstrip("/")
+    ]
+    if parseresult.port:
+      command[1:1] = ["-P", str(parseresult.port)]
+    result = subprocess.Popen(command, stdout=subprocess.PIPE).stdout
+  elif parseresult.scheme in ("http", "https"):
+    result = urllib.urlopen(path)
   elif os.path.exists(path):
     result = open(path, "rb")
   else:
     raise IOError("Path '%s' not recognized" % path)
 
   if path.endswith(".gz"):
-    result = gzip.GzipFile(fileobj=result)
+    # Built-in gzip module doesn't support streaming (fixed in Python 3.2)
+    result = subprocess.Popen(["gzip", "-cd"], stdin=result, stdout=subprocess.PIPE).stdout
   return result
 
 def get_stats_files():
   config = get_config()
 
   prefix = "mirror_"
   options = filter(lambda o: o.startswith(prefix), config.options("stats"))
   for option in options:
     if config.has_option("stats", option):
       value = config.get("stats", option)
@@ skipping 128 matching lines @@
   if ua == "Adblock Plus":
     return "ABP", ""
 
   return "Other", ""
 
 def process_ip(ip, geo, geov6):
   match = re.search(r"^::ffff:(\d+\.\d+\.\d+\.\d+)$", ip)
   if match:
     ip = match.group(1)
 
-  if ":" in ip:
-    country = geov6.country_code_by_addr(ip)
-  else:
-    country = geo.country_code_by_addr(ip)
+  try:
+    if ":" in ip:
+      country = geov6.country_code_by_addr(ip)
+    else:
+      country = geo.country_code_by_addr(ip)
+  except:
+    traceback.print_exc()
+    country = ""
+
   if country in (None, "", "--"):
     country = "unknown"
   country = country.lower()
 
   return ip, country
 
 @cache_last
 def parse_time(timestr, tz_hours, tz_minutes):
   result = datetime.strptime(timestr, "%d/%b/%Y:%H:%M:%S")
   result -= timedelta(hours = tz_hours, minutes = math.copysign(tz_minutes, tz_hours))
@@ skipping 217 matching lines @@
       data[info["month"]] = {}
     section = data[info["month"]]
 
     if info["file"] not in section:
       section[info["file"]] = {}
     section = section[info["file"]]
 
     add_record(info, section)
   return data
 
-def merge_objects(object1, object2):
+def merge_objects(object1, object2, factor=1):
   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:
-      object1[key] = value
-
-def save_stats(server_type, data):
+    try:
+      key = unicode(key)
+    except UnicodeDecodeError:
+      key = unicode(key, encoding="latin-1")
+    if isinstance(value, numbers.Number):
+      object1[key] = object1.get(key, 0) + factor * value
+    else:
+      merge_objects(object1.setdefault(key, {}), value, factor)
+
+def save_stats(server_type, data, factor=1):
   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 fileobj:
           existing = json.load(fileobj)
       else:
         existing = {}
 
-      merge_objects(existing, file_data)
+      merge_objects(existing, file_data, factor)
 
       dir = os.path.dirname(path)
       try:
         os.makedirs(dir)
       except OSError, e:
         if e.errno != errno.EEXIST:
           raise
 
       with codecs.open(path, "wb", encoding="utf-8") as fileobj:
         json.dump(existing, fileobj, indent=2, sort_keys=True)
 
-def parse_file(mirror_name, server_type, log_file, geo, geov6, verbose):
-  ignored = set()
-  fileobj = open_stats_file(log_file)
+def parse_source((mirror_name, server_type, log_file)):
   try:
-    data = parse_fileobj(mirror_name, fileobj, geo, geov6, ignored)
+    geo = pygeoip.GeoIP(get_config().get("stats", "geoip_db"), pygeoip.MEMORY_CACHE)
+    geov6 = pygeoip.GeoIP(get_config().get("stats", "geoipv6_db"), pygeoip.MEMORY_CACHE)
+
+    ignored = set()
+    fileobj = open_stats_file(log_file)
+    try:
+      data = parse_fileobj(mirror_name, fileobj, geo, geov6, ignored)
+    finally:
+      fileobj.close()
+    return server_type, log_file, data, ignored
+  except:
+    print >>sys.stderr, "Unable to process log file '%s'" % log_file
+    traceback.print_exc()
+    return None, None, None, None
+
+def parse_sources(sources, factor=1, verbose=False):
+  pool = multiprocessing.Pool()
+  try:
+    for server_type, log_file, data, ignored in pool.imap(parse_source, sources, chunksize=1):
+      if server_type == None:
+        continue
+
+      save_stats(server_type, data, factor)
+      if verbose:
+        print "Ignored files for %s" % log_file
+        print "============================================================"
+        print "\n".join(sorted(ignored))
   finally:
-    fileobj.close()
-  save_stats(server_type, data)
-
-  if verbose:
-    print "Ignored files for %s" % log_file
-    print "============================================================"
-    print "\n".join(sorted(ignored))
+    pool.close()
 
 if __name__ == "__main__":
   setupStderr()
 
   parser = argparse.ArgumentParser(description="Processes log files and merges them into the stats database")
   parser.add_argument("--verbose", dest="verbose", action="store_const", const=True, default=False, help="Verbose mode, ignored requests will be listed")
+  parser.add_argument("--revert", dest="factor", action="store_const", const=-1, default=1, help="Remove log data from the database")
   parser.add_argument("mirror_name", nargs="?", help="Name of the mirror server that the file belongs to")
   parser.add_argument("server_type", nargs="?", help="Server type like download, update or subscription")
   parser.add_argument("log_file", nargs="?", help="Log file path, can be a local file path, http:// or ssh:// URL")
   args = parser.parse_args()
 
-  geo = pygeoip.GeoIP(get_config().get("stats", "geoip_db"), pygeoip.MEMORY_CACHE)
-  geov6 = pygeoip.GeoIP(get_config().get("stats", "geoipv6_db"), pygeoip.MEMORY_CACHE)
-
   if args.mirror_name and args.server_type and args.log_file:
-    parse_file(args.mirror_name, args.server_type, args.log_file, geo, geov6, args.verbose)
+    sources = [(args.mirror_name, args.server_type, args.log_file)]
   else:
-    for mirror_name, server_type, log_file in get_stats_files():
-      try:
-        parse_file(mirror_name, server_type, log_file, geo, geov6, args.verbose)
-      except:
-        print >>sys.stderr, "Unable to process log file '%s'" % log_file
-        traceback.print_exc()
-
+    sources = get_stats_files()
+  parse_sources(sources, args.factor, args.verbose)