SNOW-2241913: non-retryable error handling (#3770)

Author: sfc-gh-yuwang

CHANGELOG.md

@@ -104,6 +104,7 @@
 
 #### Improvements
 
+- Improved `DataFrameReader.dbapi`(PuPr) that dbapi will not retry on non-retryable error such as SQL syntax error on external data source query.
 - Removed unnecessary warnings about local package version mismatch when using `session.read.option('rowTag', <tag_name>).xml(<stage_file_path>)` or `xpath` functions.
 - Improved `DataFrameReader.dbapi` (PuPr) reading performance by setting the default `fetch_size` parameter value to 100000.
 - Improved error message for XSD validation failure when reading XML files using `session.read.option('rowValidationXSDPath', <xsd_path>).xml(<stage_file_path>)`.

src/snowflake/snowpark/_internal/data_source/datasource_reader.py

@@ -85,6 +85,11 @@ def read(self, partition: str) -> Iterator[List[Any]]:
                         batch = []
             else:
                 raise ValueError("fetch size cannot be smaller than 0")
+        except Exception as exc:
+            if self.driver.non_retryable_error_checker(exc):
+                raise SnowparkDataframeReaderException(message=str(exc))
+            else:
+                raise
         finally:
             try:
                 cursor.close()

src/snowflake/snowpark/_internal/data_source/drivers/base_driver.py

@@ -55,6 +55,9 @@ def to_snow_type(self, schema: List[Any]) -> StructType:
             f"{self.__class__.__name__} has not implemented to_snow_type function"
         )
 
+    def non_retryable_error_checker(self, error: Exception) -> bool:
+        return False
+
     @staticmethod
     def prepare_connection(
         conn: "Connection",

src/snowflake/snowpark/_internal/data_source/drivers/databricks_driver.py

@@ -68,6 +68,18 @@ def to_snow_type(self, schema: List[Any]) -> StructType:
             all_columns.append(StructField(column_name, data_type, True))
         return StructType(all_columns)
 
+    def non_retryable_error_checker(self, error: Exception) -> bool:
+        import databricks.sql
+
+        if isinstance(error, databricks.sql.ServerOperationError):
+            syntax_error_codes = [
+                "PARSE_SYNTAX_ERROR",  # syntax error
+            ]
+            for error_code in syntax_error_codes:
+                if error_code in str(error):
+                    return True
+        return False
+
     def udtf_class_builder(
         self,
         fetch_size: int = 1000,

src/snowflake/snowpark/_internal/data_source/drivers/oracledb_driver.py

@@ -111,6 +111,30 @@ def prepare_connection(
             conn.outputtypehandler = output_type_handler
         return conn
 
+    def non_retryable_error_checker(self, error: Exception) -> bool:
+        import oracledb
+
+        if isinstance(error, oracledb.DatabaseError):
+            syntax_error_codes = [
+                "ORA-00900",  # invalid SQL statement
+                "ORA-00901",  # invalid CREATE command
+                "ORA-00904",  # invalid identifier
+                "ORA-00905",  # missing keyword
+                "ORA-00906",  # missing left parenthesis
+                "ORA-00907",  # missing right parenthesis
+                "ORA-00911",  # invalid character
+                "ORA-00920",  # invalid relational operator
+                "ORA-00921",  # unexpected end of SQL command
+                "ORA-00923",  # FROM keyword not found where expected
+                "ORA-00933",  # SQL command not properly ended
+                "ORA-00936",  # missing expression
+                "ORA-00942",  # table or view does not exist
+            ]
+            for error_code in syntax_error_codes:
+                if error_code in str(error):
+                    return True
+        return False
+
     def udtf_class_builder(
         self,
         fetch_size: int = 1000,

src/snowflake/snowpark/_internal/data_source/drivers/psycopg2_driver.py

@@ -211,6 +211,18 @@ def to_snow_type(self, schema: List[Any]) -> StructType:
             fields.append(StructField(name, data_type, True))
         return StructType(fields)
 
+    def non_retryable_error_checker(self, error: Exception) -> bool:
+        import psycopg2
+
+        if isinstance(error, psycopg2.errors.SyntaxError):
+            syntax_error_codes = [
+                "42601",  # syntax error
+            ]
+            for error_code in syntax_error_codes:
+                if error_code == str(error.pgcode):
+                    return True
+        return False
+
     @staticmethod
     def to_result_snowpark_df(
         session: "Session", table_name, schema, _emit_ast: bool = True

src/snowflake/snowpark/_internal/data_source/drivers/pymsql_driver.py

@@ -183,6 +183,18 @@ def to_snow_type(self, schema: List[Any]) -> StructType:
             fields.append(StructField(name, data_type, null_ok))
         return StructType(fields)
 
+    def non_retryable_error_checker(self, error: Exception) -> bool:
+        import pymysql
+
+        if isinstance(error, pymysql.err.ProgrammingError):
+            syntax_error_codes = [
+                "1064",  # syntax error
+            ]
+            for error_code in syntax_error_codes:
+                if error_code in str(error):
+                    return True
+        return False
+
     def udtf_class_builder(
         self,
         fetch_size: int = 1000,

tests/integ/datasource/test_databricks.py

@@ -265,6 +265,18 @@ def test_unsupported_type():
     assert schema == StructType([StructField("TEST_COL", StringType(), nullable=True)])
 
 
+def test_databricks_non_retryable_error(session):
+    with pytest.raises(
+        SnowparkDataframeReaderException,
+        match="PARSE_SYNTAX_ERROR",
+    ):
+        session.read.dbapi(
+            create_databricks_connection,
+            table=TEST_TABLE_NAME,
+            predicates=["invalid syntax"],
+        )
+
+
 def test_session_init(session):
     with pytest.raises(
         SnowparkDataframeReaderException,

tests/integ/datasource/test_mysql.py

@@ -305,6 +305,18 @@ def test_unsupported_type():
     assert schema == StructType([StructField("TEST_COL", StringType(), nullable=True)])
 
 
+def test_mysql_non_retryable_error(session):
+    with pytest.raises(
+        SnowparkDataframeReaderException,
+        match="You have an error in your SQL syntax",
+    ):
+        session.read.dbapi(
+            create_connection_mysql,
+            table=TEST_TABLE_NAME,
+            predicates=["invalid syntax"],
+        )
+
+
 def test_session_init(session):
     with pytest.raises(
         SnowparkDataframeReaderException,

tests/integ/datasource/test_oracledb.py

@@ -255,6 +255,18 @@ def test_unsupported_type():
     assert schema == StructType([StructField("TEST_COL", StringType(), nullable=True)])
 
 
+def test_oracledb_non_retryable_error(session):
+    with pytest.raises(
+        SnowparkDataframeReaderException,
+        match="ORA-00920: invalid relational operator",
+    ):
+        session.read.dbapi(
+            create_connection_oracledb,
+            table=ORACLEDB_TABLE_NAME,
+            predicates=["invalid syntax"],
+        ).collect()
+
+
 def test_query_timeout_and_session_init(session):
     statement = """
     BEGIN