Hide
Don't get spied on – We respect your privacy and provide numerous options to protect it. Join Siafoo Now or Learn More

Utility for loading GeoIP data into a PostgreSQL database

Revision 4 vs. Revision 5

Changelog: Added IP4r support

Legend:

Unmodified
Added
Removed
  • Code

    r4 r5  
    77from optparse import OptionGroup, OptionParser  
    88from StringIO import StringIO  
    99  
    1010class GeoIPDataLoader(object):  
    1111  
    12     def __init__(self, dsn, blocks='GeoLiteCity-Blocks.csv', locations='GeoLiteCity-Location.csv'):  
     12    def __init__(self, dsn, blocks='GeoLiteCity-Blocks.csv', locations='GeoLiteCity-Location.csv', schema='public'):  
    1313        self.con = psycopg2.connect(dsn)  
    1414        # We don't need transactions... right?  
    1515        self.con.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)  
    1616        # The data is in ISO8859_15 encoding  
    1717        self.con.set_client_encoding('iso8859_15')  
    1818        self.cur = self.con.cursor()  
    1919  
    2020        self.blocks_csv = blocks  
    2121        self.location_csv = locations  
     22        self.schema = schema  
     23   
     24    def close(self):  
     25        self.con.close()  
    2226  
    2327    def create_tables(self):  
    2428        print 'Creating structure...',  
    2529        self.db_execute(  
    2630            '''  
     
    4044                 
    4145                CREATE TABLE blocks  
    4246                (  
    4347                  start_ip bigint NOT NULL,  
    4448                  end_ip bigint NOT NULL,  
    45                   location_id bigint NOT NULL,  
    46                   CONSTRAINT blocks_pkey PRIMARY KEY (start_ip, end_ip)  
     49                  location_id bigint NOT NULL  
    4750                );  
    4851                 
    4952            '''  
    5053            )  
    5154        print '\033[1;32mDone\033[1;m'  
    5255                 
    53     def create_indexes(self):  
     56    def create_indexes(self, ip4=False):  
    5457        print 'Adding Indexes...',  
    55         self.db_execute('''  
     58        sys.stdout.flush()  
     59        if not ip4:  
     60            self.db_execute('''  
    5661             CREATE INDEX ix_start_end_ip ON blocks  
    5762                USING btree (start_ip, end_ip) WITH (FILLFACTOR=100);  
    5863             CREATE INDEX ix_end_start_ip ON blocks  
    5964                USING btree (end_ip, start_ip) WITH (FILLFACTOR=100);  
    60             ''')  
     65                ''')  
     66        else:  
     67            self.db_execute('''  
     68                 CREATE INDEX ix_ip_range ON blocks  
     69                   USING gist (ip_range) WITH (FILLFACTOR=100);  
     70                ''')  
    6171        print '\033[1;32mDone\033[1;m'  
    6272         
    63     def create_functions(self):  
     73    def create_functions(self, ip4=False):  
    6474        print 'Adding utility functions...',  
    65         self.db_execute('''  
    66             CREATE OR REPLACE FUNCTION get_location(inet) RETURNS bigint AS $$  
    67               SELECT location_id FROM blocks  
    68               WHERE $1 - inet '0.0.0.0' BETWEEN start_ip AND end_ip  
    69             $$ LANGUAGE SQL;  
    70              
    71             CREATE OR REPLACE FUNCTION inet_to_block(inet) RETURNS bigint AS $$  
    72                 SELECT $1 - inet '0.0.0.0'  
    73             $$ LANGUAGE SQL;  
    74             ''')  
    75         print '\033[1;32mDone\033[1;m'  
    76        
    77     def close(self):  
    78         self.con.close()         
     75        sys.stdout.flush()  
     76        if ip4:  
     77            self.db_execute('''  
     78                CREATE OR REPLACE FUNCTION get_location(inet) RETURNS bigint AS $$  
     79                  SELECT location_id FROM %s.blocks  
     80                  WHERE ip_range >>= ip4($1)  
     81                $$ LANGUAGE SQL;  
     82                ''' % self.schema)  
     83        else:  
     84            self.db_execute('''  
     85                CREATE OR REPLACE FUNCTION inet_to_bigint(inet) RETURNS bigint AS $$  
     86                    SELECT $1 - inet '0.0.0.0'  
     87                $$ LANGUAGE SQL;  
     88                ''' % self.schema)  
     89        print '\033[1;32mDone\033[1;m'  
     90     
     91    def create_schema(self):  
     92        try:  
     93            self.db_execute('''CREATE SCHEMA %s;''' % self.schema)  
     94        except psycopg2.ProgrammingError:  
     95          pass     
     96  
     97        self.db_execute('SET search_path TO %s,public;' % self.schema)  
    7998         
    8099    def db_execute(self, ddl):  
    81100        self.cur.execute(ddl)  
    82101#        self.con.commit()  
    83102     
     
    87106        # Load Blocks  
    88107        self.load_table(self.blocks_csv, 'blocks')  
    89108     
    90109    def load_table(self, file_name, table_name):  
    91110        print 'Loading table \033[1;34m%s\033[1;m from file \033[1;34m%s\033[1;m...' % (table_name, file_name),  
     111        sys.stdout.flush()  
    92112        geo_file = open(file_name)  
    93113        # Skip the copyright header  
    94114        geo_file.readline()  
    95115        geo_file.readline()  
    96116        #Remove quotes... psycopg2's `copy` errors on them  
    97117        string_data = geo_file.read().replace('"', '')  
    98118        self.cur.copy_from(StringIO(string_data), table_name,  sep=',', null='')  
    99119        print '\033[1;32mDone\033[1;m'  
     120     
     121    def migrate_to_ip4(self):  
     122        print 'Adding ip_range column'         
     123        self.db_execute('''  
     124                        ALTER TABLE blocks ADD COLUMN ip_range ip4r;  
     125                        ALTER TABLE blocks ALTER COLUMN ip_range SET STORAGE PLAIN;  
     126                        ''')  
    100127         
     128        print 'Migrating data to ip4...',  
     129        sys.stdout.flush()  
     130        self.db_execute('''UPDATE blocks SET ip_range = ip4r(start_ip::ip4, end_ip::ip4)''')  
     131        print '\033[1;32mDone\033[1;m'  
     132  
     133        print 'Dropping unneeded columns'  
     134        self.db_execute('''  
     135                        ALTER TABLE blocks DROP COLUMN start_ip;  
     136                        ALTER TABLE blocks DROP COLUMN end_ip;  
     137                        ''')  
    101138    def vacuum(self):  
    102139        print 'Vaccuming database...',  
     140        sys.stdout.flush()  
    103141        self.db_execute('VACUUM FULL ANALYZE')  
    104142        print '\033[1;32mDone\033[1;m'  
    105143  
    106144def main():  
    107145    DSN = "dbname='%s' user='%s' host='%s'"  
    108146  
    109147    parser = OptionParser()  
    110148    # Operational options  
    111     parser.add_option('-c', '--create-structure', dest='load_ddl', default=False,  
     149    parser.add_option('-c', '--load-ddl', dest='load_ddl', default=False,  
    112150                      action='store_true', help='Create database structure')  
    113151     
    114     parser.add_option('-n', '--no-load', dest='no_load', default=False,  
    115                       action='store_true', help='Do not load the GeoIP data')  
    116      
     152    parser.add_option('-g', '--load-data', dest='load', default=False,  
     153                      action='store_true', help='Load the GeoIP data')  
     154  
    117155    parser.add_option('-b', '--blocks-file', dest='blocks_csv', default='GeoLiteCity-Blocks.csv',  
    118156                      action='store', help='GeoIP Blocks CSV file [default: %default]', metavar='BLOCKS_FILE')  
    119157    parser.add_option('-l', '--locations-file', dest='locations_csv', default='GeoLiteCity-Location.csv',  
    120158                      action='store', help='GoeIP Locations CSV file [default: %default]', metavar='LOCATIONS_FILE')  
    121159  
    122     db_group = OptionGroup(parser, 'Database Connection Options')  
     160    db_group = OptionGroup(parser, 'Database Options')  
    123161    # Database options  
    124162    db_group.add_option('-H', '--host', dest='db_host', default='localhost',  
    125163                      action='store', help='Database host [default: %default]', metavar='DB_HOST')  
    126164    db_group.add_option('-d', '--database', dest='db_name', default='geoip_db',  
    127165                      action='store', help='Database name [default: %default]', metavar='DATABASE_NAME')  
    128166    db_group.add_option('-U', '--user', dest='db_user', default='geoip',  
    129167                      action='store', help='User [default: %default]', metavar='USER_NAME')  
     168    db_group.add_option('-s', '--schema', dest='schema', default='public',  
     169                      action='store', help='Database Schema [default: %default]', metavar='SCHEMA')  
     170  
     171    db_group.add_option('--ip4r', dest='ip4', default=False,  
     172                      action='store_true', help='Use IP4r module [default: %default]')  
    130173  
    131174    parser.add_option_group(db_group)  
    132175     
    133176    (options, args) = parser.parse_args()  
    134177  
    135178    data_loader = GeoIPDataLoader("dbname='%s' user='%s' host='%s'" % (options.db_name, options.db_user, options.db_host),  
    136                                   blocks=options.blocks_csv, locations=options.locations_csv)  
     179                                  blocks=options.blocks_csv, locations=options.locations_csv, schema=options.schema)  
     180  
     181    if not options.load_ddl and not options.load:  
     182        parser.print_help()  
     183        return  
    137184  
    138185    if options.load_ddl:  
     186        if options.schema != 'public':  
     187            data_loader.create_schema()  
    139188        data_loader.create_tables()  
    140         data_loader.create_functions()  
    141   
    142     if not options.no_load:  
     189   
     190    if options.load:  
    143191        data_loader.load_data()  
    144         data_loader.create_indexes()  
    145         data_loader.vacuum()  
     192  
     193    if options.ip4:  
     194        data_loader.migrate_to_ip4()  
     195  
     196    if options.load:  
     197        data_loader.create_indexes(options.ip4 is True)  
     198  
     199    if options.load_ddl:  
     200        data_loader.create_functions(options.ip4 is True)  
     201  
     202    data_loader.vacuum()  
    146203  
    147204if __name__ == "__main__":  
    148205    main()  
    149206