From 166abfe2fb523921c99ab33522bbe10850a85ec7 Mon Sep 17 00:00:00 2001
From: Brandon Chinn <brandon.chinn@snowflake.com>
Date: Wed, 1 Nov 2023 11:18:22 -0700
Subject: [PATCH] Store snowflake-connector-python error in
 SnowflakeSQLException

---
 CHANGELOG.md                                   |  4 ++++
 .../snowpark/_internal/error_message.py        | 18 ++----------------
 src/snowflake/snowpark/exceptions.py           | 12 ++++++++----
 3 files changed, 14 insertions(+), 20 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index eea76a2ca68..cf1b089248f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## 1.10.1 (TBD)
 
+### New Features
+
+- Add the `conn_error` attribute to `SnowflakeSQLException` that stores the whole underlying exception from `snowflake-connector-python`
+
 ### Bug Fixes
 
 - DataFrame column names qouting check now supports newline characters.
diff --git a/src/snowflake/snowpark/_internal/error_message.py b/src/snowflake/snowpark/_internal/error_message.py
index 8db3b32dda4..39d3d93078c 100644
--- a/src/snowflake/snowpark/_internal/error_message.py
+++ b/src/snowflake/snowpark/_internal/error_message.py
@@ -316,27 +316,13 @@ def SQL_PYTHON_REPORT_JOIN_AMBIGUOUS(
     def SQL_EXCEPTION_FROM_PROGRAMMING_ERROR(
         pe: ProgrammingError,
     ) -> SnowparkSQLException:
-        query = getattr(pe, "query", None)
-        return SnowparkSQLException(
-            pe.msg,
-            error_code="1304",
-            sfqid=pe.sfqid,
-            query=query,
-            sql_error_code=pe.errno,
-            raw_message=pe.raw_msg,
-        )
+        return SnowparkSQLException(pe.msg, error_code="1304", conn_error=pe)
 
     @staticmethod
     def SQL_EXCEPTION_FROM_OPERATIONAL_ERROR(
         oe: OperationalError,
     ) -> SnowparkSQLException:
-        return SnowparkSQLException(
-            oe.msg,
-            error_code="1305",
-            sfqid=oe.sfqid,
-            sql_error_code=oe.errno,
-            raw_message=oe.raw_msg,
-        )
+        return SnowparkSQLException(oe.msg, error_code="1305", conn_error=oe)
 
     # Server Error Messages 04XX
 
diff --git a/src/snowflake/snowpark/exceptions.py b/src/snowflake/snowpark/exceptions.py
index ffeb2a22d1a..ed0e776b3f1 100644
--- a/src/snowflake/snowpark/exceptions.py
+++ b/src/snowflake/snowpark/exceptions.py
@@ -7,6 +7,8 @@
 import logging
 from typing import Optional
 
+from snowflake.connector.errors import Error as ConnectorError
+
 _logger = logging.getLogger(__name__)
 
 
@@ -79,6 +81,7 @@ def __init__(
         message: str,
         *,
         error_code: Optional[str] = None,
+        conn_error: Optional[ConnectorError] = None,
         sfqid: Optional[str] = None,
         query: Optional[str] = None,
         sql_error_code: Optional[int] = None,
@@ -86,10 +89,11 @@ def __init__(
     ) -> None:
         super().__init__(message, error_code=error_code)
 
-        self.sfqid: Optional[str] = sfqid
-        self.query: Optional[str] = query
-        self.sql_error_code = sql_error_code
-        self.raw_message = raw_message
+        self.conn_error = conn_error
+        self.sfqid = sfqid or getattr(self.conn_error, "sfqid", None)
+        self.query = query or getattr(self.conn_error, "query", None)
+        self.sql_error_code = sql_error_code or getattr(self.conn_error, "errno", None)
+        self.raw_message = raw_message or getattr(self.conn_error, "raw_msg", None)
 
         pretty_error_code = f"({self.error_code}): " if self.error_code else ""
         pretty_sfqid = f"{self.sfqid}: " if self.sfqid else ""