Skip to content

Commit 9eb024c

Browse files
authored
feat(feature-flags): generalize driver-side cache (#974)
* Generalize driver feature flag cache and typed reads Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> * refactor(feature-flags): use standard integer types for validation Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> * fix(feature-flags): decouple cache setup from telemetry Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> * chore(kernel): bump to merged feature flag cache Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> * fix: initialize session feature flags lazily Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> * docs: clarify process-wide feature flag cache lifetime Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> * chore(feature-flags): annotate typed getter defaults Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> * fix(feature-flags): preserve case-insensitive boolean values Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> --------- Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com>
1 parent 61b9a7f commit 9eb024c

7 files changed

Lines changed: 392 additions & 75 deletions

File tree

Lines changed: 135 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,15 @@
11
import json
2+
import math
23
import threading
34
import time
5+
from ctypes import c_int32, c_int64
46
from 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

810
from databricks.sql.common.http import HttpMethod
911
from 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
1615
class FeatureFlagEntry:
@@ -43,77 +42,141 @@ def from_dict(cls, data: Dict[str, Any]) -> "FeatureFlagsResponse":
4342
REFRESH_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+
4664
class 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

146209
class 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

‎src/databricks/sql/session.py‎

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import logging
22
import re
3+
from functools import cached_property
34
from typing import Dict, Tuple, List, Optional, Any, Type, TYPE_CHECKING
45

56
from databricks.sql.types import SSLOptions
@@ -13,6 +14,10 @@
1314
from databricks.sql.backend.types import SessionId, BackendType
1415
from databricks.sql.common.unified_http_client import UnifiedHttpClient
1516
from databricks.sql.common.agent import detect as detect_agent
17+
from databricks.sql.common.feature_flag import (
18+
FeatureFlagsContext,
19+
FeatureFlagsContextFactory,
20+
)
1621
from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
1722

1823
if TYPE_CHECKING:
@@ -134,7 +139,8 @@ def __init__(
134139
# provider when an ``access_token`` is present, and ``None``
135140
# otherwise (OAuth M2M/U2M resolve purely from the raw kwargs
136141
# the bridge reads). The Thrift / SEA backends are unchanged.
137-
if kwargs.get("use_kernel", False):
142+
self.use_kernel = kwargs.get("use_kernel", False)
143+
if self.use_kernel:
138144
access_token = kwargs.get("access_token")
139145
self.auth_provider = (
140146
AccessTokenAuthProvider(access_token) if access_token else None
@@ -155,6 +161,19 @@ def __init__(
155161

156162
self.protocol_version = None
157163

164+
@cached_property
165+
def feature_flags(self) -> Optional[FeatureFlagsContext]:
166+
"""Attach this session's transport to the shared cache only when needed."""
167+
if self.use_kernel:
168+
return None
169+
return FeatureFlagsContextFactory.get_instance(
170+
self.host,
171+
self.http_client,
172+
self.auth_provider,
173+
self.useragent_header,
174+
self.get_spog_headers(),
175+
)
176+
158177
def _create_backend(
159178
self,
160179
server_hostname: str,
@@ -166,7 +185,6 @@ def _create_backend(
166185
) -> DatabricksClient:
167186
"""Create and return the appropriate backend client."""
168187
self.use_sea = kwargs.get("use_sea", False)
169-
self.use_kernel = kwargs.get("use_kernel", False)
170188

171189
if self.use_kernel and self.use_sea:
172190
raise ValueError(

‎src/databricks/sql/telemetry/telemetry_client.py‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,6 @@
4040
import uuid
4141
import locale
4242
from databricks.sql.telemetry.utils import BaseTelemetryClient
43-
from databricks.sql.common.feature_flag import FeatureFlagsContextFactory
4443
from databricks.sql.common.unified_http_client import UnifiedHttpClient
4544
from databricks.sql.common.http import HttpMethod
4645
from databricks.sql.exc import RequestError
@@ -133,12 +132,12 @@ def is_telemetry_enabled(connection: "Connection") -> bool:
133132
if not connection.enable_telemetry:
134133
return False
135134

136-
# Only fetch feature flags when enable_telemetry=True and not forced
137-
context = FeatureFlagsContextFactory.get_instance(connection)
138-
flag_value = context.get_flag_value(
135+
context = connection.session.feature_flags
136+
if context is None:
137+
return False
138+
return context.get_bool(
139139
TelemetryHelper.TELEMETRY_FEATURE_FLAG_NAME, default_value=False
140140
)
141-
return str(flag_value).lower() == "true"
142141

143142

144143
class NoopTelemetryClient(BaseTelemetryClient):

‎tests/unit/conftest.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
from unittest.mock import patch
2+
3+
import pytest
4+
5+
6+
@pytest.fixture(autouse=True)
7+
def session_feature_flags():
8+
# Unit connections must not contact connector-service. Cache/lifecycle tests
9+
# exercise the real reader with a mocked HTTP client.
10+
with patch("databricks.sql.session.FeatureFlagsContextFactory") as factory:
11+
factory.get_instance.return_value.get_bool.return_value = False
12+
yield factory

0 commit comments

Comments
 (0)