[dg] add --fix-env-requirements flag to dg check

Author: benpankow

python_modules/libraries/dagster-dg/dagster_dg/check.py

@@ -4,6 +4,7 @@
 from typing import Any, NamedTuple, Optional
 
 import click
+import yaml
 from dagster_shared.serdes.objects import LibraryObjectKey
 from dagster_shared.yaml_utils import parse_yaml_with_source_positions
 from dagster_shared.yaml_utils.source_position import (
@@ -62,12 +63,13 @@ class ErrorInput(NamedTuple):
 def check_yaml(
     dg_context: DgContext,
     resolved_paths: Sequence[Path],
+    fix_env_requirements: bool,
 ) -> bool:
     top_level_component_validator = Draft202012Validator(schema=COMPONENT_FILE_SCHEMA)
 
     validation_errors: list[ErrorInput] = []
     all_specified_env_var_deps = set()
-
+    updated_files = set()
     component_contents_by_key: dict[LibraryObjectKey, Any] = {}
     modules_to_fetch = set()
     for component_dir in dg_context.defs_path.iterdir():
@@ -103,21 +105,38 @@ def check_yaml(
             all_specified_env_var_deps.update(specified_env_var_deps)
 
             if used_env_vars - specified_env_var_deps:
-                msg = (
-                    "Component uses environment variables that are not specified in the component file: "
-                    + ", ".join(used_env_vars - specified_env_var_deps)
-                )
+                if fix_env_requirements:
+                    component_doc_updated = component_doc_tree.value
+                    if "requires" not in component_doc_updated:
+                        component_doc_updated["requires"] = {}
+                    if "env" not in component_doc_updated["requires"]:
+                        component_doc_updated["requires"]["env"] = []
+
+                    component_doc_updated["requires"]["env"].extend(
+                        used_env_vars - specified_env_var_deps
+                    )
 
-                validation_errors.append(
-                    ErrorInput(
-                        None,
-                        ValidationError(
-                            msg,
-                            path=["requires", "env"],
-                        ),
-                        component_doc_tree,
+                    component_path.write_text(yaml.dump(component_doc_updated, sort_keys=False))
+                    click.echo(f"Updated {component_path}")
+                    updated_files.add(component_path)
+                    all_specified_env_var_deps.update(used_env_vars)
+                else:
+                    msg = (
+                        "Component uses environment variables that are not specified in the component file: "
+                        + ", ".join(used_env_vars - specified_env_var_deps)
+                        + "\nTo automatically add these environment variables to the component file, run `dg check yaml --fix-env-requirements`"
+                    )
+
+                    validation_errors.append(
+                        ErrorInput(
+                            None,
+                            ValidationError(
+                                msg,
+                                path=["requires", "env"],
+                            ),
+                            component_doc_tree,
+                        )
                     )
-                )
 
             # First, validate the top-level structure of the component file
             # (type and params keys) before we try to validate the params themselves.
@@ -175,11 +194,22 @@ def check_yaml(
                     prefix=["attributes"] if key else [],
                 )
             )
+        if updated_files:
+            click.echo(
+                "The following component files were updated to fix environment variable requirements:\n"
+                + "\n".join(updated_files)
+            )
         return False
     else:
         missing_env_vars = (
             all_specified_env_var_deps - ProjectEnvVars.from_ctx(dg_context).values.keys()
         ) - os.environ.keys()
+        if updated_files:
+            click.echo(
+                "The following component files were updated to fix environment variable requirements:\n"
+                + "\n".join(str(file) for file in updated_files)
+            )
+
         if missing_env_vars:
             click.echo(
                 "The following environment variables are used in components but not specified in the .env file or the current shell environment:\n"

python_modules/libraries/dagster-dg/dagster_dg/cli/check.py

@@ -32,11 +32,17 @@ def check_group():
 @click.option(
     "--watch", is_flag=True, help="Watch for changes to the component files and re-validate them."
 )
+@click.option(
+    "--fix-env-requirements",
+    is_flag=True,
+    help="Automatically add environment variable requirements to component yaml files.",
+)
 @dg_global_options
 @cli_telemetry_wrapper
 def check_yaml_command(
     paths: Sequence[str],
     watch: bool,
+    fix_env_requirements: bool,
     **global_options: object,
 ) -> None:
     """Check component.yaml files against their schemas, showing validation errors."""
@@ -45,7 +51,7 @@ def check_yaml_command(
     resolved_paths = [Path(path).absolute() for path in paths]
 
     def run_check(_: Any = None) -> bool:
-        return check_yaml_fn(dg_context, resolved_paths)
+        return check_yaml_fn(dg_context, resolved_paths, fix_env_requirements=fix_env_requirements)
 
     if watch:
         watched_paths = (
@@ -145,6 +151,7 @@ def check_definitions_command(
                 check_result = check_yaml_fn(
                     dg_context.for_project_environment(project_dir, cli_config),
                     [],
+                    fix_env_requirements=False,
                 )
                 overall_check_result = overall_check_result and check_result
             if not overall_check_result:

python_modules/libraries/dagster-dg/dagster_dg/cli/env.py

@@ -1,10 +1,15 @@
+from collections.abc import Sequence
 from pathlib import Path
 
 import click
+import yaml
+from dagster_shared.yaml_utils import parse_yaml_with_source_positions
 from rich.console import Console
 from rich.table import Table
+from yaml.scanner import ScannerError
 
 from dagster_dg.cli.shared_options import dg_global_options
+from dagster_dg.component import get_specified_env_var_deps, get_used_env_vars
 from dagster_dg.config import normalize_cli_config
 from dagster_dg.context import DgContext
 from dagster_dg.env import ProjectEnvVars
@@ -42,3 +47,61 @@ def list_env_command(**global_options: object) -> None:
         table.add_row(key, value)
     console = Console()
     console.print(table)
+
+
+# ########################
+# ##### FIX COMPONENT YAML REQUIREMENTS
+# ########################
+
+
+@env_group.command(name="fix-component-requirements", cls=DgClickCommand)
+@click.argument("paths", nargs=-1, type=click.Path(exists=True))
+@dg_global_options
+def fix_component_requirements(paths: Sequence[str], **global_options: object) -> None:
+    """Automatically add environment variable requirements to component yaml files."""
+    cli_config = normalize_cli_config(global_options, click.get_current_context())
+    dg_context = DgContext.for_project_environment(Path.cwd(), cli_config)
+
+    resolved_paths = [Path(path).absolute() for path in paths]
+
+    updated_files = set()
+
+    for component_dir in dg_context.defs_path.iterdir():
+        if resolved_paths and not any(
+            path == component_dir or path in component_dir.parents for path in resolved_paths
+        ):
+            continue
+
+        component_path = component_dir / "component.yaml"
+
+        if component_path.exists():
+            text = component_path.read_text()
+            try:
+                component_doc_tree = parse_yaml_with_source_positions(
+                    text, filename=str(component_path)
+                )
+            except ScannerError:
+                continue
+
+            specified_env_var_deps = get_specified_env_var_deps(component_doc_tree.value)
+            used_env_vars = get_used_env_vars(component_doc_tree.value)
+
+            if used_env_vars - specified_env_var_deps:
+                component_doc_updated = component_doc_tree.value
+                if "requires" not in component_doc_updated:
+                    component_doc_updated["requires"] = {}
+                if "env" not in component_doc_updated["requires"]:
+                    component_doc_updated["requires"]["env"] = []
+
+                component_doc_updated["requires"]["env"].extend(
+                    used_env_vars - specified_env_var_deps
+                )
+
+                component_path.write_text(yaml.dump(component_doc_updated, sort_keys=False))
+                click.echo(f"Updated {component_path}")
+                updated_files.add(component_path)
+
+    if updated_files:
+        click.echo(f"Updated {len(updated_files)} component yaml files.")
+    else:
+        click.echo("No component yaml files were updated.")

python_modules/libraries/dagster-dg/dagster_dg/component.py

@@ -166,6 +166,6 @@ def get_used_env_vars(data_structure: Union[Mapping[str, Any], Sequence[Any], An
     elif isinstance(data_structure, str):
         return set(env_var_regex.findall(data_structure))
     elif isinstance(data_structure, Sequence):
-        return set.union(*(get_used_env_vars(item) for item in data_structure))
+        return set.union(set(), *(get_used_env_vars(item) for item in data_structure))
     else:
         return set()

python_modules/libraries/dagster-dg/dagster_dg_tests/cli_tests/test_check_yaml_command.py

@@ -29,6 +29,17 @@
     ProxyRunner,
     assert_runner_result,
     create_project_from_components,
+    set_env_var,
+)
+
+NON_SPECIFIED_ENV_VAR_TEST_CASE = ComponentValidationTestCase(
+    component_path="validation/basic_component_missing_declared_env",
+    component_type_filepath=BASIC_COMPONENT_TYPE_FILEPATH,
+    should_error=True,
+    check_error_msg=msg_includes_all_of(
+        "component.yaml:1",
+        "Component uses environment variables that are not specified in the component file: A_STRING",
+    ),
 )
 
 CLI_TEST_CASES = [
@@ -51,15 +62,7 @@
             "'an_extra_top_level_value' was unexpected",
         ),
     ),
-    ComponentValidationTestCase(
-        component_path="validation/basic_component_missing_declared_env",
-        component_type_filepath=BASIC_COMPONENT_TYPE_FILEPATH,
-        should_error=True,
-        check_error_msg=msg_includes_all_of(
-            "component.yaml:1",
-            "Component uses environment variables that are not specified in the component file: A_STRING",
-        ),
-    ),
+    NON_SPECIFIED_ENV_VAR_TEST_CASE,
     ComponentValidationTestCase(
         component_path="validation/basic_component_with_env",
         component_type_filepath=BASIC_COMPONENT_TYPE_FILEPATH,
@@ -291,3 +294,40 @@ def test_check_yaml_local_component_cache() -> None:
                 r"CACHE \[write\].*basic_component_invalid_value.*local_component_registry",
                 result.stdout,
             )
+
+
+def test_check_yaml_fix_env_requirements() -> None:
+    with (
+        ProxyRunner.test() as runner,
+        create_project_from_components(
+            runner,
+            NON_SPECIFIED_ENV_VAR_TEST_CASE.component_path,
+            local_component_defn_to_inject=NON_SPECIFIED_ENV_VAR_TEST_CASE.component_type_filepath,
+        ) as tmpdir,
+    ):
+        with pushd(tmpdir):
+            result = runner.invoke(
+                "check",
+                "yaml",
+            )
+            assert result.exit_code != 0, str(result.stdout)
+            assert NON_SPECIFIED_ENV_VAR_TEST_CASE.check_error_msg
+            NON_SPECIFIED_ENV_VAR_TEST_CASE.check_error_msg(str(result.stdout))
+
+            result = runner.invoke(
+                "check",
+                "yaml",
+                "--fix-env-requirements",
+            )
+            assert result.exit_code == 1, str(result.stdout)
+            assert (
+                "The following component files were updated to fix environment variable requirements:"
+                in str(result.stdout)
+            )
+
+            with set_env_var("A_STRING", "foo"):
+                result = runner.invoke(
+                    "check",
+                    "yaml",
+                )
+                assert result.exit_code == 0, str(result.stdout)