diff --git a/docs/docs/guides/labs/components/building-pipelines-with-components/adding-components.md b/docs/docs/guides/labs/components/building-pipelines-with-components/adding-components.md index 80add7b380812..b569d470d9021 100644 --- a/docs/docs/guides/labs/components/building-pipelines-with-components/adding-components.md +++ b/docs/docs/guides/labs/components/building-pipelines-with-components/adding-components.md @@ -50,7 +50,7 @@ This will create a new directory inside your `components/` folder that contains The `component.yaml` is the primary configuration file for a component. It contains two top-level fields: - `type`: The type of the component defined in this directory -- `attributes`: A dictionary of attributes that are specific to this component type. The schema for these attributes is defined by the `get_schema` method on the component class. +- `attributes`: A dictionary of attributes that are specific to this component type. The schema for these attributes is defined by attributes on the `Component` and totally customized by overriding `get_model_cls` method on the component class. To see a sample `component.yaml` file for your specific component, you can run: diff --git a/python_modules/dagster-test/dagster_test/components/simple_pipes_script_asset.py b/python_modules/dagster-test/dagster_test/components/simple_pipes_script_asset.py index 24d0025aad086..0540f92d7dab0 100644 --- a/python_modules/dagster-test/dagster_test/components/simple_pipes_script_asset.py +++ b/python_modules/dagster-test/dagster_test/components/simple_pipes_script_asset.py @@ -6,7 +6,7 @@ from dagster._core.definitions.definitions_class import Definitions from dagster._core.execution.context.asset_execution_context import AssetExecutionContext from dagster._core.pipes.subprocess import PipesSubprocessClient -from dagster.components import Component, ComponentLoadContext +from dagster.components import Component, ComponentLoadContext, Model from dagster.components.component_scaffolding import scaffold_component from dagster.components.scaffold.scaffold import Scaffolder, ScaffoldRequest, scaffold_with from pydantic import BaseModel @@ -18,6 +18,12 @@ class SimplePipesScriptScaffoldParams(BaseModel): filename: str +# Same schema used for file generation and defs generation +class SimplePipesScriptComponentModel(Model): + asset_key: str + filename: str + + class SimplePipesScriptScaffolder(Scaffolder): @classmethod def get_scaffold_params(cls): @@ -48,8 +54,8 @@ class SimplePipesScriptComponent(Component): """ @classmethod - def get_schema(cls): - return SimplePipesScriptScaffoldParams + def get_model_cls(cls): + return SimplePipesScriptComponentModel def __init__(self, asset_key: AssetKey, script_path: Path): self._asset_key = asset_key diff --git a/python_modules/dagster/dagster/components/cli/list.py b/python_modules/dagster/dagster/components/cli/list.py index 563e3f2303f1a..d9a4fe9d7269a 100644 --- a/python_modules/dagster/dagster/components/cli/list.py +++ b/python_modules/dagster/dagster/components/cli/list.py @@ -58,22 +58,22 @@ def list_all_components_schema_command(entry_points: bool, extra_modules: tuple[ """ component_types = _load_component_types(entry_points, extra_modules) - schemas = [] + model_cls_list = [] for key in sorted(component_types.keys(), key=lambda k: k.to_typename()): component_type = component_types[key] # Create ComponentFileModel schema for each type - schema_type = component_type.get_schema() + model_cls = component_type.get_model_cls() key_string = key.to_typename() - if schema_type: - schemas.append( + if model_cls: + model_cls_list.append( create_model( key.name, type=(Literal[key_string], key_string), - attributes=(schema_type, None), + attributes=(model_cls, None), __config__=ConfigDict(extra="forbid"), ) ) - union_type = Union[tuple(schemas)] # type: ignore + union_type = Union[tuple(model_cls_list)] # type: ignore click.echo(json.dumps(TypeAdapter(union_type).json_schema())) diff --git a/python_modules/dagster/dagster/components/component/component.py b/python_modules/dagster/dagster/components/component/component.py index 778654412fdf0..b65a0f45bc2c8 100644 --- a/python_modules/dagster/dagster/components/component/component.py +++ b/python_modules/dagster/dagster/components/component/component.py @@ -22,9 +22,18 @@ def __dg_library_object__(cls) -> None: ... @classmethod def get_schema(cls) -> Optional[type[BaseModel]]: + return None + + @classmethod + def get_model_cls(cls) -> Optional[type[BaseModel]]: if issubclass(cls, Resolvable): return cls.model() + # handle existing overrides for backwards compatibility + cls_from_get_schema = cls.get_schema() + if cls_from_get_schema: + return cls_from_get_schema + return None @classmethod diff --git a/python_modules/dagster/dagster/components/core/defs_module.py b/python_modules/dagster/dagster/components/core/defs_module.py index d320e3c725c48..320d0a5abb69b 100644 --- a/python_modules/dagster/dagster/components/core/defs_module.py +++ b/python_modules/dagster/dagster/components/core/defs_module.py @@ -77,7 +77,7 @@ class DefsFolderComponent(DefsModuleComponent): asset_post_processors: Optional[Sequence[AssetPostProcessor]] @classmethod - def get_schema(cls): + def get_model_cls(cls): return DefsFolderComponentYamlSchema.model() @classmethod @@ -226,26 +226,24 @@ def get_component(cls, context: ComponentLoadContext) -> Component: f"Component type {type_str} is of type {type(obj)}, but must be a subclass of dagster.Component" ) - component_schema = obj.get_schema() + model_cls = obj.get_model_cls() context = context.with_rendering_scope( obj.get_additional_scope() ).with_source_position_tree(source_tree.source_position_tree) # grab the attributes from the yaml file with pushd(str(context.path)): - if component_schema is None: + if model_cls is None: attributes = None elif source_tree: attributes_position_tree = source_tree.source_position_tree.children["attributes"] with enrich_validation_errors_with_source_position( attributes_position_tree, ["attributes"] ): - attributes = TypeAdapter(component_schema).validate_python( + attributes = TypeAdapter(model_cls).validate_python( component_file_model.attributes ) else: - attributes = TypeAdapter(component_schema).validate_python( - component_file_model.attributes - ) + attributes = TypeAdapter(model_cls).validate_python(component_file_model.attributes) return obj.load(attributes, context) diff --git a/python_modules/dagster/dagster/components/core/snapshot.py b/python_modules/dagster/dagster/components/core/snapshot.py index 6932a92636cc6..558cfed112431 100644 --- a/python_modules/dagster/dagster/components/core/snapshot.py +++ b/python_modules/dagster/dagster/components/core/snapshot.py @@ -42,12 +42,12 @@ def _get_summary_and_description(obj: object) -> tuple[Optional[str], Optional[s def _get_component_type_snap(key: LibraryObjectKey, obj: type[Component]) -> ComponentTypeSnap: summary, description = _get_summary_and_description(obj) - component_schema = obj.get_schema() + model_cls = obj.get_model_cls() return ComponentTypeSnap( key=key, summary=summary, description=description, - schema=component_schema.model_json_schema() if component_schema else None, + schema=model_cls.model_json_schema() if model_cls else None, scaffolder=_get_scaffolder_snap(obj), ) diff --git a/python_modules/dagster/dagster/components/test/basic_components.py b/python_modules/dagster/dagster/components/test/basic_components.py index c7cf28a833525..e1d53a053a051 100644 --- a/python_modules/dagster/dagster/components/test/basic_components.py +++ b/python_modules/dagster/dagster/components/test/basic_components.py @@ -52,7 +52,7 @@ class MyNestedModel(BaseModel): model_config = ConfigDict(extra="forbid") -class MyNestedComponentSchema(BaseModel): +class MyNestedComponentModel(BaseModel): nested: dict[str, MyNestedModel] model_config = ConfigDict(extra="forbid") @@ -60,8 +60,8 @@ class MyNestedComponentSchema(BaseModel): class MyNestedComponent(Component): @classmethod - def get_schema(cls) -> type[MyNestedComponentSchema]: - return MyNestedComponentSchema + def get_model_cls(cls) -> type[MyNestedComponentModel]: + return MyNestedComponentModel def build_defs(self, context: ComponentLoadContext) -> Definitions: return Definitions() diff --git a/python_modules/dagster/dagster_tests/components_tests/cli_tests/test_commands.py b/python_modules/dagster/dagster_tests/components_tests/cli_tests/test_commands.py index d935a88c0e2a8..22e5335fae565 100644 --- a/python_modules/dagster/dagster_tests/components_tests/cli_tests/test_commands.py +++ b/python_modules/dagster/dagster_tests/components_tests/cli_tests/test_commands.py @@ -67,24 +67,34 @@ def test_list_library_objects_from_module(): scaffolder=ScaffolderSnap(schema=None), ) - pipes_script_params_schema = { + pipes_script_component_model_schema = { + "additionalProperties": False, "properties": { "asset_key": {"title": "Asset Key", "type": "string"}, "filename": {"title": "Filename", "type": "string"}, }, "required": ["asset_key", "filename"], - "title": "SimplePipesScriptScaffoldParams", + "title": "SimplePipesScriptComponentModel", "type": "object", } + pipes_script_component_scaffold_params_schema = { + "properties": { + "asset_key": {"title": "Asset Key", "type": "string"}, + "filename": {"title": "Filename", "type": "string"}, + }, + "required": ["asset_key", "filename"], + "title": "SimplePipesScriptScaffoldParams", + "type": "object", + } assert result[3] == ComponentTypeSnap( key=LibraryObjectKey( namespace="dagster_test.components", name="SimplePipesScriptComponent" ), - schema=pipes_script_params_schema, + schema=pipes_script_component_model_schema, description="A simple asset that runs a Python script with the Pipes subprocess client.\n\nBecause it is a pipes asset, no value is returned.", summary="A simple asset that runs a Python script with the Pipes subprocess client.", - scaffolder=ScaffolderSnap(schema=pipes_script_params_schema), + scaffolder=ScaffolderSnap(schema=pipes_script_component_scaffold_params_schema), ) diff --git a/python_modules/dagster/dagster_tests/components_tests/integration_tests/integration_test_defs/definitions/other_local_component_sample/__init__.py b/python_modules/dagster/dagster_tests/components_tests/integration_tests/integration_test_defs/definitions/other_local_component_sample/__init__.py index 2e072d476ad9b..6ac26365b489a 100644 --- a/python_modules/dagster/dagster_tests/components_tests/integration_tests/integration_test_defs/definitions/other_local_component_sample/__init__.py +++ b/python_modules/dagster/dagster_tests/components_tests/integration_tests/integration_test_defs/definitions/other_local_component_sample/__init__.py @@ -9,10 +9,8 @@ class MyNewComponentSchema(BaseModel): class MyNewComponent(Component): - name = "my_new_component" - @classmethod - def get_schema(cls): + def get_model_cls(cls): return MyNewComponentSchema def build_defs(self, context: ComponentLoadContext) -> Definitions: diff --git a/python_modules/dagster/dagster_tests/components_tests/utils.py b/python_modules/dagster/dagster_tests/components_tests/utils.py index 7ac3ba3e24770..2da97e76f09ac 100644 --- a/python_modules/dagster/dagster_tests/components_tests/utils.py +++ b/python_modules/dagster/dagster_tests/components_tests/utils.py @@ -31,17 +31,17 @@ def load_context_and_component_for_test( ) -> tuple[ComponentLoadContext, T_Component]: context = ComponentLoadContext.for_test() context = context.with_rendering_scope(component_type.get_additional_scope()) - schema = check.not_none( - component_type.get_schema(), "Component must have schema for direct test" + model_cls = check.not_none( + component_type.get_model_cls(), "Component must have schema for direct test" ) if isinstance(attrs, str): source_positions = parse_yaml_with_source_positions(attrs) with enrich_validation_errors_with_source_position( source_positions.source_position_tree, [] ): - attributes = TypeAdapter(schema).validate_python(source_positions.value) + attributes = TypeAdapter(model_cls).validate_python(source_positions.value) else: - attributes = TypeAdapter(schema).validate_python(attrs) + attributes = TypeAdapter(model_cls).validate_python(attrs) component = component_type.load(attributes, context) return context, component diff --git a/python_modules/libraries/dagster-dg/dagster_dg_tests/cli_tests/test_utils_inspect_command.py b/python_modules/libraries/dagster-dg/dagster_dg_tests/cli_tests/test_utils_inspect_command.py index 81f695e973ba0..bd538de0133fb 100644 --- a/python_modules/libraries/dagster-dg/dagster_dg_tests/cli_tests/test_utils_inspect_command.py +++ b/python_modules/libraries/dagster-dg/dagster_dg_tests/cli_tests/test_utils_inspect_command.py @@ -43,6 +43,7 @@ Component schema: { + "additionalProperties": false, "properties": { "asset_key": { "title": "Asset Key", @@ -57,7 +58,7 @@ "asset_key", "filename" ], - "title": "SimplePipesScriptScaffoldParams", + "title": "SimplePipesScriptComponentModel", "type": "object" } """).strip() @@ -155,6 +156,7 @@ def test_utils_inspect_component_type_flag_fields_success(): assert result.output.strip().endswith( textwrap.dedent(""" { + "additionalProperties": false, "properties": { "asset_key": { "title": "Asset Key", @@ -169,7 +171,7 @@ def test_utils_inspect_component_type_flag_fields_success(): "asset_key", "filename" ], - "title": "SimplePipesScriptScaffoldParams", + "title": "SimplePipesScriptComponentModel", "type": "object" } """).strip()