From b6d8cd8f59f77a7ffaa5572acb72c78688766bd7 Mon Sep 17 00:00:00 2001 From: Rodrigo Lopez Date: Wed, 6 Nov 2024 12:49:12 +0100 Subject: [PATCH 1/2] Add deletion strategy option when creating classes Signed-off-by: Rodrigo Lopez --- test/integration/test_integration.py | 5 +++++ .../test_managers/test_collection_manager.py | 3 +++ weaviate_cli/commands/create.py | 8 ++++++++ weaviate_cli/managers/collection_manager.py | 14 +++++++++++++- 4 files changed, 29 insertions(+), 1 deletion(-) diff --git a/test/integration/test_integration.py b/test/integration/test_integration.py index a6e43de..49816ec 100644 --- a/test/integration/test_integration.py +++ b/test/integration/test_integration.py @@ -50,6 +50,7 @@ def test_collection_lifecycle(collection_manager: CollectionManager): auto_tenant_activation=False, vectorizer="contextionary", force_auto_schema=True, + replication_deletion_strategy=None, ) # Verify collection exists @@ -98,6 +99,7 @@ def test_multiple_collections(collection_manager: CollectionManager): auto_tenant_activation=False, vectorizer="contextionary", force_auto_schema=True, + replication_deletion_strategy=None, ) assert collection_manager.client.collections.exists(col) finally: @@ -129,6 +131,7 @@ def test_shard_operations( auto_tenant_activation=False, vectorizer="transformers", force_auto_schema=True, + replication_deletion_strategy=None, ) # Get shard info @@ -170,6 +173,7 @@ def test_error_handling(collection_manager: CollectionManager): auto_tenant_activation=False, vectorizer="transformers", force_auto_schema=True, + replication_deletion_strategy=None, ) with pytest.raises(Exception): @@ -186,6 +190,7 @@ def test_error_handling(collection_manager: CollectionManager): auto_tenant_activation=False, vectorizer="transformers", force_auto_schema=True, + replication_deletion_strategy=None, ) finally: # Test deleting non-existent collection diff --git a/test/unittests/test_managers/test_collection_manager.py b/test/unittests/test_managers/test_collection_manager.py index f6b69f1..95ee884 100644 --- a/test/unittests/test_managers/test_collection_manager.py +++ b/test/unittests/test_managers/test_collection_manager.py @@ -30,6 +30,7 @@ def test_create_collection(mock_client): force_auto_schema=False, shards=1, vectorizer=None, + replication_deletion_strategy=None, ) # Verify the collection creation was called with correct parameters @@ -68,6 +69,7 @@ def test_create_existing_collection(mock_client): force_auto_schema=False, shards=1, vectorizer=None, + replication_deletion_strategy=None, ) # Verify the error message @@ -104,6 +106,7 @@ def test_create_collection_failure(mock_client): force_auto_schema=False, shards=1, vectorizer=None, + replication_deletion_strategy=None, ) # Verify the error message diff --git a/weaviate_cli/commands/create.py b/weaviate_cli/commands/create.py index 01dcd3e..3608b93 100644 --- a/weaviate_cli/commands/create.py +++ b/weaviate_cli/commands/create.py @@ -69,6 +69,12 @@ def create() -> None: type=click.Choice(["contextionary", "transformers", "openai", "ollama"]), help="Vectorizer to use.", ) +@click.option( + "--replication_deletion_strategy", + default="delete_on_conflict", + type=click.Choice(["delete_on_conflict", "no_automated_resolution"]), + help="Replication deletion strategy (default: 'delete_on_conflict').", +) @click.pass_context def create_collection_cli( ctx: click.Context, @@ -84,6 +90,7 @@ def create_collection_cli( force_auto_schema: bool, shards: int, vectorizer: Optional[str], + replication_deletion_strategy: str, ) -> None: """Create a collection in Weaviate.""" @@ -105,6 +112,7 @@ def create_collection_cli( force_auto_schema=force_auto_schema, shards=shards, vectorizer=vectorizer, + 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 index 307c6c9..3072b73 100644 --- a/weaviate_cli/managers/collection_manager.py +++ b/weaviate_cli/managers/collection_manager.py @@ -58,6 +58,7 @@ def create_collection( force_auto_schema: bool, shards: int, vectorizer: Optional[str], + replication_deletion_strategy: Optional[str], ) -> None: if self.client.collections.exists(collection): @@ -133,6 +134,11 @@ def create_collection( wvc.Property(name="status", data_type=wvc.DataType.TEXT), ] + rds_map = { + "delete_on_conflict": wvc.ReplicationDeletionStrategy.DELETE_ON_CONFLICT, + "no_automated_resolution": wvc.ReplicationDeletionStrategy.NO_AUTOMATED_RESOLUTION, + } + try: self.client.collections.create( name=collection, @@ -141,7 +147,13 @@ def create_collection( inverted_index_map[inverted_index] if inverted_index else None ), replication_config=wvc.Configure.replication( - factor=replication_factor, async_enabled=async_enabled + factor=replication_factor, + async_enabled=async_enabled, + deletion_strategy=( + rds_map[replication_deletion_strategy] + if replication_deletion_strategy + else None + ), ), sharding_config=( wvc.Configure.sharding(desired_count=shards) if shards > 1 else None From 0f16511811924768ad053ac15a227a2a521a1ea7 Mon Sep 17 00:00:00 2001 From: Rodrigo Lopez Date: Thu, 7 Nov 2024 10:09:37 +0100 Subject: [PATCH 2/2] Add option to update deletion strategy --- test/integration/test_integration.py | 2 ++ .../test_managers/test_collection_manager.py | 2 ++ weaviate_cli/commands/update.py | 8 ++++++++ weaviate_cli/managers/collection_manager.py | 18 ++++++++++++++---- 4 files changed, 26 insertions(+), 4 deletions(-) diff --git a/test/integration/test_integration.py b/test/integration/test_integration.py index 49816ec..4714055 100644 --- 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 index 95ee884..195029a 100644 --- 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 index 8f134e3..50b15fe 100644 --- 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 index 3072b73..aab407f 100644 --- 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(