[dagster-airlift][components 2/n] scaffolder

Author: dpeng817

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

@@ -4,9 +4,12 @@
 
 from dagster._core.definitions.definitions_class import Definitions
 from dagster.components import Component, ComponentLoadContext, Resolvable
+from dagster.components.component_scaffolding import scaffold_component
 from dagster.components.resolved.context import ResolutionContext
 from dagster.components.resolved.core_models import AssetPostProcessor
 from dagster.components.resolved.model import Resolver
+from dagster.components.scaffold.scaffold import Scaffolder, ScaffoldRequest, scaffold_with
+from pydantic import BaseModel
 from typing_extensions import TypeAlias
 
 import dagster_airlift.core as dg_airlift_core
@@ -28,6 +31,32 @@ class AirflowMwaaAuthBackendModel(Resolvable):
     type: Literal["mwaa"]
 
 
+class AirflowInstanceScaffolderParams(BaseModel):
+    name: str
+    auth_type: Literal["basic_auth", "mwaa"]
+
+
+class AirflowInstanceScaffolder(Scaffolder):
+    @classmethod
+    def get_scaffold_params(cls) -> Optional[type[BaseModel]]:
+        return AirflowInstanceScaffolderParams
+
+    def scaffold(self, request: ScaffoldRequest, params: AirflowInstanceScaffolderParams) -> None:
+        full_params = {
+            "name": params.name,
+        }
+        if params.auth_type == "basic_auth":
+            full_params["auth"] = {
+                "type": "basic_auth",
+                "webserver_url": '{{ env("AIRFLOW_WEBSERVER_URL") }}',
+                "username": '{{ env("AIRFLOW_USERNAME") }}',
+                "password": '{{ env("AIRFLOW_PASSWORD") }}',
+            }
+        else:
+            raise ValueError(f"Unsupported auth type: {params.auth_type}")
+        scaffold_component(request, full_params)
+
+
 def resolve_auth(context: ResolutionContext, model) -> AirflowAuthBackend:
     if model.auth.type == "basic_auth":
         return AirflowBasicAuthBackend(
@@ -48,6 +77,7 @@ def resolve_auth(context: ResolutionContext, model) -> AirflowAuthBackend:
 ]
 
 
+@scaffold_with(AirflowInstanceScaffolder)
 @dataclass
 class AirflowInstanceComponent(Component, Resolvable):
     auth: ResolvedAirflowAuthBackend

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

@@ -1,10 +1,19 @@
+import json
+from pathlib import Path
+
 import dagster_airlift.core as dg_airlift_core
 import pytest
+import yaml
+from click.testing import CliRunner
 from dagster._core.test_utils import ensure_dagster_tests_import
+from dagster.components.cli import cli
 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
+from dagster_tests.components_tests.utils import (
+    build_component_defs_for_test,
+    temp_code_location_bar,
+)
 
 ensure_dagster_tests_import()
 
@@ -68,3 +77,56 @@ def test_load_dags_basic(component_for_test: type[AirflowInstanceComponent]) ->
         assert keyed_spec.metadata["foo"] == "bar"
 
     assert len(defs.jobs) == 3  # monitoring job + 2 dag jobs.
+
+
+def _scaffold_airlift(scaffold_format: str):
+    runner = CliRunner()
+    result = runner.invoke(
+        cli,
+        [
+            "scaffold",
+            "object",
+            "dagster_airlift.core.components.airflow_instance.component.AirflowInstanceComponent",
+            "bar/components/qux",
+            "--json-params",
+            json.dumps({"name": "qux", "auth_type": "basic_auth"}),
+            "--scaffold-format",
+            scaffold_format,
+        ],
+    )
+    assert result.exit_code == 0
+
+
+def test_scaffold_airlift_yaml():
+    with temp_code_location_bar():
+        _scaffold_airlift("yaml")
+        assert Path("bar/components/qux/component.yaml").exists()
+        with open("bar/components/qux/component.yaml") as f:
+            assert yaml.safe_load(f) == {
+                "type": "dagster_airlift.core.components.airflow_instance.component.AirflowInstanceComponent",
+                "attributes": {
+                    "name": "qux",
+                    "auth": {
+                        "type": "basic_auth",
+                        "webserver_url": '{{ env("AIRFLOW_WEBSERVER_URL") }}',
+                        "username": '{{ env("AIRFLOW_USERNAME") }}',
+                        "password": '{{ env("AIRFLOW_PASSWORD") }}',
+                    },
+                },
+            }
+
+
+def test_scaffold_airlift_python():
+    with temp_code_location_bar():
+        _scaffold_airlift("python")
+        assert Path("bar/components/qux/component.py").exists()
+        with open("bar/components/qux/component.py") as f:
+            file_contents = f.read()
+            assert file_contents == (
+                """from dagster.components import component, ComponentLoadContext
+from dagster_airlift.core.components.airflow_instance.component import AirflowInstanceComponent
+
+@component
+def load(context: ComponentLoadContext) -> AirflowInstanceComponent: ...
+"""
+            )