Replace weaviate provider with apache-airflow-providers-weaviate==1.3.0

airflow/dags/ingestion/ask-astro-forum-load.py

@@ -1,17 +1,13 @@
 import datetime
 import os
 
-from include.tasks import split
-from include.tasks.extract.astro_forum_docs import get_forum_df
-from include.tasks.extract.utils.weaviate.ask_astro_weaviate_hook import AskAstroWeaviateHook
-
 from airflow.decorators import dag, task
+from airflow.providers.weaviate.operators.weaviate import WeaviateDocumentIngestOperator
 
 ask_astro_env = os.environ.get("ASK_ASTRO_ENV", "dev")
 
 _WEAVIATE_CONN_ID = f"weaviate_{ask_astro_env}"
 WEAVIATE_CLASS = os.environ.get("WEAVIATE_CLASS", "DocsDev")
-ask_astro_weaviate_hook = AskAstroWeaviateHook(_WEAVIATE_CONN_ID)
 
 blog_cutoff_date = datetime.date(2022, 1, 1)
 
@@ -22,6 +18,8 @@
 
 @task
 def get_astro_forum_content():
+    from include.tasks.extract.astro_forum_docs import get_forum_df
+
     return get_forum_df()
 
 
@@ -33,19 +31,19 @@ def get_astro_forum_content():
     default_args=default_args,
 )
 def ask_astro_load_astro_forum():
+    from include.tasks import split
+
     split_docs = task(split.split_html).expand(dfs=[get_astro_forum_content()])
 
-    _import_data = (
-        task(ask_astro_weaviate_hook.ingest_data, retries=10)
-        .partial(
-            class_name=WEAVIATE_CLASS,
-            existing="upsert",
-            doc_key="docLink",
-            batch_params={"batch_size": 1000},
-            verbose=True,
-        )
-        .expand(dfs=[split_docs])
-    )
+    _import_data = WeaviateDocumentIngestOperator.partial(
+        class_name=WEAVIATE_CLASS,
+        existing="replace",
+        document_column="docLink",
+        batch_config_params={"batch_size": 1000},
+        verbose=True,
+        conn_id=_WEAVIATE_CONN_ID,
+        task_id="WeaviateDocumentIngestOperator",
+    ).expand(input_data=[split_docs])
 
 
 ask_astro_load_astro_forum()

airflow/dags/ingestion/ask-astro-load-airflow-docs.py

@@ -1,18 +1,14 @@
 import os
 from datetime import datetime
 
-from include.tasks import split
-from include.tasks.extract import airflow_docs
-from include.tasks.extract.utils.weaviate.ask_astro_weaviate_hook import AskAstroWeaviateHook
-
 from airflow.decorators import dag, task
+from airflow.providers.weaviate.operators.weaviate import WeaviateDocumentIngestOperator
 
 ask_astro_env = os.environ.get("ASK_ASTRO_ENV", "dev")
 
 _WEAVIATE_CONN_ID = f"weaviate_{ask_astro_env}"
 WEAVIATE_CLASS = os.environ.get("WEAVIATE_CLASS", "DocsDev")
 
-ask_astro_weaviate_hook = AskAstroWeaviateHook(_WEAVIATE_CONN_ID)
 
 airflow_docs_base_url = "https://airflow.apache.org/docs/"
 
@@ -34,22 +30,22 @@ def ask_astro_load_airflow_docs():
     data from a point-in-time data capture. By using the upsert logic of the weaviate_import decorator
     any existing documents that have been updated will be removed and re-added.
     """
+    from include.tasks import split
+    from include.tasks.extract import airflow_docs
 
     extracted_airflow_docs = task(airflow_docs.extract_airflow_docs)(docs_base_url=airflow_docs_base_url)
 
     split_md_docs = task(split.split_html).expand(dfs=[extracted_airflow_docs])
 
-    _import_data = (
-        task(ask_astro_weaviate_hook.ingest_data, retries=10)
-        .partial(
-            class_name=WEAVIATE_CLASS,
-            existing="upsert",
-            doc_key="docLink",
-            batch_params={"batch_size": 1000},
-            verbose=True,
-        )
-        .expand(dfs=[split_md_docs])
-    )
+    _import_data = WeaviateDocumentIngestOperator.partial(
+        class_name=WEAVIATE_CLASS,
+        existing="replace",
+        document_column="docLink",
+        batch_config_params={"batch_size": 1000},
+        verbose=True,
+        conn_id=_WEAVIATE_CONN_ID,
+        task_id="WeaviateDocumentIngestOperator",
+    ).expand(input_data=[split_md_docs])
 
 
 ask_astro_load_airflow_docs()

airflow/dags/ingestion/ask-astro-load-astro-cli.py

@@ -1,17 +1,13 @@
 import datetime
 import os
 
-from include.tasks import split
-from include.tasks.extract import astro_cli_docs
-from include.tasks.extract.utils.weaviate.ask_astro_weaviate_hook import AskAstroWeaviateHook
-
 from airflow.decorators import dag, task
+from airflow.providers.weaviate.operators.weaviate import WeaviateDocumentIngestOperator
 
 ask_astro_env = os.environ.get("ASK_ASTRO_ENV", "dev")
 
 _WEAVIATE_CONN_ID = f"weaviate_{ask_astro_env}"
 WEAVIATE_CLASS = os.environ.get("WEAVIATE_CLASS", "DocsDev")
-ask_astro_weaviate_hook = AskAstroWeaviateHook(_WEAVIATE_CONN_ID)
 
 default_args = {"retries": 3, "retry_delay": 30}
 
@@ -31,21 +27,21 @@ def ask_astro_load_astro_cli_docs():
     data from a point-in-time data capture. By using the upsert logic of the weaviate_import decorator
     any existing documents that have been updated will be removed and re-added.
     """
+    from include.tasks import split
+    from include.tasks.extract import astro_cli_docs
 
     extract_astro_cli_docs = task(astro_cli_docs.extract_astro_cli_docs)()
     split_md_docs = task(split.split_html).expand(dfs=[extract_astro_cli_docs])
 
-    _import_data = (
-        task(ask_astro_weaviate_hook.ingest_data, retries=10)
-        .partial(
-            class_name=WEAVIATE_CLASS,
-            existing="upsert",
-            doc_key="docLink",
-            batch_params={"batch_size": 1000},
-            verbose=True,
-        )
-        .expand(dfs=[split_md_docs])
-    )
+    _import_data = WeaviateDocumentIngestOperator.partial(
+        class_name=WEAVIATE_CLASS,
+        existing="replace",
+        document_column="docLink",
+        batch_config_params={"batch_size": 1000},
+        verbose=True,
+        conn_id=_WEAVIATE_CONN_ID,
+        task_id="WeaviateDocumentIngestOperator",
+    ).expand(input_data=[split_md_docs])
 
 
 ask_astro_load_astro_cli_docs()

airflow/dags/ingestion/ask-astro-load-astro-sdk.py

@@ -1,16 +1,13 @@
 import datetime
 import os
 
-from include.tasks.extract.astro_sdk_docs import extract_astro_sdk_docs
-from include.tasks.extract.utils.weaviate.ask_astro_weaviate_hook import AskAstroWeaviateHook
-
 from airflow.decorators import dag, task
+from airflow.providers.weaviate.operators.weaviate import WeaviateDocumentIngestOperator
 
 ask_astro_env = os.environ.get("ASK_ASTRO_ENV", "dev")
 
 _WEAVIATE_CONN_ID = f"weaviate_{ask_astro_env}"
 WEAVIATE_CLASS = os.environ.get("WEAVIATE_CLASS", "DocsDev")
-ask_astro_weaviate_hook = AskAstroWeaviateHook(_WEAVIATE_CONN_ID)
 
 blog_cutoff_date = datetime.date(2023, 1, 19)
 
@@ -21,6 +18,8 @@
 
 @task
 def get_astro_sdk_content():
+    from include.tasks.extract.astro_sdk_docs import extract_astro_sdk_docs
+
     dfs = extract_astro_sdk_docs()
     return dfs
 
@@ -33,17 +32,15 @@ def get_astro_sdk_content():
     default_args=default_args,
 )
 def ask_astro_load_astro_sdk():
-    _import_data = (
-        task(ask_astro_weaviate_hook.ingest_data, retries=10)
-        .partial(
-            class_name=WEAVIATE_CLASS,
-            existing="upsert",
-            doc_key="docLink",
-            batch_params={"batch_size": 1000},
-            verbose=True,
-        )
-        .expand(dfs=[get_astro_sdk_content()])
-    )
+    _import_data = WeaviateDocumentIngestOperator.partial(
+        class_name=WEAVIATE_CLASS,
+        existing="replace",
+        document_column="docLink",
+        batch_config_params={"batch_size": 1000},
+        verbose=True,
+        conn_id=_WEAVIATE_CONN_ID,
+        task_id="WeaviateDocumentIngestOperator",
+    ).expand(input_data=[get_astro_sdk_content()])
 
 
 ask_astro_load_astro_sdk()

airflow/dags/ingestion/ask-astro-load-astronomer-docs.py

@@ -2,6 +2,7 @@
 import os
 
 from airflow.decorators import dag, task
+from airflow.providers.weaviate.operators.weaviate import WeaviateDocumentIngestOperator
 
 ask_astro_env = os.environ.get("ASK_ASTRO_ENV", "dev")
 
@@ -27,24 +28,20 @@ def ask_astro_load_astronomer_docs():
     """
     from include.tasks import split
     from include.tasks.extract.astro_docs import extract_astro_docs
-    from include.tasks.extract.utils.weaviate.ask_astro_weaviate_hook import AskAstroWeaviateHook
 
-    ask_astro_weaviate_hook = AskAstroWeaviateHook(_WEAVIATE_CONN_ID)
     astro_docs = task(extract_astro_docs)()
 
     split_md_docs = task(split.split_markdown).expand(dfs=[astro_docs])
 
-    _import_data = (
-        task(ask_astro_weaviate_hook.ingest_data, retries=10)
-        .partial(
-            class_name=WEAVIATE_CLASS,
-            existing="upsert",
-            doc_key="docLink",
-            batch_params={"batch_size": 1000},
-            verbose=True,
-        )
-        .expand(dfs=[split_md_docs])
-    )
+    _import_data = WeaviateDocumentIngestOperator.partial(
+        class_name=WEAVIATE_CLASS,
+        existing="replace",
+        document_column="docLink",
+        batch_config_params={"batch_size": 1000},
+        verbose=True,
+        conn_id=_WEAVIATE_CONN_ID,
+        task_id="WeaviateDocumentIngestOperator",
+    ).expand(input_data=[split_md_docs])
 
 
 ask_astro_load_astronomer_docs()

airflow/dags/ingestion/ask-astro-load-astronomer-provider.py

@@ -1,16 +1,13 @@
 import datetime
 import os
 
-from include.tasks.extract.astronomer_providers_docs import extract_provider_docs
-from include.tasks.extract.utils.weaviate.ask_astro_weaviate_hook import AskAstroWeaviateHook
-
 from airflow.decorators import dag, task
+from airflow.providers.weaviate.operators.weaviate import WeaviateDocumentIngestOperator
 
 ask_astro_env = os.environ.get("ASK_ASTRO_ENV", "dev")
 
 _WEAVIATE_CONN_ID = f"weaviate_{ask_astro_env}"
 WEAVIATE_CLASS = os.environ.get("WEAVIATE_CLASS", "DocsDev")
-ask_astro_weaviate_hook = AskAstroWeaviateHook(_WEAVIATE_CONN_ID)
 
 blog_cutoff_date = datetime.date(2023, 1, 19)
 
@@ -21,6 +18,8 @@
 
 @task
 def get_provider_content():
+    from include.tasks.extract.astronomer_providers_docs import extract_provider_docs
+
     dfs = extract_provider_docs()
     return dfs
 
@@ -39,17 +38,15 @@ def ask_astro_load_astronomer_providers():
     any existing documents that have been updated will be removed and re-added.
     """
 
-    _import_data = (
-        task(ask_astro_weaviate_hook.ingest_data, retries=10)
-        .partial(
-            class_name=WEAVIATE_CLASS,
-            existing="upsert",
-            doc_key="docLink",
-            batch_params={"batch_size": 1000},
-            verbose=True,
-        )
-        .expand(dfs=[get_provider_content()])
-    )
+    _import_data = WeaviateDocumentIngestOperator.partial(
+        class_name=WEAVIATE_CLASS,
+        existing="replace",
+        document_column="docLink",
+        batch_config_params={"batch_size": 1000},
+        verbose=True,
+        conn_id=_WEAVIATE_CONN_ID,
+        task_id="WeaviateDocumentIngestOperator",
+    ).expand(input_data=[get_provider_content()])
 
 
 ask_astro_load_astronomer_providers()

airflow/dags/ingestion/ask-astro-load-blogs.py

@@ -1,17 +1,13 @@
 import datetime
 import os
 
-from include.tasks import split
-from include.tasks.extract import blogs
-from include.tasks.extract.utils.weaviate.ask_astro_weaviate_hook import AskAstroWeaviateHook
-
 from airflow.decorators import dag, task
+from airflow.providers.weaviate.operators.weaviate import WeaviateDocumentIngestOperator
 
 ask_astro_env = os.environ.get("ASK_ASTRO_ENV", "dev")
 
 _WEAVIATE_CONN_ID = f"weaviate_{ask_astro_env}"
 WEAVIATE_CLASS = os.environ.get("WEAVIATE_CLASS", "DocsDev")
-ask_astro_weaviate_hook = AskAstroWeaviateHook(_WEAVIATE_CONN_ID)
 
 blog_cutoff_date = datetime.date(2023, 1, 19)
 
@@ -33,22 +29,22 @@ def ask_astro_load_blogs():
     data from a point-in-time data capture. By using the upsert logic of the weaviate_import decorator
     any existing documents that have been updated will be removed and re-added.
     """
+    from include.tasks import split
+    from include.tasks.extract import blogs
 
     blogs_docs = task(blogs.extract_astro_blogs)(blog_cutoff_date=blog_cutoff_date)
 
     split_md_docs = task(split.split_markdown).expand(dfs=[blogs_docs])
 
-    _import_data = (
-        task(ask_astro_weaviate_hook.ingest_data, retries=10)
-        .partial(
-            class_name=WEAVIATE_CLASS,
-            existing="upsert",
-            doc_key="docLink",
-            batch_params={"batch_size": 1000},
-            verbose=True,
-        )
-        .expand(dfs=[split_md_docs])
-    )
+    _import_data = WeaviateDocumentIngestOperator.partial(
+        class_name=WEAVIATE_CLASS,
+        existing="replace",
+        document_column="docLink",
+        batch_config_params={"batch_size": 1000},
+        verbose=True,
+        conn_id=_WEAVIATE_CONN_ID,
+        task_id="WeaviateDocumentIngestOperator",
+    ).expand(input_data=[split_md_docs])
 
 
 ask_astro_load_blogs()