diff --git a/datasette/views/table_create_alter.py b/datasette/views/table_create_alter.py index ffeb7f14..94895473 100644 --- a/datasette/views/table_create_alter.py +++ b/datasette/views/table_create_alter.py @@ -772,6 +772,14 @@ def _primary_key_value(columns): return tuple(columns) +def _primary_key_columns(pk, pks): + if pks: + return list(pks) + if pk: + return [pk] + return [] + + def _default_expression_sql(default_expr): return DEFAULT_EXPR_SQL[default_expr] @@ -871,6 +879,8 @@ class TableCreateView(BaseView): return _error(["pk cannot be changed for existing table"]) pks = actual_pks + pk_columns = _primary_key_columns(pk, pks) + initial_schema = None if table_exists: initial_schema = await db.execute_fn( @@ -881,11 +891,41 @@ class TableCreateView(BaseView): db_for_write = sqlite_utils.Database(conn) table = db_for_write[table_name] if rows: + row_columns = set() + pk_columns_set = set(pk_columns) + insert_kwargs = { + "pk": pks or pk, + "ignore": ignore, + "replace": replace, + "alter": alter, + } + if not table_exists and pk_columns: + row_columns = {key for row in rows for key in row} + if pk_columns_set.issubset(row_columns): + insert_kwargs["not_null"] = pk_columns table.insert_all( - rows, pk=pks or pk, ignore=ignore, replace=replace, alter=alter + rows, + **insert_kwargs, ) + if ( + not table_exists + and pk_columns + and not pk_columns_set.issubset(row_columns) + ): + table.transform( + not_null={column: True for column in pk_columns}, + ) else: - not_null = [column.name for column in columns if column.not_null] + column_definitions = {column.name: column.type for column in columns} + if pk and pk not in column_definitions: + column_definitions = {pk: "integer", **column_definitions} + not_null = [ + column.name + for column in columns + if column.not_null or column.name in pk_columns + ] + if pk and pk not in not_null: + not_null.insert(0, pk) defaults = {} for column in columns: if "default_expr" in column.model_fields_set: @@ -900,7 +940,7 @@ class TableCreateView(BaseView): db_for_write, column.default ) table.create( - {column.name: column.type for column in columns}, + column_definitions, pk=pks or pk, foreign_keys=create_request.foreign_keys, not_null=not_null or None, @@ -1187,6 +1227,7 @@ class TableAlterView(BaseView): defaults = {} column_order = None pk = SQLITE_UTILS_DEFAULT + pk_columns = [] add_foreign_keys = [] drop_foreign_keys = [] foreign_keys = None @@ -1239,6 +1280,7 @@ class TableAlterView(BaseView): drop.add(args.name) elif operation.op == "set_primary_key": pk = _primary_key_value(args.columns) + pk_columns = args.columns elif operation.op == "reorder_columns": column_order = args.columns elif operation.op == "add_foreign_key": @@ -1281,6 +1323,8 @@ class TableAlterView(BaseView): ) ) if should_transform: + for column in pk_columns: + not_null[column] = True table.transform( types=types or None, rename=rename or None, diff --git a/tests/test_api_write.py b/tests/test_api_write.py index 76797742..563ca21e 100644 --- a/tests/test_api_write.py +++ b/tests/test_api_write.py @@ -909,6 +909,43 @@ async def test_alter_table_operations(ds_write): assert event.after_schema == data["schema"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + "table,pk_columns", + ( + ("alter_single_pk", ["id"]), + ("alter_compound_pk", ["tenant", "id"]), + ), +) +async def test_alter_table_primary_keys_are_not_null(ds_write, table, pk_columns): + token = write_token(ds_write, permissions=["at"]) + db = ds_write.get_database("data") + await db.execute_write( + "create table {} (id integer, tenant text, title text)".format( + escape_sqlite(table) + ) + ) + + response = await ds_write.client.post( + "/data/{}/-/alter".format(table), + json={ + "operations": [ + {"op": "set_primary_key", "args": {"columns": pk_columns}}, + ] + }, + headers=_headers(token), + ) + + assert response.status_code == 200, response.text + columns = ( + await db.execute( + "select * from pragma_table_info(?) where pk > 0 order by pk", [table] + ) + ).dicts() + assert [column["name"] for column in columns] == pk_columns + assert [column["notnull"] for column in columns] == [1] * len(pk_columns) + + @pytest.mark.asyncio @pytest.mark.parametrize( "default_expr,minimum_value,expected_schema", @@ -1624,7 +1661,7 @@ async def test_drop_table(ds_write, scenario): "table_api_url": "http://localhost/data/one.json", "schema": ( "CREATE TABLE [one] (\n" - " [id] INTEGER PRIMARY KEY,\n" + " [id] INTEGER PRIMARY KEY NOT NULL,\n" " [title] TEXT,\n" " [score] INTEGER,\n" " [weight] FLOAT,\n" @@ -1661,7 +1698,7 @@ async def test_drop_table(ds_write, scenario): "table_api_url": "http://localhost/data/two.json", "schema": ( "CREATE TABLE [two] (\n" - " [id] INTEGER PRIMARY KEY,\n" + " [id] INTEGER PRIMARY KEY NOT NULL,\n" " [title] TEXT,\n" " [score] FLOAT\n" ")" @@ -1690,7 +1727,7 @@ async def test_drop_table(ds_write, scenario): "table_api_url": "http://localhost/data/three.json", "schema": ( "CREATE TABLE [three] (\n" - " [id] INTEGER PRIMARY KEY,\n" + " [id] INTEGER PRIMARY KEY NOT NULL,\n" " [title] TEXT,\n" " [score] FLOAT\n" ")" @@ -1734,8 +1771,9 @@ async def test_drop_table(ds_write, scenario): "table_url": "http://localhost/data/five", "table_api_url": "http://localhost/data/five.json", "schema": ( - "CREATE TABLE [five] (\n [type] TEXT,\n [key] INTEGER,\n" - " [title] TEXT,\n PRIMARY KEY ([type], [key])\n)" + "CREATE TABLE [five] (\n [type] TEXT NOT NULL,\n" + " [key] INTEGER NOT NULL,\n [title] TEXT,\n" + " PRIMARY KEY ([type], [key])\n)" ), "row_count": 1, }, @@ -2027,6 +2065,81 @@ async def test_create_table( assert [e.name for e in events] == expected_events +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body,pk_columns", + ( + ( + { + "table": "pk_from_columns", + "columns": [ + {"name": "id", "type": "integer"}, + {"name": "title", "type": "text"}, + ], + "pk": "id", + }, + ["id"], + ), + ( + { + "table": "compound_pk_from_columns", + "columns": [ + {"name": "tenant", "type": "text"}, + {"name": "id", "type": "integer"}, + {"name": "title", "type": "text"}, + ], + "pks": ["tenant", "id"], + }, + ["tenant", "id"], + ), + ( + { + "table": "pk_omitted_from_columns", + "columns": [ + {"name": "title", "type": "text"}, + ], + "pk": "id", + }, + ["id"], + ), + ( + { + "table": "pk_from_rows", + "rows": [{"id": 1, "title": "Row 1"}], + "pk": "id", + }, + ["id"], + ), + ( + { + "table": "compound_pk_from_rows", + "row": {"tenant": "datasette", "id": 1, "title": "Row 1"}, + "pks": ["tenant", "id"], + }, + ["tenant", "id"], + ), + ), +) +async def test_create_table_primary_keys_are_not_null(ds_write, body, pk_columns): + token = write_token(ds_write) + response = await ds_write.client.post( + "/data/-/create", + json=body, + headers=_headers(token), + ) + + assert response.status_code == 201, response.text + db = ds_write.get_database("data") + columns = ( + await db.execute( + "select * from pragma_table_info(?) where pk > 0 order by pk", + [body["table"]], + ) + ).dicts() + assert [column["name"] for column in columns] == pk_columns + assert [column["notnull"] for column in columns] == [1] * len(pk_columns) + + @pytest.mark.asyncio async def test_create_table_with_foreign_key(ds_write): token = write_token(ds_write)