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 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
  65 def add_tiger_data(dsn, data_dir, threads, config, sqllib_dir):
  66     """ Import tiger data from directory or tar file
  67     """
  68
  69     sql_files, tar = handle_tarfile_or_directory(data_dir)
  70
  71     if not sql_files:
  72         return
  73
  74     with connect(dsn) as conn:
  75         sql = SQLPreprocessor(conn, config, sqllib_dir)
  76         sql.run_sql_file(conn, 'tiger_import_start.sql')
  77
  78     # Reading sql_files and then for each file line handling
  79     # sql_query in <threads - 1> chunks.
  80     sel = selectors.DefaultSelector()
  81     place_threads = max(1, threads - 1)
  82
  83     # Creates a pool of database connections
  84     for _ in range(place_threads):
  85         conn = DBConnection(dsn)
  86         conn.connect()
  87         sel.register(conn, selectors.EVENT_WRITE, conn)
  88
  89     for sql_file in sql_files:
  90         if not tar:
  91             file = open(sql_file)
  92         else:
  93             file = tar.extractfile(sql_file)
  94
  95         handle_threaded_sql_statements(sel, file)
  96
  97     # Unregistering pool of database connections
  98     while place_threads > 0:
  99         for key, _ in sel.select(1):
 100             conn = key.data
 101             sel.unregister(conn)
 102             conn.wait()
 103             conn.close()
 104             place_threads -= 1
 105
 106     if tar:
 107         tar.close()
 108     print('\n')
 109     LOG.warning("Creating indexes on Tiger data")
 110     with connect(dsn) as conn:
 111         sql = SQLPreprocessor(conn, config, sqllib_dir)
 112         sql.run_sql_file(conn, 'tiger_import_finish.sql')