diff --git a/docs/settings.md b/docs/settings.md index 90927e8..2f23b72 100644 --- a/docs/settings.md +++ b/docs/settings.md @@ -70,6 +70,22 @@ PAT can be provided: - Directly as a string - As a file path prefixed with '@' (e.g., "@~/tokens/dremio.token") +#### Basic auth (Dremio Software without PAT support) + +Dremio Software deployments that cannot issue PATs (e.g. Community edition) +can authenticate with username/password instead. The server exchanges the +credentials for a session token via `POST /apiv2/login` and refreshes it +automatically before it expires. Ignored for Dremio Cloud; if both `pat` and +`basic_auth` are configured, the PAT wins. + +```yaml +dremio: + uri: https://your-dremio-instance:9047 + basic_auth: + username: # Dremio username + password: # Direct value or '@' file reference, e.g. "@~/tokens/dremio.password" +``` + ### Tools Settings ```yaml diff --git a/src/dremioai/api/basic_auth.py b/src/dremioai/api/basic_auth.py new file mode 100644 index 0000000..8897710 --- /dev/null +++ b/src/dremioai/api/basic_auth.py @@ -0,0 +1,72 @@ +# +# Copyright (C) 2017-2025 Dremio Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Session-token login for Dremio Software deployments. + +Dremio Software Community edition cannot issue PATs, so the only credential it +offers is username/password. This module exchanges those credentials for a +session token via ``POST /apiv2/login`` and installs it as the effective token +for API calls (the REST API accepts session tokens as ``Bearer`` tokens). + +Mirrors the shape of :mod:`dremioai.api.oauth2`: ``get_session_token()`` +returns an object whose ``update_settings()`` mutates the live settings. +""" + +from dataclasses import dataclass +from datetime import datetime, timedelta +from json import dumps, loads +from typing import Optional +from urllib.request import Request, urlopen + +from dremioai.config import settings +from dremioai.log import logger + +# Re-login this long before the server-reported expiry, so an in-flight +# request never straddles the expiration boundary. +_EXPIRY_SAFETY_MARGIN = timedelta(minutes=5) + + +@dataclass +class SessionToken: + token: str + expiry: Optional[datetime] + + def update_settings(self): + dremio = settings.instance().dremio + dremio.pat = self.token + dremio.basic_auth.expiry = self.expiry + + +def get_session_token(timeout: float = 30.0) -> SessionToken: + """Exchange the configured username/password for a Dremio session token.""" + dremio = settings.instance().dremio + basic_auth = dremio.basic_auth + body = dumps( + {"userName": basic_auth.username, "password": basic_auth.password} + ).encode("utf-8") + request = Request( + f"{dremio.uri}/apiv2/login", + data=body, + headers={"Content-Type": "application/json"}, + ) + with urlopen(request, timeout=timeout) as response: + payload = loads(response.read().decode("utf-8")) + + expiry = None + if expires_ms := payload.get("expires"): + expiry = datetime.fromtimestamp(expires_ms / 1000) - _EXPIRY_SAFETY_MARGIN + + logger().info(f"Obtained Dremio session token (expires {expiry})") + return SessionToken(token=payload["token"], expiry=expiry) diff --git a/src/dremioai/api/transport.py b/src/dremioai/api/transport.py index ed1cc40..010c172 100644 --- a/src/dremioai/api/transport.py +++ b/src/dremioai/api/transport.py @@ -36,6 +36,7 @@ from dremioai.config import settings from dremioai.api.oauth2 import get_oauth2_tokens +from dremioai.api.basic_auth import get_session_token DeserializationStrategy: TypeAlias = Union[Callable, BaseModel] @@ -222,6 +223,12 @@ def __init__(self): oauth = get_oauth2_tokens() oauth.update_settings() + if dremio.basic_auth_configured and ( + dremio.pat is None or dremio.basic_auth.has_expired + ): + session = get_session_token() + session.update_settings() + uri = dremio.uri pat = dremio.pat diff --git a/src/dremioai/config/settings.py b/src/dremioai/config/settings.py index e60d193..fb28e7a 100644 --- a/src/dremioai/config/settings.py +++ b/src/dremioai/config/settings.py @@ -86,8 +86,10 @@ class NoFlag: pass + class RuntimeMutable: """Mark a field as safe to update from a runtime config reload.""" + pass @@ -95,6 +97,7 @@ class RuntimeMutable: class FlagName: name: str = None + def _has_no_flag(model_cls: type[BaseModel], field_name: str) -> bool: """Check if a field has the NoFlag annotation marker.""" info = model_cls.model_fields.get(field_name) @@ -213,6 +216,33 @@ def has_expired(self) -> bool: return self.expiry is not None and self.expiry < datetime.now() +class BasicAuth(BaseModel): + """Username/password login for Dremio Software deployments that cannot + issue PATs (e.g. Community edition). The credentials are exchanged for a + session token via /apiv2/login, which the REST API also accepts as a + Bearer token. Not applicable to Dremio Cloud.""" + + username: str + raw_password: Annotated[str, NoFlag()] = Field( + alias="password", + description="Password for basic login (can be a file path with @ prefix or direct value)", + ) + expiry: Optional[datetime] = None + model_config = ConfigDict(validate_assignment=True, populate_by_name=True) + + @property + def password(self) -> str: + return _resolve_token_file(self.raw_password) + + @field_serializer("raw_password") + def serialize_password(self, password: str): + return self.raw_password if password != self.raw_password else password + + @property + def has_expired(self) -> bool: + return self.expiry is not None and self.expiry < datetime.now() + + class Wlm(FlagAwareModel): engine_name: Optional[str] = None @@ -276,6 +306,7 @@ class Dremio(FlagAwareModel): description="enable experimental tools", ) oauth2: Optional[OAuth2] = None + basic_auth: Optional[BasicAuth] = None allow_dml: Annotated[Optional[bool], RuntimeMutable()] = Field(default=False) extract_org_id_from_jwt: Optional[bool] = Field( default=False, @@ -343,6 +374,10 @@ def serialize_pat(self, pat: str): def oauth_configured(self) -> bool: return self.oauth2 is not None + @property + def basic_auth_configured(self) -> bool: + return self.basic_auth is not None and not self.is_cloud + @property def oauth_supported(self) -> bool: return self.project_id is not None @@ -560,6 +595,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) class ConfigFingerprint: path: str diff --git a/tests/api/test_basic_auth.py b/tests/api/test_basic_auth.py new file mode 100644 index 0000000..d9d83cb --- /dev/null +++ b/tests/api/test_basic_auth.py @@ -0,0 +1,111 @@ +# +# Copyright (C) 2017-2025 Dremio Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +from datetime import datetime, timedelta +from io import BytesIO +from json import dumps, loads +from unittest.mock import patch + +import pytest + +from dremioai.api.basic_auth import get_session_token +from dremioai.config import settings + + +@pytest.fixture +def software_settings_with_basic_auth(mock_config_dir): + settings.configure(force=True) + settings.instance().dremio = settings.Dremio.model_validate( + { + "uri": "http://dremio.example.com:9047", + "basic_auth": {"username": "alice", "password": "s3cret"}, + } + ) + yield settings.instance() + settings.configure(force=True) + + +class _FakeLoginResponse: + def __init__(self, payload: dict): + self._body = BytesIO(dumps(payload).encode("utf-8")) + + def read(self): + return self._body.read() + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + +def test_basic_auth_settings_parse(software_settings_with_basic_auth): + dremio = settings.instance().dremio + assert dremio.basic_auth_configured + assert dremio.basic_auth.username == "alice" + assert dremio.basic_auth.password == "s3cret" + assert not dremio.basic_auth.has_expired + assert dremio.pat is None + + +def test_basic_auth_password_file_resolution( + software_settings_with_basic_auth, tmp_path +): + password_file = tmp_path / "dremio-password" + password_file.write_text("from-a-file\n") + dremio = settings.instance().dremio + dremio.basic_auth = settings.BasicAuth.model_validate( + {"username": "alice", "password": f"@{password_file}"} + ) + assert dremio.basic_auth.password == "from-a-file" + # serialization must keep the reference, not the resolved secret + assert dremio.basic_auth.model_dump()["raw_password"] == f"@{password_file}" + + +def test_basic_auth_not_configured_for_cloud(software_settings_with_basic_auth): + dremio = settings.instance().dremio + dremio.project_id = "01234567-89ab-cdef-0123-456789abcdef" + assert dremio.is_cloud + assert not dremio.basic_auth_configured + + +def test_get_session_token_logs_in_and_updates_settings( + software_settings_with_basic_auth, +): + expires_ms = int((datetime.now() + timedelta(hours=30)).timestamp() * 1000) + captured = {} + + def fake_urlopen(request, timeout): + captured["url"] = request.full_url + captured["body"] = loads(request.data.decode("utf-8")) + return _FakeLoginResponse({"token": "session-token-123", "expires": expires_ms}) + + with patch("dremioai.api.basic_auth.urlopen", side_effect=fake_urlopen): + session = get_session_token() + session.update_settings() + + assert captured["url"] == "http://dremio.example.com:9047/apiv2/login" + assert captured["body"] == {"userName": "alice", "password": "s3cret"} + + dremio = settings.instance().dremio + assert dremio.pat == "session-token-123" + assert dremio.basic_auth.expiry is not None + assert not dremio.basic_auth.has_expired + + +def test_expired_session_is_detected(software_settings_with_basic_auth): + dremio = settings.instance().dremio + dremio.basic_auth.expiry = datetime.now() - timedelta(minutes=1) + assert dremio.basic_auth.has_expired