Skip to content

Commit bf779cb

Browse files
authored
Merge branch 'GoogleCloudDataproc:main' into zwl-branch
2 parents 289941d + 1fc26f5 commit bf779cb

1 file changed

Lines changed: 54 additions & 1 deletion

File tree

‎google/cloud/dataproc_spark_connect/exceptions.py‎

Lines changed: 54 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,59 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
import sys
16+
17+
18+
def _setup_ipython_exception_handler():
19+
"""Setup custom exception handler for IPython environments to ensure minimal traceback display."""
20+
try:
21+
from IPython import get_ipython
22+
except ImportError:
23+
return
24+
25+
ipython = get_ipython()
26+
if ipython is None:
27+
return
28+
29+
# Store original method if not already stored
30+
if hasattr(ipython, "_dataproc_spark_connect_original_showtraceback"):
31+
return # Already patched
32+
33+
ipython._dataproc_spark_connect_original_showtraceback = (
34+
ipython.showtraceback
35+
)
36+
37+
def custom_showtraceback(
38+
shell,
39+
exc_tuple=None,
40+
filename=None,
41+
tb_offset=None,
42+
exception_only=False,
43+
running_compiled_code=False,
44+
):
45+
# Get the current exception info
46+
_, value, _ = sys.exc_info() if exc_tuple is None else exc_tuple
47+
48+
# If it's our custom exception, show only the message
49+
if isinstance(value, DataprocSparkConnectException):
50+
print(f"Error: {value.message}", file=sys.stderr)
51+
else:
52+
# Use original behavior for other exceptions
53+
shell._dataproc_spark_connect_original_showtraceback(
54+
exc_tuple,
55+
filename,
56+
tb_offset,
57+
exception_only,
58+
running_compiled_code,
59+
)
60+
61+
# Override the method
62+
ipython.showtraceback = custom_showtraceback
63+
64+
65+
# Setup the handler once at module import time
66+
_setup_ipython_exception_handler()
67+
1568

1669
class DataprocSparkConnectException(Exception):
1770
"""A custom exception class to only print the error messages.
@@ -24,4 +77,4 @@ def __init__(self, message):
2477
super().__init__(message)
2578

2679
def _render_traceback_(self):
27-
return self.message
80+
return [self.message]

0 commit comments

Comments
 (0)