nominatim/tools/tiger_data.py

   1 """
   2 Functions for importing tiger data and handling tarbar and directory files
   3 """
   4 import csv
   5 import io
   6 import logging
   7 import os
   8 import tarfile
   9
  10 from nominatim.db.connection import connect
  11 from nominatim.db.async_connection import WorkerPool
  12 from nominatim.db.sql_preprocessor import SQLPreprocessor
  13 from nominatim.errors import UsageError
  14 from nominatim.indexer.place_info import PlaceInfo
  15
  16 LOG = logging.getLogger()
  17
  18
  19 def handle_tarfile_or_directory(data_dir):
  20     """ Handles tarfile or directory for importing tiger data
  21     """
  22
  23     tar = None
  24     if data_dir.endswith('.tar.gz'):
  25         try:
  26             tar = tarfile.open(data_dir)
  27         except tarfile.ReadError as err:
  28             LOG.fatal("Cannot open '%s'. Is this a tar file?", data_dir)
  29             raise UsageError("Cannot open Tiger data file.") from err
  30
  31         csv_files = [i for i in tar.getmembers() if i.name.endswith('.csv')]
  32         LOG.warning("Found %d CSV files in tarfile with path %s", len(csv_files), data_dir)
  33         if not csv_files:
  34             LOG.warning("Tiger data import selected but no files in tarfile's path %s", data_dir)
  35             return None, None
  36     else:
  37         files = os.listdir(data_dir)
  38         csv_files = [os.path.join(data_dir, i) for i in files if i.endswith('.csv')]
  39         LOG.warning("Found %d CSV files in path %s", len(csv_files), data_dir)
  40         if not csv_files:
  41             LOG.warning("Tiger data import selected but no files found in path %s", data_dir)
  42             return None, None
  43
  44     return csv_files, tar
  45
  46
  47 def handle_threaded_sql_statements(pool, fd, analyzer):
  48     """ Handles sql statement with multiplexing
  49     """
  50     lines = 0
  51     # Using pool of database connections to execute sql statements
  52
  53     sql = "SELECT tiger_line_import(%s, %s, %s, %s, %s, %s)"
  54
  55     for row in csv.DictReader(fd, delimiter=';'):
  56         try:
  57             address = dict(street=row['street'], postcode=row['postcode'])
  58             args = ('SRID=4326;' + row['geometry'],
  59                     int(row['from']), int(row['to']), row['interpolation'],
  60                     PlaceInfo({'address': address}).analyze(analyzer),
  61                     analyzer.normalize_postcode(row['postcode']))
  62         except ValueError:
  63             continue
  64         pool.next_free_worker().perform(sql, args=args)
  65
  66         lines += 1
  67         if lines == 1000:
  68             print('.', end='', flush=True)
  69             lines = 0
  70
  71
  72 def add_tiger_data(data_dir, config, threads, tokenizer):
  73     """ Import tiger data from directory or tar file `data dir`.
  74     """
  75     dsn = config.get_libpq_dsn()
  76     files, tar = handle_tarfile_or_directory(data_dir)
  77
  78     if not files:
  79         return
  80
  81     with connect(dsn) as conn:
  82         sql = SQLPreprocessor(conn, config)
  83         sql.run_sql_file(conn, 'tiger_import_start.sql')
  84
  85     # Reading files and then for each file line handling
  86     # sql_query in <threads - 1> chunks.
  87     place_threads = max(1, threads - 1)
  88
  89     with WorkerPool(dsn, place_threads, ignore_sql_errors=True) as pool:
  90         with tokenizer.name_analyzer() as analyzer:
  91             for fname in files:
  92                 if not tar:
  93                     fd = open(fname)
  94                 else:
  95                     fd = io.TextIOWrapper(tar.extractfile(fname))
  96
  97                 handle_threaded_sql_statements(pool, fd, analyzer)
  98
  99                 fd.close()
 100
 101     if tar:
 102         tar.close()
 103     print('\n')
 104     LOG.warning("Creating indexes on Tiger data")
 105     with connect(dsn) as conn:
 106         sql = SQLPreprocessor(conn, config)
 107         sql.run_sql_file(conn, 'tiger_import_finish.sql')