Files
libpostal/scripts/geodata/geonames/geonames_sqlite.py

338 lines
10 KiB
Python

import os
import shutil
import sqlite3
import tempfile
import urlparse
import urllib2
import subprocess
import logging
import argparse
import csv
import sys
this_dir = os.path.realpath(os.path.dirname(__file__))
sys.path.append(os.path.realpath(os.path.join(os.pardir, os.pardir)))
from geodata.encoding import safe_decode
from geodata.geonames.paths import *
from geodata.file_utils import *
from geodata.log import *
from itertools import islice, chain
log_to_file(sys.stderr)
logger = logging.getLogger('geonames.sqlite')
BASE_URL = 'http://download.geonames.org/export/'
DUMP_URL = urlparse.urljoin(BASE_URL, 'dump/')
ALL_COUNTRIES_ZIP_FILE = 'allCountries.zip'
HIERARCHY_ZIP_FILE = 'hierarchy.zip'
ALTERNATE_NAMES_ZIP_FILE = 'alternateNames.zip'
ZIP_URL = urlparse.urljoin(BASE_URL, 'zip/')
GEONAMES_DUMP_FILES = (ALL_COUNTRIES_ZIP_FILE,
HIERARCHY_ZIP_FILE,
ALTERNATE_NAMES_ZIP_FILE)
# base_url, local_dir, is_gzipped, local_filename
GEONAMES_FILES = [(DUMP_URL, '', True, ALL_COUNTRIES_ZIP_FILE),
(DUMP_URL, '', True, HIERARCHY_ZIP_FILE),
(DUMP_URL, '', True, ALTERNATE_NAMES_ZIP_FILE),
(ZIP_URL, 'zip', True, ALL_COUNTRIES_ZIP_FILE),
]
def download_file(url, dest):
logger.info('Downloading file from {}'.format(url))
subprocess.check_call(['wget', url, '-O', dest])
def admin_ddl(admin_level):
columns = ['country_code TEXT'] + \
['admin{}_code TEXT'.format(i)
for i in xrange(1, admin_level)]
create = '''
CREATE TABLE admin{level}_codes (
geonames_id INT,
code TEXT,
name TEXT,
{fields}
)'''.format(level=admin_level,
fields=''',
'''.join(columns))
indices = (
'''CREATE INDEX admin{}_code_index ON
admin{}_codes (code)'''.format(admin_level, admin_level),
'''CREATE INDEX admin{}_gn_id_index ON
admin{}_codes (geonames_id)'''.format(admin_level, admin_level),
)
return (create, ) + indices
geonames_ddl = {
'geonames': (
'''CREATE TABLE geonames (
geonames_id INT PRIMARY KEY,
name TEXT,
ascii_name TEXT,
alternate_names TEXT,
latitude DOUBLE,
longitude DOUBLE,
feature_class TEXT,
feature_code TEXT,
country_code TEXT,
cc2 TEXT,
admin1_code TEXT,
admin2_code TEXT,
admin3_code TEXT,
admin4_code TEXT,
population LONG DEFAULT 0,
elevation INT,
dem INT,
timezone TEXT,
modification_date TEXT)''',
'''CREATE INDEX feature_code ON
geonames (feature_code)''',
'''CREATE INDEX country_code ON
geonames (country_code)''',
'''CREATE INDEX admin1_code ON
geonames (admin1_code)''',
'''CREATE INDEX admin2_code ON
geonames (admin2_code)''',
'''CREATE INDEX admin3_code ON
geonames (admin3_code)''',
'''CREATE INDEX admin4_code ON
geonames (admin4_code)''',
),
'alternate_names': (
'''CREATE TABLE alternate_names (
alternate_name_id INT PRIMARY KEY,
geonames_id INT,
iso_language TEXT,
alternate_name TEXT,
is_preferred_name BOOLEAN DEFAULT 0,
is_short_name BOOLEAN DEFAULT 0,
is_colloquial BOOLEAN DEFAULT 0,
is_historic BOOLEAN DEFAULT 0)''',
'''CREATE INDEX geonames_id_index ON
alternate_names (geonames_id)''',
'''CREATE INDEX geonames_id_alt_name_index ON
alternate_names(geonames_id, alternate_name)''',
),
'hierarchy': (
'''CREATE TABLE hierarchy (
parent_id INT,
child_id INT,
type TEXT
);''',
'''CREATE INDEX parent_child_index ON
hierarchy (parent_id, child_id)''',
'''CREATE INDEX child_parent_index ON
hierarchy (child_id, parent_id)''',
),
'postal_codes': (
'''CREATE TABLE postal_codes (
country_code TEXT,
postal_code TEXT,
place_name TEXT,
admin1 TEXT,
admin1_code TEXT,
admin2 TEXT,
admin2_code TEXT,
admin3 TEXT,
admin3_code TEXT,
latitude DOUBLE,
longitude DOUBLE,
accuracy INT
)''',
'''CREATE INDEX post_code_index ON
postal_codes (country_code, postal_code)''',
),
'admin1_codes': admin_ddl(1),
'admin2_codes': admin_ddl(2),
'admin3_codes': admin_ddl(3),
'admin4_codes': admin_ddl(4),
}
geonames_file_table_map = {
('', ALL_COUNTRIES_ZIP_FILE): 'geonames',
('', ALTERNATE_NAMES_ZIP_FILE): 'alternate_names',
('', HIERARCHY_ZIP_FILE): 'hierarchy',
('zip', ALL_COUNTRIES_ZIP_FILE): 'postal_codes',
}
country_codes_create_table = (
'drop table if exists country_codes',
'''
create table country_codes as
select distinct country_code from geonames
where feature_code in ('PCL', 'PCLI', 'PCLIX', 'PCLD', 'PCLF', 'PCLS', 'TERR')
''',
)
proper_countries_create_table = (
'drop table if exists proper_countries',
'''
create table proper_countries as
select * from geonames
where feature_code in ('PCL', 'PCLI', 'PCLIX', 'PCLD', 'PCLF', 'PCLS')
and country_code in (select country_code from country_codes)
''',
)
territories_create_table = (
'drop table if exists territories',
'''
create table territories as
select * from geonames where feature_code = 'TERR'
and country_code not in (select country_code from proper_countries);
''',
)
countries_create_table = (
'drop table if exists countries',
'''
create table countries as
select * from proper_countries
union
select * from territories;
''',
'create index country_geonames_id on countries (geonames_id)',
'create index conntry_country_code on countries (country_code)',
)
country_alises_create_table = (
'drop table if exists country_aliases',
'''
create table country_aliases as
select name, country_code
from countries
union
select alternate_name, country_code
from alternate_names an
join countries c
on c.geonames_id = an.geonames_id
where alternate_name != ''
and iso_language not in ('doi','faac','iata',
'icao','link','post','tcid')
'''
)
country_table_create_statements = list(chain(country_codes_create_table,
proper_countries_create_table,
territories_create_table,
countries_create_table,
country_alises_create_table))
def create_table(conn, table):
cursor = conn.cursor()
create_statements = geonames_ddl[table]
cursor.execute('DROP TABLE IF EXISTS {}'.format(table))
for statement in create_statements:
cursor.execute(statement)
conn.commit()
def batch_iter(iterable, batch_size):
source_iter = iter(iterable)
while True:
batch = list(islice(source_iter, batch_size))
if len(batch) > 0:
yield batch
else:
return
def populate_admin_table(conn, admin_level):
logging.info('Doing admin level {}'.format(admin_level))
columns = ['geonames_id',
'admin{}_code'.format(admin_level),
'name',
'country_code']
columns.extend(['admin{}_code'.format(i)
for i in xrange(1, admin_level)])
admin_insert_statement = '''
insert into "admin{}_codes"
select {}
from geonames
where feature_code = "ADM{}"
'''.format(admin_level, ','.join(columns), admin_level)
conn.execute(admin_insert_statement)
conn.commit()
logging.info('Done with admin level {}'.format(admin_level))
def import_geonames_table(conn, table, f, batch_size=2000):
# escape the brackets around the values format string so we can use later
statement = 'INSERT INTO "{}" VALUES ({{}})'.format(table)
cursor = conn.cursor()
for i, batch in enumerate(batch_iter(f, batch_size)):
num_cols = len(batch[0])
cursor.executemany(statement.format(','.join(['?'] * num_cols)), batch)
conn.commit()
cursor = conn.cursor()
logging.info('imported {} batches ({} records)'.format(i + 1, (i + 1) * batch_size))
cursor.close()
def create_geonames_sqlite_db(temp_dir, db_file=DEFAULT_GEONAMES_DB_PATH):
conn = sqlite3.connect(db_file)
logging.info('Created database at {}'.format(db_file))
for url, directory, is_gzipped, filename in GEONAMES_FILES:
table = geonames_file_table_map[(directory, filename)]
create_table(conn, table)
full_url = urlparse.urljoin(url, filename)
dest_dir = os.path.join(temp_dir, directory)
ensure_dir(dest_dir)
dest_file = os.path.join(dest_dir, filename)
download_file(full_url, dest_file)
if is_gzipped:
unzip_file(dest_file, dest_dir)
filename = dest_file.replace('.zip', '.txt')
reader = csv.reader(open(filename), delimiter='\t', quotechar=None)
lines = (map(safe_decode, line) for line in reader)
import_geonames_table(conn, table, lines)
logging.info('Creating countries tables')
for statement in country_table_create_statements:
conn.execute(statement)
conn.commit()
logging.info('Creating admin tables')
for admin_level in xrange(1, 5):
create_table(conn, 'admin{}_codes'.format(admin_level))
populate_admin_table(conn, admin_level)
conn.close()
if __name__ == '__main__':
# Handle argument parsing here
parser = argparse.ArgumentParser()
parser.add_argument('-t', '--temp-dir',
default=tempfile.gettempdir(),
help='Temporary work directory')
parser.add_argument('-o', '--out',
default=DEFAULT_GEONAMES_DB_PATH,
help='SQLite3 db filename')
args = parser.parse_args()
create_geonames_sqlite_db(args.temp_dir, args.out)