diff --git a/python_modules/libraries/dagster-components/dagster_components/core/component.py b/python_modules/libraries/dagster-components/dagster_components/core/component.py index dfb0cac23938a..dc8ca5e0dab20 100644 --- a/python_modules/libraries/dagster-components/dagster_components/core/component.py +++ b/python_modules/libraries/dagster-components/dagster_components/core/component.py @@ -210,7 +210,7 @@ def for_test( resources=resources or {}, registry=registry or ComponentRegistry.empty(), decl_node=decl_node, - templated_value_resolver=TemplatedValueResolver(context={}), + templated_value_resolver=TemplatedValueResolver.default(), ) @property diff --git a/python_modules/libraries/dagster-components/dagster_components/core/component_defs_builder.py b/python_modules/libraries/dagster-components/dagster_components/core/component_defs_builder.py index 047e458fab79f..04a1ce93ed03a 100644 --- a/python_modules/libraries/dagster-components/dagster_components/core/component_defs_builder.py +++ b/python_modules/libraries/dagster-components/dagster_components/core/component_defs_builder.py @@ -11,6 +11,7 @@ Component, ComponentLoadContext, ComponentRegistry, + TemplatedValueResolver, get_component_name, is_registered_component, ) @@ -97,7 +98,12 @@ def build_defs_from_component_path( if not decl_node: raise Exception(f"No component found at path {path}") - context = ComponentLoadContext(resources=resources, registry=registry, decl_node=decl_node) + context = ComponentLoadContext( + resources=resources, + registry=registry, + decl_node=decl_node, + templated_value_resolver=TemplatedValueResolver.default(), + ) components = load_components_from_context(context) return defs_from_components(resources=resources, context=context, components=components) diff --git a/python_modules/libraries/dagster-components/dagster_components/core/component_rendering.py b/python_modules/libraries/dagster-components/dagster_components/core/component_rendering.py index 30f002c1a3157..c7ef86e60ba32 100644 --- a/python_modules/libraries/dagster-components/dagster_components/core/component_rendering.py +++ b/python_modules/libraries/dagster-components/dagster_components/core/component_rendering.py @@ -1,4 +1,5 @@ import json +import os from typing import AbstractSet, Any, Mapping, Optional, Sequence, Type, TypeVar, Union import dagster._check as check @@ -15,7 +16,7 @@ CONTEXT_KEY = "required_rendering_context" -def add_required_rendering_context(field: FieldInfo, context: AbstractSet[str]) -> Any: +def add_required_rendering_context(field: Any, context: AbstractSet[str]) -> Any: return FieldInfo.merge_field_infos( field, FieldInfo(json_schema_extra={CONTEXT_KEY: json.dumps(list(context))}) ) @@ -26,10 +27,18 @@ def get_required_rendering_context(subschema: Mapping[str, Any]) -> Optional[Abs return set(json.loads(raw)) if raw else None +def _env(key: str) -> Optional[str]: + return os.environ.get(key) + + @record class TemplatedValueResolver: context: Mapping[str, Any] + @staticmethod + def default() -> "TemplatedValueResolver": + return TemplatedValueResolver(context={"env": _env}) + def with_context(self, **additional_context) -> "TemplatedValueResolver": return TemplatedValueResolver(context={**self.context, **additional_context}) diff --git a/python_modules/libraries/dagster-components/dagster_components_tests/integration_tests/test_templated_custom_keys_dbt_project.py b/python_modules/libraries/dagster-components/dagster_components_tests/integration_tests/test_templated_custom_keys_dbt_project.py index 60f89a8d3efd6..0c7d7554f8c88 100644 --- a/python_modules/libraries/dagster-components/dagster_components_tests/integration_tests/test_templated_custom_keys_dbt_project.py +++ b/python_modules/libraries/dagster-components/dagster_components_tests/integration_tests/test_templated_custom_keys_dbt_project.py @@ -6,6 +6,7 @@ import pytest from dagster import AssetKey +from dagster._utils.env import environ from dagster_components.core.component_decl_builder import ComponentFileModel from dagster_components.core.component_defs_builder import ( YamlComponentDecl, @@ -116,3 +117,42 @@ def test_load_from_path(dbt_path: Path) -> None: ) assert defs.get_asset_graph().get_all_asset_keys() == JAFFLE_SHOP_KEYS_WITH_PREFIX + + +def test_render_vars_root(dbt_path: Path) -> None: + with environ({"GROUP_AS_ENV": "group_in_env"}): + decl_node = YamlComponentDecl( + path=dbt_path / COMPONENT_RELPATH, + component_file_model=ComponentFileModel( + type="dbt_project", + params={ + "dbt": {"project_dir": "jaffle_shop"}, + "translator": { + "group": "{{ env('GROUP_AS_ENV') }}", + }, + }, + ), + ) + comp = DbtProjectComponent.load(context=script_load_context(decl_node)) + assert get_asset_keys(comp) == JAFFLE_SHOP_KEYS + defs: Definitions = comp.build_defs(script_load_context()) + for key in get_asset_keys(comp): + assert defs.get_assets_def(key).get_asset_spec(key).group_name == "group_in_env" + + +def test_render_vars_asset_key(dbt_path: Path) -> None: + with environ({"ASSET_KEY_PREFIX": "some_prefix"}): + decl_node = YamlComponentDecl( + path=dbt_path / COMPONENT_RELPATH, + component_file_model=ComponentFileModel( + type="dbt_project", + params={ + "dbt": {"project_dir": "jaffle_shop"}, + "translator": { + "key": "{{ env('ASSET_KEY_PREFIX') }}/{{ node.name }}", + }, + }, + ), + ) + comp = DbtProjectComponent.load(context=script_load_context(decl_node)) + assert get_asset_keys(comp) == JAFFLE_SHOP_KEYS_WITH_PREFIX diff --git a/python_modules/libraries/dagster-components/dagster_components_tests/utils.py b/python_modules/libraries/dagster-components/dagster_components_tests/utils.py index 87a9e94f1e266..9dd4669a4fe8f 100644 --- a/python_modules/libraries/dagster-components/dagster_components_tests/utils.py +++ b/python_modules/libraries/dagster-components/dagster_components_tests/utils.py @@ -19,7 +19,7 @@ def registry() -> ComponentRegistry: def script_load_context(decl_node: Optional[ComponentDeclNode] = None) -> ComponentLoadContext: - return ComponentLoadContext(registry=registry(), resources={}, decl_node=decl_node) + return ComponentLoadContext.for_test(registry=registry(), decl_node=decl_node) def get_asset_keys(component: Component) -> AbstractSet[AssetKey]: