Merge pull request #104 from ooni/fix/analysis

Fix analysis generation in airflow
ooni · Jan 9, 2025 · a49cb5f · a49cb5f
2 parents 7c14a71 + e1502a6
commit a49cb5f
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 8 deletions.
diff --git a/dags/pipeline.py b/dags/pipeline.py
@@ -31,6 +31,7 @@ def run_make_observations(
 
 
 def run_make_analysis(
+    clickhouse_url: str,
     probe_cc: List[str],
     test_name: List[str],
     day: str,
@@ -40,7 +41,9 @@ def run_make_analysis(
         make_analysis_in_a_day,
     )
 
-    params = MakeAnalysisParams(probe_cc=probe_cc, test_name=test_name, day=day)
+    params = MakeAnalysisParams(
+        probe_cc=probe_cc, test_name=test_name, day=day, clickhouse_url=clickhouse_url
+    )
     make_analysis_in_a_day(params)
 
 
@@ -59,14 +62,14 @@ def run_make_analysis(
     start_date=datetime.datetime(2012, 12, 4),
     schedule="@daily",
     catchup=False,
-) as dag:
+) as dag_full:
     start_day = "{{ ds }}"
     op_make_observations = PythonVirtualenvOperator(
         task_id="make_observations",
         python_callable=run_make_observations,
         op_kwargs={
-            "probe_cc": dag.params["probe_cc"],
-            "test_name": dag.params["test_name"],
+            "probe_cc": dag_full.params["probe_cc"],
+            "test_name": dag_full.params["test_name"],
             "clickhouse_url": Variable.get("clickhouse_url", default_var=""),
             "data_dir": Variable.get("data_dir", default_var=""),
             "bucket_date": start_day,
@@ -79,12 +82,43 @@ def run_make_analysis(
         task_id="make_analysis",
         python_callable=run_make_analysis,
         op_kwargs={
-            "probe_cc": dag.params["probe_cc"],
-            "test_name": dag.params["test_name"],
+            "probe_cc": dag_full.params["probe_cc"],
+            "test_name": dag_full.params["test_name"],
+            "clickhouse_url": Variable.get("clickhouse_url", default_var=""),
             "day": start_day,
         },
         requirements=REQUIREMENTS,
         system_site_packages=False,
     )
 
     op_make_observations >> op_make_analysis
+
+with DAG(
+    dag_id="batch_analysis_only",
+    default_args={
+        "retries": 3,
+        "retry_delay": datetime.timedelta(minutes=30),
+    },
+    params={
+        "probe_cc": Param(default=[], type=["null", "array"]),
+        "test_name": Param(default=[], type=["null", "array"]),
+    },
+    start_date=datetime.datetime(2012, 12, 4),
+    schedule=None,
+    catchup=False,
+) as dag_analysis:
+    start_day = "{{ ds }}"
+    op_make_analysis_only = PythonVirtualenvOperator(
+        task_id="make_analysis",
+        python_callable=run_make_analysis,
+        op_kwargs={
+            "probe_cc": dag_analysis.params["probe_cc"],
+            "test_name": dag_analysis.params["test_name"],
+            "clickhouse_url": Variable.get("clickhouse_url", default_var=""),
+            "day": start_day,
+        },
+        requirements=REQUIREMENTS,
+        system_site_packages=False,
+    )
+
+    op_make_analysis_only
diff --git a/oonipipeline/src/oonipipeline/tasks/analysis.py b/oonipipeline/src/oonipipeline/tasks/analysis.py
@@ -11,6 +11,7 @@
 
 @dataclass
 class MakeAnalysisParams:
+    clickhouse_url: str
     probe_cc: List[str]
     test_name: List[str]
     day: str
@@ -23,7 +24,7 @@ def make_analysis_in_a_day(params: MakeAnalysisParams):
 
     probe_cc = params.probe_cc
     test_name = params.test_name
-    db = ClickhouseConnection(config.clickhouse_url)
+    db = ClickhouseConnection(params.clickhouse_url)
 
     write_analysis_web_fuzzy_logic(
         db=db,

diff --git a/oonipipeline/tests/test_cli.py b/oonipipeline/tests/test_cli.py
@@ -32,7 +32,7 @@ def test_full_workflow(
     result = cli_runner.invoke(
         cli,
         [
-            "backfill",
+            "run",
             "--start-at",
             "2022-10-21",
             "--end-at",

diff --git a/oonipipeline/tests/test_workflows.py b/oonipipeline/tests/test_workflows.py
@@ -66,6 +66,7 @@ def test_make_file_entry_batch(datadir, db):
     make_analysis_in_a_day(
         MakeAnalysisParams(
             probe_cc=["IR"],
+            clickhouse_url=db.clickhouse_url,
             test_name=["webconnectivity"],
             day=date(2023, 10, 31).strftime("%Y-%m-%d"),
         ),