diff --git a/datasette/views/table_create_alter.py b/datasette/views/table_create_alter.py index e8264a6f..20dcc03e 100644 --- a/datasette/views/table_create_alter.py +++ b/datasette/views/table_create_alter.py @@ -292,6 +292,40 @@ class ReorderColumnsArgs(_StrictPydanticModel): columns: list[str] = Field(min_length=1) +class ForeignKeyArgs(_StrictPydanticModel): + column: str + fk_table: str | None = None + fk_column: str | None = None + + @model_validator(mode="after") + def validate_foreign_key(self): + if self.fk_column and not self.fk_table: + raise PydanticCustomError( + "alter_table_foreign_key", + "fk_column requires fk_table", + ) + if not self.fk_table: + raise PydanticCustomError( + "alter_table_foreign_key", + "fk_table is required", + ) + return self + + @property + def tuple(self): + if self.fk_column: + return (self.column, self.fk_table, self.fk_column) + return (self.column, self.fk_table) + + +class DropForeignKeyArgs(_StrictPydanticModel): + column: str + + +class SetForeignKeysArgs(_StrictPydanticModel): + foreign_keys: list[ForeignKeyArgs] + + class AddColumnOperation(_StrictPydanticModel): op: Literal["add_column"] args: AddColumnArgs @@ -322,6 +356,21 @@ class ReorderColumnsOperation(_StrictPydanticModel): args: ReorderColumnsArgs +class AddForeignKeyOperation(_StrictPydanticModel): + op: Literal["add_foreign_key"] + args: ForeignKeyArgs + + +class DropForeignKeyOperation(_StrictPydanticModel): + op: Literal["drop_foreign_key"] + args: DropForeignKeyArgs + + +class SetForeignKeysOperation(_StrictPydanticModel): + op: Literal["set_foreign_keys"] + args: SetForeignKeysArgs + + AlterTableOperation = Annotated[ Union[ AddColumnOperation, @@ -330,6 +379,9 @@ AlterTableOperation = Annotated[ DropColumnOperation, SetPrimaryKeyOperation, ReorderColumnsOperation, + AddForeignKeyOperation, + DropForeignKeyOperation, + SetForeignKeysOperation, ], Field(discriminator="op"), ] @@ -615,6 +667,9 @@ class TableAlterView(BaseView): defaults = {} column_order = None pk = SQLITE_UTILS_DEFAULT + add_foreign_keys = [] + drop_foreign_keys = [] + foreign_keys = None for operation in alter_request.operations: args = operation.args @@ -664,6 +719,12 @@ class TableAlterView(BaseView): pk = _primary_key_value(args.columns) elif operation.op == "reorder_columns": column_order = args.columns + elif operation.op == "add_foreign_key": + add_foreign_keys.append(args.tuple) + elif operation.op == "drop_foreign_key": + drop_foreign_keys.append(args.column) + elif operation.op == "set_foreign_keys": + foreign_keys = [fk.tuple for fk in args.foreign_keys] with operation_conn: for column in add_columns: @@ -692,6 +753,9 @@ class TableAlterView(BaseView): defaults, column_order is not None, pk is not SQLITE_UTILS_DEFAULT, + add_foreign_keys, + drop_foreign_keys, + foreign_keys is not None, ) ) if should_transform: @@ -703,6 +767,9 @@ class TableAlterView(BaseView): not_null=not_null or None, defaults=defaults or None, column_order=column_order, + add_foreign_keys=add_foreign_keys or None, + drop_foreign_keys=drop_foreign_keys or None, + foreign_keys=foreign_keys, ) return _table_schema_from_conn(operation_conn, table_name) diff --git a/docs/json_api.rst b/docs/json_api.rst index 1b4a196e..af16626f 100644 --- a/docs/json_api.rst +++ b/docs/json_api.rst @@ -2159,6 +2159,31 @@ The request body should include an ``operations`` array. Each operation has the "columns": ["id"] } }, + { + "op": "add_foreign_key", + "args": { + "column": "owner_id", + "fk_table": "owners" + } + }, + { + "op": "drop_foreign_key", + "args": { + "column": "old_owner_id" + } + }, + { + "op": "set_foreign_keys", + "args": { + "foreign_keys": [ + { + "column": "owner_id", + "fk_table": "owners", + "fk_column": "id" + } + ] + } + }, { "op": "reorder_columns", "args": { @@ -2177,10 +2202,15 @@ Supported operations: * ``alter_column`` changes column properties. ``args`` accepts ``name`` and at least one of ``type``, ``not_null``, literal ``default`` or ``default_expr``. Passing ``"default": null`` removes an existing default. * ``drop_column`` drops a column. ``args`` accepts ``name``. * ``set_primary_key`` changes the table primary key. ``args`` accepts ``columns``, a list of one or more column names. +* ``add_foreign_key`` adds a single-column foreign key constraint. ``args`` accepts ``column``, ``fk_table`` and optional ``fk_column``. If ``fk_column`` is omitted, Datasette will use the single primary key of ``fk_table``. +* ``drop_foreign_key`` removes the foreign key constraint for a column. ``args`` accepts ``column``. +* ``set_foreign_keys`` replaces all foreign key constraints on the table. ``args`` accepts ``foreign_keys``, a list of objects that each have ``column``, ``fk_table`` and optional ``fk_column``. An empty list removes all foreign key constraints. * ``reorder_columns`` reorders columns. ``args`` accepts ``columns``, a list of one or more column names. Columns omitted from this list will appear afterwards in their existing order. ``default`` is always treated as a literal value. ``default_expr`` accepts one of ``current_timestamp``, ``current_date`` or ``current_time`` and is rendered as the corresponding SQLite default expression. +For foreign key operations that omit ``fk_column``, the referenced ``fk_table`` must have a single-column primary key. Datasette will return an error if it cannot identify a single primary key column for that table. + A successful response returns the new schema and the previous schema: .. code-block:: json diff --git a/tests/test_api_write.py b/tests/test_api_write.py index 627b1ac1..046cf695 100644 --- a/tests/test_api_write.py +++ b/tests/test_api_write.py @@ -926,6 +926,124 @@ async def test_alter_table_dry_run(ds_write): assert last_event(ds_write) is None +@pytest.mark.asyncio +async def test_alter_table_foreign_key_operations(ds_write): + token = write_token(ds_write, permissions=["at"]) + db = ds_write.get_database("data") + await db.execute_write("create table owners (id integer primary key)") + await db.execute_write("create table categories (id integer primary key)") + + response = await ds_write.client.post( + "/data/docs/-/alter", + json={ + "operations": [ + {"op": "add_column", "args": {"name": "owner_id", "type": "integer"}}, + { + "op": "add_foreign_key", + "args": {"column": "owner_id", "fk_table": "owners"}, + }, + ] + }, + headers=_headers(token), + ) + assert response.status_code == 200, response.text + data = response.json() + assert data["operations_applied"] == 2 + assert "[owner_id] INTEGER REFERENCES [owners]([id])" in data["schema"] + + response = await ds_write.client.post( + "/data/docs/-/alter", + json={ + "operations": [{"op": "drop_foreign_key", "args": {"column": "owner_id"}}] + }, + headers=_headers(token), + ) + assert response.status_code == 200, response.text + data = response.json() + assert "[owner_id] INTEGER REFERENCES" not in data["schema"] + + response = await ds_write.client.post( + "/data/docs/-/alter", + json={ + "operations": [ + { + "op": "set_foreign_keys", + "args": { + "foreign_keys": [ + { + "column": "owner_id", + "fk_table": "categories", + "fk_column": "id", + } + ] + }, + } + ] + }, + headers=_headers(token), + ) + assert response.status_code == 200, response.text + data = response.json() + assert "[owner_id] INTEGER REFERENCES [categories]([id])" in data["schema"] + + response = await ds_write.client.post( + "/data/docs/-/alter", + json={"operations": [{"op": "set_foreign_keys", "args": {"foreign_keys": []}}]}, + headers=_headers(token), + ) + assert response.status_code == 200, response.text + data = response.json() + assert "[owner_id] INTEGER REFERENCES" not in data["schema"] + + +@pytest.mark.asyncio +async def test_alter_table_foreign_key_requires_fk_table_for_fk_column(ds_write): + response = await ds_write.client.post( + "/data/docs/-/alter", + json={ + "operations": [ + { + "op": "add_foreign_key", + "args": {"column": "age", "fk_column": "id"}, + } + ] + }, + headers=_headers(write_token(ds_write, permissions=["at"])), + ) + assert response.status_code == 400 + assert response.json() == { + "ok": False, + "errors": ["operations.0.add_foreign_key.args: fk_column requires fk_table"], + } + + +@pytest.mark.asyncio +async def test_alter_table_foreign_key_without_fk_column_requires_single_pk(ds_write): + token = write_token(ds_write, permissions=["at"]) + db = ds_write.get_database("data") + await db.execute_write( + "create table accounts (tenant_id integer, id integer, primary key (tenant_id, id))" + ) + + response = await ds_write.client.post( + "/data/docs/-/alter", + json={ + "operations": [ + { + "op": "add_foreign_key", + "args": {"column": "age", "fk_table": "accounts"}, + } + ] + }, + headers=_headers(token), + ) + assert response.status_code == 400 + assert response.json() == { + "ok": False, + "errors": ["Could not detect single primary key for table 'accounts'"], + } + + @pytest.mark.asyncio async def test_alter_table_permission_denied(ds_write): token = write_token(ds_write, permissions=["ir"])