add query to all errors generater from ProgrammingError (snowflakedb#743

)

src/snowflake/snowpark/_internal/analyzer/snowflake_plan.py

@@ -97,10 +97,13 @@ def wrap(*args, **kwargs):
                 try:
                     return func(*args, **kwargs)
                 except snowflake.connector.errors.ProgrammingError as e:
+                    query = None
+                    if "query" in e.__dict__:
+                        query = e.__getattribute__("query")
                     tb = sys.exc_info()[2]
                     if "unexpected 'as'" in e.msg.lower():
-                        ne = (
-                            SnowparkClientExceptionMessages.SQL_PYTHON_REPORT_UNEXPECTED_ALIAS()
+                        ne = SnowparkClientExceptionMessages.SQL_PYTHON_REPORT_UNEXPECTED_ALIAS(
+                            query
                         )
                         raise ne.with_traceback(tb) from None
                     elif e.sqlstate == "42000" and "invalid identifier" in e.msg:
@@ -133,7 +136,7 @@ def wrap(*args, **kwargs):
                                 unaliased_cols[0] if unaliased_cols else "<colname>"
                             )
                             ne = SnowparkClientExceptionMessages.SQL_PYTHON_REPORT_INVALID_ID(
-                                orig_col_name
+                                orig_col_name, query
                             )
                             raise ne.with_traceback(tb) from None
                         elif (
@@ -150,7 +153,7 @@ def wrap(*args, **kwargs):
                             > 1
                         ):
                             ne = SnowparkClientExceptionMessages.SQL_PYTHON_REPORT_JOIN_AMBIGUOUS(
-                                col, col
+                                col, col, query
                             )
                             raise ne.with_traceback(tb) from None
                         else:

src/snowflake/snowpark/_internal/error_message.py

@@ -3,6 +3,8 @@
 # Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
 #
 
+from typing import Optional
+
 from snowflake.connector import OperationalError, ProgrammingError
 from snowflake.snowpark.exceptions import (
     SnowparkColumnException,
@@ -236,24 +238,34 @@ def SQL_LAST_QUERY_RETURN_RESULTSET() -> SnowparkSQLException:
         )
 
     @staticmethod
-    def SQL_PYTHON_REPORT_UNEXPECTED_ALIAS() -> SnowparkSQLUnexpectedAliasException:
+    def SQL_PYTHON_REPORT_UNEXPECTED_ALIAS(
+        query: Optional[str] = None,
+    ) -> SnowparkSQLUnexpectedAliasException:
         return SnowparkSQLUnexpectedAliasException(
             "You can only define aliases for the root Columns in a DataFrame returned by "
             "select() and agg(). You cannot use aliases for Columns in expressions.",
             "1301",
+            None,  # sfqid
+            query,
         )
 
     @staticmethod
-    def SQL_PYTHON_REPORT_INVALID_ID(name: str) -> SnowparkSQLInvalidIdException:
+    def SQL_PYTHON_REPORT_INVALID_ID(
+        name: str, query: Optional[str] = None
+    ) -> SnowparkSQLInvalidIdException:
         return SnowparkSQLInvalidIdException(
             f'The column specified in df("{name}") '
             f"is not present in the output of the DataFrame.",
             "1302",
+            None,  # sfqid,
+            query,
         )
 
     @staticmethod
     def SQL_PYTHON_REPORT_JOIN_AMBIGUOUS(
-        c1: str, c2: str
+        c1: str,
+        c2: str,
+        query: Optional[str] = None,
     ) -> SnowparkSQLAmbiguousJoinException:
         return SnowparkSQLAmbiguousJoinException(
             f"The reference to the column '{c1}' is ambiguous. The column is "
@@ -264,6 +276,8 @@ def SQL_PYTHON_REPORT_JOIN_AMBIGUOUS(
             f"either DataFrame for disambiguation. See the API documentation of "
             f"the DataFrame.join() method for more details.",
             "1303",
+            None,  # sfqid
+            query,
         )
 
     @staticmethod

tests/unit/scala/test_error_message.py

@@ -266,32 +266,38 @@ def test_sql_last_query_return_resultset():
 
 
 def test_sql_python_report_unexpected_alias():
-    ex = SnowparkClientExceptionMessages.SQL_PYTHON_REPORT_UNEXPECTED_ALIAS()
+    ex = SnowparkClientExceptionMessages.SQL_PYTHON_REPORT_UNEXPECTED_ALIAS(
+        "test query"
+    )
     assert type(ex) == SnowparkSQLUnexpectedAliasException
     assert ex.error_code == "1301"
     assert (
         ex.message
         == "You can only define aliases for the root Columns in a DataFrame returned by "
         "select() and agg(). You cannot use aliases for Columns in expressions."
     )
+    assert ex.query == "test query"
 
 
 def test_sql_python_report_invalid_id():
     name = "C1"
-    ex = SnowparkClientExceptionMessages.SQL_PYTHON_REPORT_INVALID_ID(name)
+    query = "test query"
+    ex = SnowparkClientExceptionMessages.SQL_PYTHON_REPORT_INVALID_ID(name, query)
     assert type(ex) == SnowparkSQLInvalidIdException
     assert ex.error_code == "1302"
     assert (
         ex.message
         == f'The column specified in df("{name}") is not present in the output of the DataFrame.'
     )
+    assert ex.query == query
 
 
 def test_sql_report_join_ambiguous():
     column = "A"
     c1 = column
     c2 = column
-    ex = SnowparkClientExceptionMessages.SQL_PYTHON_REPORT_JOIN_AMBIGUOUS(c1, c2)
+    query = "test query"
+    ex = SnowparkClientExceptionMessages.SQL_PYTHON_REPORT_JOIN_AMBIGUOUS(c1, c2, query)
     assert type(ex) == SnowparkSQLAmbiguousJoinException
     assert ex.error_code == "1303"
     assert (
@@ -303,6 +309,7 @@ def test_sql_report_join_ambiguous():
         f"either DataFrame for disambiguation. See the API documentation of "
         f"the DataFrame.join() method for more details."
     )
+    assert ex.query == query
 
 
 def test_server_cannot_find_current_db_or_schema():