feat: Redis key_prefix field (#437)

* Redis uploader key_prefix field

* Changelog and version bump: Redis uploader key_prefix field

* Make dev version

* Revert "Make dev version"

This reverts commit 2b67e23.

Author: mpolomdeepsense

CHANGELOG.md

@@ -1,3 +1,9 @@
+## 0.5.19
+
+### Features
+
+* **Add `key_prefix` field to Redis Uploader** - Allow users to input custom prefix for keys saved inside Redis connector
+
 ## 0.5.18
 
 ### Fixes

test/integration/connectors/test_redis.py

@@ -23,20 +23,22 @@
 )
 
 
-async def delete_record(client: Redis, element_id: str) -> None:
-    await client.delete(element_id)
+async def delete_record(client: Redis, element_id: str, key_prefix: str) -> None:
+    key_with_prefix = f"{key_prefix}{element_id}"
+    await client.delete(key_with_prefix)
 
 
-async def validate_upload(client: Redis, first_element: dict):
+async def validate_upload(client: Redis, first_element: dict, key_prefix: str) -> None:
     element_id = first_element["element_id"]
+    key_with_prefix = f"{key_prefix}{element_id}"
     expected_text = first_element["text"]
     expected_embeddings = first_element["embeddings"]
     async with client.pipeline(transaction=True) as pipe:
         try:
-            response = await pipe.json().get(element_id, "$").execute()
+            response = await pipe.json().get(key_with_prefix, "$").execute()
             response = response[0][0]
         except redis_exceptions.ResponseError:
-            response = await pipe.get(element_id).execute()
+            response = await pipe.get(key_with_prefix).execute()
             response = json.loads(response[0])
 
     embedding_similarity = np.linalg.norm(
@@ -53,15 +55,17 @@ async def redis_destination_test(
     upload_file: Path,
     tmp_path: Path,
     connection_kwargs: dict,
+    uploader_config: dict,
     uri: Optional[str] = None,
     password: Optional[str] = None,
 ):
     uploader = RedisUploader(
         connection_config=RedisConnectionConfig(
             **connection_kwargs, access_config=RedisAccessConfig(uri=uri, password=password)
         ),
-        upload_config=RedisUploaderConfig(batch_size=10),
+        upload_config=RedisUploaderConfig(batch_size=10, **uploader_config),
     )
+    key_prefix = uploader.upload_config.key_prefix
 
     file_data = FileData(
         source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
@@ -78,20 +82,32 @@ async def redis_destination_test(
 
         if uri:
             async with from_url(uri) as client:
-                await validate_upload(client=client, first_element=first_element)
+                await validate_upload(
+                    client=client,
+                    first_element=first_element,
+                    key_prefix=key_prefix,
+                )
         else:
             async with Redis(**connection_kwargs, password=password) as client:
-                await validate_upload(client=client, first_element=first_element)
+                await validate_upload(
+                    client=client,
+                    first_element=first_element,
+                    key_prefix=key_prefix,
+                )
     except Exception as e:
         raise e
     finally:
         if uri:
             async with from_url(uri) as client:
-                tasks = [delete_record(client, element["element_id"]) for element in elements]
+                tasks = [
+                    delete_record(client, element["element_id"], key_prefix) for element in elements
+                ]
                 await asyncio.gather(*tasks)
         else:
             async with Redis(**connection_kwargs, password=password) as client:
-                tasks = [delete_record(client, element["element_id"]) for element in elements]
+                tasks = [
+                    delete_record(client, element["element_id"], key_prefix) for element in elements
+                ]
                 await asyncio.gather(*tasks)
 
 
@@ -105,15 +121,23 @@ async def test_redis_destination_azure_with_password(upload_file: Path, tmp_path
         "db": 0,
         "ssl": True,
     }
+    uploader_config = {
+        "key_prefix": "test_ingest:",
+    }
     redis_pw = os.environ["AZURE_REDIS_INGEST_TEST_PASSWORD"]
-    await redis_destination_test(upload_file, tmp_path, connection_kwargs, password=redis_pw)
+    await redis_destination_test(
+        upload_file, tmp_path, connection_kwargs, uploader_config, password=redis_pw
+    )
 
 
 @pytest.mark.asyncio
 @pytest.mark.tags(REDIS_CONNECTOR_TYPE, DESTINATION_TAG, "redis", NOSQL_TAG)
 @requires_env("AZURE_REDIS_INGEST_TEST_PASSWORD")
 async def test_redis_destination_azure_with_uri(upload_file: Path, tmp_path: Path):
     connection_kwargs = {}
+    uploader_config = {
+        "key_prefix": "test_ingest:",
+    }
     redis_pw = os.environ["AZURE_REDIS_INGEST_TEST_PASSWORD"]
     uri = f"rediss://:{redis_pw}@utic-dashboard-dev.redis.cache.windows.net:6380/0"
-    await redis_destination_test(upload_file, tmp_path, connection_kwargs, uri=uri)
+    await redis_destination_test(upload_file, tmp_path, connection_kwargs, uploader_config, uri=uri)

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "0.5.18"  # pragma: no cover
+__version__ = "0.5.19"  # pragma: no cover

unstructured_ingest/v2/processes/connectors/redisdb.py

@@ -110,6 +110,7 @@ def create_client(self) -> Generator["Redis", None, None]:
 
 class RedisUploaderConfig(UploaderConfig):
     batch_size: int = Field(default=100, description="Number of records per batch")
+    key_prefix: str = Field(default="", description="Prefix for Redis keys")
 
 
 @dataclass
@@ -145,11 +146,11 @@ async def _write_batch(self, batch: list[dict], redis_stack: bool) -> None:
         async with self.connection_config.create_async_client() as async_client:
             async with async_client.pipeline(transaction=True) as pipe:
                 for element in batch:
-                    element_id = element["element_id"]
+                    key_with_prefix = f"{self.upload_config.key_prefix}{element['element_id']}"
                     if redis_stack:
-                        pipe.json().set(element_id, "$", element)
+                        pipe.json().set(key_with_prefix, "$", element)
                     else:
-                        pipe.set(element_id, json.dumps(element))
+                        pipe.set(key_with_prefix, json.dumps(element))
                 await pipe.execute()
 
     @requires_dependencies(["redis"], extras="redis")
@@ -159,16 +160,16 @@ async def _check_redis_stack(self, element: dict) -> bool:
         redis_stack = True
         async with self.connection_config.create_async_client() as async_client:
             async with async_client.pipeline(transaction=True) as pipe:
-                element_id = element["element_id"]
+                key_with_prefix = f"{self.upload_config.key_prefix}{element['element_id']}"
                 try:
                     # Redis with stack extension supports JSON type
-                    await pipe.json().set(element_id, "$", element).execute()
+                    await pipe.json().set(key_with_prefix, "$", element).execute()
                 except redis_exceptions.ResponseError as e:
                     message = str(e)
                     if "unknown command `JSON.SET`" in message:
                         # if this error occurs, Redis server doesn't support JSON type,
                         # so save as string type instead
-                        await pipe.set(element_id, json.dumps(element)).execute()
+                        await pipe.set(key_with_prefix, json.dumps(element)).execute()
                         redis_stack = False
                     else:
                         raise e