nominatim/tools/tiger_data.py

   1 """
   2 Functions for setting up and importing a new Nominatim database.
   3 """
   4 import logging
   5 import os
   6 import time
   7 import tarfile
   8 import selectors
   9
  10 from ..db.connection import connect
  11 from ..db.async_connection import DBConnection
  12 from ..db.sql_preprocessor import SQLPreprocessor
  13
  14 LOG = logging.getLogger()
  15
  16
  17 def add_tiger_data(dsn, data_dir, threads, config, sqllib_dir):
  18     """ Import tiger data from directory or tar file
  19     """
  20     # Handling directory or tarball file.
  21     is_tarfile = False
  22     if(data_dir.endswith('.tar.gz')):
  23         is_tarfile = True
  24         tar = tarfile.open(data_dir)
  25         sql_files = [i for i in tar.getmembers() if i.name.endswith('.sql')]
  26         LOG.warning(f'Found {len(sql_files)} SQL files in tarfile with path {data_dir}')
  27         if(not len(sql_files)):
  28             LOG.warning(f'Tiger data import selected but no files found in tarfile with path {data_dir}')
  29             return
  30     else:
  31         files = os.listdir(data_dir)
  32         sql_files = [i for i in files if i.endswith('.sql')]
  33         LOG.warning(f'Found {len(sql_files)} SQL files in path {data_dir}')
  34         if(not len(sql_files)):
  35             LOG.warning(f'Tiger data import selected but no files found in path {data_dir}')
  36             return
  37
  38     with connect(dsn) as conn:
  39         sql = SQLPreprocessor(conn, config, sqllib_dir)
  40         sql.run_sql_file(conn, 'tiger_import_start.sql')
  41
  42     # Reading sql_files and then for each file line handling
  43     # sql_query in <threads - 1> chunks.
  44     sel = selectors.DefaultSelector()
  45     place_threads = max(1, threads - 1)
  46     for sql_file in sql_files:
  47         if(not is_tarfile):
  48             file_path = os.path.join(data_dir, sql_file)
  49             file = open(file_path)
  50         else:
  51             file = tar.extractfile(sql_file)
  52         lines = 0
  53         end_of_file = False
  54         total_used_threads = place_threads
  55         while(True):
  56             if(end_of_file):
  57                 break
  58             for imod in range(place_threads):
  59                 conn = DBConnection(dsn)
  60                 conn.connect()
  61
  62                 sql_query = file.readline()
  63                 lines+=1
  64
  65                 if(not sql_query):
  66                     end_of_file = True
  67                     total_used_threads = imod
  68                     break
  69
  70                 conn.perform(sql_query)
  71                 sel.register(conn, selectors.EVENT_READ, conn)
  72
  73                 if(lines==1000):
  74                     print('. ', end='', flush=True)
  75                     lines=0
  76
  77             todo = min(place_threads,total_used_threads)
  78             while todo > 0:
  79                 for key, _ in sel.select(1):
  80                     try:
  81                         conn = key.data
  82                         sel.unregister(conn)
  83                         conn.wait()
  84                         conn.close()
  85                         todo -= 1
  86                     except:
  87                         todo -=1
  88
  89     if(is_tarfile):
  90         tar.close()
  91     print('\n')
  92     LOG.warning("Creating indexes on Tiger data")
  93     with connect(dsn) as conn:
  94         sql = SQLPreprocessor(conn, config, sqllib_dir)
  95         sql.run_sql_file(conn, 'tiger_import_finish.sql')
  96