-
Notifications
You must be signed in to change notification settings - Fork 1.2k
Expand file tree
/
Copy patharrow_error_handler.py
More file actions
52 lines (38 loc) · 1.42 KB
/
arrow_error_handler.py
File metadata and controls
52 lines (38 loc) · 1.42 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
import logging
from functools import wraps
import pyarrow.flight as fl
from feast.errors import FeastError
logger = logging.getLogger(__name__)
def arrow_client_error_handling_decorator(func):
@wraps(func)
def wrapper(*args, **kwargs):
try:
return func(*args, **kwargs)
except Exception as e:
mapped_error = FeastError.from_error_detail(_get_exception_data(e.args[0]))
if mapped_error is not None:
raise mapped_error
raise e
return wrapper
def arrow_server_error_handling_decorator(func):
@wraps(func)
def wrapper(*args, **kwargs):
try:
return func(*args, **kwargs)
except Exception as e:
if isinstance(e, FeastError):
raise fl.FlightError(e.to_error_detail())
# Re-raise non-Feast exceptions so Arrow Flight returns a proper error
# instead of allowing the server method to return None.
raise e
return wrapper
def _get_exception_data(except_str) -> str:
substring = "Flight error: "
# Find the starting index of the substring
position = except_str.find(substring)
end_json_index = except_str.find("}")
if position != -1 and end_json_index != -1:
# Extract the part of the string after the substring
result = except_str[position + len(substring) : end_json_index + 1]
return result
return ""