33from __future__ import annotations
44
55import threading
6+ import weakref
7+ from typing import Any
8+
9+ from dbt .adapters .events .logging import AdapterLogger
10+
11+ LOGGER = AdapterLogger (__name__ )
612
713_patches_applied = False
814_patch_lock = threading .Lock ()
@@ -17,6 +23,14 @@ def apply_pyspark_workarounds() -> None:
1723 return
1824 _neutralize_release_thread_pool_shutdown ()
1925 _silence_release_all_warning ()
26+ # Athena AuthToken refresh across pyspark's reattach cycle. These
27+ # three patches cooperate as one feature:
28+ # 1. Stash the ChannelBuilder on the gRPC stub.
29+ # 2. Refresh metadata via the builder before each ReattachExecute.
30+ # 3. Allow one PERMISSION_DENIED retry so a 403 still recovers.
31+ _stash_channel_builder_on_stub ()
32+ _refresh_reattach_iterator_metadata ()
33+ _retry_permission_denied_in_spark_client ()
2034 _patches_applied = True
2135
2236
@@ -58,3 +72,103 @@ def _silence_release_all_warning() -> None:
5872 "ignore" ,
5973 message = r"ReleaseExecute failed with exception:.*" ,
6074 )
75+
76+
77+ # weakref so a reused worker thread does not pin a stale iterator.
78+ # Assumes one in-flight iterator per thread (pyspark consumes synchronously);
79+ # concurrent iterators would need a stack here.
80+ _CURRENT_ITERATOR_THREAD_LOCAL = threading .local ()
81+
82+
83+ def _stash_channel_builder_on_stub () -> None :
84+ """Cache the ChannelBuilder on the gRPC stub so the reattach iterator can find it."""
85+ from pyspark .sql .connect .client .core import SparkConnectClient
86+
87+ original_init = SparkConnectClient .__init__
88+
89+ def _patched_init (self : Any , * args : Any , ** kwargs : Any ) -> None :
90+ original_init (self , * args , ** kwargs )
91+ builder = getattr (self , "_builder" , None )
92+ stub = getattr (self , "_stub" , None )
93+ if (
94+ stub is not None
95+ and builder is not None
96+ and callable (getattr (builder , "metadata" , None ))
97+ ):
98+ stub ._dbt_athena_builder = builder
99+ LOGGER .debug (
100+ "Stashed AthenaChannelBuilder on Spark Connect stub for metadata refresh."
101+ )
102+
103+ SparkConnectClient .__init__ = _patched_init
104+
105+
106+ def _refresh_reattach_iterator_metadata () -> None :
107+ """Refresh metadata before each ReattachExecute so the AuthToken can rotate mid-stream.
108+
109+ pyspark captures ``metadata`` once at ``__init__`` and reuses the same
110+ list forever, which keeps Athena's 30-min ``x-aws-proxy-auth`` token
111+ pinned to its initial value.
112+ """
113+ from pyspark .sql .connect .client .reattach import (
114+ ExecutePlanResponseReattachableIterator ,
115+ )
116+
117+ original_init = ExecutePlanResponseReattachableIterator .__init__
118+ original_call_iter = ExecutePlanResponseReattachableIterator ._call_iter
119+
120+ def _patched_init (self : Any , * args : Any , ** kwargs : Any ) -> None :
121+ original_init (self , * args , ** kwargs )
122+ self ._dbt_athena_channel_builder = getattr (self ._stub , "_dbt_athena_builder" , None )
123+ self ._dbt_athena_pd_retried = False
124+ _CURRENT_ITERATOR_THREAD_LOCAL .iterator_ref = weakref .ref (self )
125+
126+ def _patched_call_iter (self : Any , iter_fun : Any ) -> Any :
127+ if self ._iterator is None :
128+ builder = getattr (self , "_dbt_athena_channel_builder" , None )
129+ if builder is not None :
130+ old_token = getattr (builder , "_auth_token" , None )
131+ try :
132+ self ._metadata = builder .metadata ()
133+ except Exception as e : # noqa: BLE001 - refresh is best-effort
134+ LOGGER .warning (f"Metadata refresh on reattach failed: { e } )
135+ else :
136+ new_token = getattr (builder , "_auth_token" , None )
137+ if new_token is not None and new_token != old_token :
138+ LOGGER .debug ("Reattach metadata refreshed: AuthToken rotated." )
139+ return original_call_iter (self , iter_fun )
140+
141+ ExecutePlanResponseReattachableIterator .__init__ = _patched_init
142+ ExecutePlanResponseReattachableIterator ._call_iter = _patched_call_iter
143+
144+
145+ def _retry_permission_denied_in_spark_client () -> None :
146+ """Treat PERMISSION_DENIED as retryable so a 403 from token expiry can recover.
147+
148+ pyspark's default ``retry_exception`` only retries UNAVAILABLE (and one
149+ ``INTERNAL`` cursor case), so a 403 propagates out before the reattach
150+ iterator can re-issue ``ReattachExecute``. Allowing exactly one retry
151+ per iterator pairs with ``_refresh_reattach_iterator_metadata`` so the
152+ next reattach goes out with the rotated token; a second 403 means the
153+ failure is genuine and we propagate.
154+ """
155+ import grpc
156+ from pyspark .sql .connect .client .core import SparkConnectClient
157+
158+ original = SparkConnectClient .retry_exception .__func__
159+
160+ def _patched (cls : Any , e : BaseException ) -> bool :
161+ if original (cls , e ):
162+ return True
163+ if not (isinstance (e , grpc .RpcError ) and e .code () == grpc .StatusCode .PERMISSION_DENIED ):
164+ return False
165+ iterator_ref = getattr (_CURRENT_ITERATOR_THREAD_LOCAL , "iterator_ref" , None )
166+ iterator = iterator_ref () if iterator_ref is not None else None
167+ if iterator is None or getattr (iterator , "_dbt_athena_pd_retried" , False ):
168+ LOGGER .warning ("PERMISSION_DENIED retry budget exhausted; propagating." )
169+ return False
170+ iterator ._dbt_athena_pd_retried = True
171+ LOGGER .debug ("PERMISSION_DENIED detected; allowing one reattach with refreshed metadata." )
172+ return True
173+
174+ SparkConnectClient .retry_exception = classmethod (_patched )
0 commit comments