✨Add sa_foreign_key_args and sa_foreign_key_kwargs arguments to Field…

sqlmodel/main.py

@@ -163,6 +163,10 @@ def __init__(self, default: Any = Undefined, **kwargs: Any) -> None:
         if ondelete is not Undefined:
             if foreign_key is Undefined:
                 raise RuntimeError("ondelete can only be used with foreign_key")
+            if not isinstance(foreign_key, str):
+                raise RuntimeError(
+                    "ondelete can only be used with foreign_key given as a string"
+                )
         super().__init__(default=default, **kwargs)
         self.primary_key = primary_key
         self.nullable = nullable
@@ -730,14 +734,18 @@ def get_column_from_field(field: Any) -> Column:  # type: ignore
     if unique is Undefined:
         unique = False
     if foreign_key:
-        if field_info.ondelete == "SET NULL" and not nullable:
-            raise RuntimeError('ondelete="SET NULL" requires nullable=True')
-        assert isinstance(foreign_key, str)
-        ondelete = getattr(field_info, "ondelete", Undefined)
-        if ondelete is Undefined:
-            ondelete = None
-        assert isinstance(ondelete, (str, type(None)))  # for typing
-        args.append(ForeignKey(foreign_key, ondelete=ondelete))
+        if isinstance(foreign_key, str):
+            if field_info.ondelete == "SET NULL" and not nullable:
+                raise RuntimeError('ondelete="SET NULL" requires nullable=True')
+            assert isinstance(foreign_key, str)
+            ondelete = getattr(field_info, "ondelete", Undefined)
+            if ondelete is Undefined:
+                ondelete = None
+            assert isinstance(ondelete, (str, type(None)))
+            args.append(ForeignKey(foreign_key, ondelete=ondelete))
+        else:
+            assert isinstance(foreign_key, ForeignKey)
+            args.append(foreign_key.copy())
     kwargs = {
         "primary_key": primary_key,
         "nullable": nullable,

tests/test_field_sa_fk_args_kwargs.py

@@ -0,0 +1,76 @@
+import contextlib
+import re
+from typing import Optional
+
+import pytest
+import sqlalchemy.exc
+from sqlalchemy import ForeignKey, create_engine
+from sqlmodel import Field, SQLModel
+from sqlmodel._compat import IS_PYDANTIC_V2
+
+
+def test_base_model_fk(clear_sqlmodel, caplog) -> None:
+    class User(SQLModel, table=True):
+        id: Optional[int] = Field(default=None, primary_key=True)
+
+    class Base(SQLModel):
+        owner_id: Optional[int] = Field(
+            default=None, sa_column_args=(ForeignKey("user.id", ondelete="SET NULL"),)
+        )
+
+    class Asset(Base, table=True):
+        id: Optional[int] = Field(default=None, primary_key=True)
+
+    # Fails in Pydantic v2, but not v1
+    with (
+        pytest.raises(sqlalchemy.exc.InvalidRequestError)
+        if IS_PYDANTIC_V2
+        else contextlib.nullcontext()
+    ) as e:
+
+        class Document(Base, table=True):
+            id: Optional[int] = Field(default=None, primary_key=True)
+
+    if e:
+        assert "This ForeignKey already has a parent" in str(e.errisinstance)
+
+    engine = create_engine("sqlite://", echo=True)
+    SQLModel.metadata.create_all(engine)
+
+    fk_log = [
+        message
+        for message in caplog.messages
+        if re.search(
+            r"FOREIGN KEY\s*\(owner_id\)\s*REFERENCES\s*user\s*\(id\)", message
+        )
+    ][0]
+    assert "ON DELETE SET NULL" in fk_log
+
+
+def test_base_model_fk_args(clear_sqlmodel, caplog) -> None:
+    class User(SQLModel, table=True):
+        id: Optional[int] = Field(default=None, primary_key=True)
+
+    class Base(SQLModel):
+        owner_id: Optional[int] = Field(
+            default=None,
+            foreign_key=ForeignKey("user.id", ondelete="SET NULL"),
+        )
+
+    class Asset(Base, table=True):
+        id: Optional[int] = Field(default=None, primary_key=True)
+
+    class Document(Base, table=True):
+        id: Optional[int] = Field(default=None, primary_key=True)
+
+    engine = create_engine("sqlite://", echo=True)
+    SQLModel.metadata.create_all(engine)
+
+    fk_log = [
+        message
+        for message in caplog.messages
+        if re.search(
+            r"FOREIGN KEY\s*\(owner_id\)\s*REFERENCES\s*user\s*\(id\)", message
+        )
+    ][0]
+    assert "ON DELETE SET NULL" in fk_log