Alternative implementation: remove sa_column_fk_{args,kwargs}, allow the use to pass foreign_key=ForeignKey(...)

Author: earshinov

sqlmodel/main.py

@@ -121,8 +121,6 @@ def __init__(self, default: Any = Undefined, **kwargs: Any) -> None:
         sa_column = kwargs.pop("sa_column", Undefined)
         sa_column_args = kwargs.pop("sa_column_args", Undefined)
         sa_column_kwargs = kwargs.pop("sa_column_kwargs", Undefined)
-        sa_foreign_key_args = kwargs.pop("sa_foreign_key_args", Undefined)
-        sa_foreign_key_kwargs = kwargs.pop("sa_foreign_key_kwargs", Undefined)
         if sa_column is not Undefined:
             if sa_column_args is not Undefined:
                 raise RuntimeError(
@@ -165,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
@@ -176,8 +178,6 @@ def __init__(self, default: Any = Undefined, **kwargs: Any) -> None:
         self.sa_column = sa_column
         self.sa_column_args = sa_column_args
         self.sa_column_kwargs = sa_column_kwargs
-        self.sa_foreign_key_args = sa_foreign_key_args
-        self.sa_foreign_key_kwargs = sa_foreign_key_kwargs
 
 
 class RelationshipInfo(Representation):
@@ -252,8 +252,6 @@ def Field(
     sa_type: Union[Type[Any], UndefinedType] = Undefined,
     sa_column_args: Union[Sequence[Any], UndefinedType] = Undefined,
     sa_column_kwargs: Union[Mapping[str, Any], UndefinedType] = Undefined,
-    sa_foreign_key_args: Union[Sequence[Any], UndefinedType] = Undefined,
-    sa_foreign_key_kwargs: Union[Mapping[str, Any], UndefinedType] = Undefined,
     schema_extra: Optional[Dict[str, Any]] = None,
 ) -> Any: ...
 
@@ -390,8 +388,6 @@ def Field(
     sa_column: Union[Column, UndefinedType] = Undefined,  # type: ignore
     sa_column_args: Union[Sequence[Any], UndefinedType] = Undefined,
     sa_column_kwargs: Union[Mapping[str, Any], UndefinedType] = Undefined,
-    sa_foreign_key_args: Union[Sequence[Any], UndefinedType] = Undefined,
-    sa_foreign_key_kwargs: Union[Mapping[str, Any], UndefinedType] = Undefined,
     schema_extra: Optional[Dict[str, Any]] = None,
 ) -> Any:
     current_schema_extra = schema_extra or {}
@@ -430,8 +426,6 @@ def Field(
         sa_column=sa_column,
         sa_column_args=sa_column_args,
         sa_column_kwargs=sa_column_kwargs,
-        sa_foreign_key_args=sa_foreign_key_args,
-        sa_foreign_key_kwargs=sa_foreign_key_kwargs,
         **current_schema_extra,
     )
     post_init_field_info(field_info)
@@ -740,22 +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)
-        fk_args = []
-        fk_kwargs = {}
-        ondelete = getattr(field_info, "ondelete", Undefined)
-        if ondelete is not Undefined:
+        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)))
-            fk_kwargs["ondelete"] = ondelete
-        sa_foreign_key_args = getattr(field_info, "sa_foreign_key_args", Undefined)
-        if sa_foreign_key_args is not Undefined:
-            fk_args.extend(cast(Sequence[Any], sa_foreign_key_args))
-        sa_foreign_key_kwargs = getattr(field_info, "sa_foreign_key_kwargs", Undefined)
-        if sa_foreign_key_kwargs is not Undefined:
-            fk_kwargs.update(cast(Dict[Any, Any], sa_foreign_key_kwargs))
-        args.append(ForeignKey(foreign_key, *fk_args, **fk_kwargs))
+            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

@@ -22,9 +22,11 @@ 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:
+    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)
@@ -52,8 +54,7 @@ class User(SQLModel, table=True):
     class Base(SQLModel):
         owner_id: Optional[int] = Field(
             default=None,
-            foreign_key="user.id",
-            sa_foreign_key_kwargs={"ondelete": "SET NULL"},
+            foreign_key=ForeignKey("user.id", ondelete="SET NULL"),
         )
 
     class Asset(Base, table=True):