utkarsharma2 commented on code in PR #35919: URL: https://github.com/apache/airflow/pull/35919#discussion_r1412111940
########## airflow/providers/weaviate/hooks/weaviate.py: ########## @@ -94,22 +99,117 @@ def test_connection(self) -> tuple[bool, str]: self.log.error("Error testing Weaviate connection: %s", e) return False, str(e) + @retry(reraise=True, stop=stop_after_attempt(3), retry=retry_if_exception_type(requests.ConnectionError)) def create_class(self, class_json: dict[str, Any]) -> None: """Create a new class.""" client = self.get_client() client.schema.create_class(class_json) - def create_schema(self, schema_json: dict[str, Any]) -> None: + @retry(reraise=True, stop=stop_after_attempt(3), retry=retry_if_exception_type(requests.ConnectionError)) + def create_schema(self, schema_json: dict[str, Any] | str) -> None: """ Create a new Schema. Instead of adding classes one by one , you can upload a full schema in JSON format at once. - :param schema_json: The schema to create + :param schema_json: The schema to create or path to the json file holding the schema """ client = self.get_client() client.schema.create(schema_json) + @retry(reraise=True, stop=stop_after_attempt(3), retry=retry_if_exception_type(requests.ConnectionError)) + def get_schema(self, class_name: str | None = None): + """Get the schema from Weaviate. + + get schema of a class or all classes. + """ + client = self.get_client() + return client.schema.get(class_name) + + @retry(reraise=True, stop=stop_after_attempt(3), retry=retry_if_exception_type(requests.ConnectionError)) + def delete_class(self, class_name: str): + """Delete a schema class from Weaviate. This deletes all associated data.""" + client = self.get_client() + client.schema.delete_class(class_name) + + def delete_schema( + self, class_names: list[str] | str | None = None + ) -> list[weaviate.UnexpectedStatusCodeException] | None: + """ + Deletes all or specific class if class_names are provided. + + If no class_name is given remove the entire schema from the Weaviate instance and all data associated + with it. If class_names are given, delete schema classes from Weaviate. This deletes all associated data. + :return: list of error object if any + """ + client = self.get_client() + class_names = [class_names] if class_names and isinstance(class_names, str) else class_names + if class_names: + error_list = [] + for class_name in class_names: + try: + self.delete_class(class_name=class_name) + except weaviate.UnexpectedStatusCodeException as e: + error_list.append(e) + return error_list Review Comment: > I think we should accept an argument here which tells what to do. e.g in the case when multiple classes are given for deletion, based on the argument value, while iterating it should either continue or raise exception when deletion of a class fails and maybe return a json saying how many failed or which ones failed deletion, how many successfully deleted. @ephraimbuddy @pankajkoti Yup, I too was thinking along the same lines, we can expose params. if_error="continue/stop". But for `continue` we still need to communicate that some of the classes didn't get deleted and also with the reason so that the user can debug and the reason is within the exceptions. We can just log them, instead of returning the list? and return a class count of successful deletion. WDYT? For `stop` we re-raise the exception immediately. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org