IOManager

MultiPartitionsDefinition

from dagster import StaticPartitionsDefinition, ResourceDefinition, asset, AssetIn, materialize, io_manager, IOManager, MultiPartitionKey, MultiPartitionsDefinition

dummy_multi_partition = MultiPartitionsDefinition(
    {
        "part_a": StaticPartitionsDefinition(["abc", "def"]),
        "part_b": StaticPartitionsDefinition(["123", "456"]),
    }
)

@io_manager()
def mock_delta_io_manager(_):
    class TestIOManager(IOManager):
        def handle_output(self, context, obj):
            print(f"TYPE: {type(context.asset_partition_key)}")
            print(context.asset_partition_key)
            assert isinstance(context.asset_partition_key, MultiPartitionKey)

        def load_input(self, context):
            pass

    return TestIOManager()


@asset(
    name="dummy_asset_upstream",
    metadata={"dagster_partition_cols": ["foo", "date"], "additional_partition_cols": ["test_col"]},
    io_manager_def=mock_delta_io_manager,
    partitions_def=dummy_multi_partition
)
def dummy_asset_upstream(context):
    return 1234

materialize(
    [dummy_asset_upstream], partition_key=MultiPartitionKey({"part_a": "abc", "part_b": "123"})
)