-
Notifications
You must be signed in to change notification settings - Fork 1.2k
Expand file tree
/
Copy pathrest_error_handler.py
More file actions
57 lines (45 loc) · 2.15 KB
/
rest_error_handler.py
File metadata and controls
57 lines (45 loc) · 2.15 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
53
54
55
56
57
import logging
from functools import wraps
import requests
from feast import RepoConfig
from feast.errors import FeastError
from feast.permissions.client.http_auth_requests_wrapper import (
get_http_auth_requests_session,
)
logger = logging.getLogger(__name__)
def rest_error_handling_decorator(func):
@wraps(func)
def wrapper(config: RepoConfig, *args, **kwargs):
assert isinstance(config, RepoConfig)
# Get a Session object
with get_http_auth_requests_session(config.auth_config) as session:
# Define a wrapper for session methods
def method_wrapper(method_name):
original_method = getattr(session, method_name)
@wraps(original_method)
def wrapped_method(*args, **kwargs):
logger.debug(
f"Calling {method_name} with args: {args}, kwargs: {kwargs}"
)
response = original_method(*args, **kwargs)
logger.debug(
f"{method_name} response status code: {response.status_code}"
)
try:
response.raise_for_status()
except requests.RequestException:
logger.debug(f"response.json() = {response.json()}")
mapped_error = FeastError.from_error_detail(response.json())
logger.debug(f"mapped_error = {str(mapped_error)}")
if mapped_error is not None:
raise mapped_error
return response
return wrapped_method
# Enhance session methods
session.get = method_wrapper("get") # type: ignore[method-assign]
session.post = method_wrapper("post") # type: ignore[method-assign]
session.put = method_wrapper("put") # type: ignore[method-assign]
session.delete = method_wrapper("delete") # type: ignore[method-assign]
# Pass the enhanced session object to the decorated function
return func(session, config, *args, **kwargs)
return wrapper