[dagster-airlift] log retrieval API

Author: dpeng817

python_modules/libraries/dagster-airlift/dagster_airlift/core/airflow_instance.py

@@ -542,3 +542,38 @@ def get_all_datasets(
             offset += batch_size
 
         return datasets
+
+    def get_task_instance_logs(
+        self, dag_id: str, task_id: str, run_id: str, try_number: int
+    ) -> str:
+        continuation_token = None
+        logs = []
+        while True:
+            response = self.auth_backend.get_session().get(
+                f"{self.get_api_url()}/dags/{dag_id}/dagRuns/{run_id}/taskInstances/{task_id}/logs/{try_number}",
+                headers={"Accept": "application/json"},
+                params={"token": continuation_token} if continuation_token else None,
+                timeout=5,
+            )
+            if response.status_code != 200:
+                raise DagsterError(
+                    f"Failed to fetch task instance logs for {dag_id}/{task_id}/{run_id}/{try_number}. Status code: {response.status_code}, Message: {response.text}"
+                )
+            data = response.json()
+            # Love how it's different in the two cases lol.
+            continuation_token = data.get("continuation_token")
+            log = parse_af_log_response(data["content"])
+            logs.append(log)
+            if not continuation_token or log == "":
+                break
+        return "".join(logs)
+
+
+def parse_af_log_response(logs: str) -> str:
+    import ast
+
+    parsed_data: list = ast.literal_eval(logs)
+    strs = []
+    for log_item in parsed_data:
+        strs.append(log_item[1])
+    return "".join(strs)

python_modules/libraries/dagster-airlift/kitchen-sink/kitchen_sink/airflow_dags/dataset_dags.py

@@ -16,7 +16,11 @@ def print_fn() -> None:
     import os
 
     os.environ["NO_PROXY"] = "*"
+    import json
+
     print("Hello")  # noqa: T201
+    data = json.dumps({"foo": "bar"})
+    print(f"DAGSTER_START{data}DAGSTER_END")  # noqa: T201
 
 
 # Inter-dag structure as follows:

python_modules/libraries/dagster-airlift/kitchen-sink/kitchen_sink_tests/integration_tests/test_api.py

@@ -205,3 +205,15 @@ def test_datasets(airflow_instance: None) -> None:
     )
     assert not asset_spec("example1", defs)
     assert not asset_spec("example2", defs)
+
+
+def test_log_retrieval(airflow_instance: None) -> None:
+    af_instance = local_airflow_instance()
+    run_id = af_instance.trigger_dag("dataset_producer")
+    af_instance.wait_for_run_completion(dag_id="dataset_producer", run_id=run_id)
+    logs = af_instance.get_task_instance_logs(
+        dag_id="dataset_producer", task_id="print_task", run_id=run_id, try_number=1
+    )
+    assert logs
+    assert "DAGSTER_START" in logs
+    assert "DAGSTER_END" in logs