From ce7f07ec7761d0e96bd6d41cec9ef9b7ee124000 Mon Sep 17 00:00:00 2001 From: Rehan Durrani Date: Thu, 19 Sep 2024 14:21:43 -0700 Subject: [PATCH 1/3] [SNOW-1669128]: Ensure SnowflakeConnection atexit function is registered on import --- src/snowflake/connector/connection.py | 27 ++++++++++++++++++++------- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/src/snowflake/connector/connection.py b/src/snowflake/connector/connection.py index 1d68838a3..85c45e6f0 100644 --- a/src/snowflake/connector/connection.py +++ b/src/snowflake/connector/connection.py @@ -314,6 +314,25 @@ class TypeAndBinding(NamedTuple): binding: str | None +# SNOW-1669128 previously, the close_at_exit function for the SnowflakeConnection object +# was registered with the `atexit` module during initialization of a SnowflakeConnection object. +# This caused the Snowpark Session object's atexit function to be registered before the +# SnowflakeConnection's object when the Session was made without an existing SnowflakeConnection +# object, because the SnowflakeConnection object's atexit was only registered after it was initialized +# whereas the Snowpark Session object's atexit was registered when the file was imported. During the +# atexit for the Snowpark Session, it attempts to send telemetry regarding cleanup operations, but since +# it's connection is already closed, it fails to send the telemetry. Moving the registration of the atexit +# function here means that it will happen when the SnowflakeConnection object is imported - and so the +# Snowpark Session's atexit function will be registered after, and the order will be correct. +def _close_connection_at_exit(snowflake_connection: SnowflakeConnection): + with suppress(Exception): + snowflake_connection.close(retry=False) + + +# check SNOW-1218851 for long term improvement plan to refactor ocsp code +atexit.register(_close_connection_at_exit) + + class SnowflakeConnection: """Implementation of the connection object for the Snowflake Database. @@ -459,8 +478,6 @@ def __init__( # get the imported modules from sys.modules self._log_telemetry_imported_packages() - # check SNOW-1218851 for long term improvement plan to refactor ocsp code - atexit.register(self._close_at_exit) @property def insecure_mode(self) -> bool: @@ -773,7 +790,7 @@ def connect(self, **kwargs) -> None: def close(self, retry: bool = True) -> None: """Closes the connection.""" # unregister to dereference connection object as it's already closed after the execution - atexit.unregister(self._close_at_exit) + atexit.unregister(_close_connection_at_exit) try: if not self.rest: logger.debug("Rest object has been destroyed, cannot close session") @@ -1813,10 +1830,6 @@ def _cache_query_status(self, sf_qid: str, status_ret: QueryStatus) -> None: ) # Prevent KeyError when multiple threads try to remove the same query id self._done_async_sfqids[sf_qid] = None - def _close_at_exit(self): - with suppress(Exception): - self.close(retry=False) - def _process_error_query_status( self, sf_qid: str, From 890fa57c0358e8870ef2e46daee2951a0323944c Mon Sep 17 00:00:00 2001 From: Rehan Durrani Date: Thu, 19 Sep 2024 14:27:22 -0700 Subject: [PATCH 2/3] Update description --- DESCRIPTION.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/DESCRIPTION.md b/DESCRIPTION.md index 71d77406b..d44ea2bdd 100644 --- a/DESCRIPTION.md +++ b/DESCRIPTION.md @@ -8,6 +8,9 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne # Release Notes +- v3.12.3(DATE TBD) + - Ensure `snowflake.connector.SnowflakeConnection`'s `atexit` function is registered on import. + - v3.12.2(September 11,2024) - Improved error handling for asynchronous queries, providing more detailed and informative error messages when an async query fails. - Improved inference of top-level domains for accounts specifying a region in China, now defaulting to snowflakecomputing.cn. From 306f5812cc4222213d187da5380f8e19c5fd508f Mon Sep 17 00:00:00 2001 From: Rehan Durrani Date: Thu, 19 Sep 2024 14:39:49 -0700 Subject: [PATCH 3/3] Add reference to active connections --- src/snowflake/connector/connection.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/src/snowflake/connector/connection.py b/src/snowflake/connector/connection.py index 85c45e6f0..304968070 100644 --- a/src/snowflake/connector/connection.py +++ b/src/snowflake/connector/connection.py @@ -314,6 +314,9 @@ class TypeAndBinding(NamedTuple): binding: str | None +active_connections: list[SnowflakeConnection] = [] + + # SNOW-1669128 previously, the close_at_exit function for the SnowflakeConnection object # was registered with the `atexit` module during initialization of a SnowflakeConnection object. # This caused the Snowpark Session object's atexit function to be registered before the @@ -324,9 +327,10 @@ class TypeAndBinding(NamedTuple): # it's connection is already closed, it fails to send the telemetry. Moving the registration of the atexit # function here means that it will happen when the SnowflakeConnection object is imported - and so the # Snowpark Session's atexit function will be registered after, and the order will be correct. -def _close_connection_at_exit(snowflake_connection: SnowflakeConnection): +def _close_connection_at_exit(): with suppress(Exception): - snowflake_connection.close(retry=False) + for connection in active_connections: + connection.close(retry=False) # check SNOW-1218851 for long term improvement plan to refactor ocsp code @@ -475,6 +479,7 @@ def __init__( self.connect(**kwargs) self._telemetry = TelemetryClient(self._rest) self.expired = False + active_connections.append(self) # get the imported modules from sys.modules self._log_telemetry_imported_packages() @@ -789,8 +794,7 @@ def connect(self, **kwargs) -> None: def close(self, retry: bool = True) -> None: """Closes the connection.""" - # unregister to dereference connection object as it's already closed after the execution - atexit.unregister(_close_connection_at_exit) + active_connections.remove(self) try: if not self.rest: logger.debug("Rest object has been destroyed, cannot close session")