Adding retries and pool to Weaviate dags to handle connection latency issues

Author: Atul Singh

providers/weaviate/tests/system/weaviate/example_weaviate_cohere.py

@@ -16,6 +16,8 @@
 # under the License.
 from __future__ import annotations
 
+from datetime import timedelta
+
 import pendulum
 
 try:
@@ -28,10 +30,17 @@
 
 COLLECTION_NAME = "weaviate_cohere_example_collection"
 
+default_args = {
+    "retries": 5,
+    "retry_delay": timedelta(seconds=15),
+    "pool": "weaviate_pool",
+}
+
 
 @dag(
     schedule=None,
     start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
+    default_args=default_args,
     catchup=False,
     tags=["example", "weaviate", "cohere"],
 )

providers/weaviate/tests/system/weaviate/example_weaviate_dynamic_mapping_dag.py

@@ -16,6 +16,8 @@
 # under the License.
 from __future__ import annotations
 
+from datetime import timedelta
+
 import pendulum
 from weaviate.collections.classes.config import Configure
 
@@ -28,10 +30,17 @@
 
 COLLECTION_NAMES = ["Weaviate_DTM_example_collection_1", "Weaviate_DTM_example_collection_2"]
 
+default_args = {
+    "retries": 5,
+    "retry_delay": timedelta(seconds=15),
+    "pool": "weaviate_pool",
+}
+
 
 @dag(
     schedule=None,
     start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
+    default_args=default_args,
     catchup=False,
     tags=["example", "weaviate"],
 )

providers/weaviate/tests/system/weaviate/example_weaviate_openai.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import json
+from datetime import timedelta
 from pathlib import Path
 
 import pendulum
@@ -32,10 +33,17 @@
 
 COLLECTION_NAME = "Weaviate_openai_example_collection"
 
+default_args = {
+    "retries": 5,
+    "retry_delay": timedelta(seconds=15),
+    "pool": "weaviate_pool",
+}
+
 
 @dag(
     schedule=None,
     start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
+    default_args=default_args,
     catchup=False,
     tags=["example", "weaviate"],
 )

providers/weaviate/tests/system/weaviate/example_weaviate_operator.py

@@ -16,6 +16,8 @@
 # under the License.
 from __future__ import annotations
 
+from datetime import timedelta
+
 import pendulum
 from weaviate.classes.config import DataType, Property
 from weaviate.collections.classes.config import Configure
@@ -92,9 +94,17 @@ def get_data_without_vectors(*args, **kwargs):
     return sample_data_without_vector
 
 
+default_args = {
+    "retries": 5,
+    "retry_delay": timedelta(seconds=15),
+    "pool": "weaviate_pool",
+}
+
+
 @dag(
     schedule=None,
     start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
+    default_args=default_args,
     catchup=False,
     tags=["example", "weaviate"],
 )

providers/weaviate/tests/system/weaviate/example_weaviate_using_hook.py

@@ -16,6 +16,8 @@
 # under the License.
 from __future__ import annotations
 
+from datetime import timedelta
+
 import pendulum
 from weaviate.classes.config import DataType, Property
 from weaviate.collections.classes.config import Configure
@@ -28,10 +30,17 @@
 
 COLLECTION_NAME = "QuestionWithOpenAIVectorizerUsingHook"
 
+default_args = {
+    "retries": 5,
+    "retry_delay": timedelta(seconds=15),
+    "pool": "weaviate_pool",
+}
+
 
 @dag(
     schedule=None,
     start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
+    default_args=default_args,
     catchup=False,
     tags=["example", "weaviate"],
 )

providers/weaviate/tests/system/weaviate/example_weaviate_vectorizer_dag.py

@@ -16,6 +16,8 @@
 # under the License.
 from __future__ import annotations
 
+from datetime import timedelta
+
 import pendulum
 from weaviate.collections.classes.config import Configure
 
@@ -28,10 +30,17 @@
 
 COLLECTION_NAME = "Weaviate_with_vectorizer_example_collection"
 
+default_args = {
+    "retries": 5,
+    "retry_delay": timedelta(seconds=15),
+    "pool": "weaviate_pool",
+}
+
 
 @dag(
     schedule=None,
     start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
+    default_args=default_args,
     catchup=False,
     tags=["example", "weaviate"],
 )

providers/weaviate/tests/system/weaviate/example_weaviate_without_vectorizer_dag.py

@@ -16,6 +16,8 @@
 # under the License.
 from __future__ import annotations
 
+from datetime import timedelta
+
 import pendulum
 
 try:
@@ -28,10 +30,17 @@
 
 COLLECTION_NAME = "Weaviate_example_without_vectorizer_collection"
 
+default_args = {
+    "retries": 5,
+    "retry_delay": timedelta(seconds=15),
+    "pool": "weaviate_pool",
+}
+
 
 @dag(
     schedule=None,
     start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
+    default_args=default_args,
     catchup=False,
     tags=["example", "weaviate"],
 )