nominatim/tools/tiger_data.py

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