diff --git a/redisvl/index/storage.py b/redisvl/index/storage.py index da9290b4a..43ecdc1b0 100644 --- a/redisvl/index/storage.py +++ b/redisvl/index/storage.py @@ -375,7 +375,9 @@ def write( ValueError: If the length of provided keys does not match the length of objects, or if validation fails. """ - if keys and len(keys) != len(objects): # type: ignore + objects = list(objects) + keys = list(keys) if keys is not None else None + if keys is not None and len(keys) != len(objects): raise ValueError("Length of keys does not match the length of objects") if batch_size is None: @@ -386,7 +388,7 @@ def write( # Pass 1: Preprocess and validate all objects prepared_objects = self._preprocess_and_validate_objects( - list(objects), # Convert Iterable to List + objects, id_field=id_field, keys=keys, preprocess=preprocess, @@ -454,7 +456,9 @@ async def awrite( ValueError: If the length of provided keys does not match the length of objects, or if validation fails. """ - if keys and len(keys) != len(objects): # type: ignore + objects = list(objects) + keys = list(keys) if keys is not None else None + if keys is not None and len(keys) != len(objects): raise ValueError("Length of keys does not match the length of objects") if batch_size is None: @@ -465,7 +469,7 @@ async def awrite( # Pass 1: Preprocess and validate all objects prepared_objects = self._preprocess_and_validate_objects( - list(objects), # Convert Iterable to List + objects, id_field=id_field, keys=keys, preprocess=preprocess, diff --git a/tests/unit/test_storage.py b/tests/unit/test_storage.py index e79c9650d..c5fbe4ff8 100644 --- a/tests/unit/test_storage.py +++ b/tests/unit/test_storage.py @@ -1,4 +1,4 @@ -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pytest from pydantic import ValidationError @@ -181,3 +181,31 @@ def fn(d): preprocessed_data = storage_instance._preprocess(obj=data, preprocess=fn) assert "foo" in preprocessed_data assert preprocessed_data["foo"] == "bar" + + +ITERABLE_SCHEMA = IndexSchema.from_dict( + { + "index": {"name": "t", "prefix": "t"}, + "fields": [{"name": "title", "type": "text"}], + } +) +ITERABLE_OBJS = [{"title": "a"}, {"title": "b"}] + + +async def test_write_accepts_single_pass_iterables(): + storage = HashStorage(index_schema=ITERABLE_SCHEMA) + keys = ["t:1", "t:2"] + assert storage.write(MagicMock(), iter(ITERABLE_OBJS), keys=iter(keys)) == keys + client = MagicMock() + client.pipeline.return_value.__aenter__.return_value = AsyncMock() + assert await storage.awrite(client, iter(ITERABLE_OBJS), keys=iter(keys)) == keys + + +@pytest.mark.parametrize("keys", [[], ["t:1"], ["t:1", "t:2", "t:3"]]) +async def test_write_rejects_mismatched_keys_before_writing(keys): + storage, client = HashStorage(index_schema=ITERABLE_SCHEMA), MagicMock() + with pytest.raises(ValueError, match="Length of keys"): + storage.write(client, iter(ITERABLE_OBJS), keys=iter(keys)) + with pytest.raises(ValueError, match="Length of keys"): + await storage.awrite(client, iter(ITERABLE_OBJS), keys=iter(keys)) + client.pipeline.assert_not_called()