import json

import logging

import time

from functools import wraps

import pyarrow.flight as fl

from feast.errors import FeastError

logger = logging.getLogger(__name__)

BACKOFF_FACTOR = 0.5

def arrow_client_error_handling_decorator(func):

@wraps(func)

def wrapper(*args, **kwargs):

# Retry only applies to FeastFlightClient methods where args[0] (self)

# carries _connection_retries from RemoteOfflineStoreConfig.

# Standalone stream functions (write_table, read_all) get 0 retries:

# broken streams can't be reused and retrying risks duplicate writes.

max_retries = max(0, getattr(args[0], "_connection_retries", 0)) if args else 0

for attempt in range(max_retries + 1):

try:

return func(*args, **kwargs)

except fl.FlightUnavailableError as e:

if attempt < max_retries:

wait_time = BACKOFF_FACTOR * (2**attempt)

logger.warning(

"Transient Arrow Flight error on attempt %d/%d, "

"retrying in %.1fs: %s",

attempt + 1,

max_retries + 1,

wait_time,

e,

)

time.sleep(wait_time)

continue

mapped_error = FeastError.from_error_detail(_get_exception_data(str(e)))

if mapped_error is not None:

raise mapped_error

raise e

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:

if not isinstance(except_str, str):

return ""

substring = "Flight error: "

position = except_str.find(substring)

if position == -1:

return ""

search_start = position + len(substring)

json_start = except_str.find("{", search_start)

if json_start == -1:

return ""

search_pos = json_start

while True:

json_end = except_str.find("}", search_pos + 1)

if json_end == -1:

return ""

candidate = except_str[json_start : json_end + 1]

try:

json.loads(candidate)

return candidate

except (json.JSONDecodeError, ValueError):

search_pos = json_end