diff --git a/src/snowflake/snowpark/mock/_connection.py b/src/snowflake/snowpark/mock/_connection.py index 92fea43002..4c09e656f3 100644 --- a/src/snowflake/snowpark/mock/_connection.py +++ b/src/snowflake/snowpark/mock/_connection.py @@ -31,6 +31,7 @@ from snowflake.snowpark._internal.error_message import SnowparkClientExceptionMessages from snowflake.snowpark._internal.server_connection import DEFAULT_STRING_SIZE from snowflake.snowpark._internal.utils import ( + IS_V5_DRIVER, is_in_stored_procedure, result_set_to_rows, ) @@ -78,6 +79,10 @@ def connect(self, **kwargs) -> None: } self._rest = Mock(**attrs) + if IS_V5_DRIVER: + # UD Connection.__init__ calls _connect(); legacy called connect(). + _connect = connect + def close(self, retry: bool = True) -> None: self._rest = None