59 lines
2.4 KiB
Python
59 lines
2.4 KiB
Python
"""
|
|
Patch for Thrift-generated exception classes to make them compatible with Airflow's secret masking.
|
|
"""
|
|
|
|
import logging
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional, Tuple, Union
|
|
|
|
# --- Python Path Setup ---
|
|
project_root = Path(__file__).parent.absolute()
|
|
# Add project root to sys.path (needed for the 'pangramia' symlink)
|
|
if str(project_root) not in sys.path: sys.path.insert(0, str(project_root))
|
|
# --- End Python Path Setup ---
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def patch_thrift_exceptions():
|
|
"""
|
|
Patch Thrift-generated exception classes to make them compatible with Airflow's secret masking.
|
|
"""
|
|
try:
|
|
from pangramia.yt.exceptions.ttypes import PBServiceException, PBUserException
|
|
|
|
# Save original __setattr__ methods
|
|
original_service_setattr = PBServiceException.__setattr__
|
|
original_user_setattr = PBUserException.__setattr__
|
|
|
|
# Define a new __setattr__ method that allows modifying any attribute
|
|
def new_service_setattr(self, name, value):
|
|
logger.debug(f"Setting attribute {name} on PBServiceException")
|
|
object.__setattr__(self, name, value)
|
|
|
|
def new_user_setattr(self, name, value):
|
|
logger.debug(f"Setting attribute {name} on PBUserException")
|
|
object.__setattr__(self, name, value)
|
|
|
|
# Apply the patch to both exception classes
|
|
PBServiceException.__setattr__ = new_service_setattr
|
|
PBUserException.__setattr__ = new_user_setattr
|
|
|
|
logger.info("Successfully patched Thrift exception classes for Airflow compatibility")
|
|
|
|
# Verify the patch
|
|
try:
|
|
test_exception = PBServiceException(message="Test")
|
|
test_exception.args = ("Test",) # Try to modify an attribute
|
|
logger.info("Verified Thrift exception patch is working correctly")
|
|
except Exception as e:
|
|
logger.error(f"Thrift exception patch verification failed: {e}")
|
|
except ImportError as e:
|
|
logger.warning(f"Could not import Thrift exception classes: {e}")
|
|
logger.warning("Airflow error handling may not work properly with Thrift exceptions")
|
|
except Exception as e:
|
|
logger.error(f"Error patching Thrift exception classes: {e}")
|
|
|
|
# Apply the patch when this module is imported
|
|
patch_thrift_exceptions()
|