From: Darkshredder <srivastavayash58@gmail.com>
Date: Tue, 13 Apr 2021 09:36:02 +0000 (+0530)
Subject: Fix: Removed error if endstatement is wrong and improved tests
X-Git-Tag: v3.7.1~19
X-Git-Url: https://git.openstreetmap.org/nominatim.git/commitdiff_plain/6ed495bfc2dc935935a3e40966776ef968cd7424

Fix: Removed error if endstatement is wrong and improved tests
---

diff --git a/nominatim/tools/tiger_data.py b/nominatim/tools/tiger_data.py
index 9b960e2d..c655f91d 100644
--- a/nominatim/tools/tiger_data.py
+++ b/nominatim/tools/tiger_data.py
@@ -61,6 +61,20 @@ def handle_threaded_sql_statements(sel, file):
             except Exception as exc: # pylint: disable=broad-except
                 LOG.info('Wrong SQL statement: %s', exc)
 
+def handle_unregister_connection_pool(sel, place_threads):
+    """ Handles unregistering pool of connections
+    """
+
+    while place_threads > 0:
+        for key, _ in sel.select(1):
+            conn = key.data
+            sel.unregister(conn)
+            try:
+                conn.wait()
+            except Exception as exc: # pylint: disable=broad-except
+                LOG.info('Wrong SQL statement: %s', exc)
+            conn.close()
+            place_threads -= 1
 
 def add_tiger_data(dsn, data_dir, threads, config, sqllib_dir):
     """ Import tiger data from directory or tar file
@@ -95,13 +109,7 @@ def add_tiger_data(dsn, data_dir, threads, config, sqllib_dir):
         handle_threaded_sql_statements(sel, file)
 
     # Unregistering pool of database connections
-    while place_threads > 0:
-        for key, _ in sel.select(1):
-            conn = key.data
-            sel.unregister(conn)
-            conn.wait()
-            conn.close()
-            place_threads -= 1
+    handle_unregister_connection_pool(sel, place_threads)
 
     if tar:
         tar.close()
diff --git a/test/python/test_tools_tiger_data.py b/test/python/test_tools_tiger_data.py
index 5029a132..1366fe3e 100644
--- a/test/python/test_tools_tiger_data.py
+++ b/test/python/test_tools_tiger_data.py
@@ -16,11 +16,24 @@ def test_add_tiger_data(dsn, src_dir, def_config, tmp_path, sql_preprocessor,
     temp_db_cursor.execute('CREATE EXTENSION postgis')
     temp_db_cursor.execute('CREATE TABLE place (id INT)')
     sqlfile = tmp_path / '1010.sql'
-    sqlfile.write_text("""INSERT INTO place values (1)""")
+    sqlfile.write_text("""INSERT INTO place values (1);
+                          INSERT INTO non_existant_table values (1);""")
     tiger_data.add_tiger_data(dsn, str(tmp_path), threads, def_config, src_dir / 'lib-sql')
 
     assert temp_db_cursor.table_rows('place') == 1
 
+@pytest.mark.parametrize("threads", (1, 5))
+def test_add_tiger_data_bad_file(dsn, src_dir, def_config, tmp_path, sql_preprocessor,
+                        temp_db_cursor, threads, temp_db):
+    temp_db_cursor.execute('CREATE EXTENSION hstore')
+    temp_db_cursor.execute('CREATE EXTENSION postgis')
+    temp_db_cursor.execute('CREATE TABLE place (id INT)')
+    sqlfile = tmp_path / '1010.txt'
+    sqlfile.write_text("""Random text""")
+    tiger_data.add_tiger_data(dsn, str(tmp_path), threads, def_config, src_dir / 'lib-sql')
+
+    assert temp_db_cursor.table_rows('place') == 0
+
 @pytest.mark.parametrize("threads", (1, 5))
 def test_add_tiger_data_tarfile(dsn, src_dir, def_config, tmp_path,
                         temp_db_cursor, threads, temp_db, sql_preprocessor):
@@ -28,10 +41,26 @@ def test_add_tiger_data_tarfile(dsn, src_dir, def_config, tmp_path,
     temp_db_cursor.execute('CREATE EXTENSION postgis')
     temp_db_cursor.execute('CREATE TABLE place (id INT)')
     sqlfile = tmp_path / '1010.sql'
-    sqlfile.write_text("""INSERT INTO place values (1)""")
+    sqlfile.write_text("""INSERT INTO place values (1);
+                          INSERT INTO non_existant_table values (1);""")
+    tar = tarfile.open("sample.tar.gz", "w:gz")
+    tar.add(sqlfile)
+    tar.close()
+    tiger_data.add_tiger_data(dsn, str(src_dir / 'sample.tar.gz'), threads, def_config, src_dir / 'lib-sql')
+    
+    assert temp_db_cursor.table_rows('place') == 1
+
+@pytest.mark.parametrize("threads", (1, 5))
+def test_add_tiger_data_bad_tarfile(dsn, src_dir, def_config, tmp_path,
+                        temp_db_cursor, threads, temp_db, sql_preprocessor):
+    temp_db_cursor.execute('CREATE EXTENSION hstore')
+    temp_db_cursor.execute('CREATE EXTENSION postgis')
+    temp_db_cursor.execute('CREATE TABLE place (id INT)')
+    sqlfile = tmp_path / '1010.txt'
+    sqlfile.write_text("""Random text""")
     tar = tarfile.open("sample.tar.gz", "w:gz")
     tar.add(sqlfile)
     tar.close()
     tiger_data.add_tiger_data(dsn, str(src_dir / 'sample.tar.gz'), threads, def_config, src_dir / 'lib-sql')
     
-    assert temp_db_cursor.table_rows('place') == 1
\ No newline at end of file
+    assert temp_db_cursor.table_rows('place') == 0
\ No newline at end of file