From 01be2b809a48817a71133dfcdfb9ced449882199 Mon Sep 17 00:00:00 2001 From: Aniket Kulkarni Date: Thu, 14 May 2026 11:37:58 -0400 Subject: [PATCH 01/12] DX-119909: support runtime reload of mutable MCP settings --- src/dremioai/api/dremio/sql.py | 2 +- src/dremioai/api/transport.py | 22 +- src/dremioai/config/settings.py | 254 ++++++++++++++--- src/dremioai/servers/mcp.py | 24 +- tests/api/dremio/test_sql.py | 98 ++++--- tests/api/test_transport_retry.py | 66 +++-- tests/config/test_launchdarkly_integration.py | 117 +++++++- tests/config/test_settings.py | 262 +++++++++++++++--- tests/conftest.py | 49 ++-- tests/servers/test_mcp.py | 4 +- tests/stremable_http_cli.py | 9 +- tests/test_fastmcp_basic.py | 4 +- tests/test_simple_fastmcp_server.py | 30 +- tests/tools/test_output_validation.py | 14 +- tests/tools/test_tools.py | 6 +- 15 files changed, 750 insertions(+), 211 deletions(-) diff --git a/src/dremioai/api/dremio/sql.py b/src/dremioai/api/dremio/sql.py index 02f5793..f272fa6 100644 --- a/src/dremioai/api/dremio/sql.py +++ b/src/dremioai/api/dremio/sql.py @@ -185,7 +185,7 @@ async def get_results( if client is None: client = AsyncHttpClient() - delay = settings.instance().dremio.api.polling_interval + delay = settings.instance().dremio.api.get("polling_interval") endpoint = f"/v0/projects/{project_id}" if project_id else "/api/v3" job: Job = await client.get(f"{endpoint}/job/{qs.id}", deser=Job) diff --git a/src/dremioai/api/transport.py b/src/dremioai/api/transport.py index b80d3bd..ed1cc40 100644 --- a/src/dremioai/api/transport.py +++ b/src/dremioai/api/transport.py @@ -43,18 +43,18 @@ class RetryConfig: def __init__(self): if settings.instance() and settings.instance().dremio: - self.config = settings.instance().dremio.api.http_retry + self.http_retry = settings.instance().dremio.api.http_retry else: - self.config = settings.HttpRetry() + self.http_retry = settings.HttpRetry() @property def max_retries(self) -> int: """Expose max_retries from config for convenience""" - return self.config.max_retries + return self.http_retry.get("max_retries") def get_config_delay(self, attempt_number: int = 0) -> float: - return self.config.initial_delay * ( - self.config.backoff_multiplier**attempt_number + return self.http_retry.get("initial_delay") * ( + self.http_retry.get("backoff_multiplier") ** attempt_number ) def get_delay( @@ -72,7 +72,7 @@ def get_delay( f"Invalid Retry-After header, using exponential backoff - {e}" ) - return min(delay, self.config.max_delay) + return min(delay, self.http_retry.get("max_delay")) async def retry_middleware( @@ -200,7 +200,11 @@ async def post( async with ClientSession(middlewares=(retry_middleware,)) as session: self.log_request("POST", endpoint, params) async with session.post( - f"{self.uri}{endpoint}", params=params, headers=self.headers, json=body, ssl=False + f"{self.uri}{endpoint}", + params=params, + headers=self.headers, + json=body, + ssl=False, ) as response: return await self.handle_response( response, deser, file, top_level_list=top_level_list @@ -222,5 +226,7 @@ def __init__(self): pat = dremio.pat if uri is None or pat is None: - raise RuntimeError("Dremio connection is not properly configured. Both URI and authentication token are required.") + raise RuntimeError( + "Dremio connection is not properly configured. Both URI and authentication token are required." + ) super().__init__(uri, pat) diff --git a/src/dremioai/config/settings.py b/src/dremioai/config/settings.py index 58a16d2..af3f459 100644 --- a/src/dremioai/config/settings.py +++ b/src/dremioai/config/settings.py @@ -26,7 +26,12 @@ field_serializer, AliasChoices, ) -from pydantic_settings import BaseSettings, SettingsConfigDict, PydanticBaseSettingsSource, YamlConfigSettingsSource +from pydantic_settings import ( + BaseSettings, + SettingsConfigDict, + PydanticBaseSettingsSource, + YamlConfigSettingsSource, +) from typing import ( Optional, Union, @@ -52,6 +57,7 @@ from os import environ from importlib.util import find_spec from datetime import datetime +from types import NoneType from dremioai import log from dremioai.config.feature_flags import FeatureFlagManager @@ -73,6 +79,13 @@ def get(self, field_name: str): # Usage: field: Annotated[type, NoFlag()] = Field(...) class NoFlag: """Mark a field to skip LD flag lookups in FlagAwareMixin.get().""" + + pass + + +class RuntimeMutable: + """Mark a field as safe to update from a runtime config reload.""" + pass @@ -82,6 +95,14 @@ def _has_no_flag(model_cls: type, field_name: str) -> bool: return info is not None and any(isinstance(m, NoFlag) for m in info.metadata) +def _has_runtime_mutable(model_cls: type, field_name: str) -> bool: + """Check if a field has the RuntimeMutable annotation marker.""" + info = model_cls.model_fields.get(field_name) + return info is not None and any( + isinstance(m, RuntimeMutable) for m in info.metadata + ) + + # FlagAwareMixin overrides .get() to check LaunchDarkly before returning the # config value. The LD flag key is "{_flag_prefix}.{field_name}", e.g. # "dremio.allow_dml" or "dremio.api.http_retry.max_retries". @@ -101,9 +122,7 @@ def get(self, field_name: str): if _has_no_flag(type(self), field_name): return super().get(field_name) key = f"{self._flag_prefix}.{field_name}" if self._flag_prefix else field_name - return FeatureFlagManager.instance().get_flag( - key, super().get(field_name) - ) + return FeatureFlagManager.instance().get_flag(key, super().get(field_name)) # Convenience base for sub-models that need both FlagAwareMixin and BaseModel. @@ -199,17 +218,17 @@ class Metrics(FlagAwareModel): class HttpRetry(FlagAwareModel): - max_retries: Optional[int] = Field( + max_retries: Annotated[Optional[int], RuntimeMutable()] = Field( default=20, description="Maximum number of retry attempts for rate-limited requests", ) - initial_delay: Optional[float] = Field( + initial_delay: Annotated[Optional[float], RuntimeMutable()] = Field( default=1.0, description="Initial delay in seconds before first retry" ) - max_delay: Optional[float] = Field( + max_delay: Annotated[Optional[float], RuntimeMutable()] = Field( default=60.0, description="Maximum delay in seconds between retries" ) - backoff_multiplier: Optional[float] = Field( + backoff_multiplier: Annotated[Optional[float], RuntimeMutable()] = Field( default=2.0, description="Multiplier for exponential backoff" ) @@ -217,7 +236,7 @@ class HttpRetry(FlagAwareModel): class ApiSettings(FlagAwareModel): # HTTP retry configuration http_retry: Optional[HttpRetry] = Field(default_factory=HttpRetry) - polling_interval: Optional[float] = Field( + polling_interval: Annotated[Optional[float], RuntimeMutable()] = Field( default=1, description="Polling interval for REST api in seconds" ) @@ -237,17 +256,21 @@ def enabled(self) -> bool: class Dremio(FlagAwareModel): uri: Annotated[ - Union[str, HttpUrl, DremioCloudUri], AfterValidator(_resolve_dremio_uri), NoFlag() + Union[str, HttpUrl, DremioCloudUri], + AfterValidator(_resolve_dremio_uri), + NoFlag(), ] raw_pat: Annotated[Optional[str], NoFlag()] = Field(default=None, alias="pat") - raw_project_id: Annotated[Optional[ProjectId], NoFlag()] = Field(default=None, alias="project_id") - enable_search: Optional[bool] = Field( + raw_project_id: Annotated[Optional[ProjectId], NoFlag()] = Field( + default=None, alias="project_id" + ) + enable_search: Annotated[Optional[bool], RuntimeMutable()] = Field( default=False, alias=AliasChoices("enable_search", "enable_experimental"), description="enable experimental tools", ) oauth2: Optional[OAuth2] = None - allow_dml: Optional[bool] = Field(default=False) + allow_dml: Annotated[Optional[bool], RuntimeMutable()] = Field(default=False) extract_org_id_from_jwt: Optional[bool] = Field( default=False, description="Extract org ID from JWT aud claim for LD context targeting", @@ -274,7 +297,7 @@ class Dremio(FlagAwareModel): wlm: Optional[Wlm] = None api: Optional[ApiSettings] = Field(default_factory=ApiSettings) metrics: Optional[Metrics] = None - enable_remote_tools: Optional[bool] = Field( + enable_remote_tools: Annotated[Optional[bool], RuntimeMutable()] = Field( default=False, description="Enable dynamic registration of remote tools from Dremio's Java-side tool registry", ) @@ -407,7 +430,7 @@ class BeeAI(BaseModel): class Settings(FlagAwareMixin, BaseSettings): - log_level: Optional[str] = Field(default="INFO") + log_level: Annotated[Optional[str], RuntimeMutable()] = Field(default="INFO") dremio: Optional[Dremio] = Field(default=None) tools: Optional[Tools] = Field(default_factory=Tools) launchdarkly: Optional[LaunchDarkly] = Field(default_factory=LaunchDarkly) @@ -442,8 +465,6 @@ def settings_customise_sources( def model_post_init(self, __context): _propagate_flag_prefixes(self, "") - if self.launchdarkly and self.launchdarkly.sdk_key: - FeatureFlagManager.initialize(self.launchdarkly.sdk_key) def with_overrides(self, overrides: Dict[str, Any]) -> Self: def set_values(aparts: List[str], value: Any, obj: Any): @@ -502,8 +523,157 @@ def collect_flag_keys(model_cls: type, prefix: str = "") -> list[str]: # Module-level holder so configure() can pass the YAML path to the Settings constructor _yaml_file: Path | None = None +_base_settings: Settings | None = None +_settings_override: ContextVar[Settings | None] = ContextVar( + "settings_override", default=None +) +_config_fingerprint: tuple[str, int, int, int] | None = None + + +def _initialize_launchdarkly(inst: Settings | None): + sdk_key = ( + inst.launchdarkly.sdk_key + if isinstance(inst, Settings) + and inst.launchdarkly is not None + and inst.launchdarkly.sdk_key + else None + ) + FeatureFlagManager.initialize(sdk_key) + + +def _set_base_settings( + inst: Settings, + *, + fingerprint: tuple[str, int, int, int] | None = None, + initialize_ld: bool = True, +) -> Settings: + global _base_settings, _config_fingerprint + _base_settings = inst + _config_fingerprint = fingerprint + if initialize_ld: + _initialize_launchdarkly(inst) + return inst + + +def _reset_state_for_tests(): + global _yaml_file, _base_settings, _config_fingerprint + _yaml_file = None + _base_settings = None + _config_fingerprint = None + FeatureFlagManager.reset() + + +def _build_settings_candidate() -> Settings: + return Settings() + + +def _resolved_config_fingerprint(cfg: Path | None) -> tuple[str, int, int, int] | None: + if cfg is None: + return None + resolved = cfg.resolve() + stat = resolved.stat() + return (str(resolved), stat.st_ino, stat.st_size, stat.st_mtime_ns) + -_settings: ContextVar[Settings] = ContextVar("settings", default=None) +def _unwrap_annotation(annotation: Any) -> Any: + args = [a for a in get_args(annotation) if a is not NoneType] + if len(args) == 1: + return args[0] + return annotation + + +def _copy_runtime_mutable_fields( + current_obj: Any, + candidate_obj: Any, + model_cls: type, + changed_paths: list[str], + prefix: str = "", +): + hints = get_type_hints(model_cls, include_extras=True) + for field_name in model_cls.model_fields: + value = getattr(candidate_obj, field_name, None) + current_value = getattr(current_obj, field_name, None) + field_path = f"{prefix}.{field_name}" if prefix else field_name + annotation = _unwrap_annotation(hints[field_name]) + + if _has_runtime_mutable(model_cls, field_name): + if current_value != value: + setattr(current_obj, field_name, value) + changed_paths.append(field_path) + continue + + if ( + isinstance(annotation, type) + and issubclass(annotation, BaseModel) + and value is not None + ): + if current_value is None: + current_value = value.model_copy(deep=True) + _blank_model_subtree(current_value, annotation) + setattr(current_obj, field_name, current_value) + _copy_runtime_mutable_fields( + current_value, value, annotation, changed_paths, field_path + ) + + +def _blank_model_subtree(current_obj: Any, model_cls: type): + hints = get_type_hints(model_cls, include_extras=True) + for field_name in model_cls.model_fields: + value = getattr(current_obj, field_name, None) + annotation = _unwrap_annotation(hints[field_name]) + + if ( + isinstance(annotation, type) + and issubclass(annotation, BaseModel) + and value is not None + ): + _blank_model_subtree(value, annotation) + continue + + object.__setattr__(current_obj, field_name, None) + + +def reload_mutable_settings_if_changed() -> list[str]: + global _base_settings, _config_fingerprint + _log = log.logger("settings_reload") + if _yaml_file is None: + return [] + + try: + fingerprint = _resolved_config_fingerprint(_yaml_file) + except Exception as exc: + _log.warning(f"Unable to fingerprint config file {_yaml_file}: {exc}") + return [] + + if fingerprint == _config_fingerprint: + return [] + + try: + candidate = _build_settings_candidate() + except Exception as exc: + _log.warning(f"Skipping config reload for {_yaml_file}: {exc}") + return [] + + if not isinstance(_base_settings, Settings): + _set_base_settings(candidate, fingerprint=fingerprint, initialize_ld=False) + return [] + + updated = _base_settings.model_copy(deep=True) + changed_paths: list[str] = [] + _copy_runtime_mutable_fields(updated, candidate, Settings, changed_paths) + + _base_settings = updated + _config_fingerprint = fingerprint + + if changed_paths: + _log.info( + f"Reloaded runtime-mutable settings from {_yaml_file}: {', '.join(changed_paths)}" + ) + else: + _log.info( + f"Config file {_yaml_file} changed but runtime-mutable settings were unchanged" + ) + return changed_paths # the default config is ~/.config/dremioai/config.yaml, use it if it exists @@ -520,16 +690,21 @@ def default_config() -> Path: # configures the settings using the given config file and overwrites the global # settings instance if force is True -def configure(cfg: Union[str, Path] = None, force=False) -> ContextVar[Settings]: - global _settings - if force and isinstance(_settings.get(), Settings): - old = _settings.get() +def configure(cfg: Union[str, Path] = None, force=False) -> Settings: + global _yaml_file, _base_settings, _config_fingerprint + if force and isinstance(_base_settings, Settings): + old_settings = _base_settings + old_yaml_file = _yaml_file + old_fingerprint = _config_fingerprint try: - _settings.set(None) - configure(cfg, force=False) - except: + _base_settings = None + _config_fingerprint = None + return configure(cfg, force=False) + except Exception: # don't replace the old if there is an issue setting the new value - _settings.set(old) + _base_settings = old_settings + _yaml_file = old_yaml_file + _config_fingerprint = old_fingerprint raise if isinstance(cfg, str): @@ -542,25 +717,28 @@ def configure(cfg: Union[str, Path] = None, force=False) -> ContextVar[Settings] cfg.parent.mkdir(parents=True, exist_ok=True) cfg.touch() - global _yaml_file _yaml_file = cfg - _settings.set(Settings()) - - return _settings + candidate = _build_settings_candidate() + return _set_base_settings( + candidate, fingerprint=_resolved_config_fingerprint(cfg), initialize_ld=True + ) # Get the current settings instance if one has been configured. If not try # to configure it using the default config file. If that fails, create a new # empty settings instance. def instance() -> Settings | None: - global _settings - if not isinstance(_settings.get(), Settings): + global _base_settings + override = _settings_override.get() + if isinstance(override, Settings): + return override + if not isinstance(_base_settings, Settings): try: configure() # use default config, if exists except FileNotFoundError: # no default config, create a new default one - _settings.set(Settings()) - return _settings.get() + _set_base_settings(Settings()) + return _base_settings async def run_with( @@ -569,14 +747,14 @@ async def run_with( args: Optional[List[Any]] = [], kw: Optional[Dict[str, Any]] = {}, ) -> Any: - global _settings - async def _call(): - tok = _settings.set(instance().model_copy(deep=True).with_overrides(overrides)) + tok = _settings_override.set( + instance().model_copy(deep=True).with_overrides(overrides) + ) try: return await func(*args, **kw) finally: - _settings.reset(tok) + _settings_override.reset(tok) return await _call() diff --git a/src/dremioai/servers/mcp.py b/src/dremioai/servers/mcp.py index 868f784..422816c 100644 --- a/src/dremioai/servers/mcp.py +++ b/src/dremioai/servers/mcp.py @@ -553,31 +553,32 @@ def create_metrics_server(host: str, port: int, log_level: str) -> uvicorn.Serve return server -_LOG_LEVEL_REFRESH_INTERVAL = 60 # seconds +_SETTINGS_REFRESH_INTERVAL = 60 # seconds -async def _log_level_refresh_loop(): - """Periodically sync log level from LD flags.""" - _log = log.logger("log_level_refresh") +async def _settings_refresh_loop(): + """Periodically reload runtime-mutable settings and sync log level.""" + _log = log.logger("settings_refresh") while True: - await asyncio.sleep(_LOG_LEVEL_REFRESH_INTERVAL) + await asyncio.sleep(_SETTINGS_REFRESH_INTERVAL) try: s = settings.instance() if s is None: continue - level_name = s.get("log_level") + settings.reload_mutable_settings_if_changed() + level_name = settings.instance().get("log_level") level = getattr(logging, level_name.upper(), None) if level is not None and level != log.level(): _log.info(f"Updating log level to {level_name}") log.set_level(level) except Exception as e: - _log.debug(f"Log level refresh failed: {e}") + _log.debug(f"Settings refresh failed: {e}") @contextlib.asynccontextmanager async def _server_lifespan(app: FastMCP): """Lifespan context manager that runs background tasks alongside the server.""" - task = asyncio.create_task(_log_level_refresh_loop()) + task = asyncio.create_task(_settings_refresh_loop()) try: yield finally: @@ -656,7 +657,7 @@ def main( if mock: transport = Transports.streamable_http # In mock mode, create a minimal settings instance — no Dremio config needed - settings._settings.set( + settings._set_base_settings( settings.Settings.model_validate( { "dremio": { @@ -664,14 +665,15 @@ def main( "pat": "mock-pat", } } - ) + ), + initialize_ld=True, ) else: if enable_streaming_http: transport = Transports.streamable_http else: transport = Transports.stdio - cfg = settings.configure(config_file).get() + settings.configure(config_file) dremio = settings.instance().dremio if ( dremio.oauth_supported diff --git a/tests/api/dremio/test_sql.py b/tests/api/dremio/test_sql.py index c0faf21..3ff1a56 100644 --- a/tests/api/dremio/test_sql.py +++ b/tests/api/dremio/test_sql.py @@ -14,41 +14,67 @@ # limitations under the License. # +from unittest.mock import AsyncMock, MagicMock, patch + import pytest -from dremioai.api.dremio.sql import Job - - -@pytest.mark.parametrize( - "js", - [ - pytest.param( - """ - { - "jobState": "COMPLETED", - "rowCount": 1, - "errorMessage": "", - "startedAt": "2025-06-11T15:35:11.636Z", - "endedAt": "2025-06-11T15:35:15.949Z", - "queryType": "REST", - "queueName": "SMALL", - "queueId": "SMALL", - "resourceSchedulingStartedAt": "2025-06-11T15:35:12.435Z", - "resourceSchedulingEndedAt": "2025-06-11T15:35:12.503Z", - "cancellationReason": "" - }""", - id="with rows", - ), - pytest.param( - """{ - "jobState": "METADATA_RETRIEVAL", - "errorMessage": "", - "startedAt": "2025-06-11T15:35:11.565Z", - "queryType": "REST", - "cancellationReason": "" - }""", - id="without rows", - ), - ], + +from dremioai.api.dremio.sql import ( + Job, + JobResults, + JobResultsWrapper, + JobState, + QuerySubmission, + QueryType, + get_results, ) -def test_basic_job(js: str): - j = Job.model_validate_json(js) +from dremioai.config import settings + + +@pytest.mark.asyncio +@patch("dremioai.config.feature_flags.ldclient") +async def test_get_results_uses_polling_interval_get_for_ld_precedence(mock_ldclient): + mock_client = MagicMock() + mock_client.is_initialized.return_value = True + mock_client.variation.side_effect = lambda key, ctx, default: ( + 0.25 if key == "dremio.api.polling_interval" else default + ) + mock_ldclient.get.return_value = mock_client + + settings._set_base_settings( + settings.Settings.model_validate( + { + "launchdarkly": {"sdk_key": "test-key"}, + "dremio": { + "uri": "https://test.dremio.cloud", + "pat": "test-pat", + "api": {"polling_interval": 3.0}, + }, + } + ) + ) + + running = Job(jobState=JobState.RUNNING, queryType=QueryType.REST) + complete = Job( + jobState=JobState.COMPLETED, + rowCount=1, + queryType=QueryType.REST, + ) + results = JobResults(rowCount=1, schema=[], rows=[{"value": 1}]) + client = AsyncMock() + client.get.side_effect = [running, complete] + + with ( + patch( + "dremioai.api.dremio.sql._fetch_results", + new=AsyncMock(return_value=results), + ), + patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep, + ): + response = await get_results( + project_id=None, + qs=QuerySubmission(id="job-1"), + client=client, + ) + + assert isinstance(response, JobResultsWrapper) + assert mock_sleep.await_args_list[0].args[0] == 0.25 diff --git a/tests/api/test_transport_retry.py b/tests/api/test_transport_retry.py index 87d0b2d..cd6cea4 100644 --- a/tests/api/test_transport_retry.py +++ b/tests/api/test_transport_retry.py @@ -26,6 +26,7 @@ from dremioai.api.transport import RetryConfig, retry_middleware from dremioai.config import settings +from dremioai.config.feature_flags import FeatureFlagManager class TestRetryConfig: @@ -37,20 +38,21 @@ def test_retry_config_with_default_settings(self): retry_config = RetryConfig() # Verify default values from HttpRetry model - assert retry_config.config.max_retries == 20 - assert retry_config.config.initial_delay == 1.0 - assert retry_config.config.max_delay == 60.0 - assert retry_config.config.backoff_multiplier == 2.0 + assert retry_config.max_retries == 20 + assert retry_config.get_config_delay(0) == 1.0 + assert ( + retry_config.get_delay(MagicMock(headers={"Retry-After": None}), 10) + == 60.0 + ) def test_retry_config_with_custom_settings(self, mock_settings_instance): """Test RetryConfig initialization with custom settings""" retry_config = RetryConfig() # Verify it uses settings from mock - assert retry_config.config.max_retries == 5 - assert retry_config.config.initial_delay == 2.0 - assert retry_config.config.max_delay == 120.0 - assert retry_config.config.backoff_multiplier == 3.0 + assert retry_config.max_retries == 5 + assert retry_config.get_config_delay(0) == 2.0 + assert retry_config.get_config_delay(1) == 6.0 def test_get_config_delay_exponential_backoff(self, mock_settings_instance): """Test exponential backoff calculation""" @@ -124,6 +126,22 @@ def test_get_delay_with_invalid_retry_after_header(self, mock_settings_instance) delay = retry_config.get_delay(mock_response, attempt_number=1) assert delay == 6.0 # Falls back to config delay + @patch("dremioai.config.feature_flags.ldclient") + def test_retry_config_uses_get_for_ld_precedence( + self, mock_ldclient, mock_settings_instance + ): + mock_client = MagicMock() + mock_client.is_initialized.return_value = True + mock_client.variation.side_effect = lambda key, ctx, default: ( + 9 if key == "dremio.api.http_retry.max_retries" else default + ) + mock_ldclient.get.return_value = mock_client + FeatureFlagManager.initialize("test-key") + + retry_config = RetryConfig() + + assert retry_config.max_retries == 9 + class TestRetryMiddleware: """Test the retry_middleware function""" @@ -294,19 +312,21 @@ async def test_no_retry_on_other_errors(self, mock_settings_instance): @pytest.fixture def mock_settings_instance(): """Create a mock settings instance with custom retry configuration""" - # Create actual HttpRetry instance with custom values - http_retry_config = settings.HttpRetry( - max_retries=5, initial_delay=2.0, max_delay=120.0, backoff_multiplier=3.0 + mock_settings = settings.Settings.model_validate( + { + "dremio": { + "uri": "https://test.dremio.cloud", + "pat": "test-pat", + "api": { + "http_retry": { + "max_retries": 5, + "initial_delay": 2.0, + "max_delay": 120.0, + "backoff_multiplier": 3.0, + } + }, + } + } ) - - mock_dremio = MagicMock() - mock_dremio.api = MagicMock() - mock_dremio.api.http_retry = http_retry_config - - # Create mock settings object - mock_settings = MagicMock() - mock_settings.dremio = mock_dremio - - # Patch settings.instance to return our mock - with patch.object(settings, "instance", return_value=mock_settings): - yield mock_settings + settings._set_base_settings(mock_settings) + yield mock_settings diff --git a/tests/config/test_launchdarkly_integration.py b/tests/config/test_launchdarkly_integration.py index 3f48195..f098857 100644 --- a/tests/config/test_launchdarkly_integration.py +++ b/tests/config/test_launchdarkly_integration.py @@ -27,10 +27,8 @@ @pytest.fixture(autouse=True) def reset_feature_flag_manager(): FeatureFlagManager.reset() - old = settings._settings.get() yield FeatureFlagManager.reset() - settings._settings.set(old) def _make_settings(launchdarkly=None, **dremio_overrides): @@ -40,7 +38,7 @@ def _make_settings(launchdarkly=None, **dremio_overrides): if launchdarkly is not None: cfg["launchdarkly"] = launchdarkly s = settings.Settings.model_validate(cfg) - settings._settings.set(s) + settings._set_base_settings(s) return s @@ -572,15 +570,15 @@ def test_flag_keys_match_golden(): ), "Flag keys changed! If intentional, run: uv run python scripts/generate_flag_keys.py --write" -# -- Periodic log level refresh ----------------------------------------------- +# -- Periodic settings refresh ------------------------------------------------ @pytest.mark.asyncio -async def test_log_level_refresh_updates_level(): +async def test_settings_refresh_updates_level(): """Periodic refresh picks up LD log_level changes and updates logging.""" from dremioai.servers.mcp import ( - _log_level_refresh_loop, - _LOG_LEVEL_REFRESH_INTERVAL, + _settings_refresh_loop, + _SETTINGS_REFRESH_INTERVAL, ) from dremioai import log @@ -598,8 +596,8 @@ async def test_log_level_refresh_updates_level(): FeatureFlagManager.initialize("test-sdk-key") # Run one iteration with a short interval - with patch("dremioai.servers.mcp._LOG_LEVEL_REFRESH_INTERVAL", 0): - task = asyncio.create_task(_log_level_refresh_loop()) + with patch("dremioai.servers.mcp._SETTINGS_REFRESH_INTERVAL", 0): + task = asyncio.create_task(_settings_refresh_loop()) await asyncio.sleep(0.05) task.cancel() try: @@ -614,18 +612,18 @@ async def test_log_level_refresh_updates_level(): @pytest.mark.asyncio -async def test_log_level_refresh_no_change_when_same(): +async def test_settings_refresh_no_change_when_same(): """Refresh loop does not call set_level when level hasn't changed.""" - from dremioai.servers.mcp import _log_level_refresh_loop + from dremioai.servers.mcp import _settings_refresh_loop from dremioai import log _make_settings() with ( - patch("dremioai.servers.mcp._LOG_LEVEL_REFRESH_INTERVAL", 0), + patch("dremioai.servers.mcp._SETTINGS_REFRESH_INTERVAL", 0), patch.object(log, "set_level") as mock_set_level, ): - task = asyncio.create_task(_log_level_refresh_loop()) + task = asyncio.create_task(_settings_refresh_loop()) await asyncio.sleep(0.05) task.cancel() try: @@ -637,6 +635,99 @@ async def test_log_level_refresh_no_change_when_same(): mock_set_level.assert_not_called() +@patch("dremioai.config.feature_flags.ldclient") +def test_reload_mutable_settings_preserves_ld_precedence(mock_ldclient, tmp_path): + mock_client = _make_mock_ld_client({"dremio.allow_dml": True}) + mock_ldclient.get.return_value = mock_client + + cfg = tmp_path / "config.yaml" + cfg.write_text(""" +launchdarkly: + sdk_key: test-key +dremio: + uri: https://test.dremio.cloud + pat: test-pat + allow_dml: false +""") + settings.configure(cfg) + + cfg.write_text(""" +launchdarkly: + sdk_key: test-key +dremio: + uri: https://test.dremio.cloud + pat: test-pat + allow_dml: false + api: + polling_interval: 7 +""") + settings.reload_mutable_settings_if_changed() + + assert settings.instance().dremio.allow_dml is False + assert settings.instance().dremio.get("allow_dml") is True + assert settings.instance().dremio.api.get("polling_interval") == 7 + + +@patch("dremioai.config.feature_flags.ldclient") +def test_reload_mutable_settings_does_not_reinitialize_ld_on_validation_failure( + mock_ldclient, tmp_path +): + mock_client = _make_mock_ld_client({}) + mock_ldclient.get.return_value = mock_client + + cfg = tmp_path / "config.yaml" + cfg.write_text(""" +launchdarkly: + sdk_key: test-key +dremio: + uri: https://test.dremio.cloud + pat: test-pat +""") + settings.configure(cfg) + assert mock_ldclient.set_config.call_count == 1 + + cfg.write_text("dremio: [") + settings.reload_mutable_settings_if_changed() + + assert mock_ldclient.set_config.call_count == 1 + + +@pytest.mark.asyncio +async def test_settings_refresh_reloads_yaml_before_log_level_evaluation(tmp_path): + from dremioai.servers.mcp import _settings_refresh_loop + from dremioai import log + + cfg = tmp_path / "config.yaml" + cfg.write_text(""" +log_level: INFO +dremio: + uri: https://test.dremio.cloud + pat: test-pat +""") + settings.configure(cfg) + + cfg.write_text(""" +log_level: ERROR +dremio: + uri: https://test.dremio.cloud + pat: test-pat +""") + + original_level = log.level() + try: + with patch("dremioai.servers.mcp._SETTINGS_REFRESH_INTERVAL", 0): + task = asyncio.create_task(_settings_refresh_loop()) + await asyncio.sleep(0.05) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + assert log.level() == logging.ERROR + finally: + log.set_level(original_level) + + # -- _build_context ----------------------------------------------------------- diff --git a/tests/config/test_settings.py b/tests/config/test_settings.py index 7e6a320..f79b492 100644 --- a/tests/config/test_settings.py +++ b/tests/config/test_settings.py @@ -56,7 +56,7 @@ def test_create_default_config(mock_config_dir): project_id = uuid.uuid4() mode = ToolType.FOR_DATA_PATTERNS settings.configure(force=True) - settings._settings.set( + settings._set_base_settings( settings.instance().model_validate( { "dremio": { @@ -81,6 +81,180 @@ def test_create_default_config(mock_config_dir): assert tools.server_mode == mode +async def _read_runtime_settings(): + cfg = settings.instance() + return ( + cfg.log_level, + cfg.dremio.enable_search, + cfg.dremio.api.polling_interval, + ) + + +@pytest.mark.asyncio +async def test_run_with_keeps_overrides_request_scoped(): + base = settings.Settings.model_validate( + { + "log_level": "INFO", + "dremio": { + "uri": "https://test.dremio.cloud", + "pat": "test-pat", + "enable_search": False, + "api": {"polling_interval": 1.0}, + }, + } + ) + settings._set_base_settings(base) + + original = await _read_runtime_settings() + overridden = await settings.run_with( + _read_runtime_settings, + overrides={ + "log_level": "DEBUG", + "dremio.enable_search": True, + "dremio.api.polling_interval": 3.5, + }, + ) + + assert overridden == ("DEBUG", True, 3.5) + assert await _read_runtime_settings() == original + + +def test_reload_mutable_settings_if_changed_updates_runtime_mutable_only(tmp_path): + cfg = tmp_path / "config.yaml" + cfg.write_text( + yaml.safe_dump( + { + "log_level": "INFO", + "dremio": { + "uri": "https://one.dremio.cloud", + "pat": "test-pat", + "enable_search": False, + "allow_dml": False, + "api": { + "polling_interval": 1.0, + "http_retry": {"max_retries": 5}, + }, + }, + } + ) + ) + + settings.configure(cfg) + + cfg.write_text( + yaml.safe_dump( + { + "log_level": "DEBUG", + "dremio": { + "uri": "https://two.dremio.cloud", + "pat": "changed-pat", + "enable_search": True, + "allow_dml": True, + "api": { + "polling_interval": 9.0, + "http_retry": {"max_retries": 11}, + }, + }, + } + ) + ) + + changed = settings.reload_mutable_settings_if_changed() + + assert changed == [ + "log_level", + "dremio.enable_search", + "dremio.allow_dml", + "dremio.api.http_retry.max_retries", + "dremio.api.polling_interval", + ] + assert settings.instance().log_level == "DEBUG" + assert settings.instance().dremio.enable_search is True + assert settings.instance().dremio.allow_dml is True + assert settings.instance().dremio.api.http_retry.max_retries == 11 + assert settings.instance().dremio.api.polling_interval == 9.0 + assert settings.instance().dremio.uri == "https://one.dremio.cloud" + assert settings.instance().dremio.pat == "test-pat" + + +def test_reload_mutable_settings_if_changed_preserves_base_on_invalid_yaml(tmp_path): + cfg = tmp_path / "config.yaml" + cfg.write_text( + yaml.safe_dump( + { + "log_level": "INFO", + "dremio": {"uri": "https://test.dremio.cloud", "pat": "test-pat"}, + } + ) + ) + settings.configure(cfg) + + before = settings.instance().model_copy(deep=True) + cfg.write_text("dremio: [") + + changed = settings.reload_mutable_settings_if_changed() + + assert changed == [] + assert settings.instance().model_dump() == before.model_dump() + + +def test_reload_mutable_settings_if_changed_ignores_non_mutable_changes(tmp_path): + cfg = tmp_path / "config.yaml" + cfg.write_text( + yaml.safe_dump( + { + "log_level": "INFO", + "dremio": {"uri": "https://one.dremio.cloud", "pat": "test-pat"}, + } + ) + ) + settings.configure(cfg) + + cfg.write_text( + yaml.safe_dump( + { + "log_level": "INFO", + "dremio": {"uri": "https://two.dremio.cloud", "pat": "changed-pat"}, + } + ) + ) + + changed = settings.reload_mutable_settings_if_changed() + + assert changed == [] + assert settings.instance().dremio.uri == "https://one.dremio.cloud" + assert settings.instance().dremio.pat == "test-pat" + + +def test_reload_mutable_settings_if_changed_materializes_missing_nested_subtree( + tmp_path, +): + cfg = tmp_path / "config.yaml" + cfg.write_text(yaml.safe_dump({"log_level": "INFO"})) + settings.configure(cfg) + + cfg.write_text( + yaml.safe_dump( + { + "log_level": "INFO", + "dremio": { + "uri": "https://later.dremio.cloud", + "pat": "later-pat", + "allow_dml": True, + "api": {"polling_interval": 7.0}, + }, + } + ) + ) + + changed = settings.reload_mutable_settings_if_changed() + + assert "dremio.allow_dml" in changed + assert "dremio.api.polling_interval" in changed + assert settings.instance().dremio.get("allow_dml") is True + assert settings.instance().dremio.api.get("polling_interval") == 7.0 + + @pytest.mark.parametrize( "name,value", [ @@ -177,7 +351,15 @@ def test_auth_urls( ) if iss_override: issuer = iss_override - auth = (f"{issuer}/oauth/authorize", f"{issuer}/oauth/token", f"{issuer}/oauth/register") if not error else None + auth = ( + ( + f"{issuer}/oauth/authorize", + f"{issuer}/oauth/token", + f"{issuer}/oauth/register", + ) + if not error + else None + ) issuer = issuer if not error else None assert d.auth_issuer_uri == issuer assert d.auth_endpoints == auth @@ -187,12 +369,14 @@ def test_auth_urls( def test_launchdarkly_sdk_key_from_env(monkeypatch, sdk_key): monkeypatch.setenv("DREMIOAI_LAUNCHDARKLY__SDK_KEY", sdk_key) - s = settings.Settings.model_validate({ - "dremio": { - "uri": "https://test.dremio.cloud", - "pat": "test-pat", + s = settings.Settings.model_validate( + { + "dremio": { + "uri": "https://test.dremio.cloud", + "pat": "test-pat", + } } - }) + ) assert s.launchdarkly.sdk_key == sdk_key assert s.launchdarkly.enabled is True @@ -203,11 +387,9 @@ def test_launchdarkly_sdk_key_from_file(tmp_path): sdk_key_file = tmp_path / "sdk_key.txt" sdk_key_file.write_text("sdk-file-key-abcdef") - s = settings.Settings.model_validate({ - "launchdarkly": { - "sdk_key": f"@{sdk_key_file}" - } - }) + s = settings.Settings.model_validate( + {"launchdarkly": {"sdk_key": f"@{sdk_key_file}"}} + ) assert s.launchdarkly.sdk_key == "sdk-file-key-abcdef" assert s.launchdarkly.enabled is True @@ -224,39 +406,45 @@ def test_launchdarkly_defaults(): def test_dremio_get_without_launchdarkly(): """Test that get() returns config value when LaunchDarkly is not configured.""" - s = settings.Settings.model_validate({ - "dremio": { - "uri": "https://test.dremio.cloud", - "pat": "test-pat", - "allow_dml": True, + s = settings.Settings.model_validate( + { + "dremio": { + "uri": "https://test.dremio.cloud", + "pat": "test-pat", + "allow_dml": True, + } } - }) + ) assert s.dremio.get("allow_dml") is True def test_dremio_get_with_launchdarkly_disabled(): """Test that get() returns config value when LaunchDarkly is disabled.""" - s = settings.Settings.model_validate({ - "dremio": { - "uri": "https://test.dremio.cloud", - "pat": "test-pat", - "enable_search": True, + s = settings.Settings.model_validate( + { + "dremio": { + "uri": "https://test.dremio.cloud", + "pat": "test-pat", + "enable_search": True, + } } - }) + ) assert s.dremio.get("enable_search") is True def test_dremio_enable_search_fallback(): """Test that enable_search returns config value when LD is disabled.""" - s = settings.Settings.model_validate({ - "dremio": { - "uri": "https://test.dremio.cloud", - "pat": "test-pat", - "enable_search": True + s = settings.Settings.model_validate( + { + "dremio": { + "uri": "https://test.dremio.cloud", + "pat": "test-pat", + "enable_search": True, + } } - }) + ) assert s.dremio.enable_search is True assert s.dremio.get("enable_search") is True @@ -264,13 +452,15 @@ def test_dremio_enable_search_fallback(): def test_dremio_allow_dml_fallback(): """Test that allow_dml returns config value when LD is disabled.""" - s = settings.Settings.model_validate({ - "dremio": { - "uri": "https://test.dremio.cloud", - "pat": "test-pat", - "allow_dml": True + s = settings.Settings.model_validate( + { + "dremio": { + "uri": "https://test.dremio.cloud", + "pat": "test-pat", + "allow_dml": True, + } } - }) + ) assert s.dremio.allow_dml is True assert s.dremio.get("allow_dml") is True diff --git a/tests/conftest.py b/tests/conftest.py index 2f4626a..4c28902 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -17,6 +17,7 @@ """ Global pytest fixtures for dremio-mcp tests. """ + import os import random import uuid @@ -50,6 +51,13 @@ from prometheus_client import CollectorRegistry +@pytest.fixture(autouse=True) +def reset_settings_state(): + settings.reset_state_for_tests() + yield + settings.reset_state_for_tests() + + @pytest.fixture(autouse=True) def reset_uvicorn_logger_propagation(): """Reset uvicorn logger state between tests. @@ -135,23 +143,26 @@ def mock_config_dir(temp_config_dir): @pytest.fixture def mock_settings_instance(): """Create a mock settings instance with default values""" - old_settings = settings.instance() - try: - settings._settings.set( - settings.Settings.model_validate( - { - "dremio": { - "uri": "https://test-dremio-uri.com", - "pat": "test-pat", - "project_id": uuid.uuid4(), - }, - "tools": {"server_mode": ToolType.FOR_SELF.name}, - } - ) - ) - yield settings.instance() - finally: - settings._settings.set(old_settings) + config = settings.Settings.model_validate( + { + "dremio": { + "uri": "https://test-dremio-uri.com", + "pat": "test-pat", + "project_id": uuid.uuid4(), + "api": { + "http_retry": { + "max_retries": 5, + "initial_delay": 2.0, + "max_delay": 120.0, + "backoff_multiplier": 3.0, + } + }, + }, + "tools": {"server_mode": ToolType.FOR_SELF.name}, + } + ) + settings._set_base_settings(config) + yield settings.instance() @pytest.fixture @@ -258,7 +269,7 @@ async def http_streamable_mcp_server( config["dremio"]["wlm"] = {"engine_name": wlm_engine} if dremio_overrides: config["dremio"].update(dremio_overrides) - settings._settings.set(settings.Settings.model_validate(config)) + settings._set_base_settings(settings.Settings.model_validate(config)) settings.write_settings() set_level(logging_level.upper()) @@ -294,7 +305,7 @@ async def http_streamable_mcp_server( if sf is not None: sf.close() print(f"{sf} closed") - settings._settings.set(old) + settings._set_base_settings(old) @contextlib.asynccontextmanager diff --git a/tests/servers/test_mcp.py b/tests/servers/test_mcp.py index 59d886f..f148fa9 100644 --- a/tests/servers/test_mcp.py +++ b/tests/servers/test_mcp.py @@ -39,7 +39,7 @@ def mock_settings(mode: ToolType): old = settings.instance() with TemporaryDirectory() as temp_dir: temp_dir = Path(temp_dir) - settings._settings.set( + settings._set_base_settings( settings.Settings.model_validate( { "dremio": { @@ -54,7 +54,7 @@ def mock_settings(mode: ToolType): settings.write_settings(cfg=cfg, inst=settings.instance()) yield settings.instance(), cfg finally: - settings._settings.set(old) + settings._set_base_settings(old) @asynccontextmanager diff --git a/tests/stremable_http_cli.py b/tests/stremable_http_cli.py index 7b00691..0fb5d0d 100644 --- a/tests/stremable_http_cli.py +++ b/tests/stremable_http_cli.py @@ -399,7 +399,7 @@ def _local_mcp_server(dremio_uri: str, port: int = 8989, ld_sdk_key: str | None if ld_sdk_key: overrides["launchdarkly.sdk_key"] = ld_sdk_key configured_settings = old.model_copy(deep=True).with_overrides(overrides) - settings._settings.set(configured_settings) + settings._set_base_settings(configured_settings) mcp_server = init( transport=Transports.streamable_http, port=port, @@ -409,11 +409,6 @@ def _local_mcp_server(dremio_uri: str, port: int = 8989, ld_sdk_key: str | None ) def _run(): - # Propagate settings to server thread — ContextVar doesn't - # automatically transfer across threads, so without this the - # server would fall back to the default config file. - settings._settings.set(configured_settings) - a = mcp_server.streamable_http_app() c = uvicorn.Config(app=a, host="127.0.0.1", port=port, log_level="warning") s = uvicorn.Server(c) @@ -434,7 +429,7 @@ def _run(): yield port finally: - settings._settings.set(old) + settings._set_base_settings(old) @app.command("test", help="Run a quick smoketest for a deployed MCP server") diff --git a/tests/test_fastmcp_basic.py b/tests/test_fastmcp_basic.py index 989786a..db886b9 100644 --- a/tests/test_fastmcp_basic.py +++ b/tests/test_fastmcp_basic.py @@ -42,7 +42,7 @@ def mock_settings_for_test(mode: ToolType): """Create mock settings for testing FastMCP server""" try: old = settings.instance() - settings._settings.set( + settings._set_base_settings( settings.Settings.model_validate( { "dremio": { @@ -57,7 +57,7 @@ def mock_settings_for_test(mode: ToolType): ) yield settings.instance() finally: - settings._settings.set(old) + settings._set_base_settings(old) @pytest.mark.asyncio diff --git a/tests/test_simple_fastmcp_server.py b/tests/test_simple_fastmcp_server.py index 0627df2..fc6404b 100644 --- a/tests/test_simple_fastmcp_server.py +++ b/tests/test_simple_fastmcp_server.py @@ -36,7 +36,7 @@ def mock_settings_for_fastmcp(self, mode: ToolType): """Create mock settings for testing FastMCP server""" try: old = settings.instance() - settings._settings.set( + settings._set_base_settings( settings.Settings.model_validate( { "dremio": { @@ -51,7 +51,7 @@ def mock_settings_for_fastmcp(self, mode: ToolType): ) yield settings.instance() finally: - settings._settings.set(old) + settings._set_base_settings(old) @pytest.mark.asyncio async def test_fastmcp_server_creation_and_tool_registration(self): @@ -157,7 +157,7 @@ def mock_settings_for_dynamic_tools(self, enable_remote_tools: bool = True): """Create mock settings with remote tools enabled/disabled""" try: old = settings.instance() - settings._settings.set( + settings._set_base_settings( settings.Settings.model_validate( { "dremio": { @@ -175,7 +175,7 @@ def mock_settings_for_dynamic_tools(self, enable_remote_tools: bool = True): ) yield settings.instance() finally: - settings._settings.set(old) + settings._set_base_settings(old) @pytest.mark.asyncio async def test_meta_tools_registered_when_enabled(self): @@ -206,10 +206,24 @@ async def test_meta_tools_disabled_at_invocation(self): async def test_discover_returns_tool_list(self): """DiscoverDynamicTools should return JSON with tools from Dremio""" from dremioai.api.dremio.ai_tools import AiTool - fake_response = ListToolsResponse(tools=[ - AiTool(name="JavaTool1", description="A tool from Java", inputSchema={"type": "object", "properties": {"x": {"type": "string"}}}), - AiTool(name="JavaTool2", description="Another Java tool", inputSchema={"type": "object"}), - ]) + + fake_response = ListToolsResponse( + tools=[ + AiTool( + name="JavaTool1", + description="A tool from Java", + inputSchema={ + "type": "object", + "properties": {"x": {"type": "string"}}, + }, + ), + AiTool( + name="JavaTool2", + description="Another Java tool", + inputSchema={"type": "object"}, + ), + ] + ) with self.mock_settings_for_dynamic_tools(enable_remote_tools=True): server = mcp_server.init(mode=ToolType.FOR_DATA_PATTERNS) diff --git a/tests/tools/test_output_validation.py b/tests/tools/test_output_validation.py index 609d316..4e06f19 100644 --- a/tests/tools/test_output_validation.py +++ b/tests/tools/test_output_validation.py @@ -23,7 +23,11 @@ from unittest.mock import AsyncMock, patch from mcp.server.fastmcp.utilities.func_metadata import func_metadata from dremioai.config import settings -from dremioai.tools.tools import GetUsefulSystemTableNames, GetSchemaOfTable, RunSqlQuery +from dremioai.tools.tools import ( + GetUsefulSystemTableNames, + GetSchemaOfTable, + RunSqlQuery, +) async def mock_mcp_validate_tool_output(tool, *args, **kwargs): @@ -82,15 +86,17 @@ async def test_run_sql_query_json_safe_output(): ] ) - with patch("dremioai.tools.tools.sql.run_query", new_callable=AsyncMock) as mock_run_query: + with patch( + "dremioai.tools.tools.sql.run_query", new_callable=AsyncMock + ) as mock_run_query: mock_run_query.return_value = df - token = settings._settings.set( + token = settings._settings_override.set( settings.Settings.model_validate({"dremio": {"uri": "https://test"}}) ) try: result = await tool.invoke("SELECT 1") finally: - settings._settings.reset(token) + settings._settings_override.reset(token) assert isinstance(result, dict) assert "result" in result diff --git a/tests/tools/test_tools.py b/tests/tools/test_tools.py index 6d0597d..d6e1efd 100644 --- a/tests/tools/test_tools.py +++ b/tests/tools/test_tools.py @@ -137,7 +137,7 @@ "comment": "Recursive CTE for hierarchical data traversal", }, { - "sql": "CREATE TABLE \"temp\".tbl AS SELECT 1;", + "sql": 'CREATE TABLE "temp".tbl AS SELECT 1;', "allowed": False, "comment": "Simple CREATE TABLE AS SELECT query", }, @@ -415,7 +415,7 @@ def mock_settings(dml_allowed: bool): old_settings = settings.instance() try: - settings._settings.set( + settings._set_base_settings( settings.Settings.model_validate( { "dremio": { @@ -427,7 +427,7 @@ def mock_settings(dml_allowed: bool): ) yield finally: - settings._settings.set(old_settings) + settings._set_base_settings(old_settings) @pytest.mark.parametrize( From 0ff05f3c59fa6a758d96d8d0d772e2c305ecae35 Mon Sep 17 00:00:00 2001 From: Aniket Kulkarni Date: Thu, 14 May 2026 11:58:46 -0400 Subject: [PATCH 02/12] DX-119909: address settings reload review comments --- src/dremioai/config/settings.py | 358 ++++++++++-------- src/dremioai/servers/mcp.py | 2 +- tests/api/dremio/test_sql.py | 2 +- tests/api/test_transport_retry.py | 2 +- tests/config/test_launchdarkly_integration.py | 2 +- tests/config/test_settings.py | 4 +- tests/conftest.py | 6 +- tests/servers/test_mcp.py | 4 +- tests/stremable_http_cli.py | 4 +- tests/test_fastmcp_basic.py | 4 +- tests/test_simple_fastmcp_server.py | 8 +- tests/tools/test_output_validation.py | 4 +- tests/tools/test_tools.py | 4 +- 13 files changed, 233 insertions(+), 171 deletions(-) diff --git a/src/dremioai/config/settings.py b/src/dremioai/config/settings.py index af3f459..f0c8b9e 100644 --- a/src/dremioai/config/settings.py +++ b/src/dremioai/config/settings.py @@ -14,6 +14,7 @@ # limitations under the License. # import uuid +import hashlib from uuid import UUID from urllib.parse import urlparse @@ -523,157 +524,217 @@ def collect_flag_keys(model_cls: type, prefix: str = "") -> list[str]: # Module-level holder so configure() can pass the YAML path to the Settings constructor _yaml_file: Path | None = None -_base_settings: Settings | None = None -_settings_override: ContextVar[Settings | None] = ContextVar( - "settings_override", default=None -) -_config_fingerprint: tuple[str, int, int, int] | None = None +ConfigFingerprint = tuple[str, int, int, int, str] -def _initialize_launchdarkly(inst: Settings | None): - sdk_key = ( - inst.launchdarkly.sdk_key - if isinstance(inst, Settings) - and inst.launchdarkly is not None - and inst.launchdarkly.sdk_key - else None - ) - FeatureFlagManager.initialize(sdk_key) +class SettingsReloader: + """Singleton runtime helper for base settings snapshots and config reloads. + The active settings snapshot must be process-wide so background refreshes + are visible across request contexts. Request-scoped overrides remain in a + ContextVar so `run_with()` keeps its existing isolation semantics. + """ -def _set_base_settings( - inst: Settings, - *, - fingerprint: tuple[str, int, int, int] | None = None, - initialize_ld: bool = True, -) -> Settings: - global _base_settings, _config_fingerprint - _base_settings = inst - _config_fingerprint = fingerprint - if initialize_ld: - _initialize_launchdarkly(inst) - return inst + def __init__(self): + self.base_settings: Settings | None = None + self.settings_override: ContextVar[Settings | None] = ContextVar( + "settings_override", default=None + ) + self.config_fingerprint: ConfigFingerprint | None = None + + def initialize_launchdarkly(self, inst: Settings | None): + sdk_key = ( + inst.launchdarkly.sdk_key + if isinstance(inst, Settings) + and inst.launchdarkly is not None + and inst.launchdarkly.sdk_key + else None + ) + FeatureFlagManager.initialize(sdk_key) + + def activate_base_settings( + self, + inst: Settings, + *, + fingerprint: ConfigFingerprint | None = None, + initialize_ld: bool = True, + ) -> Settings: + """Activate a base settings snapshot and optionally refresh LD wiring. + + Candidate validation during reload intentionally skips LD reinitialization + so the refresh loop can compare a scratch `Settings()` object without + reconfiguring LaunchDarkly on every tick. + """ + self.base_settings = inst + self.config_fingerprint = fingerprint + if initialize_ld: + self.initialize_launchdarkly(inst) + return inst + + def reset_for_tests(self): + global _yaml_file + _yaml_file = None + self.base_settings = None + self.config_fingerprint = None + FeatureFlagManager.reset() + + def build_candidate(self) -> Settings: + return Settings() + + def compute_fingerprint(self, cfg: Path | None) -> ConfigFingerprint | None: + """Return a resolved-path fingerprint plus a content hash. + + Resolved path + inode + size + mtime_ns handles ConfigMap swaps well. + A SHA-256 of the current file contents covers the rare edge case where + metadata stays the same across a rewrite. + """ + if cfg is None: + return None + resolved = cfg.resolve() + stat = resolved.stat() + content_hash = hashlib.sha256(resolved.read_bytes()).hexdigest() + return ( + str(resolved), + stat.st_ino, + stat.st_size, + stat.st_mtime_ns, + content_hash, + ) + def unwrap_annotation(self, annotation: Any) -> Any: + args = [a for a in get_args(annotation) if a is not NoneType] + if len(args) == 1: + return args[0] + return annotation + + def copy_runtime_mutable_fields( + self, + current_obj: Any, + candidate_obj: Any, + model_cls: type, + changed_paths: list[str], + prefix: str = "", + ): + """Copy only RuntimeMutable leaves from candidate_obj into current_obj.""" + hints = get_type_hints(model_cls, include_extras=True) + for field_name in model_cls.model_fields: + value = getattr(candidate_obj, field_name, None) + current_value = getattr(current_obj, field_name, None) + field_path = f"{prefix}.{field_name}" if prefix else field_name + annotation = self.unwrap_annotation(hints[field_name]) + + if _has_runtime_mutable(model_cls, field_name): + if current_value != value: + setattr(current_obj, field_name, value) + changed_paths.append(field_path) + continue + + if ( + isinstance(annotation, type) + and issubclass(annotation, BaseModel) + and value is not None + ): + if current_value is None: + # Simple instantiation is not sufficient here because some + # nested models have required, startup-only fields. We copy + # the candidate subtree shape, blank it, then replay only + # runtime-mutable leaves into that shell. + current_value = value.model_copy(deep=True) + self.blank_model_subtree(current_value, annotation) + setattr(current_obj, field_name, current_value) + self.copy_runtime_mutable_fields( + current_value, value, annotation, changed_paths, field_path + ) + + def blank_model_subtree(self, current_obj: Any, model_cls: type): + """Clear a copied subtree before replaying mutable leaves into it.""" + hints = get_type_hints(model_cls, include_extras=True) + for field_name in model_cls.model_fields: + value = getattr(current_obj, field_name, None) + annotation = self.unwrap_annotation(hints[field_name]) + + if ( + isinstance(annotation, type) + and issubclass(annotation, BaseModel) + and value is not None + ): + self.blank_model_subtree(value, annotation) + continue + + object.__setattr__(current_obj, field_name, None) + + def reload_if_changed(self) -> list[str]: + _log = log.logger("settings_reload") + if _yaml_file is None: + return [] + + try: + fingerprint = self.compute_fingerprint(_yaml_file) + except Exception as exc: + _log.warning(f"Unable to fingerprint config file {_yaml_file}: {exc}") + return [] -def _reset_state_for_tests(): - global _yaml_file, _base_settings, _config_fingerprint - _yaml_file = None - _base_settings = None - _config_fingerprint = None - FeatureFlagManager.reset() + if fingerprint == self.config_fingerprint: + return [] + try: + candidate = self.build_candidate() + except Exception as exc: + _log.warning(f"Skipping config reload for {_yaml_file}: {exc}") + return [] + + if not isinstance(self.base_settings, Settings): + self.activate_base_settings( + candidate, fingerprint=fingerprint, initialize_ld=False + ) + return [] -def _build_settings_candidate() -> Settings: - return Settings() + updated = self.base_settings.model_copy(deep=True) + changed_paths: list[str] = [] + self.copy_runtime_mutable_fields(updated, candidate, Settings, changed_paths) + self.base_settings = updated + self.config_fingerprint = fingerprint -def _resolved_config_fingerprint(cfg: Path | None) -> tuple[str, int, int, int] | None: - if cfg is None: - return None - resolved = cfg.resolve() - stat = resolved.stat() - return (str(resolved), stat.st_ino, stat.st_size, stat.st_mtime_ns) + if changed_paths: + _log.info( + f"Reloaded runtime-mutable settings from {_yaml_file}: {', '.join(changed_paths)}" + ) + else: + _log.info( + f"Config file {_yaml_file} changed but runtime-mutable settings were unchanged" + ) + return changed_paths -def _unwrap_annotation(annotation: Any) -> Any: - args = [a for a in get_args(annotation) if a is not NoneType] - if len(args) == 1: - return args[0] - return annotation +settings_reloader = SettingsReloader() -def _copy_runtime_mutable_fields( - current_obj: Any, - candidate_obj: Any, - model_cls: type, - changed_paths: list[str], - prefix: str = "", -): - hints = get_type_hints(model_cls, include_extras=True) - for field_name in model_cls.model_fields: - value = getattr(candidate_obj, field_name, None) - current_value = getattr(current_obj, field_name, None) - field_path = f"{prefix}.{field_name}" if prefix else field_name - annotation = _unwrap_annotation(hints[field_name]) - - if _has_runtime_mutable(model_cls, field_name): - if current_value != value: - setattr(current_obj, field_name, value) - changed_paths.append(field_path) - continue +def set_base_settings( + inst: Settings, + *, + fingerprint: ConfigFingerprint | None = None, + initialize_ld: bool = True, +) -> Settings: + return settings_reloader.activate_base_settings( + inst, fingerprint=fingerprint, initialize_ld=initialize_ld + ) - if ( - isinstance(annotation, type) - and issubclass(annotation, BaseModel) - and value is not None - ): - if current_value is None: - current_value = value.model_copy(deep=True) - _blank_model_subtree(current_value, annotation) - setattr(current_obj, field_name, current_value) - _copy_runtime_mutable_fields( - current_value, value, annotation, changed_paths, field_path - ) +def reset_state_for_tests(): + settings_reloader.reset_for_tests() -def _blank_model_subtree(current_obj: Any, model_cls: type): - hints = get_type_hints(model_cls, include_extras=True) - for field_name in model_cls.model_fields: - value = getattr(current_obj, field_name, None) - annotation = _unwrap_annotation(hints[field_name]) - - if ( - isinstance(annotation, type) - and issubclass(annotation, BaseModel) - and value is not None - ): - _blank_model_subtree(value, annotation) - continue - object.__setattr__(current_obj, field_name, None) +def push_settings_override(inst: Settings): + return settings_reloader.settings_override.set(inst) + + +def pop_settings_override(token): + settings_reloader.settings_override.reset(token) def reload_mutable_settings_if_changed() -> list[str]: - global _base_settings, _config_fingerprint - _log = log.logger("settings_reload") - if _yaml_file is None: - return [] - - try: - fingerprint = _resolved_config_fingerprint(_yaml_file) - except Exception as exc: - _log.warning(f"Unable to fingerprint config file {_yaml_file}: {exc}") - return [] - - if fingerprint == _config_fingerprint: - return [] - - try: - candidate = _build_settings_candidate() - except Exception as exc: - _log.warning(f"Skipping config reload for {_yaml_file}: {exc}") - return [] - - if not isinstance(_base_settings, Settings): - _set_base_settings(candidate, fingerprint=fingerprint, initialize_ld=False) - return [] - - updated = _base_settings.model_copy(deep=True) - changed_paths: list[str] = [] - _copy_runtime_mutable_fields(updated, candidate, Settings, changed_paths) - - _base_settings = updated - _config_fingerprint = fingerprint - - if changed_paths: - _log.info( - f"Reloaded runtime-mutable settings from {_yaml_file}: {', '.join(changed_paths)}" - ) - else: - _log.info( - f"Config file {_yaml_file} changed but runtime-mutable settings were unchanged" - ) - return changed_paths + return settings_reloader.reload_if_changed() # the default config is ~/.config/dremioai/config.yaml, use it if it exists @@ -691,20 +752,20 @@ def default_config() -> Path: # configures the settings using the given config file and overwrites the global # settings instance if force is True def configure(cfg: Union[str, Path] = None, force=False) -> Settings: - global _yaml_file, _base_settings, _config_fingerprint - if force and isinstance(_base_settings, Settings): - old_settings = _base_settings + global _yaml_file + if force and isinstance(settings_reloader.base_settings, Settings): + old_settings = settings_reloader.base_settings old_yaml_file = _yaml_file - old_fingerprint = _config_fingerprint + old_fingerprint = settings_reloader.config_fingerprint try: - _base_settings = None - _config_fingerprint = None + settings_reloader.base_settings = None + settings_reloader.config_fingerprint = None return configure(cfg, force=False) except Exception: # don't replace the old if there is an issue setting the new value - _base_settings = old_settings + settings_reloader.base_settings = old_settings _yaml_file = old_yaml_file - _config_fingerprint = old_fingerprint + settings_reloader.config_fingerprint = old_fingerprint raise if isinstance(cfg, str): @@ -718,9 +779,11 @@ def configure(cfg: Union[str, Path] = None, force=False) -> Settings: cfg.touch() _yaml_file = cfg - candidate = _build_settings_candidate() - return _set_base_settings( - candidate, fingerprint=_resolved_config_fingerprint(cfg), initialize_ld=True + candidate = settings_reloader.build_candidate() + return set_base_settings( + candidate, + fingerprint=settings_reloader.compute_fingerprint(cfg), + initialize_ld=True, ) @@ -728,17 +791,16 @@ def configure(cfg: Union[str, Path] = None, force=False) -> Settings: # to configure it using the default config file. If that fails, create a new # empty settings instance. def instance() -> Settings | None: - global _base_settings - override = _settings_override.get() + override = settings_reloader.settings_override.get() if isinstance(override, Settings): return override - if not isinstance(_base_settings, Settings): + if not isinstance(settings_reloader.base_settings, Settings): try: configure() # use default config, if exists except FileNotFoundError: # no default config, create a new default one - _set_base_settings(Settings()) - return _base_settings + set_base_settings(Settings()) + return settings_reloader.base_settings async def run_with( @@ -748,13 +810,13 @@ async def run_with( kw: Optional[Dict[str, Any]] = {}, ) -> Any: async def _call(): - tok = _settings_override.set( + tok = push_settings_override( instance().model_copy(deep=True).with_overrides(overrides) ) try: return await func(*args, **kw) finally: - _settings_override.reset(tok) + pop_settings_override(tok) return await _call() diff --git a/src/dremioai/servers/mcp.py b/src/dremioai/servers/mcp.py index 422816c..7a28cab 100644 --- a/src/dremioai/servers/mcp.py +++ b/src/dremioai/servers/mcp.py @@ -657,7 +657,7 @@ def main( if mock: transport = Transports.streamable_http # In mock mode, create a minimal settings instance — no Dremio config needed - settings._set_base_settings( + settings.set_base_settings( settings.Settings.model_validate( { "dremio": { diff --git a/tests/api/dremio/test_sql.py b/tests/api/dremio/test_sql.py index 3ff1a56..9793eef 100644 --- a/tests/api/dremio/test_sql.py +++ b/tests/api/dremio/test_sql.py @@ -40,7 +40,7 @@ async def test_get_results_uses_polling_interval_get_for_ld_precedence(mock_ldcl ) mock_ldclient.get.return_value = mock_client - settings._set_base_settings( + settings.set_base_settings( settings.Settings.model_validate( { "launchdarkly": {"sdk_key": "test-key"}, diff --git a/tests/api/test_transport_retry.py b/tests/api/test_transport_retry.py index cd6cea4..d6f2d75 100644 --- a/tests/api/test_transport_retry.py +++ b/tests/api/test_transport_retry.py @@ -328,5 +328,5 @@ def mock_settings_instance(): } } ) - settings._set_base_settings(mock_settings) + settings.set_base_settings(mock_settings) yield mock_settings diff --git a/tests/config/test_launchdarkly_integration.py b/tests/config/test_launchdarkly_integration.py index f098857..3e7b0bd 100644 --- a/tests/config/test_launchdarkly_integration.py +++ b/tests/config/test_launchdarkly_integration.py @@ -38,7 +38,7 @@ def _make_settings(launchdarkly=None, **dremio_overrides): if launchdarkly is not None: cfg["launchdarkly"] = launchdarkly s = settings.Settings.model_validate(cfg) - settings._set_base_settings(s) + settings.set_base_settings(s) return s diff --git a/tests/config/test_settings.py b/tests/config/test_settings.py index f79b492..a5171af 100644 --- a/tests/config/test_settings.py +++ b/tests/config/test_settings.py @@ -56,7 +56,7 @@ def test_create_default_config(mock_config_dir): project_id = uuid.uuid4() mode = ToolType.FOR_DATA_PATTERNS settings.configure(force=True) - settings._set_base_settings( + settings.set_base_settings( settings.instance().model_validate( { "dremio": { @@ -103,7 +103,7 @@ async def test_run_with_keeps_overrides_request_scoped(): }, } ) - settings._set_base_settings(base) + settings.set_base_settings(base) original = await _read_runtime_settings() overridden = await settings.run_with( diff --git a/tests/conftest.py b/tests/conftest.py index 4c28902..8111610 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -161,7 +161,7 @@ def mock_settings_instance(): "tools": {"server_mode": ToolType.FOR_SELF.name}, } ) - settings._set_base_settings(config) + settings.set_base_settings(config) yield settings.instance() @@ -269,7 +269,7 @@ async def http_streamable_mcp_server( config["dremio"]["wlm"] = {"engine_name": wlm_engine} if dremio_overrides: config["dremio"].update(dremio_overrides) - settings._set_base_settings(settings.Settings.model_validate(config)) + settings.set_base_settings(settings.Settings.model_validate(config)) settings.write_settings() set_level(logging_level.upper()) @@ -305,7 +305,7 @@ async def http_streamable_mcp_server( if sf is not None: sf.close() print(f"{sf} closed") - settings._set_base_settings(old) + settings.set_base_settings(old) @contextlib.asynccontextmanager diff --git a/tests/servers/test_mcp.py b/tests/servers/test_mcp.py index f148fa9..e94746c 100644 --- a/tests/servers/test_mcp.py +++ b/tests/servers/test_mcp.py @@ -39,7 +39,7 @@ def mock_settings(mode: ToolType): old = settings.instance() with TemporaryDirectory() as temp_dir: temp_dir = Path(temp_dir) - settings._set_base_settings( + settings.set_base_settings( settings.Settings.model_validate( { "dremio": { @@ -54,7 +54,7 @@ def mock_settings(mode: ToolType): settings.write_settings(cfg=cfg, inst=settings.instance()) yield settings.instance(), cfg finally: - settings._set_base_settings(old) + settings.set_base_settings(old) @asynccontextmanager diff --git a/tests/stremable_http_cli.py b/tests/stremable_http_cli.py index 0fb5d0d..0fb7ce0 100644 --- a/tests/stremable_http_cli.py +++ b/tests/stremable_http_cli.py @@ -399,7 +399,7 @@ def _local_mcp_server(dremio_uri: str, port: int = 8989, ld_sdk_key: str | None if ld_sdk_key: overrides["launchdarkly.sdk_key"] = ld_sdk_key configured_settings = old.model_copy(deep=True).with_overrides(overrides) - settings._set_base_settings(configured_settings) + settings.set_base_settings(configured_settings) mcp_server = init( transport=Transports.streamable_http, port=port, @@ -429,7 +429,7 @@ def _run(): yield port finally: - settings._set_base_settings(old) + settings.set_base_settings(old) @app.command("test", help="Run a quick smoketest for a deployed MCP server") diff --git a/tests/test_fastmcp_basic.py b/tests/test_fastmcp_basic.py index db886b9..de9ee72 100644 --- a/tests/test_fastmcp_basic.py +++ b/tests/test_fastmcp_basic.py @@ -42,7 +42,7 @@ def mock_settings_for_test(mode: ToolType): """Create mock settings for testing FastMCP server""" try: old = settings.instance() - settings._set_base_settings( + settings.set_base_settings( settings.Settings.model_validate( { "dremio": { @@ -57,7 +57,7 @@ def mock_settings_for_test(mode: ToolType): ) yield settings.instance() finally: - settings._set_base_settings(old) + settings.set_base_settings(old) @pytest.mark.asyncio diff --git a/tests/test_simple_fastmcp_server.py b/tests/test_simple_fastmcp_server.py index fc6404b..3902b9a 100644 --- a/tests/test_simple_fastmcp_server.py +++ b/tests/test_simple_fastmcp_server.py @@ -36,7 +36,7 @@ def mock_settings_for_fastmcp(self, mode: ToolType): """Create mock settings for testing FastMCP server""" try: old = settings.instance() - settings._set_base_settings( + settings.set_base_settings( settings.Settings.model_validate( { "dremio": { @@ -51,7 +51,7 @@ def mock_settings_for_fastmcp(self, mode: ToolType): ) yield settings.instance() finally: - settings._set_base_settings(old) + settings.set_base_settings(old) @pytest.mark.asyncio async def test_fastmcp_server_creation_and_tool_registration(self): @@ -157,7 +157,7 @@ def mock_settings_for_dynamic_tools(self, enable_remote_tools: bool = True): """Create mock settings with remote tools enabled/disabled""" try: old = settings.instance() - settings._set_base_settings( + settings.set_base_settings( settings.Settings.model_validate( { "dremio": { @@ -175,7 +175,7 @@ def mock_settings_for_dynamic_tools(self, enable_remote_tools: bool = True): ) yield settings.instance() finally: - settings._set_base_settings(old) + settings.set_base_settings(old) @pytest.mark.asyncio async def test_meta_tools_registered_when_enabled(self): diff --git a/tests/tools/test_output_validation.py b/tests/tools/test_output_validation.py index 4e06f19..a945e6a 100644 --- a/tests/tools/test_output_validation.py +++ b/tests/tools/test_output_validation.py @@ -90,13 +90,13 @@ async def test_run_sql_query_json_safe_output(): "dremioai.tools.tools.sql.run_query", new_callable=AsyncMock ) as mock_run_query: mock_run_query.return_value = df - token = settings._settings_override.set( + token = settings.push_settings_override( settings.Settings.model_validate({"dremio": {"uri": "https://test"}}) ) try: result = await tool.invoke("SELECT 1") finally: - settings._settings_override.reset(token) + settings.pop_settings_override(token) assert isinstance(result, dict) assert "result" in result diff --git a/tests/tools/test_tools.py b/tests/tools/test_tools.py index d6e1efd..2b929fb 100644 --- a/tests/tools/test_tools.py +++ b/tests/tools/test_tools.py @@ -415,7 +415,7 @@ def mock_settings(dml_allowed: bool): old_settings = settings.instance() try: - settings._set_base_settings( + settings.set_base_settings( settings.Settings.model_validate( { "dremio": { @@ -427,7 +427,7 @@ def mock_settings(dml_allowed: bool): ) yield finally: - settings._set_base_settings(old_settings) + settings.set_base_settings(old_settings) @pytest.mark.parametrize( From 524b55e30b1cb23dbf20f153fbda77bf878ef538 Mon Sep 17 00:00:00 2001 From: Aniket Kulkarni Date: Thu, 14 May 2026 12:05:04 -0400 Subject: [PATCH 03/12] DX-119909: tighten settings reload synchronization --- src/dremioai/config/settings.py | 63 ++++++++++++++++++--------------- 1 file changed, 34 insertions(+), 29 deletions(-) diff --git a/src/dremioai/config/settings.py b/src/dremioai/config/settings.py index f0c8b9e..edb61e8 100644 --- a/src/dremioai/config/settings.py +++ b/src/dremioai/config/settings.py @@ -15,6 +15,7 @@ # import uuid import hashlib +import threading from uuid import UUID from urllib.parse import urlparse @@ -541,6 +542,7 @@ def __init__(self): "settings_override", default=None ) self.config_fingerprint: ConfigFingerprint | None = None + self.lock = threading.RLock() def initialize_launchdarkly(self, inst: Settings | None): sdk_key = ( @@ -565,18 +567,20 @@ def activate_base_settings( so the refresh loop can compare a scratch `Settings()` object without reconfiguring LaunchDarkly on every tick. """ - self.base_settings = inst - self.config_fingerprint = fingerprint - if initialize_ld: - self.initialize_launchdarkly(inst) + with self.lock: + self.base_settings = inst + self.config_fingerprint = fingerprint + if initialize_ld: + self.initialize_launchdarkly(inst) return inst def reset_for_tests(self): global _yaml_file - _yaml_file = None - self.base_settings = None - self.config_fingerprint = None - FeatureFlagManager.reset() + with self.lock: + _yaml_file = None + self.base_settings = None + self.config_fingerprint = None + FeatureFlagManager.reset() def build_candidate(self) -> Settings: return Settings() @@ -670,31 +674,28 @@ def reload_if_changed(self) -> list[str]: try: fingerprint = self.compute_fingerprint(_yaml_file) - except Exception as exc: - _log.warning(f"Unable to fingerprint config file {_yaml_file}: {exc}") + candidate = self.build_candidate() + except Exception: + _log.exception(f"Skipping config reload for {_yaml_file}") return [] - if fingerprint == self.config_fingerprint: - return [] + with self.lock: + if fingerprint == self.config_fingerprint: + return [] - try: - candidate = self.build_candidate() - except Exception as exc: - _log.warning(f"Skipping config reload for {_yaml_file}: {exc}") - return [] + if not isinstance(self.base_settings, Settings): + self.base_settings = candidate + self.config_fingerprint = fingerprint + return [] - if not isinstance(self.base_settings, Settings): - self.activate_base_settings( - candidate, fingerprint=fingerprint, initialize_ld=False + updated = self.base_settings.model_copy(deep=True) + changed_paths: list[str] = [] + self.copy_runtime_mutable_fields( + updated, candidate, Settings, changed_paths ) - return [] - - updated = self.base_settings.model_copy(deep=True) - changed_paths: list[str] = [] - self.copy_runtime_mutable_fields(updated, candidate, Settings, changed_paths) - self.base_settings = updated - self.config_fingerprint = fingerprint + self.base_settings = updated + self.config_fingerprint = fingerprint if changed_paths: _log.info( @@ -794,13 +795,17 @@ def instance() -> Settings | None: override = settings_reloader.settings_override.get() if isinstance(override, Settings): return override - if not isinstance(settings_reloader.base_settings, Settings): + with settings_reloader.lock: + base_settings = settings_reloader.base_settings + if not isinstance(base_settings, Settings): try: configure() # use default config, if exists except FileNotFoundError: # no default config, create a new default one set_base_settings(Settings()) - return settings_reloader.base_settings + with settings_reloader.lock: + return settings_reloader.base_settings + return base_settings async def run_with( From 3f84ddca57665eac44e7c5da0e461444490772fd Mon Sep 17 00:00:00 2001 From: Aniket Kulkarni Date: Thu, 14 May 2026 12:10:26 -0400 Subject: [PATCH 04/12] DX-119909: add scoped logger refresh support --- src/dremioai/config/settings.py | 4 ++ src/dremioai/log.py | 52 ++++++++++++++++--- src/dremioai/servers/mcp.py | 29 +++++++++-- tests/config/test_launchdarkly_integration.py | 35 +++++++++++++ tests/test_log.py | 27 ++++++++++ 5 files changed, 138 insertions(+), 9 deletions(-) diff --git a/src/dremioai/config/settings.py b/src/dremioai/config/settings.py index edb61e8..02c866e 100644 --- a/src/dremioai/config/settings.py +++ b/src/dremioai/config/settings.py @@ -433,6 +433,10 @@ class BeeAI(BaseModel): class Settings(FlagAwareMixin, BaseSettings): log_level: Annotated[Optional[str], RuntimeMutable()] = Field(default="INFO") + loggers: Annotated[Optional[List[str]], RuntimeMutable(), NoFlag()] = Field( + default=None, + description="Optional logger names to scope log_level to; when unset, log_level applies globally.", + ) dremio: Optional[Dremio] = Field(default=None) tools: Optional[Tools] = Field(default_factory=Tools) launchdarkly: Optional[LaunchDarkly] = Field(default_factory=LaunchDarkly) diff --git a/src/dremioai/log.py b/src/dremioai/log.py index 43c70dc..66c9804 100644 --- a/src/dremioai/log.py +++ b/src/dremioai/log.py @@ -56,6 +56,8 @@ def logger(name=None): _level = None +_scoped_level = None +_scoped_logger_names = set() def _rename_exception_field(_logger, _name, event_dict): @@ -77,13 +79,52 @@ def level(): return getattr(logging, os.environ.get("LOG_LEVEL", "INFO"), logging.INFO) -def set_level(l): - global _level +def scoped_level(): + return _scoped_level + + +def scoped_loggers(): + return sorted(_scoped_logger_names) + + +def _normalize_level(l): + if isinstance(l, str): + return getattr(logging, l.upper(), logging.INFO) + return l + + +def _set_handler_level(l): + for handler in logging.getLogger().handlers: + handler.setLevel(l) + + +def set_level(l, logger_names=None): + global _level, _scoped_level, _scoped_logger_names + l = _normalize_level(l) + + if logger_names: + global_level = level() + logger_names = set(logger_names) + + for name in _scoped_logger_names - logger_names: + logging.getLogger(name).setLevel(global_level) + + for name in logger_names: + logging.getLogger(name).setLevel(l) + + _scoped_logger_names = logger_names + _scoped_level = l + logging.getLogger().setLevel(global_level) + _set_handler_level(min(global_level, l)) + return + _level = l - # propagate to all loggers + _scoped_level = None + _scoped_logger_names.clear() logging.getLogger().setLevel(l) for name in logging.getLogger().manager.loggerDict: logging.getLogger(name).setLevel(l) + _set_handler_level(l) def configure(enable_json_logging=None, to_file=False): @@ -101,7 +142,6 @@ def configure(enable_json_logging=None, to_file=False): else: handler = logging.StreamHandler(sys.stderr) - handler.setLevel(level()) logging.getLogger().handlers.clear() logging.getLogger().addHandler(handler) @@ -133,8 +173,8 @@ def configure(enable_json_logging=None, to_file=False): formatter_processors.append(_rename_exception_field) formatter_processors.extend( [ - structlog.processors.EventRenamer("message"), - renderer, + structlog.processors.EventRenamer("message"), + renderer, ] ) handler.setFormatter( diff --git a/src/dremioai/servers/mcp.py b/src/dremioai/servers/mcp.py index 7a28cab..ccd6b21 100644 --- a/src/dremioai/servers/mcp.py +++ b/src/dremioai/servers/mcp.py @@ -566,10 +566,25 @@ async def _settings_refresh_loop(): if s is None: continue settings.reload_mutable_settings_if_changed() - level_name = settings.instance().get("log_level") + current_settings = settings.instance() + level_name = current_settings.get("log_level") + logger_names = current_settings.loggers or [] level = getattr(logging, level_name.upper(), None) - if level is not None and level != log.level(): - _log.info(f"Updating log level to {level_name}") + if level is None: + continue + + current_scoped_loggers = log.scoped_loggers() + if logger_names: + if ( + level != log.scoped_level() + or logger_names != current_scoped_loggers + ): + _log.info( + f"Updating log level to {level_name} for loggers {', '.join(logger_names)}" + ) + log.set_level(level, logger_names=logger_names) + elif level != log.level() or current_scoped_loggers: + _log.info(f"Updating global log level to {level_name}") log.set_level(level) except Exception as e: _log.debug(f"Settings refresh failed: {e}") @@ -674,6 +689,14 @@ def main( else: transport = Transports.stdio settings.configure(config_file) + if settings.instance().loggers: + configured_level = getattr( + logging, settings.instance().get("log_level").upper(), None + ) + if configured_level is not None: + log.set_level( + configured_level, logger_names=settings.instance().loggers + ) dremio = settings.instance().dremio if ( dremio.oauth_supported diff --git a/tests/config/test_launchdarkly_integration.py b/tests/config/test_launchdarkly_integration.py index 3e7b0bd..5fee48d 100644 --- a/tests/config/test_launchdarkly_integration.py +++ b/tests/config/test_launchdarkly_integration.py @@ -728,6 +728,41 @@ async def test_settings_refresh_reloads_yaml_before_log_level_evaluation(tmp_pat log.set_level(original_level) +@pytest.mark.asyncio +async def test_settings_refresh_scopes_log_level_to_configured_loggers(tmp_path): + from dremioai.servers.mcp import _settings_refresh_loop + from dremioai import log + + cfg = tmp_path / "config.yaml" + cfg.write_text(""" +log_level: DEBUG +loggers: + - scoped.logger +dremio: + uri: https://test.dremio.cloud + pat: test-pat +""") + settings.configure(cfg) + + original_level = log.level() + try: + with ( + patch("dremioai.servers.mcp._SETTINGS_REFRESH_INTERVAL", 0), + patch.object(log, "set_level") as mock_set_level, + ): + task = asyncio.create_task(_settings_refresh_loop()) + await asyncio.sleep(0.05) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + mock_set_level.assert_called_with(logging.DEBUG, logger_names=["scoped.logger"]) + finally: + log.set_level(original_level) + + # -- _build_context ----------------------------------------------------------- diff --git a/tests/test_log.py b/tests/test_log.py index a4b0d05..27ccd76 100644 --- a/tests/test_log.py +++ b/tests/test_log.py @@ -49,6 +49,8 @@ def reset_structlog(): structlog.reset_defaults() # Reset the global level log._level = None + log._scoped_level = None + log._scoped_logger_names.clear() # Clear all handlers from root logger root_logger = logging.getLogger() for handler in root_logger.handlers[:]: @@ -60,6 +62,8 @@ def reset_structlog(): root_logger = logging.getLogger() for handler in root_logger.handlers[:]: root_logger.removeHandler(handler) + log._scoped_level = None + log._scoped_logger_names.clear() class TestGetLogDirectory: @@ -355,3 +359,26 @@ def test_log_level_filtering(self, mock_home_dir): assert "Info message" not in content assert "Warning message" in content assert "Error message" in content + + def test_scoped_log_level_filtering(self, mock_home_dir): + """Scoped log level should only apply to the named loggers.""" + with patch("sys.platform", "linux"): + structlog.reset_defaults() + log.set_level(logging.INFO) + log.configure(to_file=True) + log.set_level(logging.DEBUG, logger_names=["scoped.logger"]) + + scoped_logger = log.logger("scoped.logger") + other_logger = log.logger("other.logger") + + scoped_logger.debug("Scoped debug message") + other_logger.debug("Other debug message") + other_logger.info("Other info message") + + expected_dir = mock_home_dir / ".local" / "share" / "dremioai" / "logs" + log_file = expected_dir / "dremioai.log" + content = log_file.read_text() + + assert "Scoped debug message" in content + assert "Other debug message" not in content + assert "Other info message" in content From 283ec5f17fe3fb19712cbce497ed4441daffe245 Mon Sep 17 00:00:00 2001 From: Aniket Kulkarni Date: Thu, 14 May 2026 16:43:48 -0400 Subject: [PATCH 05/12] DX-119909: offload settings reload from refresh loop --- src/dremioai/servers/mcp.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/dremioai/servers/mcp.py b/src/dremioai/servers/mcp.py index ccd6b21..0bf8ac3 100644 --- a/src/dremioai/servers/mcp.py +++ b/src/dremioai/servers/mcp.py @@ -565,7 +565,11 @@ async def _settings_refresh_loop(): s = settings.instance() if s is None: continue - settings.reload_mutable_settings_if_changed() + # Config reload is typically a small local file read, so blocking the + # event loop here would be an edge case. We still offload it to a + # worker thread to keep the refresh loop non-blocking if config I/O + # ever becomes unexpectedly slow. + await asyncio.to_thread(settings.reload_mutable_settings_if_changed) current_settings = settings.instance() level_name = current_settings.get("log_level") logger_names = current_settings.loggers or [] From 0088c8c9cdd8721970be8655dbe685d1cf7777a3 Mon Sep 17 00:00:00 2001 From: Aniket Kulkarni Date: Mon, 18 May 2026 11:09:27 -0400 Subject: [PATCH 06/12] DX-119909: harden auth warning logging --- src/dremioai/servers/mcp.py | 8 +++++++- tests/servers/test_jwks_verifier.py | 2 ++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/src/dremioai/servers/mcp.py b/src/dremioai/servers/mcp.py index 0bf8ac3..c63429b 100644 --- a/src/dremioai/servers/mcp.py +++ b/src/dremioai/servers/mcp.py @@ -183,12 +183,18 @@ async def dispatch(self, request: Request, call_next: RequestResponseEndpoint): and request.url.path.startswith("/mcp") ): client_host = request.client.host if request.client else "unknown" + inst = settings.instance() + endpoint = ( + str(inst.dremio.uri) + if inst is not None and inst.dremio is not None and inst.dremio.uri + else None + ) self.logger.warning( "Unauthorized request rejected", path=request.url.path, client=client_host, project_id=ProjectIdMiddleware.get_project_id(), - endpoint=str(settings.instance().dremio.uri), + endpoint=endpoint, ) # Return 401 with WWW-Authenticate header return StarletteResponse( diff --git a/tests/servers/test_jwks_verifier.py b/tests/servers/test_jwks_verifier.py index ace2019..a17e0cc 100644 --- a/tests/servers/test_jwks_verifier.py +++ b/tests/servers/test_jwks_verifier.py @@ -36,6 +36,7 @@ from starlette.responses import Response from dremioai import log +from dremioai.config import settings from dremioai.servers.jwks_verifier import JWKSVerifier, VerifiedClaims, TokenExpiredError from dremioai.servers.mcp import ( FastMCPServerWithAuthToken, @@ -293,6 +294,7 @@ class TestDispatchWarning: @pytest.mark.asyncio async def test_dispatch_logs_warning_on_401(self, caplog): + settings.set_base_settings(settings.Settings()) middleware = RequireAuthWithWWWAuthenticateMiddleware(app=MagicMock()) mock_user = MagicMock() From db1a9934f509395ceb6f5931e3d3e75f376ace83 Mon Sep 17 00:00:00 2001 From: Aniket Kulkarni Date: Mon, 18 May 2026 12:21:32 -0400 Subject: [PATCH 07/12] DX-119909: guard tool filtering without dremio config --- src/dremioai/tools/tools.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/dremioai/tools/tools.py b/src/dremioai/tools/tools.py index 6c56608..b073e31 100644 --- a/src/dremioai/tools/tools.py +++ b/src/dremioai/tools/tools.py @@ -219,15 +219,18 @@ def _get_class_var_hints(tool: Tools, name: str) -> bool: def is_tool_for( tool: Tools, tool_type: ToolType, dremio: settings.Dremio = None ) -> bool: - if dremio is None and settings.instance().dremio: - dremio = settings.instance().dremio + inst = settings.instance() + if dremio is None and inst is not None and inst.dremio is not None: + dremio = inst.dremio if project_id_required := get_project_id_required(tool): if dremio is not None and dremio.project_id is None: return False if (For := get_for(tool)) is not None: - if For & ToolType.EXPERIMENTAL and not dremio.get("enable_search"): + if For & ToolType.EXPERIMENTAL and ( + dremio is None or not dremio.get("enable_search") + ): return False return (For & tool_type) != 0 # == tool_type return False From 6ac1b394ecbb624cf29ef7bd97bbd5f299b3800b Mon Sep 17 00:00:00 2001 From: Aniket Kulkarni Date: Mon, 18 May 2026 17:12:35 -0400 Subject: [PATCH 08/12] DX-119909: address review feedback on reload boundaries --- src/dremioai/config/settings.py | 28 +---- tests/config/test_launchdarkly_integration.py | 38 ++++++ tests/config/test_settings.py | 109 +++++++++++++++++- 3 files changed, 144 insertions(+), 31 deletions(-) diff --git a/src/dremioai/config/settings.py b/src/dremioai/config/settings.py index 02c866e..bf6cb47 100644 --- a/src/dremioai/config/settings.py +++ b/src/dremioai/config/settings.py @@ -433,7 +433,7 @@ class BeeAI(BaseModel): class Settings(FlagAwareMixin, BaseSettings): log_level: Annotated[Optional[str], RuntimeMutable()] = Field(default="INFO") - loggers: Annotated[Optional[List[str]], RuntimeMutable(), NoFlag()] = Field( + loggers: Annotated[Optional[List[str]], NoFlag()] = Field( default=None, description="Optional logger names to scope log_level to; when unset, log_level applies globally.", ) @@ -640,37 +640,13 @@ def copy_runtime_mutable_fields( if ( isinstance(annotation, type) and issubclass(annotation, BaseModel) + and current_value is not None and value is not None ): - if current_value is None: - # Simple instantiation is not sufficient here because some - # nested models have required, startup-only fields. We copy - # the candidate subtree shape, blank it, then replay only - # runtime-mutable leaves into that shell. - current_value = value.model_copy(deep=True) - self.blank_model_subtree(current_value, annotation) - setattr(current_obj, field_name, current_value) self.copy_runtime_mutable_fields( current_value, value, annotation, changed_paths, field_path ) - def blank_model_subtree(self, current_obj: Any, model_cls: type): - """Clear a copied subtree before replaying mutable leaves into it.""" - hints = get_type_hints(model_cls, include_extras=True) - for field_name in model_cls.model_fields: - value = getattr(current_obj, field_name, None) - annotation = self.unwrap_annotation(hints[field_name]) - - if ( - isinstance(annotation, type) - and issubclass(annotation, BaseModel) - and value is not None - ): - self.blank_model_subtree(value, annotation) - continue - - object.__setattr__(current_obj, field_name, None) - def reload_if_changed(self) -> list[str]: _log = log.logger("settings_reload") if _yaml_file is None: diff --git a/tests/config/test_launchdarkly_integration.py b/tests/config/test_launchdarkly_integration.py index 5fee48d..cdab3ef 100644 --- a/tests/config/test_launchdarkly_integration.py +++ b/tests/config/test_launchdarkly_integration.py @@ -692,6 +692,44 @@ def test_reload_mutable_settings_does_not_reinitialize_ld_on_validation_failure( assert mock_ldclient.set_config.call_count == 1 +@patch("dremioai.config.feature_flags.ldclient") +def test_reload_mutable_settings_preserves_snapshot_on_schema_validation_failure( + mock_ldclient, tmp_path +): + mock_client = _make_mock_ld_client({}) + mock_ldclient.get.return_value = mock_client + + cfg = tmp_path / "config.yaml" + cfg.write_text(""" +launchdarkly: + sdk_key: test-key +dremio: + uri: https://test.dremio.cloud + pat: test-pat + api: + polling_interval: 3 +""") + settings.configure(cfg) + before = settings.instance().model_dump() + assert mock_ldclient.set_config.call_count == 1 + + cfg.write_text(""" +launchdarkly: + sdk_key: test-key +dremio: + uri: https://test.dremio.cloud + pat: test-pat + api: + http_retry: + max_retries: not-a-number +""") + changed = settings.reload_mutable_settings_if_changed() + + assert changed == [] + assert settings.instance().model_dump() == before + assert mock_ldclient.set_config.call_count == 1 + + @pytest.mark.asyncio async def test_settings_refresh_reloads_yaml_before_log_level_evaluation(tmp_path): from dremioai.servers.mcp import _settings_refresh_loop diff --git a/tests/config/test_settings.py b/tests/config/test_settings.py index a5171af..5a3c5f9 100644 --- a/tests/config/test_settings.py +++ b/tests/config/test_settings.py @@ -14,6 +14,7 @@ # limitations under the License. # +import asyncio import os import uuid @@ -28,6 +29,7 @@ from dremioai.config import settings from dremioai.config.tools import ToolType +from dremioai.tools.tools import get_tools def test_configure_with_no_file_works(mock_config_dir): @@ -90,6 +92,14 @@ async def _read_runtime_settings(): ) +async def _read_runtime_settings_after_event( + started: asyncio.Event, release: asyncio.Event +): + started.set() + await release.wait() + return await _read_runtime_settings() + + @pytest.mark.asyncio async def test_run_with_keeps_overrides_request_scoped(): base = settings.Settings.model_validate( @@ -119,6 +129,60 @@ async def test_run_with_keeps_overrides_request_scoped(): assert await _read_runtime_settings() == original +@pytest.mark.asyncio +async def test_run_with_keeps_overrides_request_scoped_under_concurrency(): + base = settings.Settings.model_validate( + { + "log_level": "INFO", + "dremio": { + "uri": "https://test.dremio.cloud", + "pat": "test-pat", + "enable_search": False, + "api": {"polling_interval": 1.0}, + }, + } + ) + settings.set_base_settings(base) + + started_one = asyncio.Event() + started_two = asyncio.Event() + release = asyncio.Event() + + override_one = asyncio.create_task( + settings.run_with( + _read_runtime_settings_after_event, + overrides={ + "log_level": "DEBUG", + "dremio.enable_search": True, + "dremio.api.polling_interval": 3.5, + }, + args=[started_one, release], + ) + ) + override_two = asyncio.create_task( + settings.run_with( + _read_runtime_settings_after_event, + overrides={ + "log_level": "ERROR", + "dremio.enable_search": False, + "dremio.api.polling_interval": 9.0, + }, + args=[started_two, release], + ) + ) + + await asyncio.gather(started_one.wait(), started_two.wait()) + base_read = await _read_runtime_settings() + release.set() + + result_one, result_two = await asyncio.gather(override_one, override_two) + + assert base_read == ("INFO", False, 1.0) + assert result_one == ("DEBUG", True, 3.5) + assert result_two == ("ERROR", False, 9.0) + assert await _read_runtime_settings() == ("INFO", False, 1.0) + + def test_reload_mutable_settings_if_changed_updates_runtime_mutable_only(tmp_path): cfg = tmp_path / "config.yaml" cfg.write_text( @@ -226,7 +290,7 @@ def test_reload_mutable_settings_if_changed_ignores_non_mutable_changes(tmp_path assert settings.instance().dremio.pat == "test-pat" -def test_reload_mutable_settings_if_changed_materializes_missing_nested_subtree( +def test_reload_mutable_settings_if_changed_does_not_materialize_missing_subtree( tmp_path, ): cfg = tmp_path / "config.yaml" @@ -249,10 +313,45 @@ def test_reload_mutable_settings_if_changed_materializes_missing_nested_subtree( changed = settings.reload_mutable_settings_if_changed() - assert "dremio.allow_dml" in changed - assert "dremio.api.polling_interval" in changed - assert settings.instance().dremio.get("allow_dml") is True - assert settings.instance().dremio.api.get("polling_interval") == 7.0 + assert changed == [] + assert settings.instance().dremio is None + + +def test_reload_mutable_settings_if_changed_keeps_tools_server_mode_startup_only( + tmp_path, +): + cfg = tmp_path / "config.yaml" + cfg.write_text( + yaml.safe_dump( + { + "log_level": "INFO", + "dremio": {"uri": "https://one.dremio.cloud", "pat": "test-pat"}, + "tools": {"server_mode": ToolType.FOR_SELF.name}, + } + ) + ) + settings.configure(cfg) + + before_mode = settings.instance().tools.server_mode + before_tools = {tool.__name__ for tool in get_tools(For=before_mode)} + + cfg.write_text( + yaml.safe_dump( + { + "log_level": "DEBUG", + "dremio": {"uri": "https://one.dremio.cloud", "pat": "test-pat"}, + "tools": {"server_mode": ToolType.FOR_DATA_PATTERNS.name}, + } + ) + ) + + changed = settings.reload_mutable_settings_if_changed() + after_mode = settings.instance().tools.server_mode + after_tools = {tool.__name__ for tool in get_tools(For=after_mode)} + + assert changed == ["log_level"] + assert after_mode == before_mode + assert after_tools == before_tools @pytest.mark.parametrize( From 8a3daff6561742d19adee82c3f083280908570bd Mon Sep 17 00:00:00 2001 From: Aniket Kulkarni Date: Tue, 19 May 2026 09:55:59 -0400 Subject: [PATCH 09/12] DX-119909: clean up settings reload helpers --- src/dremioai/config/settings.py | 63 ++++++++++++++++++++------------- 1 file changed, 38 insertions(+), 25 deletions(-) diff --git a/src/dremioai/config/settings.py b/src/dremioai/config/settings.py index bf6cb47..3505970 100644 --- a/src/dremioai/config/settings.py +++ b/src/dremioai/config/settings.py @@ -16,6 +16,7 @@ import uuid import hashlib import threading +from dataclasses import dataclass from uuid import UUID from urllib.parse import urlparse @@ -529,7 +530,25 @@ def collect_flag_keys(model_cls: type, prefix: str = "") -> list[str]: # Module-level holder so configure() can pass the YAML path to the Settings constructor _yaml_file: Path | None = None -ConfigFingerprint = tuple[str, int, int, int, str] +@dataclass(frozen=True, init=False) +class ConfigFingerprint: + path: str + inode: int + size: int + mtime_ns: int + content_hash: str + + def __init__(self, resolved: Path): + stat = resolved.stat() + object.__setattr__(self, "path", str(resolved)) + object.__setattr__(self, "inode", stat.st_ino) + object.__setattr__(self, "size", stat.st_size) + object.__setattr__(self, "mtime_ns", stat.st_mtime_ns) + object.__setattr__( + self, + "content_hash", + hashlib.sha256(resolved.read_bytes()).hexdigest(), + ) class SettingsReloader: @@ -590,27 +609,14 @@ def build_candidate(self) -> Settings: return Settings() def compute_fingerprint(self, cfg: Path | None) -> ConfigFingerprint | None: - """Return a resolved-path fingerprint plus a content hash. - - Resolved path + inode + size + mtime_ns handles ConfigMap swaps well. - A SHA-256 of the current file contents covers the rare edge case where - metadata stays the same across a rewrite. - """ + """Return a config fingerprint for the resolved config path.""" if cfg is None: return None - resolved = cfg.resolve() - stat = resolved.stat() - content_hash = hashlib.sha256(resolved.read_bytes()).hexdigest() - return ( - str(resolved), - stat.st_ino, - stat.st_size, - stat.st_mtime_ns, - content_hash, - ) + return ConfigFingerprint(cfg.resolve()) - def unwrap_annotation(self, annotation: Any) -> Any: - args = [a for a in get_args(annotation) if a is not NoneType] + def unwrap_optional_annotation(self, annotation: Any) -> Any: + """Unwrap Optional[T] / T | None, but leave broader unions unchanged.""" + args = tuple(a for a in get_args(annotation) if a is not NoneType) if len(args) == 1: return args[0] return annotation @@ -629,7 +635,7 @@ def copy_runtime_mutable_fields( value = getattr(candidate_obj, field_name, None) current_value = getattr(current_obj, field_name, None) field_path = f"{prefix}.{field_name}" if prefix else field_name - annotation = self.unwrap_annotation(hints[field_name]) + annotation = self.unwrap_optional_annotation(hints[field_name]) if _has_runtime_mutable(model_cls, field_name): if current_value != value: @@ -668,12 +674,19 @@ def reload_if_changed(self) -> list[str]: self.config_fingerprint = fingerprint return [] - updated = self.base_settings.model_copy(deep=True) - changed_paths: list[str] = [] - self.copy_runtime_mutable_fields( - updated, candidate, Settings, changed_paths - ) + current_base = self.base_settings + current_fingerprint = self.config_fingerprint + updated = current_base.model_copy(deep=True) + changed_paths: list[str] = [] + self.copy_runtime_mutable_fields(updated, candidate, Settings, changed_paths) + + with self.lock: + if ( + self.base_settings is not current_base + or self.config_fingerprint != current_fingerprint + ): + return [] self.base_settings = updated self.config_fingerprint = fingerprint From 9d5237c17a4628722cfb9283e26311ba6e821439 Mon Sep 17 00:00:00 2001 From: Aniket Kulkarni Date: Tue, 19 May 2026 10:06:17 -0400 Subject: [PATCH 10/12] DX-119909: allow mutable scoped logger config --- src/dremioai/config/settings.py | 2 +- tests/config/test_launchdarkly_integration.py | 9 ++++++++- tests/config/test_settings.py | 4 ++++ 3 files changed, 13 insertions(+), 2 deletions(-) diff --git a/src/dremioai/config/settings.py b/src/dremioai/config/settings.py index 3505970..5f3b93f 100644 --- a/src/dremioai/config/settings.py +++ b/src/dremioai/config/settings.py @@ -434,7 +434,7 @@ class BeeAI(BaseModel): class Settings(FlagAwareMixin, BaseSettings): log_level: Annotated[Optional[str], RuntimeMutable()] = Field(default="INFO") - loggers: Annotated[Optional[List[str]], NoFlag()] = Field( + loggers: Annotated[Optional[List[str]], RuntimeMutable(), NoFlag()] = Field( default=None, description="Optional logger names to scope log_level to; when unset, log_level applies globally.", ) diff --git a/tests/config/test_launchdarkly_integration.py b/tests/config/test_launchdarkly_integration.py index cdab3ef..32436ac 100644 --- a/tests/config/test_launchdarkly_integration.py +++ b/tests/config/test_launchdarkly_integration.py @@ -772,6 +772,14 @@ async def test_settings_refresh_scopes_log_level_to_configured_loggers(tmp_path) from dremioai import log cfg = tmp_path / "config.yaml" + cfg.write_text(""" +log_level: INFO +dremio: + uri: https://test.dremio.cloud + pat: test-pat +""") + settings.configure(cfg) + cfg.write_text(""" log_level: DEBUG loggers: @@ -780,7 +788,6 @@ async def test_settings_refresh_scopes_log_level_to_configured_loggers(tmp_path) uri: https://test.dremio.cloud pat: test-pat """) - settings.configure(cfg) original_level = log.level() try: diff --git a/tests/config/test_settings.py b/tests/config/test_settings.py index 5a3c5f9..8204eb1 100644 --- a/tests/config/test_settings.py +++ b/tests/config/test_settings.py @@ -189,6 +189,7 @@ def test_reload_mutable_settings_if_changed_updates_runtime_mutable_only(tmp_pat yaml.safe_dump( { "log_level": "INFO", + "loggers": ["initial.logger"], "dremio": { "uri": "https://one.dremio.cloud", "pat": "test-pat", @@ -209,6 +210,7 @@ def test_reload_mutable_settings_if_changed_updates_runtime_mutable_only(tmp_pat yaml.safe_dump( { "log_level": "DEBUG", + "loggers": ["updated.logger"], "dremio": { "uri": "https://two.dremio.cloud", "pat": "changed-pat", @@ -227,12 +229,14 @@ def test_reload_mutable_settings_if_changed_updates_runtime_mutable_only(tmp_pat assert changed == [ "log_level", + "loggers", "dremio.enable_search", "dremio.allow_dml", "dremio.api.http_retry.max_retries", "dremio.api.polling_interval", ] assert settings.instance().log_level == "DEBUG" + assert settings.instance().loggers == ["updated.logger"] assert settings.instance().dremio.enable_search is True assert settings.instance().dremio.allow_dml is True assert settings.instance().dremio.api.http_retry.max_retries == 11 From e2674499969dc28dabf4cd77da60fac951af17d5 Mon Sep 17 00:00:00 2001 From: Aniket Kulkarni Date: Tue, 19 May 2026 10:11:23 -0400 Subject: [PATCH 11/12] DX-119909: simplify config fingerprint construction --- src/dremioai/config/settings.py | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/src/dremioai/config/settings.py b/src/dremioai/config/settings.py index 5f3b93f..2e27cae 100644 --- a/src/dremioai/config/settings.py +++ b/src/dremioai/config/settings.py @@ -530,7 +530,7 @@ def collect_flag_keys(model_cls: type, prefix: str = "") -> list[str]: # Module-level holder so configure() can pass the YAML path to the Settings constructor _yaml_file: Path | None = None -@dataclass(frozen=True, init=False) +@dataclass(frozen=True) class ConfigFingerprint: path: str inode: int @@ -538,16 +538,15 @@ class ConfigFingerprint: mtime_ns: int content_hash: str - def __init__(self, resolved: Path): + @classmethod + def from_resolved_path(cls, resolved: Path) -> Self: stat = resolved.stat() - object.__setattr__(self, "path", str(resolved)) - object.__setattr__(self, "inode", stat.st_ino) - object.__setattr__(self, "size", stat.st_size) - object.__setattr__(self, "mtime_ns", stat.st_mtime_ns) - object.__setattr__( - self, - "content_hash", - hashlib.sha256(resolved.read_bytes()).hexdigest(), + return cls( + path=str(resolved), + inode=stat.st_ino, + size=stat.st_size, + mtime_ns=stat.st_mtime_ns, + content_hash=hashlib.sha256(resolved.read_bytes()).hexdigest(), ) @@ -612,7 +611,7 @@ def compute_fingerprint(self, cfg: Path | None) -> ConfigFingerprint | None: """Return a config fingerprint for the resolved config path.""" if cfg is None: return None - return ConfigFingerprint(cfg.resolve()) + return ConfigFingerprint.from_resolved_path(cfg.resolve()) def unwrap_optional_annotation(self, annotation: Any) -> Any: """Unwrap Optional[T] / T | None, but leave broader unions unchanged.""" From 77ddee766753b39a6211178c3cf0566733c798f2 Mon Sep 17 00:00:00 2001 From: Aniket Kulkarni Date: Tue, 19 May 2026 10:25:47 -0400 Subject: [PATCH 12/12] DX-119909: avoid e2e port collisions --- tests/conftest.py | 19 ++++++++++++------- tests/mocks/http_mock.py | 8 +++++++- 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 8111610..de10a81 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -19,7 +19,7 @@ """ import os -import random +import socket import uuid from typing import AsyncGenerator, NamedTuple @@ -51,6 +51,13 @@ from prometheus_client import CollectorRegistry +def _reserve_local_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + sock.listen(1) + return sock.getsockname()[1] + + @pytest.fixture(autouse=True) def reset_settings_state(): settings.reset_state_for_tests() @@ -203,7 +210,7 @@ def _create_logging_server(log_level="warning"): ) return create_pytest_logging_server_fixture( - mock_data=mock_data, port=8000, log_level=log_level + mock_data=mock_data, log_level=log_level ) @@ -244,12 +251,10 @@ async def http_streamable_mcp_server( try: settings.configure(force=True) host = "127.0.0.1" - port = random.randrange(9000, 12000) - metrics_port = random.randrange(9000, 12000) - - # Ensure metrics port is different from main port + port = _reserve_local_port() + metrics_port = _reserve_local_port() while metrics_port == port: - metrics_port = random.randrange(9000, 12000) + metrics_port = _reserve_local_port() config = { "dremio": { diff --git a/tests/mocks/http_mock.py b/tests/mocks/http_mock.py index 7cc527f..10434b3 100644 --- a/tests/mocks/http_mock.py +++ b/tests/mocks/http_mock.py @@ -17,6 +17,7 @@ import json import re import asyncio +import socket import time from pathlib import Path from typing import Dict, Any, Optional, Union, TextIO, List, Coroutine, Callable @@ -404,9 +405,14 @@ def logs(self) -> List[LogEntry]: def create_pytest_logging_server_fixture( mock_data: Optional[OrderedDict[str, str]] = None, - port: int = 8000, + port: int | None = None, log_level="warning", ) -> LoggingServerFixture: + if port is None: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + sock.listen(1) + port = sock.getsockname()[1] log_file = io.StringIO() thread, stop_event = start_logging_server(