Skip to content
Merged
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
53 changes: 0 additions & 53 deletions google/cloud/dataproc_spark_connect/exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,59 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.

import sys


def _setup_ipython_exception_handler():
"""Setup custom exception handler for IPython environments to ensure minimal traceback display."""
try:
from IPython import get_ipython
except ImportError:
return

ipython = get_ipython()
if ipython is None:
return

# Store original method if not already stored
if hasattr(ipython, "_dataproc_spark_connect_original_showtraceback"):
return # Already patched

ipython._dataproc_spark_connect_original_showtraceback = (
ipython.showtraceback
)

def custom_showtraceback(
shell,
exc_tuple=None,
filename=None,
tb_offset=None,
exception_only=False,
running_compiled_code=False,
):
# Get the current exception info
_, value, _ = sys.exc_info() if exc_tuple is None else exc_tuple

# If it's our custom exception, show only the message
if isinstance(value, DataprocSparkConnectException):
print(f"Error: {value.message}", file=sys.stderr)
else:
# Use original behavior for other exceptions
shell._dataproc_spark_connect_original_showtraceback(
exc_tuple,
filename,
tb_offset,
exception_only,
running_compiled_code,
)

# Override the method
ipython.showtraceback = custom_showtraceback


# Setup the handler once at module import time
_setup_ipython_exception_handler()


class DataprocSparkConnectException(Exception):
"""A custom exception class to only print the error messages.
Expand Down