added maxmind.py and updated mlab_mysql_import to take advantage of it.

maxmind.py indexes one of the maxmind tables, to greatly speed up the
locId lookup process.

Author: axello

maxmind.py

@@ -0,0 +1,100 @@
+#!/usr/bin/python
+#
+# maxmind.py
+#
+# A class to represent a full maxmind database as a python hash, to make locId lookups much faster (hopefully)
+#
+# Algorithm:
+# set up two tables: 
+#   startips = list of all the startips, we binary search through this using bisect
+#   ipranges =  hash of all the ranges, where the key is the startip, and the endip is the value
+# 
+# binary search through startips to find the closest startip of the target ip
+# get the endip of the range from the hash, and compare
+# return locId to caller
+#
+# Initials:
+#               AX  Axel Roest
+#               RB  Ruben Bloemgarten
+#
+# Version history
+# 20120710      AR  first version
+#
+# ToDO: 
+#   
+
+import sys
+import re
+import os
+import math
+import bisect
+from datetime import datetime
+import MySQLdb
+
+#################################################################
+#																#
+#			settings											#
+#																#
+#################################################################
+
+# Defaults
+
+#################################################################
+#																#
+#			the meat											#
+#																#
+#################################################################
+
+class MaxMind:
+    # two tables:
+    # startips = []   # list of all the startips, we binary search through this using bisect
+    # ipranges = {}   # hash of all the ranges, where the key is the startip, and the endip is the value
+    
+    # initialise the object with the database table to instantiate with
+    def __init__(self, maxmind_db_host, maxmind_db_user, maxmind_db_passwd, maxmind_db_name, maxmind_table_name):
+        global_start_time = datetime.now()
+        try:
+            # Connect to the mysql database
+            maxmind_db = MySQLdb.connect(host = maxmind_db_host, 
+                                                     user = maxmind_db_user, 
+                                                     passwd = maxmind_db_passwd, 
+                                                     db = maxmind_db_name) 
+            maxmind_cursor = maxmind_db.cursor()
+            MaxMind.loadTable(self, maxmind_cursor, maxmind_table_name)
+        except Exception as e:
+            print('Aborting maxmind due to error: ' + str(e))
+            exit(1)
+        finally:
+            maxmind_cursor.close()
+        global_end_time = datetime.now()
+        print 'MaxMind: Read and indexed `' + maxmind_table_name + '` in ' +  str(global_end_time - global_start_time) + ' seconds.'
+    
+    def loadTable(self, maxmind_cursor, maxmind_table_name):
+        # grab everything (whole table, 3,5 million items)
+        sql = """SELECT startIpNum, endIpNum, locId FROM `{0}`""".format(maxmind_table_name)
+        maxmind_cursor.execute(sql)
+        result = maxmind_cursor.fetchall()
+        if result:
+        # fromkeys
+            self.startips = [x[0] for x in result]
+            self.startips.sort()
+            self.ipranges = {int(start) : [int(e),int(l)] for start,e,l in result}
+        
+    def find_le(self, a, x):
+        # Find rightmost value less than or equal to x
+        i = bisect.bisect_right(a, x)
+        if i:
+            return a[i-1]
+        raise ValueError
+    
+    def lookup(self, ipnumber):
+        # print "looking up: " + str(ipnumber)
+        i = self.find_le(self.startips, ipnumber)
+        # print "i="+str(i)
+        (endip, loc) = self.ipranges[i]
+        # print endip, loc
+        if ipnumber <= endip:
+            return loc
+        else:
+            return -1
+        

mlab_mysql_import.py

@@ -10,6 +10,7 @@
 # 20120629      AX  added loop over all arguments, exception handling, restructured code, moved processed files to archive or error folder
 # 20120708      AX  skip empty ip lines instead or error message
 # 20120708      RB  cleaning some names and spelling, also we don't want processed_files.log to clobber the downloaders processed_files.log. So we should use overly descriptive names
+# 20120710      AX  added locId lookup and added longip to insert query
 #
 # test: 
 # cd /DATA
@@ -21,15 +22,14 @@
 #       v move error files naar error directory
 #       v log process and errors
 #       v skip empty ip lines instead or error message
+#       v added locId lookup and added longip to insert query
 #
 #       Get the date from the filename, and look up the correct maxmind database
 #       then, insert the locId directly with the line in the mlab/{glasnost,ndt} database, preventing slow future updates
 #       on the other hand, all these updates might be extremely slow: TEST
 #
-#		todo : refactor all the utility functions in a separate file
-#		todo : refactor all the passwords in a separate file (which is NOT in the repo, AND is in the .gitignore list
-
-
+#       todo : refactor all the utility functions in a separate file
+#       todo : refactor all the passwords in a separate file (which is NOT in the repo, AND is in the .gitignore list
 
 import sys
 import re
@@ -39,6 +39,8 @@
 import dateutil.parser as dparser
 import MySQLdb
 import shutil
+from maxmind import MaxMind
+import socket, struct
 
 #################################################################
 #                                                               #
@@ -51,11 +53,12 @@
 db_user = "root" # your username
 db_passwd = "" # your password
 db_name = "mlab" # name of the database
-db_tables = {"glasnost": "glasnost", "ndt": "ndt"} # a mapping from testname to tablename
+db_tables = {"glasnost": "glasnost", "ndt": "ndt_test"} # a mapping from testname to tablename
 db_filetable = 'files'
 
 # directories
 baseDir     = '/DATA/mlab/'
+#baseDir     = '/home/axel/mlab/'
 scratchDir  = baseDir + 'scratch/'
 workDir     = baseDir + 'work/'
 archiveDir  = baseDir + 'archive/'
@@ -67,14 +70,26 @@
 errorLog    = "mlab_mysql_import_error.log"
 processLog  = "mlab_mysql_import_processed_files.log"
 
+# default tables
+maxmind_table = 'Blocks_GeoLiteCity_Last'
+ndt_import    = 'ndt_import'
 #################################################################
 #                                                               #
 #           functions                                           #
 #                                                               #
 #################################################################
 
+# Convert an IP string to long
+def ip2long(ip):
+  packedIP = socket.inet_aton(ip)
+  return struct.unpack("!L", packedIP)[0]
+
+def long2ip(l):
+  return socket.inet_ntoa(struct.pack('!L', l))
+
 def usage():
-  print "Usage: mlab_mysql_import3.py mlab_file1.csv [mlab_files.csv ...]"
+  print "Usage: mlab_mysql_import.py [ -m maxmind_Blocks_Tablename ] mlab_file1.csv [mlab_files.csv ...]"
+  print "Default: maxmind_Blocks_Tablename = `Blocks_GeoLiteCity_Last`"
   sys.exit(1)
 
 # This routine extracts the destination server of the mlab file. 
@@ -133,9 +148,13 @@ def exists_dbentry(cur, file_id, db_table, test_datetime, destination, source_ip
 
 # Insert a connection to the database without testing.
 def blunt_insert_dbentry(cur, file_id, db_table, test_datetime, destination, source_ip):
-    columns = ', '.join(['date', 'destination', 'source', 'file_id'])
-    values = '"' + '", "'.join([test_datetime.isoformat(), destination, source_ip, str(file_id)]) + '"'
+    longip = ip2long(source_ip)
+    # locid = 0
+    locid = mm.lookup(longip)                  # lookup location id from ip number
+    columns = ', '.join(['date', 'destination', 'source', 'file_id', 'longip', 'locId'])
+    values = '"' + '", "'.join([test_datetime.isoformat(), destination, source_ip, str(file_id), str(longip), str(locid)]) + '"'
     sql = "INSERT INTO  " + db_table + " (" + columns + ") VALUES(" + values + ") "
+    # print sql
     cur.execute(sql)
 
 # Insert a test connection to the database, if it not already exists
@@ -170,6 +189,26 @@ def dedup(file_id, table, test_datetime, destination, source_ip):
         deduplookup[key] = True
         return True
 
+# for the temp table, look up all the locations with the locId
+def lookup_locations(cur, destination):
+    location_table_name = maxmind_table.replace("Blocks", "Location")
+    # sql = 'UPDATE mlab.`' + destination + '` L, maxmind.`' + location_table_name + '` M SET L.country_code = M.country, L.region=M.region, L.city=M.city, L.postalCode=M.postalCode, L.latitude=M.latitude, L.longitude=M.longitude, L.metroCode=M.metroCode, L.areaCode=M.areaCode WHERE L.`locId` = M.`locId`'
+    sql = 'UPDATE mlab.`ndt_import` L, maxmind.`' + location_table_name + '` M SET L.country_code = M.country, L.region=M.region, L.city=M.city, L.postalCode=M.postalCode, L.latitude=M.latitude, L.longitude=M.longitude, L.metroCode=M.metroCode, L.areaCode=M.areaCode WHERE L.`locId` = M.`locId`'
+    updated = cur.execute(sql)
+    # update country from country_code later?
+    return updated
+
+# clear the temp table
+def clear_temp_table(cur):
+    sql = 'truncate table `' + ndt_import + '`'
+    cur.execute(sql)
+    
+# move the temp table to the real on (either ndt_test or ndt)
+def move_temp_table(cur, destination):
+    sql = 'INSERT INTO `' + destination + '` (`created_at`, `date`, `destination`, `source`, `file_id`, `country_code`, `longip`, `locId`, `country`, `region`, `city`, `postalCode`, `latitude`, `longitude`, `metroCode`, `areaCode`)  SELECT * FROM `' + ndt_import + '`'
+    updated = cur.execute(sql)
+    return updated
+    
 # returns True on error, False on correct processing
 def process_file(f, filename):
     start_time = datetime.now()
@@ -181,6 +220,7 @@ def process_file(f, filename):
                              passwd = db_passwd, 
                              db = db_name) 
         cur = db.cursor() 
+        clear_temp_table(cur)
 
         # Find the destination server by investigating the filename
         destination = extract_destination(filename)
@@ -190,15 +230,15 @@ def process_file(f, filename):
         file_id = get_file_id(cur, filename) 
         db.commit()
 
-        # Find the testsuite by investigating the filename
+        # Find the testsuite (glasnost or ndt) by investigating the filename
         try:
             test = [test for test in db_tables.keys() if test in filename][0]
         except IndexError:
             sys.stderr.write('The filename ' + filename + ' does not contain a valid testname.')
             return 1
         # print "Found test suite " + test 
 
-    	# The filetest ALONE, takes 3 seconds with a 9 million records database, without indexes
+        # The filetest ALONE, takes 3 seconds with a 9 million records database, without indexes
         # But falls back to less than half a second when indexing is turned on on the db
         filetest=True
         # Read the file line by line and import it into the database
@@ -215,9 +255,12 @@ def process_file(f, filename):
             filetest=False
           # test if we have already done it in this or last filetest
           if (dedup(file_id, db_tables[test], test_datetime, destination, source_ip)):
-              blunt_insert_dbentry(cur, file_id, db_tables[test], test_datetime, destination, source_ip)
+              # blunt_insert_dbentry(cur, file_id, db_tables[test], test_datetime, destination, source_ip)
+              blunt_insert_dbentry(cur, file_id, ndt_import, test_datetime, destination, source_ip)
         end_time = datetime.now()
         print 'File done in ' + str(end_time - start_time)
+        lookup_locations(cur, destination)
+        move_temp_table(cur, db_tables[test])
         failure = False
     except Exception as inst:
         sys.stderr.write('Exception: '+str(inst.args)  + '\n')
@@ -231,9 +274,6 @@ def process_file(f, filename):
             f.write(pathname + '\n')
             f.write('Error handling file ' + filename + ' (' + str(e.args) + ')\n')
         print
-# This bit should probably be cleaned up.        
-#    except:
-#        sys.stderr.write('Process error ' + '\n')
     finally:
         # Commit and finish up
         sys.stderr.flush()
@@ -273,7 +313,11 @@ def move_archive(pathname):
 
 parser = OptionParser()
 parser.add_option("-q", "--quiet", action="store_false", dest="verbose", default=False, help="don't print status messages to stdout")
+parser.add_option("-m", "--maxmind", dest="maxmind_table", default='', help="optional maxmind_table, if omitted we use 'Last'")
 (options, args) = parser.parse_args()
+if options.maxmind_table != '':
+    maxmind_table = options.maxmind_table
+    
 if len(args) == 0:
   usage()
 
@@ -295,17 +339,26 @@ def move_archive(pathname):
 #################################################################
 global_start_time = datetime.now()
 
+# get instance of maxmind table
+print "using " + maxmind_table
+
+mm = MaxMind(db_host, db_user, db_passwd, "maxmind",maxmind_table)
+
+if not mm:
+    sys.stderr.write('maxmind table does not exist: ' + maxmind_table + ' (' + str(e.args) + ')\n')
+    exit(1)
+
 # Iterate over ALL filenames
 for pathname in args:
-	try:
-		 with open(pathname, 'r') as f:
-        	# Extract the basename of the filename, as the path is not of interest after this point
-        	filename = os.path.basename(pathname)
+    try:
+         with open(pathname, 'r') as f:
+            # Extract the basename of the filename, as the path is not of interest after this point
+            filename = os.path.basename(pathname)
             print "processing file " + filename,
-        	if (process_file(f, filename)):
-    	        shutil.move(pathname,errorDir)
-        	else:
-        	    move_archive(pathname)
+            if (process_file(f, filename)):
+                shutil.move(pathname,errorDir)
+            else:
+                move_archive(pathname)
         # file is automatically closed if needed
     except IOError as e:
          print 'Could not open file ' + pathname + '\nError: ' + str(e.args)