Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 10 additions & 3 deletions flytekit/exceptions/user.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,12 +31,19 @@ def value(self):

@property
def error_code(self):
code = None
if hasattr(self.value, "error_code"):
return self.value.error_code
code = self.value.error_code
elif hasattr(type(self.value), "error_code"):
return type(self.value).error_code
else:
code = type(self.value).error_code

if code is None:
return self._ERROR_CODE
if isinstance(code, str):
return code
if isinstance(code, int):
return str(code)
raise _FlyteException(f"error_code must be a str or int, got {type(code).__name__}")


class FlyteTypeException(FlyteUserException, TypeError):
Expand Down
47 changes: 47 additions & 0 deletions tests/flytekit/unit/bin/test_python_entrypoint.py
Original file line number Diff line number Diff line change
Expand Up @@ -1044,3 +1044,50 @@ def verify_output(*args, **kwargs):
mock_write_to_file.side_effect = verify_output
_dispatch_execute(ctx, lambda: python_task, "inputs path", "outputs prefix")
assert mock_write_to_file.call_count == 1


@mock.patch("flytekit.core.utils.load_proto_from_file")
@mock.patch("flytekit.core.data_persistence.FileAccessProvider.get_data")
@mock.patch("flytekit.core.data_persistence.FileAccessProvider.put_data")
@mock.patch("flytekit.core.utils.write_proto_to_file")
def test_dispatch_execute_non_string_error_code_serializes_cleanly(mock_write_to_file, mock_upload_dir, mock_get_data, mock_load_proto):
"""A user exception with a non-string error_code must not cause TypeError during proto serialization.

Previously, a non-string error_code (e.g. an integer HTTP status code) was passed
directly to protobuf's string `code` field, raising TypeError inside to_flyte_idl()
and hiding the original exception in the UI.
"""
mock_get_data.return_value = True
mock_upload_dir.return_value = True

class DataFormatException(Exception):
error_code = 422 # integer, not a string

@task
def t1(a: int) -> str:
raise DataFormatException("Empty value for specimen_type")
return ""

ctx = context_manager.FlyteContext.current_context()
with context_manager.FlyteContextManager.with_context(
ctx.with_execution_state(
ctx.execution_state.with_params(mode=context_manager.ExecutionState.Mode.TASK_EXECUTION)
)
) as ctx:
input_literal_map = TypeEngine.dict_to_literal_map(ctx, {"a": 5})
mock_load_proto.return_value = input_literal_map.to_flyte_idl()

files = OrderedDict()
mock_write_to_file.side_effect = get_output_collector(files)
system_entry_point(_dispatch_execute)(ctx, lambda: t1, "inputs path", "outputs prefix")
assert len(files) == 1

k = list(files.keys())[0]
assert "error.pb" in k

v = list(files.values())[0]
ed = error_models.ErrorDocument.from_flyte_idl(v)
assert ed.error.code == "422"
assert "DataFormatException" in ed.error.message
assert "Empty value for specimen_type" in ed.error.message
assert ed.error.origin == execution_models.ExecutionError.ErrorKind.USER
35 changes: 35 additions & 0 deletions tests/flytekit/unit/exceptions/test_user.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import pytest

from flytekit.exceptions import base, user


Expand Down Expand Up @@ -116,3 +118,36 @@ def test_flyte_validation_error():
assert isinstance(e, AssertionError)
assert type(e).error_code == "USER:ValidationError"
assert isinstance(e, user.FlyteUserException)


@pytest.mark.parametrize(
"error_code_value, expected",
[
("CUSTOM:Error", "CUSTOM:Error"), # string passthrough
(422, "422"), # integer converted to string
(None, "USER:RuntimeError"), # None falls back to default
],
)
def test_flyte_user_runtime_exception_error_code_always_str(error_code_value, expected):
"""Non-string error_code on a user exception must not reach protobuf as a non-string.

Previously, a non-string error_code caused TypeError inside
ContainerError.to_flyte_idl(), hiding the original exception in the UI.
Integers are now coerced to str; None falls back to the default code.
"""
class UserException(Exception):
error_code = error_code_value

exc = user.FlyteUserRuntimeException(UserException("boom"))
assert exc.error_code == expected
assert isinstance(exc.error_code, str)


def test_flyte_user_runtime_exception_error_code_unsupported_type_raises():
"""An error_code that is neither str nor int should raise FlyteException."""
class UserException(Exception):
error_code = [1, 2, 3] # unsupported type

exc = user.FlyteUserRuntimeException(UserException("boom"))
with pytest.raises(base.FlyteException, match="error_code must be a str or int"):
_ = exc.error_code
Loading