Add exception handlers to bootstrap python code (#1258)

Port these changes to the preview branch: #1133 Co-authored-by: kartikgupta-db <[email protected]>
databricks · Jun 27, 2024 · b40f14c · b40f14c
1 parent 3264900
commit b40f14c
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/packages/databricks-vscode/resources/python/bootstrap.py b/packages/databricks-vscode/resources/python/bootstrap.py
@@ -37,8 +37,12 @@
 }
 
 # Set log level to "ERROR". See https://kb.databricks.com/notebooks/cmd-c-on-object-id-p0.html
-import logging; logger = spark._jvm.org.apache.log4j;
-logging.getLogger("py4j.java_gateway").setLevel(logging.ERROR)
+try:
+    import logging; logger = spark._jvm.org.apache.log4j;
+    logging.getLogger("py4j.java_gateway").setLevel(logging.ERROR)
+except Exception as e:
+    logging.debug("Failed to set py4j.java_gateway log level to ERROR", exc_info=True)
+    pass
 
 runpy.run_path(python_file, run_name="__main__", init_globals=user_ns)
 None
diff --git a/packages/databricks-vscode/resources/python/file.workflow-wrapper.py b/packages/databricks-vscode/resources/python/file.workflow-wrapper.py
@@ -27,9 +27,13 @@
     "sqlContext": sqlContext,
 }
 
-# Set log level to "ERROR". See https://kb.databricks.com/notebooks/cmd-c-on-object-id-p0.html
-import logging; logger = spark._jvm.org.apache.log4j;
-logging.getLogger("py4j.java_gateway").setLevel(logging.ERROR)
+try:
+    # Set log level to "ERROR". See https://kb.databricks.com/notebooks/cmd-c-on-object-id-p0.html
+    import logging; logger = spark._jvm.org.apache.log4j;
+    logging.getLogger("py4j.java_gateway").setLevel(logging.ERROR)
+except Exception as e:
+    logging.debug("Failed to set py4j.java_gateway log level to ERROR", exc_info=True)
+    pass
 
 runpy.run_path(python_file, run_name="__main__", init_globals=user_ns)
 None