Fix Weaviate tenant-aware ingestion

providers/weaviate/src/airflow/providers/weaviate/hooks/weaviate.py

@@ -411,6 +411,7 @@ def batch_data(
         uuid_col: str = "id",
         retry_attempts_per_object: int = 5,
         references: ReferenceInputs | None = None,
+        tenant: str | None = None,
     ) -> None:
         """
         Add multiple objects or object references at once into weaviate.
@@ -421,10 +422,13 @@ def batch_data(
         :param uuid_col: Name of the column containing the UUID.
         :param retry_attempts_per_object: number of time to try in case of failure before giving up.
         :param references: The references of the object to be added as a dictionary. Use `wvc.Reference.to` to create the correct values in the dict.
+        :param tenant: The tenant to which the objects will be added.
         """
         converted_data = self._convert_dataframe_to_list(data)
 
         collection = self.get_collection(collection_name)
+        if tenant:
+            collection = collection.with_tenant(tenant)
         with collection.batch.dynamic() as batch:
             # Batch import all data
             for data_obj in converted_data:
@@ -585,14 +589,17 @@ def get_all_objects(
             )
         return all_objects
 
-    def delete_object(self, collection_name: str, uuid: UUID | str) -> bool:
+    def delete_object(self, collection_name: str, uuid: UUID | str, tenant: str | None = None) -> bool:
         """
         Delete an object from weaviate.
 
         :param collection_name: Collection name associated with the object given.
         :param uuid: uuid of the object to be deleted
+        :param tenant: The tenant from which the object will be deleted.
         """
         collection = self.get_collection(collection_name)
+        if tenant:
+            collection = collection.with_tenant(tenant)
         return collection.data.delete_by_id(uuid=uuid)
 
     def update_object(
@@ -640,7 +647,11 @@ def object_exists(self, collection_name: str, uuid: str | UUID) -> bool:
         return collection.data.exists(uuid=uuid)
 
     def _delete_objects(
-        self, uuids: list[UUID], collection_name: str, retry_attempts_per_object: int = 5
+        self,
+        uuids: list[UUID],
+        collection_name: str,
+        retry_attempts_per_object: int = 5,
+        tenant: str | None = None,
     ) -> None:
         """
         Delete multiple objects.
@@ -650,6 +661,7 @@ def _delete_objects(
         :param uuids: Collection of uuids.
         :param collection_name: Name of the collection in Weaviate schema where data is to be ingested.
         :param retry_attempts_per_object: number of times to try in case of failure before giving up.
+        :param tenant: The tenant from which the objects will be deleted.
         """
         for uuid in uuids:
             for attempt in Retrying(
@@ -661,7 +673,7 @@ def _delete_objects(
             ):
                 with attempt:
                     try:
-                        self.delete_object(uuid=uuid, collection_name=collection_name)
+                        self.delete_object(uuid=uuid, collection_name=collection_name, tenant=tenant)
                         self.log.debug("Deleted object with uuid %s", uuid)
                     except weaviate.exceptions.UnexpectedStatusCodeException as e:
                         if e.status_code == 404:
@@ -728,6 +740,7 @@ def _get_documents_to_uuid_map(
         document_column: str,
         uuid_column: str,
         collection_name: str,
+        tenant: str | None = None,
         offset: int = 0,
         limit: int = 2000,
     ) -> dict[str, set]:
@@ -737,6 +750,7 @@ def _get_documents_to_uuid_map(
         :param data: A single pandas DataFrame.
         :param document_column: The name of the property to query.
         :param collection_name: The name of the collection to query.
+        :param tenant: The tenant to query.
         :param uuid_column: The name of the column containing the UUID.
         :param offset: pagination parameter to indicate the which object to start fetching data.
         :param limit: pagination param to indicate the number of records to fetch from start object.
@@ -745,6 +759,8 @@ def _get_documents_to_uuid_map(
         document_keys = set(data[document_column])
         while True:
             collection = self.get_collection(collection_name)
+            if tenant:
+                collection = collection.with_tenant(tenant)
             data_objects = collection.query.fetch_objects(
                 filters=Filter.any_of(
                     [Filter.by_property(document_column).equal(key) for key in document_keys]
@@ -791,7 +807,12 @@ def _prepare_document_to_uuid_map(
         return grouped_key_to_set
 
     def _get_segregated_documents(
-        self, data: pd.DataFrame, document_column: str, collection_name: str, uuid_column: str
+        self,
+        data: pd.DataFrame,
+        document_column: str,
+        collection_name: str,
+        uuid_column: str,
+        tenant: str | None = None,
     ) -> tuple[dict[str, set], set, set, set]:
         """
         Segregate documents into changed, unchanged and new document, when compared to Weaviate db.
@@ -800,6 +821,7 @@ def _get_segregated_documents(
         :param document_column: The name of the property to query.
         :param collection_name: The name of the collection to query.
         :param uuid_column: The name of the column containing the UUID.
+        :param tenant: The tenant to query.
         """
         changed_documents = set()
         unchanged_docs = set()
@@ -809,6 +831,7 @@ def _get_segregated_documents(
             uuid_column=uuid_column,
             document_column=document_column,
             collection_name=collection_name,
+            tenant=tenant,
         )
 
         input_documents_to_uuid = self._prepare_document_to_uuid_map(
@@ -836,6 +859,7 @@ def _delete_all_documents_objects(
         total_objects_count: int = 1,
         batch_delete_error: Sequence | None = None,
         verbose: bool = False,
+        tenant: str | None = None,
     ) -> Sequence[dict[str, UUID | str]]:
         """
         Delete all object that belong to list of documents.
@@ -847,13 +871,16 @@ def _delete_all_documents_objects(
             query is 10,000, if we have more objects to delete we need to run query multiple times.
         :param batch_delete_error: list to hold errors while inserting.
         :param verbose: Flag to enable verbose output during the ingestion process.
+        :param tenant: The tenant from which document objects will be deleted.
         """
         batch_delete_error = batch_delete_error or []
 
         # This limit is imposed by Weavaite database
         MAX_LIMIT_ON_TOTAL_DELETABLE_OBJECTS = 10000
 
         collection = self.get_collection(collection_name)
+        if tenant:
+            collection = collection.with_tenant(tenant)
         delete_many_return = collection.data.delete_many(
             where=Filter.any_of([Filter.by_property(document_column).equal(key) for key in document_keys]),
             verbose=verbose,
@@ -881,6 +908,7 @@ def create_or_replace_document_objects(
         uuid_column: str | None = None,
         vector_column: str = "Vector",
         verbose: bool = False,
+        tenant: str | None = None,
     ) -> Sequence[dict[str, UUID | str] | None]:
         """
         Create or replace objects belonging to documents.
@@ -909,6 +937,7 @@ def create_or_replace_document_objects(
         :param uuid_column: Column with pre-generated UUIDs. If not provided, UUIDs will be generated.
         :param vector_column: Column with embedding vectors for pre-embedded data.
         :param verbose: Flag to enable verbose output during the ingestion process.
+        :param tenant: The tenant to which objects will be added.
         :return: list of UUID which failed to create
         """
         if existing not in ["skip", "replace", "error"]:
@@ -960,6 +989,7 @@ def create_or_replace_document_objects(
             document_column=document_column,
             uuid_column=uuid_column,
             collection_name=collection_name,
+            tenant=tenant,
         )
         if verbose:
             self.log.info(
@@ -1001,6 +1031,7 @@ def create_or_replace_document_objects(
                     total_objects_count=total_objects_count,
                     batch_delete_error=batch_delete_error,
                     verbose=verbose,
+                    tenant=tenant,
                 )
             data = data[data[document_column].isin(new_documents.union(changed_documents))]
             self.log.info("Batch inserting %s objects for non-existing and changed documents.", data.shape[0])
@@ -1011,6 +1042,7 @@ def create_or_replace_document_objects(
                 data=data,
                 vector_col=vector_column,
                 uuid_col=uuid_column,
+                tenant=tenant,
             )
             if batch_delete_error:
                 if batch_delete_error:
@@ -1019,10 +1051,13 @@ def create_or_replace_document_objects(
                 self._delete_objects(
                     [item["uuid"] for item in batch_delete_error],
                     collection_name=collection_name,
+                    tenant=tenant,
                 )
 
         if verbose:
             collection = self.get_collection(collection_name)
+            if tenant:
+                collection = collection.with_tenant(tenant)
             self.log.info(
                 "Total objects in collection %s : %s ",
                 collection_name,

providers/weaviate/src/airflow/providers/weaviate/operators/weaviate.py

@@ -43,10 +43,12 @@ class WeaviateIngestOperator(BaseOperator):
     custom vectors and store them in the Weaviate class.
 
     :param conn_id: The Weaviate connection.
-    :param collection: The Weaviate collection to be used for storing the data objects into.
+    :param collection_name: The Weaviate collection to be used for storing the data objects into.
     :param input_data: The list of dicts or pandas dataframe representing Weaviate data objects to generate
         embeddings on (or provides custom vectors) and store them in the Weaviate class.
     :param vector_col: key/column name in which the vectors are stored.
+    :param uuid_column: Column with pre-generated UUIDs.
+    :param tenant: The tenant to which objects will be added.
     :param hook_params: Optional config params to be passed to the underlying hook.
         Should match the desired hook constructor params.
     """
@@ -88,6 +90,7 @@ def execute(self, context: Context) -> None:
             data=self.input_data,
             vector_col=self.vector_col,
             uuid_col=self.uuid_column,
+            tenant=self.tenant,
         )
 
 
@@ -118,7 +121,7 @@ class WeaviateDocumentIngestOperator(BaseOperator):
     :param document_column: Column in DataFrame that identifying source document.
     :param uuid_column: Column with pre-generated UUIDs. If not provided, UUIDs will be generated.
     :param vector_column: Column with embedding vectors for pre-embedded data.
-    :param tenant: The tenant to which the object will be added.
+    :param tenant: The tenant to which objects will be added.
     :param verbose: Flag to enable verbose output during the ingestion process.
     :param hook_params: Optional config params to be passed to the underlying hook.
         Should match the desired hook constructor params.
@@ -172,5 +175,6 @@ def execute(self, context: Context) -> Sequence[dict[str, UUID | str] | None]:
             uuid_column=self.uuid_column,
             vector_column=self.vector_col,
             verbose=self.verbose,
+            tenant=self.tenant,
         )
         return batch_delete_error