From 0e797033f96a1c61b173a3d8af2ff36905687a2e Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sat, 19 Jun 2021 08:28:26 -0700 Subject: [PATCH] .transform() on rowid (non-pk) tables bug fix, closes #284 --- sqlite_utils/db.py | 4 ++- tests/test_cli.py | 4 +-- tests/test_cli_memory.py | 17 ++++----- tests/test_extract.py | 15 +++++++- tests/test_transform.py | 76 ++++++++++++++++++++++++++++++++++++++-- 5 files changed, 100 insertions(+), 16 deletions(-) diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 2d8fb24..9da1b97 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -1016,7 +1016,9 @@ class Table(Queryable): sqls = [] if pk is DEFAULT: - pks_renamed = tuple(rename.get(p) or p for p in self.pks) + pks_renamed = tuple( + rename.get(p.name) or p.name for p in self.columns if p.is_pk + ) if len(pks_renamed) == 1: pk = pks_renamed[0] else: diff --git a/tests/test_cli.py b/tests/test_cli.py index 8ddf71b..8d59c0c 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -2088,8 +2088,8 @@ def test_insert_detect_types(tmpdir, option_or_env_var): assert result.exit_code == 0 db = Database(db_path) assert list(db["creatures"].rows) == [ - {"rowid": 1, "name": "Cleo", "age": 6, "weight": 45.5}, - {"rowid": 2, "name": "Dori", "age": 1, "weight": 3.5}, + {"name": "Cleo", "age": 6, "weight": 45.5}, + {"name": "Dori", "age": 1, "weight": 3.5}, ] if option_or_env_var is None: diff --git a/tests/test_cli_memory.py b/tests/test_cli_memory.py index 2a1fb85..6966d05 100644 --- a/tests/test_cli_memory.py +++ b/tests/test_cli_memory.py @@ -32,8 +32,7 @@ def test_memory_csv(tmpdir, sql_from, use_stdin): ) assert result.exit_code == 0 assert ( - result.output.strip() - == '{"rowid": 1, "id": 1, "name": "Cleo"}\n{"rowid": 2, "id": 2, "name": "Bants"}' + result.output.strip() == '{"id": 1, "name": "Cleo"}\n{"id": 2, "name": "Bants"}' ) @@ -57,8 +56,8 @@ def test_memory_tsv(tmpdir, use_stdin): ) assert result.exit_code == 0, result.output assert json.loads(result.output.strip()) == [ - {"rowid": 1, "id": 1, "name": "Cleo"}, - {"rowid": 2, "id": 2, "name": "Bants"}, + {"id": 1, "name": "Cleo"}, + {"id": 2, "name": "Bants"}, ] @@ -146,7 +145,6 @@ def test_memory_csv_encoding(tmpdir, use_stdin): ) assert result.exit_code == 0, result.output assert json.loads(result.output.strip()) == { - "rowid": 1, "date": "2020-03-04", "name": "São Paulo", "latitude": -23.561, @@ -165,12 +163,11 @@ def test_memory_dump(extra_args): assert result.output.strip() == ( "BEGIN TRANSACTION;\n" 'CREATE TABLE "stdin" (\n' - " [rowid] INTEGER PRIMARY KEY,\n" " [id] INTEGER,\n" " [name] TEXT\n" ");\n" - "INSERT INTO \"stdin\" VALUES(1,1,'Cleo');\n" - "INSERT INTO \"stdin\" VALUES(2,2,'Bants');\n" + "INSERT INTO \"stdin\" VALUES(1,'Cleo');\n" + "INSERT INTO \"stdin\" VALUES(2,'Bants');\n" "CREATE VIEW t1 AS select * from [stdin];\n" "CREATE VIEW t AS select * from [stdin];\n" "COMMIT;" @@ -188,8 +185,8 @@ def test_memory_save(tmpdir, extra_args): assert result.exit_code == 0 db = Database(save_to) assert list(db["stdin"].rows) == [ - {"rowid": 1, "id": 1, "name": "Cleo"}, - {"rowid": 2, "id": 2, "name": "Bants"}, + {"id": 1, "name": "Cleo"}, + {"id": 2, "name": "Bants"}, ] diff --git a/tests/test_extract.py b/tests/test_extract.py index 9eae704..280c59c 100644 --- a/tests/test_extract.py +++ b/tests/test_extract.py @@ -126,12 +126,25 @@ def test_extract_rowid_table(fresh_db): fresh_db["tree"].extract(["common_name", "latin_name"]) assert fresh_db["tree"].schema == ( 'CREATE TABLE "tree" (\n' - " [rowid] INTEGER PRIMARY KEY,\n" " [name] TEXT,\n" " [common_name_latin_name_id] INTEGER,\n" " FOREIGN KEY([common_name_latin_name_id]) REFERENCES [common_name_latin_name]([id])\n" ")" ) + assert ( + fresh_db.execute( + """ + select + tree.name, + common_name_latin_name.common_name, + common_name_latin_name.latin_name + from tree + join common_name_latin_name + on tree.common_name_latin_name_id = common_name_latin_name.id + """ + ).fetchall() + == [("Tree 1", "Palm", "Arecaceae")] + ) def test_reuse_lookup_table(fresh_db): diff --git a/tests/test_transform.py b/tests/test_transform.py index b3ef009..19e447e 100644 --- a/tests/test_transform.py +++ b/tests/test_transform.py @@ -89,7 +89,9 @@ import pytest ], ) @pytest.mark.parametrize("use_pragma_foreign_keys", [False, True]) -def test_transform_sql(fresh_db, params, expected_sql, use_pragma_foreign_keys): +def test_transform_sql_table_with_primary_key( + fresh_db, params, expected_sql, use_pragma_foreign_keys +): captured = [] tracer = lambda sql, params: captured.append((sql, params)) dogs = fresh_db["dogs"] @@ -111,7 +113,77 @@ def test_transform_sql(fresh_db, params, expected_sql, use_pragma_foreign_keys): assert ("PRAGMA foreign_keys=1;", None) not in captured -def test_transform_sql_rowid_to_id(fresh_db): +@pytest.mark.parametrize( + "params,expected_sql", + [ + # Identity transform - nothing changes + ( + {}, + [ + "CREATE TABLE [dogs_new_suffix] (\n [id] INTEGER,\n [name] TEXT,\n [age] TEXT\n);", + "INSERT INTO [dogs_new_suffix] ([id], [name], [age])\n SELECT [id], [name], [age] FROM [dogs];", + "DROP TABLE [dogs];", + "ALTER TABLE [dogs_new_suffix] RENAME TO [dogs];", + ], + ), + # Change column type + ( + {"types": {"age": int}}, + [ + "CREATE TABLE [dogs_new_suffix] (\n [id] INTEGER,\n [name] TEXT,\n [age] INTEGER\n);", + "INSERT INTO [dogs_new_suffix] ([id], [name], [age])\n SELECT [id], [name], [age] FROM [dogs];", + "DROP TABLE [dogs];", + "ALTER TABLE [dogs_new_suffix] RENAME TO [dogs];", + ], + ), + # Rename a column + ( + {"rename": {"age": "dog_age"}}, + [ + "CREATE TABLE [dogs_new_suffix] (\n [id] INTEGER,\n [name] TEXT,\n [dog_age] TEXT\n);", + "INSERT INTO [dogs_new_suffix] ([id], [name], [dog_age])\n SELECT [id], [name], [age] FROM [dogs];", + "DROP TABLE [dogs];", + "ALTER TABLE [dogs_new_suffix] RENAME TO [dogs];", + ], + ), + # Make ID a primary key + ( + {"pk": "id"}, + [ + "CREATE TABLE [dogs_new_suffix] (\n [id] INTEGER PRIMARY KEY,\n [name] TEXT,\n [age] TEXT\n);", + "INSERT INTO [dogs_new_suffix] ([id], [name], [age])\n SELECT [id], [name], [age] FROM [dogs];", + "DROP TABLE [dogs];", + "ALTER TABLE [dogs_new_suffix] RENAME TO [dogs];", + ], + ), + ], +) +@pytest.mark.parametrize("use_pragma_foreign_keys", [False, True]) +def test_transform_sql_table_with_no_primary_key( + fresh_db, params, expected_sql, use_pragma_foreign_keys +): + captured = [] + tracer = lambda sql, params: captured.append((sql, params)) + dogs = fresh_db["dogs"] + if use_pragma_foreign_keys: + fresh_db.conn.execute("PRAGMA foreign_keys=ON") + dogs.insert({"id": 1, "name": "Cleo", "age": "5"}) + sql = dogs.transform_sql(**{**params, **{"tmp_suffix": "suffix"}}) + assert sql == expected_sql + # Check that .transform() runs without exceptions: + with fresh_db.tracer(tracer): + dogs.transform(**params) + # If use_pragma_foreign_keys, check that we did the right thing + if use_pragma_foreign_keys: + assert ("PRAGMA foreign_keys=0;", None) in captured + assert captured[-2] == ("PRAGMA foreign_key_check;", None) + assert captured[-1] == ("PRAGMA foreign_keys=1;", None) + else: + assert ("PRAGMA foreign_keys=0;", None) not in captured + assert ("PRAGMA foreign_keys=1;", None) not in captured + + +def test_transform_sql_with_no_primary_key_to_primary_key_of_id(fresh_db): dogs = fresh_db["dogs"] dogs.insert({"id": 1, "name": "Cleo", "age": "5"}) assert (