Preserve primary key when loading models

aredis_om/model/model.py

@@ -3232,7 +3232,8 @@ async def get(cls: Type["Model"], pk: Any) -> "Model":
             document = convert_base64_to_bytes(document, cls.model_fields)
             # Convert bytes back to list[float] for vector fields
             document = convert_bytes_to_vector(document, cls.model_fields)
-            result = cls.model_validate(document)
+            document_with_pk = {**document, cls._meta.primary_key.name: pk}
+            result = cls.model_validate(document_with_pk)
         except TypeError as e:
             log.warning(
                 f'Could not parse Redis response. Error was: "{e}". Probably, the '
@@ -3249,6 +3250,7 @@ async def get(cls: Type["Model"], pk: Any) -> "Model":
             document = convert_base64_to_bytes(document, cls.model_fields)
             # Convert bytes back to list[float] for vector fields
             document = convert_bytes_to_vector(document, cls.model_fields)
+            document[cls._meta.primary_key.name] = pk
             result = cls.model_validate(document)
         return result
 
@@ -3587,6 +3589,7 @@ async def get(cls: Type["Model"], pk: Any) -> "Model":
         document_data = await cls.db().json().get(cls.make_key(pk))
         if document_data is None:
             raise NotFoundError
+        document_data[cls._meta.primary_key.name] = pk
         # Convert timestamps back to datetime objects before validation
         document_data = convert_timestamp_to_datetime(document_data, cls.model_fields)
         # Convert base64 strings back to bytes for bytes fields

tests/test_primary_key_loading.py

@@ -0,0 +1,75 @@
+from unittest import mock
+
+from aredis_om import Field, HashModel, JsonModel
+
+from .conftest import py_test_mark_asyncio
+
+
+@py_test_mark_asyncio
+async def test_hash_model_get_uses_requested_pk_when_document_pk_is_missing():
+    fake_db = mock.AsyncMock()
+    fake_db.hgetall.return_value = {"name": "legacy"}
+
+    class LegacyModel(HashModel, index=True):
+        name: str
+
+        class Meta:
+            database = fake_db
+
+    model = await LegacyModel.get("existing-pk")
+
+    assert model.pk == "existing-pk"
+
+
+@py_test_mark_asyncio
+async def test_hash_model_get_decodes_bytes_before_adding_requested_pk():
+    fake_db = mock.AsyncMock()
+    fake_db.hgetall.return_value = {b"name": b"legacy"}
+
+    class LegacyModel(HashModel, index=True):
+        name: str
+
+        class Meta:
+            database = fake_db
+
+    model = await LegacyModel.get("existing-pk")
+
+    assert model.pk == "existing-pk"
+    assert model.name == "legacy"
+
+
+@py_test_mark_asyncio
+async def test_json_model_get_uses_requested_pk_when_document_pk_is_missing():
+    json_client = mock.Mock()
+    json_client.get = mock.AsyncMock(return_value={"name": "legacy"})
+    fake_db = mock.Mock()
+    fake_db.json.return_value = json_client
+    fake_db.execute_command.return_value = [object()]
+
+    class LegacyModel(JsonModel, index=True):
+        name: str
+
+        class Meta:
+            database = fake_db
+
+    model = await LegacyModel.get("existing-pk")
+
+    assert model.pk == "existing-pk"
+
+
+@py_test_mark_asyncio
+async def test_get_uses_configured_custom_primary_key():
+    fake_db = mock.AsyncMock()
+    fake_db.hgetall.return_value = {"name": "legacy"}
+
+    class LegacyModel(HashModel, index=True):
+        external_id: int = Field(primary_key=True)
+        name: str
+
+        class Meta:
+            database = fake_db
+
+    model = await LegacyModel.get(42)
+
+    assert model.pk == 42
+    assert model.external_id == 42