feat(airflow): add astro cli bulk load

astronomer · Dec 5, 2023 · 236162a · 236162a
1 parent 1aeb78a
commit 236162a
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/airflow/dags/ingestion/ask-astro-load.py b/airflow/dags/ingestion/ask-astro-load.py
@@ -8,7 +8,7 @@
 
 import pandas as pd
 from include.tasks import split
-from include.tasks.extract import airflow_docs, blogs, github, registry, stack_overflow
+from include.tasks.extract import airflow_docs, astro_cli_docs, blogs, github, registry, stack_overflow
 from include.tasks.extract.utils.weaviate.ask_astro_weaviate_hook import AskAstroWeaviateHook
 
 from airflow.decorators import dag, task
@@ -177,6 +177,17 @@ def extract_airflow_docs():
 
         return [df]
 
+    @task(trigger_rule="none_failed")
+    def extract_astro_cli_docs():
+        astro_cli_parquet_path = "include/data/astronomer/docs/astro-cli.parquet"
+        try:
+            df = pd.read_parquet(astro_cli_parquet_path)
+        except Exception:
+            df = astro_cli_docs.extract_astro_cli_docs()[0]
+            df.to_parquet(astro_cli_parquet_path)
+
+        return [df]
+
     @task(trigger_rule="none_failed")
     def extract_stack_overflow(tag: str, stackoverflow_cutoff_date: str = stackoverflow_cutoff_date):
         try:
@@ -237,6 +248,7 @@ def extract_astro_blogs():
     registry_dags_docs = extract_astro_registry_dags()
     code_samples = extract_github_python.expand(source=code_samples_sources)
     _airflow_docs = extract_airflow_docs()
+    _astro_cli_docs = extract_astro_cli_docs()
 
     _get_schema = get_schema_and_process(schema_file="include/data/schema.json")
     _check_schema = check_schema(class_objects=_get_schema)
@@ -251,7 +263,7 @@ def extract_astro_blogs():
         registry_cells_docs,
     ]
 
-    html_tasks = [_airflow_docs]
+    html_tasks = [_airflow_docs, _astro_cli_docs]
 
     python_code_tasks = [registry_dags_docs, code_samples]
 

diff --git a/airflow/include/tasks/extract/astro_cli_docs.py b/airflow/include/tasks/extract/astro_cli_docs.py
@@ -27,7 +27,7 @@ def extract_astro_cli_docs() -> list[pd.DataFrame]:
     df["content"] = df["content"].apply(lambda x: re.sub("¶", "", x))
 
     df["sha"] = df["content"].apply(generate_uuid5)
-    df["docSource"] = "apache/airflow/docs"
+    df["docSource"] = "astronomer/docs/astro-cli"
     df.reset_index(drop=True, inplace=True)
 
     # column order matters for uuid generation