X-Git-Url: https://git.openstreetmap.org./nominatim.git/blobdiff_plain/2af82975cd968ec09683ae5b16a9aa157a7f2176..36a15601174ab238888fcc19621bb647a94d819c:/nominatim/tools/tiger_data.py diff --git a/nominatim/tools/tiger_data.py b/nominatim/tools/tiger_data.py index 521d11c4..8610880f 100644 --- a/nominatim/tools/tiger_data.py +++ b/nominatim/tools/tiger_data.py @@ -1,96 +1,113 @@ +# SPDX-License-Identifier: GPL-2.0-only +# +# This file is part of Nominatim. (https://nominatim.org) +# +# Copyright (C) 2022 by the Nominatim developer community. +# For a full list of authors see the git log. """ -Functions for setting up and importing a new Nominatim database. +Functions for importing tiger data and handling tarbar and directory files """ +import csv +import io import logging import os -import time import tarfile -import selectors -from ..db.connection import connect -from ..db.async_connection import DBConnection -from ..db.sql_preprocessor import SQLPreprocessor +from nominatim.db.connection import connect +from nominatim.db.async_connection import WorkerPool +from nominatim.db.sql_preprocessor import SQLPreprocessor +from nominatim.errors import UsageError +from nominatim.indexer.place_info import PlaceInfo LOG = logging.getLogger() -def add_tiger_data(dsn, data_dir, threads, config, sqllib_dir): - """ Import tiger data from directory or tar file +def handle_tarfile_or_directory(data_dir): + """ Handles tarfile or directory for importing tiger data """ - # Handling directory or tarball file. - is_tarfile = False - if(data_dir.endswith('.tar.gz')): - is_tarfile = True - tar = tarfile.open(data_dir) - sql_files = [i for i in tar.getmembers() if i.name.endswith('.sql')] - LOG.warning(f'Found {len(sql_files)} SQL files in tarfile with path {data_dir}') - if(not len(sql_files)): - LOG.warning(f'Tiger data import selected but no files found in tarfile with path {data_dir}') - return + + tar = None + if data_dir.endswith('.tar.gz'): + try: + tar = tarfile.open(data_dir) + except tarfile.ReadError as err: + LOG.fatal("Cannot open '%s'. Is this a tar file?", data_dir) + raise UsageError("Cannot open Tiger data file.") from err + + csv_files = [i for i in tar.getmembers() if i.name.endswith('.csv')] + LOG.warning("Found %d CSV files in tarfile with path %s", len(csv_files), data_dir) + if not csv_files: + LOG.warning("Tiger data import selected but no files in tarfile's path %s", data_dir) + return None, None else: files = os.listdir(data_dir) - sql_files = [i for i in files if i.endswith('.sql')] - LOG.warning(f'Found {len(sql_files)} SQL files in path {data_dir}') - if(not len(sql_files)): - LOG.warning(f'Tiger data import selected but no files found in path {data_dir}') - return - + csv_files = [os.path.join(data_dir, i) for i in files if i.endswith('.csv')] + LOG.warning("Found %d CSV files in path %s", len(csv_files), data_dir) + if not csv_files: + LOG.warning("Tiger data import selected but no files found in path %s", data_dir) + return None, None + + return csv_files, tar + + +def handle_threaded_sql_statements(pool, fd, analyzer): + """ Handles sql statement with multiplexing + """ + lines = 0 + # Using pool of database connections to execute sql statements + + sql = "SELECT tiger_line_import(%s, %s, %s, %s, %s, %s)" + + for row in csv.DictReader(fd, delimiter=';'): + try: + address = dict(street=row['street'], postcode=row['postcode']) + args = ('SRID=4326;' + row['geometry'], + int(row['from']), int(row['to']), row['interpolation'], + PlaceInfo({'address': address}).analyze(analyzer), + analyzer.normalize_postcode(row['postcode'])) + except ValueError: + continue + pool.next_free_worker().perform(sql, args=args) + + lines += 1 + if lines == 1000: + print('.', end='', flush=True) + lines = 0 + + +def add_tiger_data(data_dir, config, threads, tokenizer): + """ Import tiger data from directory or tar file `data dir`. + """ + dsn = config.get_libpq_dsn() + files, tar = handle_tarfile_or_directory(data_dir) + + if not files: + return + with connect(dsn) as conn: - sql = SQLPreprocessor(conn, config, sqllib_dir) + sql = SQLPreprocessor(conn, config) sql.run_sql_file(conn, 'tiger_import_start.sql') - # Reading sql_files and then for each file line handling + # Reading files and then for each file line handling # sql_query in chunks. - sel = selectors.DefaultSelector() place_threads = max(1, threads - 1) - for sql_file in sql_files: - if(not is_tarfile): - file_path = os.path.join(data_dir, sql_file) - file = open(file_path) - else: - file = tar.extractfile(sql_file) - lines = 0 - end_of_file = False - total_used_threads = place_threads - while(True): - if(end_of_file): - break - for imod in range(place_threads): - conn = DBConnection(dsn) - conn.connect() - - sql_query = file.readline() - lines+=1 - - if(not sql_query): - end_of_file = True - total_used_threads = imod - break - - conn.perform(sql_query) - sel.register(conn, selectors.EVENT_READ, conn) - - if(lines==1000): - print('. ', end='', flush=True) - lines=0 - - todo = min(place_threads,total_used_threads) - while todo > 0: - for key, _ in sel.select(1): - try: - conn = key.data - sel.unregister(conn) - conn.wait() - conn.close() - todo -= 1 - except: - todo -=1 - - if(is_tarfile): + + with WorkerPool(dsn, place_threads, ignore_sql_errors=True) as pool: + with tokenizer.name_analyzer() as analyzer: + for fname in files: + if not tar: + fd = open(fname) + else: + fd = io.TextIOWrapper(tar.extractfile(fname)) + + handle_threaded_sql_statements(pool, fd, analyzer) + + fd.close() + + if tar: tar.close() print('\n') LOG.warning("Creating indexes on Tiger data") with connect(dsn) as conn: - sql = SQLPreprocessor(conn, config, sqllib_dir) + sql = SQLPreprocessor(conn, config) sql.run_sql_file(conn, 'tiger_import_finish.sql') - \ No newline at end of file