Simplify collection alias code.

Now it's possible to query schema using the alias
so we don't need to keep the alias and collection
referenced to do either one operation or the other.

Also, add support for the override_alias when restoring
the backup.

Author: jfrancoa

setup.cfg

@@ -37,7 +37,7 @@ classifiers =
 include_package_data = True
 python_requires = >=3.9
 install_requires =
-    weaviate-client>=4.16.0
+    weaviate-client==v4.17.0-rc1
     click==8.1.7
     semver>=3.0.2
     numpy>=1.24.0

weaviate_cli/commands/restore.py

@@ -38,8 +38,13 @@ def restore() -> None:
     default=RestoreBackupDefaults.exclude,
     help="Collection to exclude in backup (default: None).",
 )
+@click.option(
+    "--override-alias",
+    is_flag=True,
+    help="Override the alias of the collection (default: False).",
+)
 @click.pass_context
-def restore_backup_cli(ctx, backend, include, exclude, backup_id, wait):
+def restore_backup_cli(ctx, backend, include, exclude, backup_id, wait, override_alias):
     """Restore a backup in Weaviate."""
 
     client = None
@@ -52,6 +57,7 @@ def restore_backup_cli(ctx, backend, include, exclude, backup_id, wait):
             include=include,
             exclude=exclude,
             wait=wait,
+            override_alias=override_alias,
         )
     except Exception as e:
         click.echo(f"Error: {e}")

weaviate_cli/defaults.py

@@ -200,6 +200,7 @@ class RestoreBackupDefaults:
     wait: bool = False
     include: Optional[str] = None
     exclude: Optional[str] = None
+    override_alias: bool = False
 
 
 @dataclass

weaviate_cli/managers/backup_manager.py

@@ -13,7 +13,7 @@
 
 class BackupManager:
     def __init__(self, client: WeaviateClient) -> None:
-        self.client = client
+        self.client: WeaviateClient = client
 
     def create_backup(
         self,
@@ -68,13 +68,15 @@ def restore_backup(
         include: Optional[str] = RestoreBackupDefaults.include,
         exclude: Optional[str] = RestoreBackupDefaults.exclude,
         wait: bool = RestoreBackupDefaults.wait,
+        override_alias: bool = RestoreBackupDefaults.override_alias,
     ) -> None:
 
         result = self.client.backup.restore(
             backup_id=backup_id,
             backend=backend,
             include_collections=include.split(",") if include else None,
             exclude_collections=exclude.split(",") if exclude else None,
+            overwrite_alias=override_alias,
             wait_for_completion=wait,
         )
 

weaviate_cli/managers/data_manager.py

@@ -237,18 +237,10 @@ def __import_json(
         file_name: str,
         cl: wvc.ConsistencyLevel,
         num_objects: Optional[int] = None,
-        alias: Optional[str] = None,
     ) -> int:
         counter = 0
-        if alias is None:
-            properties: List[wvc.Property] = collection.config.get().properties
-        else:
-            collection_from_alias = self.client.alias.get(alias_name=alias).collection
-            properties: List[wvc.Property] = (
-                self.client.collections.get(collection_from_alias)
-                .config.get()
-                .properties
-            )
+
+        properties: List[wvc.Property] = collection.config.get().properties
 
         try:
             with (
@@ -451,20 +443,13 @@ def __ingest_data(
         uuid: Optional[str] = None,
         verbose: bool = False,
         multi_vector: bool = False,
-        alias: Optional[str] = None,
     ) -> Collection:
         if randomize:
             click.echo(f"Generating {num_objects} objects")
             start_time = time.time()
 
             # Determine vector dimensions based on vectorizer
-            if alias is None:
-                config = collection.config.get()
-            else:
-                collection_from_alias = self.client.alias.get(
-                    alias_name=alias
-                ).collection
-                config = self.client.collections.get(collection_from_alias).config.get()
+            config = collection.config.get()
 
             if not config.vectorizer and config.vector_config:
                 # Named vectors
@@ -599,11 +584,9 @@ def __ingest_data(
             )
             return cl_collection
         else:
-            click.echo(
-                f"Importing {num_objects} objects from Movies dataset {'using alias' if alias else ''}"
-            )
+            click.echo(f"Importing {num_objects} objects from Movies dataset")
             num_objects_inserted = self.__import_json(
-                collection, "movies.json", cl, num_objects, alias
+                collection, "movies.json", cl, num_objects
             )
             print(
                 f"Inserted {num_objects_inserted} objects into class '{collection.name}'"
@@ -627,16 +610,12 @@ def create_data(
         multi_vector: bool = CreateDataDefaults.multi_vector,
     ) -> Collection:
 
-        alias = None
         if not self.client.collections.exists(collection):
             alias_list = self.client.alias.list_all()
             if collection not in alias_list.keys():
                 raise Exception(
                     f"Class '{collection}' does not exist in Weaviate. Create first using <create class> command"
                 )
-            else:
-                alias = collection
-                collection = str(alias_list[collection].collection)
 
         col: Collection = self.client.collections.get(collection)
         mt_enabled = col.config.get().multi_tenancy_config.enabled
@@ -705,9 +684,7 @@ def create_data(
             if tenant == "None":
                 initial_length = len(col)
                 collection = self.__ingest_data(
-                    collection=(
-                        col if alias is None else self.client.collections.get(alias)
-                    ),
+                    collection=col,
                     num_objects=limit,
                     cl=cl_map[consistency_level],
                     randomize=randomize,
@@ -716,7 +693,6 @@ def create_data(
                     uuid=uuid,
                     verbose=verbose,
                     multi_vector=multi_vector,
-                    alias=alias,
                 )
                 after_length = len(col)
             else:
@@ -739,11 +715,7 @@ def create_data(
                     initial_length = len(col.with_tenant(tenant))
                 click.echo(f"Processing objects for tenant '{tenant}'")
                 collection = self.__ingest_data(
-                    collection=(
-                        col.with_tenant(tenant)
-                        if alias is None
-                        else self.client.collections.get(alias).with_tenant(tenant)
-                    ),
+                    collection=col.with_tenant(tenant),
                     num_objects=limit,
                     cl=cl_map[consistency_level],
                     randomize=randomize,
@@ -752,7 +724,6 @@ def create_data(
                     uuid=uuid,
                     verbose=verbose,
                     multi_vector=multi_vector,
-                    alias=alias,
                 )
                 after_length = len(col.with_tenant(tenant))
             if wait_for_indexing:
@@ -983,16 +954,12 @@ def update_data(
         verbose: bool = UpdateDataDefaults.verbose,
     ) -> None:
 
-        alias = None
         if not self.client.collections.exists(collection):
             alias_list = self.client.alias.list_all()
             if collection not in alias_list.keys():
                 raise Exception(
                     f"Class '{collection}' does not exist in Weaviate. Create first using ./create_class.py"
                 )
-            else:
-                alias = collection
-                collection = str(alias_list[collection].collection)
 
         col: Collection = self.client.collections.get(collection)
         try:
@@ -1012,8 +979,6 @@ def update_data(
         }
 
         click.echo(f"Preparing to update {limit} objects into class '{col.name}'")
-        # Override collection if alias is provided
-        col = col if alias is None else self.client.collections.get(alias)
         for tenant in tenants:
             if tenant == "None":
                 ret = self.__update_data(
@@ -1130,18 +1095,12 @@ def delete_data(
         verbose: bool = DeleteDataDefaults.verbose,
     ) -> None:
 
-        alias = None
         if not self.client.collections.exists(collection):
             alias_list = self.client.alias.list_all()
             if collection not in alias_list.keys():
                 raise Exception(
                     f"Class '{collection}' does not exist in Weaviate. Create first using <create class> command."
                 )
-            else:
-                alias = collection
-                collection = str(alias_list[collection].collection)
-
-            return 1
 
         col: Collection = self.client.collections.get(collection)
         mt_enabled = col.config.get().multi_tenancy_config.enabled
@@ -1161,9 +1120,6 @@ def delete_data(
         else:
             tenants = existing_tenants
 
-        # Override collection if alias is provided
-        col = col if alias is None else self.client.collections.get(alias)
-
         for tenant in tenants:
             if tenant == "None":
                 ret = self.__delete_data(
@@ -1264,16 +1220,12 @@ def query_data(
         target_vector: Optional[str] = QueryDataDefaults.target_vector,
     ) -> None:
 
-        alias = None
         if not self.client.collections.exists(collection):
             alias_list = self.client.alias.list_all()
             if collection not in alias_list.keys():
                 raise Exception(
                     f"Class '{collection}' does not exist in Weaviate. Create first using <create class> command."
                 )
-            else:
-                alias = collection
-                collection = str(alias_list[collection].collection)
 
         col: Collection = self.client.collections.get(collection)
         mt_enabled = col.config.get().multi_tenancy_config.enabled
@@ -1303,8 +1255,6 @@ def query_data(
             "all": wvc.ConsistencyLevel.ALL,
             "one": wvc.ConsistencyLevel.ONE,
         }
-        # Override collection if alias is provided
-        col = col if alias is None else self.client.collections.get(alias)
 
         for tenant in existing_tenants:
             if tenant == "None":