Add foreign keys to alter table API

- Add add_foreign_key, drop_foreign_key, and set_foreign_keys operations.
- Validate flat fk_table and fk_column arguments with Pydantic.
- Document the API and cover inferred primary-key and validation cases.

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)

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

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"])