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 ..db.connection import connect
  10 from ..db.async_connection import DBConnection
  11 from ..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 True:
  48         if end_of_file:
  49             break
  50         for key, _ in sel.select(1):
  51             conn = key.data
  52             try:
  53                 if conn.is_done():
  54                     sql_query = file.readline()
  55                     lines += 1
  56                     if not sql_query:
  57                         end_of_file = True
  58                         break
  59                     conn.perform(sql_query)
  60                     if lines == 1000:
  61                         print('. ', end='', flush=True)
  62                         lines = 0
  63             except Exception as exc: # pylint: disable=broad-except
  64                 LOG.error('Wrong SQL statement: %s', exc)
  65
  66
  67 def add_tiger_data(dsn, data_dir, threads, config, sqllib_dir):
  68     """ Import tiger data from directory or tar file
  69     """
  70
  71     sql_files, tar = handle_tarfile_or_directory(data_dir)
  72
  73     if not sql_files:
  74         return
  75
  76     with connect(dsn) as conn:
  77         sql = SQLPreprocessor(conn, config, sqllib_dir)
  78         sql.run_sql_file(conn, 'tiger_import_start.sql')
  79
  80     # Reading sql_files and then for each file line handling
  81     # sql_query in <threads - 1> chunks.
  82     sel = selectors.DefaultSelector()
  83     place_threads = max(1, threads - 1)
  84
  85     # Creates a pool of database connections
  86     for _ in range(place_threads):
  87         conn = DBConnection(dsn)
  88         conn.connect()
  89         sel.register(conn, selectors.EVENT_WRITE, conn)
  90
  91     for sql_file in sql_files:
  92         if not tar:
  93             file = open(sql_file)
  94         else:
  95             file = tar.extractfile(sql_file)
  96
  97         handle_threaded_sql_statements(sel, file)
  98
  99     # Unregistering pool of database connections
 100     while place_threads > 0:
 101         for key, _ in sel.select(1):
 102             conn = key.data
 103             sel.unregister(conn)
 104             conn.wait()
 105             conn.close()
 106             place_threads -= 1
 107
 108     if tar:
 109         tar.close()
 110     print('\n')
 111     LOG.warning("Creating indexes on Tiger data")
 112     with connect(dsn) as conn:
 113         sql = SQLPreprocessor(conn, config, sqllib_dir)
 114         sql.run_sql_file(conn, 'tiger_import_finish.sql')