[Fix] AJAX fields: persist selected values after validation errors (#1039)

* * Add `format_by_pk`
* Fix `form` in context after form validation
* Now, even with override Field, the loader field will be used (for example, for custom AJAX widgets)
* If the loader field is None, an empty list will be returned instead of an error.
* If the model has not been transferred, it will be automatically loaded by the PK field.

* fix: enhance AJAX loader with improved error handling and logging and fix composite pk formatting

---------

Co-authored-by: maxim-f1 <“maxconal228@gmail.com”>
Co-authored-by: Miradil Zeynalli <miradil.zeynalli@gmail.com>

Author: 3 people

sqladmin/ajax.py

@@ -4,7 +4,11 @@
 
 from sqlalchemy import String, cast, inspect, or_, select
 
-from sqladmin.helpers import get_object_identifier, get_primary_keys
+from sqladmin.helpers import (
+    get_object_identifier,
+    get_primary_keys,
+    object_identifier_values,
+)
 
 if TYPE_CHECKING:
     from sqladmin.models import ModelView
@@ -61,6 +65,39 @@ def format(self, model: type) -> dict[str, Any]:
 
         return {"id": str(get_object_identifier(model)), "text": str(model)}
 
+    async def format_by_pk(self, pk: Any) -> dict[str, Any]:
+        if pk is None:
+            return {}
+
+        stmt = select(self.model)
+        primary_keys = tuple(inspect(self.model).primary_key)
+
+        try:
+            values = object_identifier_values(str(pk), self.model)
+        except (TypeError, ValueError):
+            return {}
+
+        if len(values) != len(primary_keys):
+            return {}
+
+        conditions = [field == value for field, value in zip(primary_keys, values)]
+        stmt = stmt.where(*conditions)
+
+        if self.order_by:
+            if isinstance(self.order_by, list):
+                for o in self.order_by:
+                    stmt = stmt.order_by(o)
+            else:
+                stmt = stmt.order_by(self.order_by)
+
+        stmt = stmt.limit(1)
+
+        result = await self.model_admin._run_query(stmt)
+        if len(result) < 1:
+            return {}
+
+        return {"id": str(get_object_identifier(result[0])), "text": str(result[0])}
+
     async def get_list(self, term: str) -> list[Any]:
         stmt = select(self.model)
 

sqladmin/application.py

@@ -704,8 +704,9 @@ async def edit(self, request: Request) -> Response:
 
         form_data = await self._handle_form_data(request, model)
         form = Form(form_data)
+        context["form"] = form
+
         if not form.validate():
-            context["form"] = form
             return await self.templates.TemplateResponse(
                 request, model_view.edit_template, context, status_code=400
             )

sqladmin/fields.py

@@ -299,7 +299,7 @@ def pre_validate(self, form: Form) -> None:
 
 
 class AjaxSelectField(fields.SelectFieldBase):
-    widget = sqladmin_widgets.AjaxSelect2Widget()
+    widget = sqladmin_widgets.AjaxSelect2Widget()  # type: ignore[assignment]
     separator = ","
 
     def __init__(

sqladmin/forms.py

@@ -314,6 +314,9 @@ async def convert(
             if not issubclass(override, Field):
                 raise TypeError("Expected Field, got %s" % type(override))
 
+            if loader:
+                kwargs.setdefault("loader", loader)
+
             return override(**kwargs)
 
         multiple = (

sqladmin/widgets.py

@@ -1,6 +1,7 @@
 # mypy: disable-error-code="override"
 
 import json
+import logging
 from typing import TYPE_CHECKING, Any
 
 from markupsafe import Markup
@@ -17,6 +18,8 @@
     "Select2TagsWidget",
 ]
 
+logger = logging.getLogger(__name__)
+
 
 class DatePickerWidget(widgets.TextInput):
     """
@@ -43,7 +46,7 @@ def __init__(self, multiple: bool = False):
         self.multiple = multiple
         self.lookup_url = ""
 
-    def __call__(self, field: "AjaxSelectField", **kwargs: Any) -> Markup:
+    async def __call__(self, field: "AjaxSelectField", **kwargs: Any) -> Markup:
         kwargs.setdefault("data-role", "select2-ajax")
         kwargs.setdefault("data-url", field.loader.model_admin.ajax_lookup_url)
 
@@ -55,14 +58,59 @@ def __call__(self, field: "AjaxSelectField", **kwargs: Any) -> Markup:
         kwargs.setdefault("type", "hidden")
 
         if self.multiple:
-            result = [field.loader.format(value) for value in field.data]
+            result = []
+            for value in field.data:
+                try:
+                    result.append(field.loader.format(value))
+                    continue
+                except Exception:
+                    logger.debug(
+                        "Fallback to format_by_pk for ajax value=%r",
+                        value,
+                        exc_info=True,
+                    )
+
+                try:
+                    result_value = await field.loader.format_by_pk(value)
+                    if result_value == {}:
+                        continue
+                    else:
+                        result.append(result_value)
+                except Exception:
+                    logger.debug(
+                        "Unable to resolve ajax value by pk for field=%s value=%r",
+                        field.name,
+                        value,
+                        exc_info=True,
+                    )
+
             kwargs.setdefault("data-json", json.dumps(result))
             kwargs.setdefault("multiple", "1")
+
         else:
             try:
                 data = field.loader.format(field.data)
             except Exception:
-                data = None
+                logger.debug(
+                    "Fallback to format_by_pk for ajax field=%s value=%r",
+                    field.name,
+                    field.data,
+                    exc_info=True,
+                )
+                try:
+                    data = await field.loader.format_by_pk(field.data)
+                    if data == {}:
+                        data = None
+                except Exception:
+                    logger.debug(
+                        "Unable to resolve ajax single value by pk for "
+                        "field=%s value=%r",
+                        field.name,
+                        field.data,
+                        exc_info=True,
+                    )
+                    data = None
+
             if data:
                 kwargs.setdefault("data-json", json.dumps([data]))
 

tests/test_ajax.py

@@ -8,7 +8,7 @@
 from starlette.applications import Starlette
 
 from sqladmin import Admin, ModelView
-from sqladmin.ajax import create_ajax_loader
+from sqladmin.ajax import QueryAjaxModelLoader, create_ajax_loader
 from tests.common import async_engine as engine
 
 pytestmark = pytest.mark.anyio
@@ -77,6 +77,40 @@ def __str__(self) -> str:
         return f"Room {self.id}"
 
 
+class Team(Base):
+    __tablename__ = "teams"
+
+    id = Column(Integer, primary_key=True)
+    name = Column(String(length=32), nullable=False)
+
+    def __str__(self) -> str:
+        return f"Team {self.id}"
+
+
+class Member(Base):
+    __tablename__ = "members"
+
+    id = Column(Integer, primary_key=True)
+    name = Column(String(length=32), nullable=False)
+    team_id = Column(Integer, ForeignKey("teams.id"))
+
+    team = relationship("Team")
+
+    def __str__(self) -> str:
+        return f"Member {self.id}"
+
+
+class CompositeTag(Base):
+    __tablename__ = "composite_tags"
+
+    key = Column(String(length=16), primary_key=True)
+    locale = Column(String(length=8), primary_key=True)
+    label = Column(String(length=32), nullable=False)
+
+    def __str__(self) -> str:
+        return f"{self.label}:{self.key}:{self.locale}"
+
+
 class UserAdmin(ModelView, model=User):
     form_ajax_refs = {
         "addresses": {
@@ -110,9 +144,18 @@ class RoomAdmin(ModelView, model=Room):
     }
 
 
+class MemberAdmin(ModelView, model=Member):
+    form_ajax_refs = {
+        "team": {
+            "fields": ("name",),
+        }
+    }
+
+
 admin.add_view(UserAdmin)
 admin.add_view(AddressAdmin)
 admin.add_view(RoomAdmin)
+admin.add_view(MemberAdmin)
 
 
 @pytest.fixture(autouse=True)
@@ -330,3 +373,73 @@ async def test_create_and_edit_forms(client: AsyncClient) -> None:
 
     user = result.scalar_one()
     assert len(user.addresses) == 2
+
+
+async def test_edit_validation_error_preserves_selected_ajax_value(
+    client: AsyncClient,
+) -> None:
+    async with session_maker() as s:
+        s.add_all([Team(name="A"), Team(name="B")])
+        await s.commit()
+
+    async with session_maker() as s:
+        member = Member(name="John", team_id=1)
+        s.add(member)
+        await s.commit()
+
+    response = await client.post(
+        "/admin/member/edit/1",
+        data={"name": "", "team": "2"},
+    )
+
+    assert response.status_code == 400
+    assert (
+        'data-json="[{"id": "2", "text": "Team 2"}]"'
+        in response.text
+    )
+
+
+async def test_format_by_pk_single_pk() -> None:
+    async with session_maker() as s:
+        user = User(name="Arya")
+        s.add(user)
+        await s.commit()
+
+    loader = QueryAjaxModelLoader(
+        name="user",
+        model=User,
+        model_admin=UserAdmin(),
+        fields=("name",),
+    )
+
+    assert await loader.format_by_pk(1) == {"id": "1", "text": "User 1"}
+
+
+async def test_format_by_pk_composite_pk_identifier() -> None:
+    async with session_maker() as s:
+        tag = CompositeTag(key="greeting", locale="en", label="Hello")
+        s.add(tag)
+        await s.commit()
+
+    loader = QueryAjaxModelLoader(
+        name="composite",
+        model=CompositeTag,
+        model_admin=UserAdmin(),
+        fields=("label",),
+    )
+
+    assert await loader.format_by_pk("greeting;en") == {
+        "id": "greeting;en",
+        "text": "Hello:greeting:en",
+    }
+
+
+async def test_format_by_pk_returns_empty_for_missing_record() -> None:
+    loader = QueryAjaxModelLoader(
+        name="user",
+        model=User,
+        model_admin=UserAdmin(),
+        fields=("name",),
+    )
+
+    assert await loader.format_by_pk("999") == {}

tests/test_forms/test_forms.py

@@ -32,6 +32,7 @@
 from wtforms.fields.core import UnboundField
 
 from sqladmin import ModelView
+from sqladmin.ajax import create_ajax_loader
 from sqladmin.fields import Select2TagsField, SelectField
 from sqladmin.forms import ModelConverter, converts, get_model_form
 from tests.common import async_engine as engine
@@ -233,6 +234,32 @@ class ExampleField(Field):
     assert not isinstance(Form()._fields["email"], ExampleField)
 
 
+async def test_model_form_override_receives_ajax_loader() -> None:
+    class LoaderAwareField(Field):
+        def __init__(self, *args: Any, loader=None, **kwargs: Any) -> None:
+            self.loader = loader
+            kwargs.pop("allow_blank", None)
+            super().__init__(*args, **kwargs)
+
+    class AddressAdmin(ModelView, model=Address):
+        form_ajax_refs = {"user": {"fields": ("name",)}}
+
+    loader = create_ajax_loader(
+        model_admin=AddressAdmin(),
+        name="user",
+        options={"fields": ("name",)},
+    )
+
+    Form = await get_model_form(
+        model=Address,
+        session_maker=session_maker,
+        form_overrides={"user": LoaderAwareField},
+        form_ajax_refs={"user": loader},
+    )
+
+    assert Form()._fields["user"].loader is loader
+
+
 @pytest.mark.skipif(engine.name != "postgresql", reason="PostgreSQL only")
 async def test_model_form_postgresql() -> None:
     class PostgresModel(Base):