diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 40b734c..5d17ad2 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -245,11 +245,7 @@ class Database: ), "defaults set {} includes items not in columns {}".format( repr(set(defaults)), repr(set(columns.keys())) ) - # Validate no columns contain '[' or ']' - #86 - for column in columns.keys(): - assert ( - "[" not in column and "]" not in column - ), "'[' and ']' cannot be used in column names" + validate_column_names(columns.keys()) column_items = list(columns.items()) if column_order is not None: column_items.sort( @@ -892,6 +888,7 @@ class Table(Queryable): args = [] sets = [] wheres = [] + validate_column_names(updates.keys()) for key, value in updates.items(): sets.append("[{}] = {}".format(key, conversions.get(key, "?"))) args.append(value) @@ -1026,8 +1023,8 @@ class Table(Queryable): all_columns = list(sorted(all_columns)) if hash_id: all_columns.insert(0, hash_id) + validate_column_names(all_columns) first = False - # values is the list of insert data that is passed to the # .execute() method - but some of them may be replaced by # new primary keys if we are extracting any columns. @@ -1310,3 +1307,11 @@ def resolve_extracts(extracts): if isinstance(extracts, (list, tuple)): extracts = {item: item for item in extracts} return extracts + + +def validate_column_names(columns): + # Validate no columns contain '[' or ']' - #86 + for column in columns: + assert ( + "[" not in column and "]" not in column + ), "'[' and ']' cannot be used in column names" diff --git a/tests/test_create.py b/tests/test_create.py index 6ad54d0..92e0f8a 100644 --- a/tests/test_create.py +++ b/tests/test_create.py @@ -77,7 +77,7 @@ def test_create_table_with_bad_defaults(fresh_db): ) -def test_create_table_with_invalid_column_charactters(fresh_db): +def test_create_table_with_invalid_column_characters(fresh_db): with pytest.raises(AssertionError): fresh_db.create_table("players", {"name[foo]": str}) @@ -449,6 +449,13 @@ def test_insert_row_alter_table( ] +def test_insert_row_alter_table_invalid_column_characters(fresh_db): + table = fresh_db["table"] + rowid = table.insert({"foo": "bar"}).last_pk + with pytest.raises(AssertionError): + table.insert({"foo": "baz", "new_col[abc]": 1.2}, alter=True) + + @pytest.mark.parametrize("use_table_factory", [True, False]) def test_insert_replace_rows_alter_table(fresh_db, use_table_factory): first_row = {"id": 1, "title": "Hedgehogs of the world", "author_id": 1} diff --git a/tests/test_update.py b/tests/test_update.py index 70f8e9e..3f4c4ba 100644 --- a/tests/test_update.py +++ b/tests/test_update.py @@ -66,6 +66,13 @@ def test_update_alter(fresh_db): ] == list(table.rows) +def test_update_alter_with_invalid_column_characters(fresh_db): + table = fresh_db["table"] + rowid = table.insert({"foo": "bar"}).last_pk + with pytest.raises(AssertionError): + table.update(rowid, {"new_col[abc]": 1.2}, alter=True) + + def test_update_with_no_values_sets_last_pk(fresh_db): table = fresh_db.table("dogs", pk="id") table.insert_all([{"id": 1, "name": "Cleo"}, {"id": 2, "name": "Pancakes"}])