diff --git a/python_modules/libraries/dagster-components/dagster_components/cli/generate.py b/python_modules/libraries/dagster-components/dagster_components/cli/generate.py index cb7324062f259..c28c755278244 100644 --- a/python_modules/libraries/dagster-components/dagster_components/cli/generate.py +++ b/python_modules/libraries/dagster-components/dagster_components/cli/generate.py @@ -62,7 +62,7 @@ def generate_component_command( raise Exception( f"Component type {component_type} does not have a generator. Reason: {generator.message}." ) - generate_params = TypeAdapter(generator.generator_params).validate_json(json_params) + generate_params = TypeAdapter(generator.get_params_schema_type()).validate_json(json_params) else: generate_params = {} 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 20cd8a9473aac..2f779c94cf1ac 100644 --- a/python_modules/libraries/dagster-components/dagster_components/core/component.py +++ b/python_modules/libraries/dagster-components/dagster_components/core/component.py @@ -28,7 +28,7 @@ from dagster._core.definitions.definitions_class import Definitions from dagster._core.errors import DagsterError from dagster._utils import pushd, snakecase -from pydantic import TypeAdapter +from pydantic import BaseModel, TypeAdapter from typing_extensions import Self from dagster_components.core.component_generator import ( @@ -44,7 +44,10 @@ class ComponentDeclNode: ... class Component(ABC): name: ClassVar[Optional[str]] = None - params_schema: ClassVar = None + + @classmethod + def get_params_schema_type(cls) -> Optional[Type[BaseModel]]: + return None @classmethod def get_generator(cls) -> Union[ComponentGenerator, ComponentGeneratorUnavailableReason]: @@ -78,13 +81,17 @@ def get_metadata(cls) -> "ComponentTypeInternalMetadata": if isinstance(generator, ComponentGeneratorUnavailableReason): raise DagsterError(f"Component {cls.__name__} is not scaffoldable: {generator.message}") + component_params = cls.get_params_schema_type() + generator_params = generator.get_params_schema_type() return { "summary": clean_docstring.split("\n\n")[0] if clean_docstring else None, "description": clean_docstring if clean_docstring else None, - "generate_params_schema": generator.generator_params.schema() - if generator.generator_params - else None, - "component_params_schema": cls.params_schema.schema() if cls.params_schema else None, + "generate_params_schema": None + if generator_params is None + else generator_params.model_json_schema(), + "component_params_schema": None + if component_params is None + else component_params.model_json_schema(), } @classmethod diff --git a/python_modules/libraries/dagster-components/dagster_components/core/component_generator.py b/python_modules/libraries/dagster-components/dagster_components/core/component_generator.py index 34b0c3fd83327..309b53bdd0475 100644 --- a/python_modules/libraries/dagster-components/dagster_components/core/component_generator.py +++ b/python_modules/libraries/dagster-components/dagster_components/core/component_generator.py @@ -1,9 +1,10 @@ from abc import abstractmethod from dataclasses import dataclass from pathlib import Path -from typing import Any, ClassVar +from typing import Any, Optional, Type from dagster._record import record +from pydantic import BaseModel @record @@ -13,7 +14,9 @@ class ComponentGenerateRequest: class ComponentGenerator: - generator_params: ClassVar = None + @classmethod + def get_params_schema_type(cls) -> Optional[Type[BaseModel]]: + return None @abstractmethod def generate_files(self, request: ComponentGenerateRequest, params: Any) -> None: ... diff --git a/python_modules/libraries/dagster-components/dagster_components/lib/dbt_project/component.py b/python_modules/libraries/dagster-components/dagster_components/lib/dbt_project/component.py index 489a201371d06..4a4c5421b2e2d 100644 --- a/python_modules/libraries/dagster-components/dagster_components/lib/dbt_project/component.py +++ b/python_modules/libraries/dagster-components/dagster_components/lib/dbt_project/component.py @@ -65,8 +65,6 @@ def get_automation_condition(self, dbt_resource_props): @component_type(name="dbt_project") class DbtProjectComponent(Component): - params_schema = DbtProjectParams - def __init__( self, dbt_resource: DbtCliResource, @@ -83,9 +81,13 @@ def __init__( def get_generator(cls) -> "DbtProjectComponentGenerator": return DbtProjectComponentGenerator() + @classmethod + def get_params_schema_type(cls): + return DbtProjectParams + @classmethod def load(cls, context: ComponentLoadContext) -> Self: - loaded_params = context.load_params(cls.params_schema) + loaded_params = context.load_params(cls.get_params_schema_type()) return cls( dbt_resource=loaded_params.dbt, diff --git a/python_modules/libraries/dagster-components/dagster_components/lib/dbt_project/generator.py b/python_modules/libraries/dagster-components/dagster_components/lib/dbt_project/generator.py index 675d280a31086..8f4e795a2c3a6 100644 --- a/python_modules/libraries/dagster-components/dagster_components/lib/dbt_project/generator.py +++ b/python_modules/libraries/dagster-components/dagster_components/lib/dbt_project/generator.py @@ -1,6 +1,6 @@ import os from pathlib import Path -from typing import Optional +from typing import Optional, Type import dagster._check as check from dbt.cli.main import dbtRunner @@ -16,7 +16,9 @@ class DbtGenerateParams(BaseModel): class DbtProjectComponentGenerator(ComponentGenerator): - generator_params = DbtGenerateParams + @classmethod + def get_params_schema_type(cls) -> Optional[Type[BaseModel]]: + return DbtGenerateParams def generate_files(self, request: ComponentGenerateRequest, params: DbtGenerateParams) -> None: cwd = os.getcwd() diff --git a/python_modules/libraries/dagster-components/dagster_components/lib/definitions_component.py b/python_modules/libraries/dagster-components/dagster_components/lib/definitions_component.py deleted file mode 100644 index e69de29bb2d1d..0000000000000 diff --git a/python_modules/libraries/dagster-components/dagster_components/lib/definitions_component/component.py b/python_modules/libraries/dagster-components/dagster_components/lib/definitions_component/component.py index 7833f0ed123fb..17b1b04dece6a 100644 --- a/python_modules/libraries/dagster-components/dagster_components/lib/definitions_component/component.py +++ b/python_modules/libraries/dagster-components/dagster_components/lib/definitions_component/component.py @@ -23,16 +23,18 @@ class DefinitionsComponent(Component): def __init__(self, definitions_path: Path): self.definitions_path = definitions_path - params_schema = DefinitionsParamSchema - @classmethod def get_generator(cls) -> DefinitionsComponentGenerator: return DefinitionsComponentGenerator() + @classmethod + def get_params_schema_type(cls): + return DefinitionsParamSchema + @classmethod def load(cls, context: ComponentLoadContext) -> Self: # all paths should be resolved relative to the directory we're in - loaded_params = context.load_params(cls.params_schema) + loaded_params = context.load_params(cls.get_params_schema_type()) return cls(definitions_path=Path(loaded_params.definitions_path or "definitions.py")) diff --git a/python_modules/libraries/dagster-components/dagster_components/lib/definitions_component/generator.py b/python_modules/libraries/dagster-components/dagster_components/lib/definitions_component/generator.py index 81594297b040e..d4ace58c55b26 100644 --- a/python_modules/libraries/dagster-components/dagster_components/lib/definitions_component/generator.py +++ b/python_modules/libraries/dagster-components/dagster_components/lib/definitions_component/generator.py @@ -14,7 +14,9 @@ class DefinitionsGenerateParams(BaseModel): class DefinitionsComponentGenerator(ComponentGenerator): - generator_params = DefinitionsGenerateParams + @classmethod + def get_params_schema_type(cls): + return DefinitionsGenerateParams def generate_files( self, request: ComponentGenerateRequest, params: DefinitionsGenerateParams diff --git a/python_modules/libraries/dagster-components/dagster_components/lib/pipes_subprocess_script_collection.py b/python_modules/libraries/dagster-components/dagster_components/lib/pipes_subprocess_script_collection.py index 084fa282698f0..1adf4996b1f3d 100644 --- a/python_modules/libraries/dagster-components/dagster_components/lib/pipes_subprocess_script_collection.py +++ b/python_modules/libraries/dagster-components/dagster_components/lib/pipes_subprocess_script_collection.py @@ -29,8 +29,6 @@ class PipesSubprocessScriptCollectionParams(BaseModel): class PipesSubprocessScriptCollection(Component): """Assets that wrap Python scripts executed with Dagster's PipesSubprocessClient.""" - params_schema = PipesSubprocessScriptCollectionParams - def __init__(self, dirpath: Path, path_specs: Mapping[Path, Sequence[AssetSpec]]): self.dirpath = dirpath # mapping from the script name (e.g. /path/to/script_abc.py -> script_abc) @@ -42,9 +40,13 @@ def introspect_from_path(path: Path) -> "PipesSubprocessScriptCollection": path_specs = {path: [AssetSpec(path.stem)] for path in list(path.rglob("*.py"))} return PipesSubprocessScriptCollection(dirpath=path, path_specs=path_specs) + @classmethod + def get_params_schema_type(cls): + return PipesSubprocessScriptCollectionParams + @classmethod def load(cls, context: ComponentLoadContext) -> "PipesSubprocessScriptCollection": - loaded_params = context.load_params(cls.params_schema) + loaded_params = context.load_params(cls.get_params_schema_type()) path_specs = {} for script in loaded_params.scripts: diff --git a/python_modules/libraries/dagster-components/dagster_components/lib/sling_replication_collection/component.py b/python_modules/libraries/dagster-components/dagster_components/lib/sling_replication_collection/component.py index 441b1d546de7a..b5ea37466cca9 100644 --- a/python_modules/libraries/dagster-components/dagster_components/lib/sling_replication_collection/component.py +++ b/python_modules/libraries/dagster-components/dagster_components/lib/sling_replication_collection/component.py @@ -78,8 +78,6 @@ def get_auto_materialize_policy( @component_type(name="sling_replication_collection") class SlingReplicationCollectionComponent(Component): - params_schema = SlingReplicationCollectionParams - def __init__( self, dirpath: Path, @@ -100,9 +98,13 @@ def get_generator(cls) -> ComponentGenerator: return SlingReplicationComponentGenerator() + @classmethod + def get_params_schema_type(cls): + return SlingReplicationCollectionParams + @classmethod def load(cls, context: ComponentLoadContext) -> Self: - loaded_params = context.load_params(cls.params_schema) + loaded_params = context.load_params(cls.get_params_schema_type()) return cls( dirpath=context.path, resource=loaded_params.sling or SlingResource(), diff --git a/python_modules/libraries/dagster-components/dagster_components/lib/test/simple_asset.py b/python_modules/libraries/dagster-components/dagster_components/lib/test/simple_asset.py index 591a0ec17fa3b..f9ef3adc81992 100644 --- a/python_modules/libraries/dagster-components/dagster_components/lib/test/simple_asset.py +++ b/python_modules/libraries/dagster-components/dagster_components/lib/test/simple_asset.py @@ -27,7 +27,9 @@ class SimpleAssetParams(BaseModel): class SimpleAsset(Component): """A simple asset that returns a constant string value.""" - params_schema = SimpleAssetParams + @classmethod + def get_params_schema_type(cls): + return SimpleAssetParams @classmethod def get_generator(cls) -> ComponentGenerator: @@ -38,7 +40,7 @@ def from_decl_node( cls, context: "ComponentLoadContext", decl_node: "ComponentDeclNode" ) -> Self: assert isinstance(decl_node, YamlComponentDecl) - loaded_params = TypeAdapter(cls.params_schema).validate_python( + loaded_params = TypeAdapter(cls.get_params_schema_type()).validate_python( decl_node.component_file_model.params ) return cls( diff --git a/python_modules/libraries/dagster-components/dagster_components/lib/test/simple_pipes_script_asset.py b/python_modules/libraries/dagster-components/dagster_components/lib/test/simple_pipes_script_asset.py index 41dd54247159f..c0ccce65e592f 100644 --- a/python_modules/libraries/dagster-components/dagster_components/lib/test/simple_pipes_script_asset.py +++ b/python_modules/libraries/dagster-components/dagster_components/lib/test/simple_pipes_script_asset.py @@ -26,7 +26,9 @@ class SimplePipesScriptAssetParams(BaseModel): class SimplePipesScriptAssetGenerator(ComponentGenerator): - generator_params = SimplePipesScriptAssetParams + @classmethod + def get_params_schema_type(cls): + return SimplePipesScriptAssetParams def generate_files( self, request: ComponentGenerateRequest, params: SimplePipesScriptAssetParams @@ -54,18 +56,20 @@ class SimplePipesScriptAsset(Component): Because it is a pipes asset, no value is returned. """ - params_schema = SimplePipesScriptAssetParams - @classmethod def get_generator(cls) -> ComponentGenerator: return SimplePipesScriptAssetGenerator() + @classmethod + def get_params_schema_type(cls): + return SimplePipesScriptAssetParams + @classmethod def from_decl_node( cls, context: "ComponentLoadContext", decl_node: "ComponentDeclNode" ) -> Self: assert isinstance(decl_node, YamlComponentDecl) - loaded_params = TypeAdapter(cls.params_schema).validate_python( + loaded_params = TypeAdapter(cls.get_params_schema_type()).validate_python( decl_node.component_file_model.params ) return cls( diff --git a/python_modules/libraries/dagster-components/dagster_components_tests/rendering_tests/custom_scope_component/component.py b/python_modules/libraries/dagster-components/dagster_components_tests/rendering_tests/custom_scope_component/component.py index a0e8be5b68fa5..1dbc0114cb16e 100644 --- a/python_modules/libraries/dagster-components/dagster_components_tests/rendering_tests/custom_scope_component/component.py +++ b/python_modules/libraries/dagster-components/dagster_components_tests/rendering_tests/custom_scope_component/component.py @@ -19,8 +19,6 @@ class CustomScopeParams(BaseModel): @component_type(name="custom_scope_component") class HasCustomScope(Component): - params_schema = CustomScopeParams - @classmethod def get_rendering_scope(cls) -> Mapping[str, Any]: return { @@ -33,9 +31,13 @@ def get_rendering_scope(cls) -> Mapping[str, Any]: def __init__(self, attributes: Mapping[str, Any]): self.attributes = attributes + @classmethod + def get_params_schema_type(cls): + return CustomScopeParams + @classmethod def load(cls, context: ComponentLoadContext): - loaded_params = context.load_params(cls.params_schema) + loaded_params = context.load_params(cls.get_params_schema_type()) return cls(attributes=loaded_params.attributes) def build_defs(self, context: ComponentLoadContext):