Make dict get_relationship_to more robust

Author: nkrishnaswami

sqlmodel/_compat.py

@@ -178,9 +178,16 @@ def get_relationship_to(
         # If a list, then also get the real field
         elif origin is list:
             use_annotation = get_args(annotation)[0]
-        # If a dict, then use the value type
-        elif origin is dict:
-            use_annotation = get_args(annotation)[1]
+        # If a dict or Mapping, then use the value (second) type argument
+        elif origin is dict or origin is Mapping:
+            args = get_args(annotation)
+            if len(args) >= 2:
+                use_annotation = args[1]
+            else:
+                raise ValueError(
+                    f"Dict/Mapping relationship field '{name}' must have both "
+                    "key and value type arguments (e.g., dict[str, Model])"
+                )
 
         return get_relationship_to(
             name=name, rel_info=rel_info, annotation=use_annotation

tests/test_attribute_keyed_dict.py

@@ -1,6 +1,8 @@
+import re
 from enum import Enum
 from typing import Dict, Optional
 
+import pytest
 from sqlalchemy.orm.collections import attribute_keyed_dict
 from sqlmodel import Field, Relationship, Session, SQLModel, create_engine
 
@@ -43,3 +45,46 @@ class Parent(SQLModel, table=True):
         assert parent.children_by_color[Color.Blue].parent_id == parent.id
         assert parent.children_by_color[Color.Blue].color == Color.Blue
         assert parent.children_by_color[Color.Blue].value == 2
+
+
+def test_dict_relationship_throws_on_missing_annotation_arg(clear_sqlmodel):
+    class Color(str, Enum):
+        Orange = "Orange"
+        Blue = "Blue"
+
+    class Child(SQLModel, table=True):
+        __tablename__ = "children"
+
+        id: Optional[int] = Field(primary_key=True, default=None)
+        parent_id: int = Field(foreign_key="parents.id")
+        color: Color
+        value: int
+
+    error_msg_re = re.escape(
+        "Dict/Mapping relationship field 'children_by_color' must have both key and value type arguments (e.g., dict[str, Model])"
+    )
+    # No type args
+    with pytest.raises(ValueError, match=error_msg_re):
+
+        class Parent(SQLModel, table=True):
+            __tablename__ = "parents"
+
+            id: Optional[int] = Field(primary_key=True, default=None)
+            children_by_color: dict[()] = Relationship(
+                sa_relationship_kwargs={
+                    "collection_class": attribute_keyed_dict("color")
+                }
+            )
+
+    # One type arg
+    with pytest.raises(ValueError, match=error_msg_re):
+
+        class Parent(SQLModel, table=True):
+            __tablename__ = "parents"
+
+            id: Optional[int] = Field(primary_key=True, default=None)
+            children_by_color: dict[Color] = Relationship(
+                sa_relationship_kwargs={
+                    "collection_class": attribute_keyed_dict("color")
+                }
+            )