diff --git a/etlhelper/etl.py b/etlhelper/etl.py index 9ea91f3..2ec922a 100644 --- a/etlhelper/etl.py +++ b/etlhelper/etl.py @@ -482,6 +482,8 @@ def copy_table_rows(table, source_conn, dest_conn, target=None, e.g. WHERE, ORDER BY, LIMIT :return processed, failed: (int, int) number of rows processed, failed """ + validate_identifier(table) + select_query = f"SELECT * FROM {table}" if not target: target = table diff --git a/test/integration/etl/test_etl_transform.py b/test/integration/etl/test_etl_transform.py index 51abad3..fec94de 100644 --- a/test/integration/etl/test_etl_transform.py +++ b/test/integration/etl/test_etl_transform.py @@ -13,6 +13,7 @@ iter_rows, ) from etlhelper.row_factories import dict_row_factory +from etlhelper.exceptions import ETLHelperBadIdentifierError def test_copy_rows_happy_path(pgtestdb_conn, pgtestdb_test_tables, @@ -74,6 +75,13 @@ def test_copy_table_rows_on_error(pgtestdb_test_tables, pgtestdb_conn, assert result[1:] == test_table_data[1:] +def test_copy_table_rows_bad_table(pgtestdb_conn, pgtestdb_test_tables, + pgtestdb_insert_sql, test_table_data): + # Arrange and act + with pytest.raises(ETLHelperBadIdentifierError): + copy_table_rows('bad; sql', pgtestdb_conn, pgtestdb_conn, target='dest') + + def transform_return_list(rows): # Simple transform function that changes size and number of rows return [(row.id, row.value) for row in rows if row.id > 1]