Add option to update deletion strategy

weaviate · Nov 7, 2024 · 0f16511 · 0f16511
1 parent b6d8cd8
commit 0f16511
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 4 deletions.
diff --git a/test/integration/test_integration.py b/test/integration/test_integration.py
@@ -65,6 +65,7 @@ def test_collection_lifecycle(collection_manager: CollectionManager):
             async_enabled=False,
             auto_tenant_creation=None,
             auto_tenant_activation=None,
+            replication_deletion_strategy=None,
         )
 
         # Get collection config to verify update
@@ -209,4 +210,5 @@ def test_error_handling(collection_manager: CollectionManager):
             async_enabled=True,
             auto_tenant_creation=None,
             auto_tenant_activation=None,
+            replication_deletion_strategy=None,
         )
diff --git a/test/unittests/test_managers/test_collection_manager.py b/test/unittests/test_managers/test_collection_manager.py
@@ -182,6 +182,7 @@ def test_update_collection(mock_client):
         async_enabled=True,
         auto_tenant_creation=True,
         auto_tenant_activation=True,
+        replication_deletion_strategy="delete_on_conflict",
     )
 
     mock_collection.config.update.assert_called_once()
@@ -221,6 +222,7 @@ def test_update_nonexistent_collection(mock_client):
             async_enabled=True,
             auto_tenant_creation=True,
             auto_tenant_activation=True,
+            replication_deletion_strategy="delete_on_conflict",
         )
 
     assert "Error: Collection 'TestCollection' does not exist in Weaviate." in str(

diff --git a/weaviate_cli/commands/update.py b/weaviate_cli/commands/update.py
@@ -49,6 +49,12 @@ def update() -> None:
     type=bool,
     help="Enable auto tenant activation (default: None).",
 )
+@click.option(
+    "--replication_deletion_strategy",
+    default=None,
+    type=click.Choice(["delete_on_conflict", "no_automated_resolution"]),
+    help="Replication deletion strategy.",
+)
 @click.pass_context
 def update_collection_cli(
     ctx: click.Context,
@@ -59,6 +65,7 @@ def update_collection_cli(
     training_limit: int,
     auto_tenant_creation: Optional[bool],
     auto_tenant_activation: Optional[bool],
+    replication_deletion_strategy: Optional[str],
 ) -> None:
     """Update a collection in Weaviate."""
 
@@ -75,6 +82,7 @@ def update_collection_cli(
             training_limit=training_limit,
             auto_tenant_creation=auto_tenant_creation,
             auto_tenant_activation=auto_tenant_activation,
+            replication_deletion_strategy=replication_deletion_strategy,
         )
     except Exception as e:
         click.echo(f"Error: {e}")

diff --git a/weaviate_cli/managers/collection_manager.py b/weaviate_cli/managers/collection_manager.py
@@ -183,6 +183,7 @@ def update_collection(
         async_enabled: Optional[bool],
         auto_tenant_creation: Optional[bool],
         auto_tenant_activation: Optional[bool],
+        replication_deletion_strategy: Optional[str],
     ) -> None:
 
         if not self.client.collections.exists(collection):
@@ -214,6 +215,10 @@ def update_collection(
 
         col_obj: Collection = self.client.collections.get(collection)
         rf = col_obj.config.get().replication_config.factor
+        rds_map = {
+            "delete_on_conflict": wvc.ReplicationDeletionStrategy.DELETE_ON_CONFLICT,
+            "no_automated_resolution": wvc.ReplicationDeletionStrategy.NO_AUTOMATED_RESOLUTION,
+        }
         mt = col_obj.config.get().multi_tenancy_config.enabled
         auto_tenant_creation = (
             auto_tenant_creation
@@ -225,16 +230,21 @@ def update_collection(
             if auto_tenant_activation is not None
             else col_obj.config.get().multi_tenancy_config.auto_tenant_activation
         )
-
         col_obj.config.update(
             description=description,
             vectorizer_config=(
                 vector_index_map[vector_index] if vector_index else None
             ),
             replication_config=(
-                wvc.Reconfigure.replication(factor=rf, async_enabled=async_enabled)
-                if async_enabled is not None
-                else None
+                wvc.Reconfigure.replication(
+                    factor=rf,
+                    async_enabled=async_enabled if async_enabled is not None else None,
+                    deletion_strategy=(
+                        rds_map[replication_deletion_strategy]
+                        if replication_deletion_strategy
+                        else None
+                    ),
+                )
             ),
             multi_tenancy_config=(
                 wvc.Reconfigure.multi_tenancy(