utkarsharma2 commented on code in PR #35919:
URL: https://github.com/apache/airflow/pull/35919#discussion_r1415613287


##########
airflow/providers/weaviate/hooks/weaviate.py:
##########
@@ -112,22 +117,170 @@ 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(

Review Comment:
   I have split the method into two methods - `delete_classes` and 
`delete_all_schema`.



-- 
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

Reply via email to