nominatim/tools/tiger_data.py

   1 """
   2 Functions for importing tiger data and handling tarbar and directory files
   3 """
   4 import logging
   5 import os
   6 import tarfile
   7 import selectors
   8
   9 from nominatim.db.connection import connect
  10 from nominatim.db.async_connection import DBConnection
  11 from nominatim.db.sql_preprocessor import SQLPreprocessor
  12
  13
  14 LOG = logging.getLogger()
  15
  16
  17 def handle_tarfile_or_directory(data_dir):
  18     """ Handles tarfile or directory for importing tiger data
  19     """
  20
  21     tar = None
  22     if data_dir.endswith('.tar.gz'):
  23         tar = tarfile.open(data_dir)
  24         sql_files = [i for i in tar.getmembers() if i.name.endswith('.sql')]
  25         LOG.warning("Found %d SQL files in tarfile with path %s", len(sql_files), data_dir)
  26         if not sql_files:
  27             LOG.warning("Tiger data import selected but no files in tarfile's path %s", data_dir)
  28             return None, None
  29     else:
  30         files = os.listdir(data_dir)
  31         sql_files = [os.path.join(data_dir, i) for i in files if i.endswith('.sql')]
  32         LOG.warning("Found %d SQL files in path %s", len(sql_files), data_dir)
  33         if not sql_files:
  34             LOG.warning("Tiger data import selected but no files found in path %s", data_dir)
  35             return None, None
  36
  37     return sql_files, tar
  38
  39
  40 def handle_threaded_sql_statements(sel, file):
  41     """ Handles sql statement with multiplexing
  42     """
  43
  44     lines = 0
  45     end_of_file = False
  46     # Using pool of database connections to execute sql statements
  47     while not end_of_file:
  48         for key, _ in sel.select(1):
  49             conn = key.data
  50             try:
  51                 if conn.is_done():
  52                     sql_query = file.readline()
  53                     lines += 1
  54                     if not sql_query:
  55                         end_of_file = True
  56                         break
  57                     conn.perform(sql_query)
  58                     if lines == 1000:
  59                         print('. ', end='', flush=True)
  60                         lines = 0
  61             except Exception as exc: # pylint: disable=broad-except
  62                 LOG.info('Wrong SQL statement: %s', exc)
  63
  64 def handle_unregister_connection_pool(sel, place_threads):
  65     """ Handles unregistering pool of connections
  66     """
  67
  68     while place_threads > 0:
  69         for key, _ in sel.select(1):
  70             conn = key.data
  71             sel.unregister(conn)
  72             try:
  73                 conn.wait()
  74             except Exception as exc: # pylint: disable=broad-except
  75                 LOG.info('Wrong SQL statement: %s', exc)
  76             conn.close()
  77             place_threads -= 1
  78
  79 def add_tiger_data(data_dir, config, threads):
  80     """ Import tiger data from directory or tar file `data dir`.
  81     """
  82     dsn = config.get_libpq_dsn()
  83     sql_files, tar = handle_tarfile_or_directory(data_dir)
  84
  85     if not sql_files:
  86         return
  87
  88     with connect(dsn) as conn:
  89         sql = SQLPreprocessor(conn, config)
  90         sql.run_sql_file(conn, 'tiger_import_start.sql')
  91
  92     # Reading sql_files and then for each file line handling
  93     # sql_query in <threads - 1> chunks.
  94     sel = selectors.DefaultSelector()
  95     place_threads = max(1, threads - 1)
  96
  97     # Creates a pool of database connections
  98     for _ in range(place_threads):
  99         conn = DBConnection(dsn)
 100         conn.connect()
 101         sel.register(conn, selectors.EVENT_WRITE, conn)
 102
 103     for sql_file in sql_files:
 104         if not tar:
 105             file = open(sql_file)
 106         else:
 107             file = tar.extractfile(sql_file)
 108
 109         handle_threaded_sql_statements(sel, file)
 110
 111     # Unregistering pool of database connections
 112     handle_unregister_connection_pool(sel, place_threads)
 113
 114     if tar:
 115         tar.close()
 116     print('\n')
 117     LOG.warning("Creating indexes on Tiger data")
 118     with connect(dsn) as conn:
 119         sql = SQLPreprocessor(conn, config)
 120         sql.run_sql_file(conn, 'tiger_import_finish.sql')