[dagster-airlift] instance component

Author: dpeng817

.mcp.json

@@ -0,0 +1,3 @@
+{
+  "mcpServers": {}
+}

python_modules/libraries/dagster-airlift/dagster_airlift/core/components/__init__.py

@@ -0,0 +1,69 @@
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Annotated, Literal, Optional, Union
+
+from dagster._core.definitions.definitions_class import Definitions
+from dagster.components import Component, ComponentLoadContext, Resolvable
+from dagster.components.resolved.context import ResolutionContext
+from dagster.components.resolved.core_models import AssetPostProcessor
+from dagster.components.resolved.model import Resolver
+from typing_extensions import TypeAlias
+
+import dagster_airlift.core as dg_airlift_core
+from dagster_airlift.core.airflow_instance import AirflowAuthBackend
+from dagster_airlift.core.basic_auth import AirflowBasicAuthBackend
+from dagster_airlift.core.load_defs import build_job_based_airflow_defs
+
+
+@dataclass
+class AirflowBasicAuthBackendModel(Resolvable):
+    type: Literal["basic_auth"]
+    webserver_url: str
+    username: str
+    password: str
+
+
+@dataclass
+class AirflowMwaaAuthBackendModel(Resolvable):
+    type: Literal["mwaa"]
+
+
+def resolve_auth(context: ResolutionContext, model) -> AirflowAuthBackend:
+    if model.auth.type == "basic_auth":
+        return AirflowBasicAuthBackend(
+            webserver_url=model.auth.webserver_url,
+            username=model.auth.username,
+            password=model.auth.password,
+        )
+    else:
+        raise ValueError(f"Unsupported auth type: {model.auth.type}")
+
+
+ResolvedAirflowAuthBackend: TypeAlias = Annotated[
+    AirflowAuthBackend,
+    Resolver.from_model(
+        resolve_auth,
+        model_field_type=Union[AirflowBasicAuthBackendModel, AirflowMwaaAuthBackendModel],
+    ),
+]
+
+
+@dataclass
+class AirflowInstanceComponent(Component, Resolvable):
+    auth: ResolvedAirflowAuthBackend
+    name: str
+    asset_post_processors: Optional[Sequence[AssetPostProcessor]] = None
+
+    def _get_instance(self) -> dg_airlift_core.AirflowInstance:
+        return dg_airlift_core.AirflowInstance(
+            auth_backend=self.auth,
+            name=self.name,
+        )
+
+    def build_defs(self, context: ComponentLoadContext) -> Definitions:
+        defs = build_job_based_airflow_defs(
+            airflow_instance=self._get_instance(),
+        )
+        for post_processor in self.asset_post_processors or []:
+            defs = post_processor(defs)
+        return defs

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

@@ -0,0 +1,67 @@
+import dagster_airlift.core as dg_airlift_core
+import pytest
+from dagster_airlift.core.components.airflow_instance.component import AirflowInstanceComponent
+from dagster_airlift.test import make_instance
+from dagster_airlift.test.test_utils import asset_spec
+from dagster_tests.components_tests.utils import build_component_defs_for_test
+
+
+@pytest.fixture
+def component_for_test():
+    airflow_instance = make_instance(
+        {"dag_1": ["dag_1_task_1", "dag_1_task_2"], "dag_2": ["dag_2_task_1", "dag_2_task_2"]},
+        dataset_construction_info=[
+            {
+                "uri": "s3://dataset-bucket/example1.csv",
+                "producing_tasks": [
+                    {"dag_id": "dag_1", "task_id": "dag_1_task_1"},
+                ],
+                "consuming_dags": ["dag_2"],
+            },
+            {
+                "uri": "s3://dataset-bucket/example2.csv",
+                "producing_tasks": [
+                    {"dag_id": "dag_2", "task_id": "dag_2_task_1"},
+                ],
+                "consuming_dags": [],
+            },
+        ],
+    )
+
+    class DebugAirflowInstanceComponent(AirflowInstanceComponent):
+        def _get_instance(self) -> dg_airlift_core.AirflowInstance:
+            return airflow_instance
+
+    return DebugAirflowInstanceComponent
+
+
+def test_load_dags_basic(component_for_test: type[AirflowInstanceComponent]) -> None:
+    defs = build_component_defs_for_test(
+        component_for_test,
+        {
+            "auth": {
+                "type": "basic_auth",
+                "webserver_url": "http://localhost:8080",
+                "username": "admin",
+                "password": "admin",
+            },
+            "name": "test_instance",
+            "asset_post_processors": [
+                {
+                    "target": "*",
+                    "attributes": {
+                        "metadata": {
+                            "foo": "bar",
+                        },
+                    },
+                }
+            ],
+        },
+    )
+
+    for asset_key in ["example1", "example2"]:
+        keyed_spec = asset_spec(asset_key, defs)
+        assert keyed_spec is not None
+        assert keyed_spec.metadata["foo"] == "bar"
+
+    assert len(defs.jobs) == 3  # monitoring job + 2 dag jobs.

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

@@ -37,10 +37,14 @@ run_observation_defs:
 run_job_based_defs:
 	dagster dev -m kitchen_sink.dagster_defs.job_based_defs -p 3333
 
+run_component_defs:
+	dagster dev -m kitchen_sink.dagster_defs.component_defs -p 3333
+
 # Command to point at a workspace.yaml
 run_dagster_multi_code_locations:
 	dagster dev -w $(MAKEFILE_DIR)/kitchen_sink/dagster_multi_code_locations/workspace.yaml -p 3333
 
+
 wipe: ## Wipe out all the files created by the Makefile
 	rm -rf $(AIRFLOW_HOME) $(DAGSTER_HOME)
 

python_modules/libraries/dagster-airlift/dagster_airlift_tests/unit_tests/core_tests/test_component_defs.py

@@ -0,0 +1,5 @@
+from dagster.components import load_defs
+
+import kitchen_sink.dagster_defs.inner_component_defs as inner_component_defs
+
+defs = load_defs(defs_root=inner_component_defs)

python_modules/libraries/dagster-airlift/kitchen-sink/Makefile

@@ -0,0 +1,9 @@
+type: dagster_airlift.core.components.airflow_instance.component.AirflowInstanceComponent
+
+attributes:
+  name: kitchen_sink_instance 
+  auth:
+    type: basic_auth
+    username: admin
+    password: admin
+    webserver_url: http://localhost:8080

python_modules/libraries/dagster-airlift/kitchen-sink/kitchen_sink/dagster_defs/component_defs.py

@@ -0,0 +1,96 @@
+from dagster._core.definitions.asset_key import AssetKey
+from dagster._core.definitions.metadata.metadata_value import TimestampMetadataValue
+from dagster._core.event_api import EventRecordsFilter
+from dagster._core.events import DagsterEventType
+from dagster._core.instance_for_test import instance_for_test
+from dagster._core.storage.dagster_run import DagsterRunStatus
+from dagster._core.storage.tags import EXTERNAL_JOB_SOURCE_TAG_KEY
+from dagster_airlift.constants import DAG_ID_TAG_KEY, DAG_RUN_ID_TAG_KEY
+from dagster_airlift.core.utils import monitoring_job_name
+from dagster_airlift.test.test_utils import asset_spec
+from kitchen_sink.airflow_instance import local_airflow_instance
+
+from kitchen_sink_tests.integration_tests.conftest import (
+    poll_for_airflow_run_existence_and_completion,
+)
+
+
+def test_component_based_defs(
+    airflow_instance: None,
+) -> None:
+    """Test that component based defs load properly."""
+    from kitchen_sink.dagster_defs.component_defs import defs
+
+    assert len(defs.jobs) == 20
+    assert len(defs.assets) == 1
+    for key in ["example1", "example2"]:
+        assert asset_spec(key, defs)
+
+    # First, execute dataset producer dag
+    af_instance = local_airflow_instance("kitchen_sink_instance")
+    af_run_id = af_instance.trigger_dag("dataset_producer")
+    poll_for_airflow_run_existence_and_completion(
+        af_instance=af_instance, af_run_id=af_run_id, dag_id="dataset_producer", duration=30
+    )
+
+    # Then, execute monitoring job
+    with instance_for_test() as instance:
+        result = defs.execute_job_in_process(
+            monitoring_job_name(af_instance.name), instance=instance
+        )
+        assert result.success
+        # There should be a run for the dataset producer dag
+        runs = instance.get_runs()
+        assert len(runs) == 2
+        producer_run = next(run for run in runs if run.job_name == "dataset_producer")
+        assert producer_run.status == DagsterRunStatus.SUCCESS
+        assert producer_run.tags[DAG_RUN_ID_TAG_KEY] == af_run_id
+        assert producer_run.tags[DAG_ID_TAG_KEY] == "dataset_producer"
+        assert producer_run.tags[EXTERNAL_JOB_SOURCE_TAG_KEY] == "airflow"
+
+        # Check that there are asset planned events for the two assets
+        planned_records = instance.get_event_records(
+            event_records_filter=EventRecordsFilter(
+                event_type=DagsterEventType.ASSET_MATERIALIZATION_PLANNED
+            ),
+        )
+        assert len(planned_records) == 2
+        assert {planned_records.asset_key for planned_records in planned_records} == {
+            AssetKey("example1"),
+            AssetKey("example2"),
+        }
+        assert {planned_records.run_id for planned_records in planned_records} == {
+            producer_run.run_id
+        }
+
+        # Check that there are asset materialized events for the two assets
+        materialized_records = instance.get_event_records(
+            event_records_filter=EventRecordsFilter(
+                event_type=DagsterEventType.ASSET_MATERIALIZATION
+            ),
+        )
+        assert len(materialized_records) == 2
+        assert {
+            materialized_records.asset_key for materialized_records in materialized_records
+        } == {
+            AssetKey("example1"),
+            AssetKey("example2"),
+        }
+        assert {materialized_records.run_id for materialized_records in materialized_records} == {
+            producer_run.run_id
+        }
+        for key in ["example1", "example2"]:
+            key_record = next(
+                materialized_records
+                for materialized_records in materialized_records
+                if materialized_records.asset_key == AssetKey(key)
+            )
+            assert key_record.asset_materialization.metadata
+            assert key_record.asset_materialization.metadata[
+                "my_timestamp"
+            ] == TimestampMetadataValue(value=111.0)
+            assert key_record.asset_materialization.metadata[
+                "my_other_timestamp"
+            ] == TimestampMetadataValue(value=113.0)
+            # It gets overridden by the second print
+            assert key_record.asset_materialization.metadata["foo"].value == "baz"