11import json
2+ import math
23import threading
34import time
5+ from ctypes import c_int32 , c_int64
46from dataclasses import dataclass , field
5- from concurrent .futures import ThreadPoolExecutor
6- from typing import Dict , Optional , List , Any , TYPE_CHECKING
7+ from concurrent .futures import Future , ThreadPoolExecutor
8+ from typing import Dict , Optional , List , Any , Type , Union
79
810from databricks .sql .common .http import HttpMethod
911from databricks .sql .common .url_utils import normalize_host_with_protocol
1012
11- if TYPE_CHECKING :
12- from databricks .sql .client import Connection
13-
1413
1514@dataclass
1615class FeatureFlagEntry :
@@ -43,77 +42,141 @@ def from_dict(cls, data: Dict[str, Any]) -> "FeatureFlagsResponse":
4342REFRESH_BEFORE_EXPIRY_SECONDS = 10 # Start proactive refresh 10s before expiry
4443
4544
45+ @dataclass
46+ class _CacheState :
47+ # Only values/coordination are shared; credentials and HTTP clients are not.
48+ flags : Optional [Dict [str , str ]] = None
49+ ttl_seconds : int = DEFAULT_TTL_SECONDS
50+ last_refresh_time : float = 0
51+ lock : Any = field (default_factory = threading .RLock )
52+ refresh : Optional [Future ] = None
53+
54+
55+ def _cache_key (host , headers ):
56+ workspace_id = (headers or {}).get ("x-databricks-org-id" )
57+ return (
58+ ("workspace" , workspace_id )
59+ if workspace_id
60+ else ("host" , normalize_host_with_protocol (host ).lower ())
61+ )
62+
63+
4664class FeatureFlagsContext :
4765 """
48- Manages fetching and caching of server-side feature flags for a connection .
66+ Authenticated flag reader usable before any session/backend is opened .
4967
5068 1. The very first check for any flag is a synchronous, BLOCKING operation.
5169 2. Subsequent refreshes (triggered near TTL expiry) are done asynchronously
5270 in the background, returning stale data until the refresh completes.
5371 """
5472
5573 def __init__ (
56- self , connection : "Connection" , executor : ThreadPoolExecutor , http_client
74+ self , host , executor , http_client , auth_provider , user_agent , headers , state
5775 ):
5876 from databricks .sql import __version__
5977
60- self ._connection = connection
6178 self ._executor = executor # Used for ASYNCHRONOUS refreshes
62- self ._lock = threading .RLock ()
63-
64- # Cache state: `None` indicates the cache has never been loaded.
65- self ._flags : Optional [Dict [str , str ]] = None
66- self ._ttl_seconds : int = DEFAULT_TTL_SECONDS
67- self ._last_refresh_time : float = 0
79+ self ._state = state
80+ self ._auth_provider = auth_provider
81+ self ._headers = {"User-Agent" : user_agent , ** headers }
6882
6983 endpoint_suffix = FEATURE_FLAGS_ENDPOINT_SUFFIX_FORMAT .format (__version__ )
7084 self ._feature_flag_endpoint = (
71- normalize_host_with_protocol (self ._connection .session .host )
72- + endpoint_suffix
85+ normalize_host_with_protocol (host ) + endpoint_suffix
7386 )
7487
7588 # Use the provided HTTP client
7689 self ._http_client = http_client
7790
7891 def _is_refresh_needed (self ) -> bool :
7992 """Checks if the cache is due for a proactive background refresh."""
80- if self ._flags is None :
93+ if self ._state . flags is None :
8194 return False # Not eligible for refresh until loaded once.
8295
83- refresh_threshold = self ._last_refresh_time + (
84- self ._ttl_seconds - REFRESH_BEFORE_EXPIRY_SECONDS
96+ refresh_threshold = self ._state . last_refresh_time + (
97+ self ._state . ttl_seconds - REFRESH_BEFORE_EXPIRY_SECONDS
8598 )
8699 return time .monotonic () > refresh_threshold
87100
88- def get_flag_value (self , name : str , default_value : Any ) -> Any :
101+ def _get_value (self , name : str ) -> Any :
89102 """
90- Checks if a feature is enabled .
103+ Reads and parses a flag's JSON value .
91104 - BLOCKS on the first call until flags are fetched.
92105 - Returns cached values on subsequent calls, triggering non-blocking refreshes if needed.
93106 """
94- with self ._lock :
107+ with self ._state . lock :
95108 # If cache has never been loaded, perform a synchronous, blocking fetch.
96- if self ._flags is None :
109+ if self ._state . flags is None :
97110 self ._refresh_flags ()
98111
99112 # If a proactive background refresh is needed, start one. This is non-blocking.
100- elif self ._is_refresh_needed ():
101- # We don't check for an in-flight refresh; the executor queues the task, which is safe.
102- self ._executor .submit (self ._refresh_flags )
113+ elif self ._is_refresh_needed () and (
114+ self ._state .refresh is None or self ._state .refresh .done ()
115+ ):
116+ self ._state .refresh = self ._executor .submit (self ._refresh_flags )
103117
104- assert self ._flags is not None
105-
106- # Now, return the value from the populated cache.
107- return self ._flags .get (name , default_value )
118+ raw = (self ._state .flags or {}).get (name )
119+ try :
120+ return json .loads (raw ) if raw is not None else None
121+ except (TypeError , ValueError ):
122+ return {"true" : True , "false" : False }.get (str (raw ).lower ())
123+
124+ def get_bool (self , name : str , default_value : bool = False ) -> bool :
125+ value = self ._get_value (name )
126+ return value if type (value ) is bool else default_value
127+
128+ def _get_int (
129+ self ,
130+ name : str ,
131+ integer_type : Type [Union [c_int32 , c_int64 ]],
132+ default_value : Optional [int ],
133+ ) -> Optional [int ]:
134+ value = self ._get_value (name )
135+ if type (value ) is int and integer_type (value ).value == value :
136+ return value
137+ return default_value
138+
139+ def get_int32 (
140+ self , name : str , default_value : Optional [int ] = None
141+ ) -> Optional [int ]:
142+ return self ._get_int (name , c_int32 , default_value )
143+
144+ def get_int64 (
145+ self , name : str , default_value : Optional [int ] = None
146+ ) -> Optional [int ]:
147+ return self ._get_int (name , c_int64 , default_value )
148+
149+ def get_double (
150+ self , name : str , default_value : Optional [float ] = None
151+ ) -> Optional [float ]:
152+ value = self ._get_value (name )
153+ if type (value ) is int :
154+ try :
155+ value = float (value )
156+ except OverflowError :
157+ return default_value
158+ return value if type (value ) is float and math .isfinite (value ) else default_value
159+
160+ def get_string (
161+ self , name : str , default_value : Optional [str ] = None
162+ ) -> Optional [str ]:
163+ value = self ._get_value (name )
164+ return value if isinstance (value , str ) else default_value
165+
166+ def get_string_list (
167+ self , name : str , default_value : Optional [List [str ]] = None
168+ ) -> Optional [List [str ]]:
169+ value = self ._get_value (name )
170+ if isinstance (value , list ) and all (isinstance (item , str ) for item in value ):
171+ return value
172+ return default_value
108173
109174 def _refresh_flags (self ):
110175 """Performs a synchronous network request to fetch and update flags."""
111- headers = {}
176+ headers = dict ( self . _headers )
112177 try :
113178 # Authenticate the request
114- self ._connection .session .auth_provider .add_headers (headers )
115- headers ["User-Agent" ] = self ._connection .session .useragent_header
116- headers .update (self ._connection .session .get_spog_headers ())
179+ self ._auth_provider .add_headers (headers )
117180
118181 response = self ._http_client .request (
119182 HttpMethod .GET , self ._feature_flag_endpoint , headers = headers , timeout = 30
@@ -126,30 +189,32 @@ def _refresh_flags(self):
126189 self ._update_cache_from_response (ff_response )
127190 else :
128191 # On failure, initialize with an empty dictionary to prevent re-blocking.
129- if self ._flags is None :
130- self ._flags = {}
192+ if self ._state . flags is None :
193+ self ._state . flags = {}
131194
132- except Exception as e :
195+ except Exception :
133196 # On exception, initialize with an empty dictionary to prevent re-blocking.
134- if self ._flags is None :
135- self ._flags = {}
197+ if self ._state . flags is None :
198+ self ._state . flags = {}
136199
137200 def _update_cache_from_response (self , ff_response : FeatureFlagsResponse ):
138201 """Atomically updates the internal cache state from a successful server response."""
139- with self ._lock :
140- self ._flags = {flag .name : flag .value for flag in ff_response .flags }
202+ with self ._state . lock :
203+ self ._state . flags = {flag .name : flag .value for flag in ff_response .flags }
141204 if ff_response .ttl_seconds is not None and ff_response .ttl_seconds > 0 :
142- self ._ttl_seconds = ff_response .ttl_seconds
143- self ._last_refresh_time = time .monotonic ()
205+ self ._state . ttl_seconds = ff_response .ttl_seconds
206+ self ._state . last_refresh_time = time .monotonic ()
144207
145208
146209class FeatureFlagsContextFactory :
147210 """
148- Manages a singleton instance of FeatureFlagsContext per connection session.
149- Also manages a shared ThreadPoolExecutor for all background refresh operations.
211+ Process-wide flag values per workspace and a shared refresh executor.
212+
213+ Both are created lazily and retained until process exit. Session close does
214+ not evict values or shut down the executor, which other readers may still use.
150215 """
151216
152- _context_map : Dict [str , FeatureFlagsContext ] = {}
217+ _context_map : Dict [tuple , _CacheState ] = {}
153218 _executor : Optional [ThreadPoolExecutor ] = None
154219 _lock = threading .Lock ()
155220
@@ -162,31 +227,41 @@ def _initialize(cls):
162227 )
163228
164229 @classmethod
165- def get_instance (cls , connection : "Connection" ) -> FeatureFlagsContext :
166- """Gets or creates a FeatureFlagsContext for the given connection."""
230+ def get_instance (
231+ cls , host , http_client , auth_provider , user_agent , headers = None
232+ ) -> FeatureFlagsContext :
233+ """Reuse the cache with this caller's authenticated transport, even pre-session."""
234+ headers = {name .lower (): value for name , value in (headers or {}).items ()}
167235 with cls ._lock :
168236 cls ._initialize ()
169237 assert cls ._executor is not None
170238
171- # Cache at HOST level - share feature flags across connections to same host
172- # Feature flags are per-host, not per-session
173- key = connection .session .host
239+ key = _cache_key (host , headers )
174240 if key not in cls ._context_map :
175- cls ._context_map [key ] = FeatureFlagsContext (
176- connection , cls ._executor , connection .session .http_client
177- )
178- return cls ._context_map [key ]
241+ cls ._context_map [key ] = _CacheState ()
242+ return FeatureFlagsContext (
243+ host ,
244+ cls ._executor ,
245+ http_client ,
246+ auth_provider ,
247+ user_agent ,
248+ headers ,
249+ cls ._context_map [key ],
250+ )
179251
180252 @classmethod
181- def remove_instance (cls , connection : "Connection" ):
182- """Removes the context for a given connection and shuts down the executor if no clients remain."""
253+ def remove_instance (cls , host , headers = None ):
254+ """Explicitly evict workspace values and stop the executor if the cache is empty.
255+
256+ Used for test/reset cleanup, not individual session teardown.
257+ """
183258 with cls ._lock :
184- # Use host as key to match get_instance
185- key = connection . session . host
259+ headers = { name . lower (): value for name , value in ( headers or {}). items ()}
260+ key = _cache_key ( host , headers )
186261 if key in cls ._context_map :
187262 cls ._context_map .pop (key , None )
188263
189- # If this was the last active context , clean up the thread pool.
264+ # If no cached workspaces remain , clean up the thread pool.
190265 if not cls ._context_map and cls ._executor is not None :
191266 cls ._executor .shutdown (wait = False )
192267 cls ._executor = None
0 commit comments