From 59743220bb9c8ecde4c758b017eae484c1eeec2f Mon Sep 17 00:00:00 2001 From: Abigail Emery Date: Thu, 8 Jan 2026 14:49:45 +0000 Subject: [PATCH] Add url option to cli interface --- src/blueapi/cli/cli.py | 264 ++++++++++++++++++++----------- src/blueapi/client/client.py | 3 + src/blueapi/client/rest.py | 5 +- src/blueapi/service/interface.py | 4 + src/blueapi/service/main.py | 18 ++- 5 files changed, 204 insertions(+), 90 deletions(-) diff --git a/src/blueapi/cli/cli.py b/src/blueapi/cli/cli.py index e70922b9c5..dedb721215 100644 --- a/src/blueapi/cli/cli.py +++ b/src/blueapi/cli/cli.py @@ -16,7 +16,7 @@ from bluesky_stomp.models import Broker from click.exceptions import ClickException from observability_utils.tracing import setup_tracing -from pydantic import ValidationError +from pydantic import HttpUrl, ValidationError from requests.exceptions import ConnectionError from blueapi import __version__, config @@ -44,23 +44,52 @@ LOGGER = logging.getLogger(__name__) +P = ParamSpec("P") +T = TypeVar("T") -@click.group( - invoke_without_command=True, context_settings={"auto_envvar_prefix": "BLUEAPI"} -) -@click.version_option(version=__version__, prog_name="blueapi") -@click.option( - "-c", "--config", type=Path, help="Path to configuration YAML file", multiple=True -) -@click.pass_context -def main(ctx: click.Context, config: Path | None | tuple[Path, ...]) -> None: - # if no command is supplied, run with the options passed - # Set umask to DLS standard - os.umask(stat.S_IWOTH) +def check_connection(func: Callable[P, T]) -> Callable[P, T]: + @wraps(func) + def wrapper(*args: P.args, **kwargs: P.kwargs) -> T: + try: + return func(*args, **kwargs) + except ConnectionError as ce: + raise ClickException( + "Failed to establish connection to blueapi server." + ) from ce + except BlueskyRemoteControlError as e: + if str(e) == "": + raise ClickException( + "Access denied. Please check your login status and try again." + ) from e + else: + raise e + return wrapper + + +def _default_config(ctx: click.Context) -> None: + ctx.ensure_object(dict) config_loader = ConfigLoader(ApplicationConfig) + + loaded_config: ApplicationConfig = config_loader.load() + + set_up_logging(loaded_config.logging) + + ctx.obj["config"] = loaded_config + + +def _load_config( + ctx: click.Context, + config: Path | None | tuple[Path, ...], +) -> None: + ctx.ensure_object(dict) + + config_loader = ConfigLoader(ApplicationConfig) + ctx.obj["custom_config"] = False + if config is not None: + ctx.obj["custom_config"] = True configs = (config,) if isinstance(config, Path) else config for path in configs: if path.exists(): @@ -68,13 +97,34 @@ def main(ctx: click.Context, config: Path | None | tuple[Path, ...]) -> None: else: raise FileNotFoundError(f"Cannot find file: {path}") - ctx.ensure_object(dict) loaded_config: ApplicationConfig = config_loader.load() - set_up_logging(loaded_config.logging) - ctx.obj["config"] = loaded_config + +@click.group( + invoke_without_command=True, context_settings={"auto_envvar_prefix": "BLUEAPI"} +) +@click.version_option(version=__version__, prog_name="blueapi") +@click.option( + "-c", + "--config", + type=Path, + help="Path to configuration YAML file", + multiple=True, +) +@click.pass_context +def main(ctx: click.Context, config: Path | None | tuple[Path, ...]) -> None: + # if no command is supplied, run with the options passed + + # Set umask to DLS standard + os.umask(stat.S_IWOTH) + + if config == (): + config = None + + _load_config(ctx, config) + if ctx.invoked_subcommand is None: print("Please invoke subcommand!") @@ -136,10 +186,10 @@ def config_schema(output: Path | None = None, update: bool = False) -> None: @main.command(name="serve") -@click.pass_obj -def start_application(obj: dict): +@click.pass_context +def start_application(ctx: click.Context): """Run a worker that accepts plans to run""" - config: ApplicationConfig = obj["config"] + config: ApplicationConfig = ctx.obj["config"] """Only import the service functions when starting the service or generating the schema, not the controller as a new FastAPI app will be started each time. @@ -154,6 +204,88 @@ def start_application(obj: dict): start(config) +@main.command(name="login") +@click.option( + "--url", + type=HttpUrl, + help="The url of the blueapi server you want to connect to.", + default=None, +) +@click.pass_obj +@check_connection +def login( + obj: dict, + url: HttpUrl | None, +) -> None: + """ + Authenticate with the blueapi using the OIDC (OpenID Connect) flow. + """ + config: ApplicationConfig = obj["config"] + + if url is not None: + if obj["custom_config"] is True: + LOGGER.warning( + "Custom config has been used. This will take precidence " + "over a provided url" + ) + else: + config.api.url = HttpUrl(url) + try: + auth: SessionManager = SessionManager.from_cache(config.auth_token_path) + access_token = auth.get_valid_access_token() + assert access_token + print("Logged in") + except Exception: + client = BlueapiClient.from_config(config) + oidc_config = client.get_oidc_config() + if oidc_config is None: + print("Server is not configured to use authentication!") + return + auth = SessionManager( + oidc_config, cache_manager=SessionCacheManager(config.auth_token_path) + ) + auth.start_device_flow() + + +@main.command(name="logout") +@click.option( + "--url", + type=HttpUrl, + help="The url of the blueapi server you want to connect to.", + default=None, +) +@click.pass_obj +def logout( + obj: dict, + url: HttpUrl | None, +) -> None: + """ + Logs out from the OIDC provider and removes the cached access token. + """ + config: ApplicationConfig = obj["config"] + + if url is not None: + if obj["custom_config"] is True: + LOGGER.warning( + "Custom config has been used. This will take precidence " + "over a provided url" + ) + else: + config.api.url = HttpUrl(url) + try: + auth: SessionManager = SessionManager.from_cache(config.auth_token_path) + auth.logout() + except FileNotFoundError: + print("Logged out") + except ValueError as e: + LOGGER.debug("Invalid login token: %s", e) + raise ClickException( + "Login token is not valid - remove before trying again" + ) from e + except Exception as e: + raise ClickException(f"Error logging out: {e}") from e + + @main.group() @click.option( "-o", @@ -161,8 +293,18 @@ def start_application(obj: dict): type=click.Choice([o.name.lower() for o in OutputFormat]), default="compact", ) +@click.option( + "--url", + type=HttpUrl, + help="The url of the blueapi server you want to connect to.", + default=None, +) @click.pass_context -def controller(ctx: click.Context, output: str) -> None: +def controller( + ctx: click.Context, + output: str, + url: HttpUrl | None, +) -> None: """Client utility for controlling and introspecting the worker""" setup_tracing("BlueAPICLI", OTLP_EXPORT_ENABLED) @@ -171,33 +313,25 @@ def controller(ctx: click.Context, output: str) -> None: return ctx.ensure_object(dict) - config: ApplicationConfig = ctx.obj["config"] ctx.obj["fmt"] = OutputFormat(output) - ctx.obj["client"] = BlueapiClient.from_config(config) + config: ApplicationConfig = ctx.obj["config"] -P = ParamSpec("P") -T = TypeVar("T") + if url is not None: + if ctx.obj["custom_config"] is True: + LOGGER.warning( + "Custom config has been used. This will take precidence " + "over a provided url" + ) + else: + config.api.url = HttpUrl(url) + tmp_client = BlueapiClient.from_config(config) + config.stomp = tmp_client.get_stomp_config() + ctx.obj["config"] = config -def check_connection(func: Callable[P, T]) -> Callable[P, T]: - @wraps(func) - def wrapper(*args: P.args, **kwargs: P.kwargs) -> T: - try: - return func(*args, **kwargs) - except ConnectionError as ce: - raise ClickException( - "Failed to establish connection to blueapi server." - ) from ce - except BlueskyRemoteControlError as e: - if str(e) == "": - raise ClickException( - "Access denied. Please check your login status and try again." - ) from e - else: - raise e - - return wrapper + set_up_logging(config.logging) + ctx.obj["client"] = BlueapiClient.from_config(config) @controller.command(name="plans") @@ -455,49 +589,3 @@ def get_python_env(obj: dict, name: str, source: SourceInfo) -> None: """ client: BlueapiClient = obj["client"] obj["fmt"].display(client.get_python_env(name=name, source=source)) - - -@main.command(name="login") -@click.pass_obj -@check_connection -def login(obj: dict) -> None: - """ - Authenticate with the blueapi using the OIDC (OpenID Connect) flow. - """ - config: ApplicationConfig = obj["config"] - try: - auth: SessionManager = SessionManager.from_cache(config.auth_token_path) - access_token = auth.get_valid_access_token() - assert access_token - print("Logged in") - except Exception: - client = BlueapiClient.from_config(config) - oidc_config = client.get_oidc_config() - if oidc_config is None: - print("Server is not configured to use authentication!") - return - auth = SessionManager( - oidc_config, cache_manager=SessionCacheManager(config.auth_token_path) - ) - auth.start_device_flow() - - -@main.command(name="logout") -@click.pass_obj -def logout(obj: dict) -> None: - """ - Logs out from the OIDC provider and removes the cached access token. - """ - config: ApplicationConfig = obj["config"] - try: - auth: SessionManager = SessionManager.from_cache(config.auth_token_path) - auth.logout() - except FileNotFoundError: - print("Logged out") - except ValueError as e: - LOGGER.debug("Invalid login token: %s", e) - raise ClickException( - "Login token is not valid - remove before trying again" - ) from e - except Exception as e: - raise ClickException(f"Error logging out: {e}") from e diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py index 0930e240a9..3a896c8e6d 100644 --- a/src/blueapi/client/client.py +++ b/src/blueapi/client/client.py @@ -71,6 +71,9 @@ def from_config(cls, config: ApplicationConfig) -> "BlueapiClient": else: return cls(rest) + def get_stomp_config(self): + return self._rest.get_stomp_config() + @start_as_current_span(TRACER) def get_plans(self) -> PlanResponse: """ diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py index 3ff119449e..33676aa4ad 100644 --- a/src/blueapi/client/rest.py +++ b/src/blueapi/client/rest.py @@ -10,7 +10,7 @@ ) from pydantic import BaseModel, TypeAdapter, ValidationError -from blueapi.config import RestConfig +from blueapi.config import RestConfig, StompConfig from blueapi.service.authentication import JWTAuth, SessionManager from blueapi.service.model import ( DeviceModel, @@ -215,6 +215,9 @@ def cancel_current_task( data={"new_state": state, "reason": reason}, ) + def get_stomp_config(self): + return self._request_and_deserialize("/config/stomp", StompConfig) + def get_environment(self) -> EnvironmentResponse: return self._request_and_deserialize("/environment", EnvironmentResponse) diff --git a/src/blueapi/service/interface.py b/src/blueapi/service/interface.py index 9bc8bcef85..4c4c2633bd 100644 --- a/src/blueapi/service/interface.py +++ b/src/blueapi/service/interface.py @@ -260,6 +260,10 @@ def get_oidc_config() -> OIDCConfig | None: return config().oidc +def get_stomp_config() -> StompConfig | None: + return config().stomp + + def get_python_env( name: str | None = None, source: SourceInfo | None = None ) -> PythonEnvironmentResponse: diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py index 9711f20053..d33533e79c 100644 --- a/src/blueapi/service/main.py +++ b/src/blueapi/service/main.py @@ -34,7 +34,7 @@ from starlette.responses import JSONResponse from super_state_machine.errors import TransitionError -from blueapi.config import ApplicationConfig, OIDCConfig +from blueapi.config import ApplicationConfig, OIDCConfig, StompConfig from blueapi.service import interface from blueapi.worker import TrackableTask, WorkerState from blueapi.worker.event import TaskStatusEnum @@ -253,6 +253,22 @@ def get_oidc_config( return config +@open_router.get( + "/config/stomp", + tags=[Tag.META], + responses={status.HTTP_204_NO_CONTENT: {"description": "No Stomp configured"}}, +) +@start_as_current_span(TRACER) +def get_stomp_config( + runner: Annotated[WorkerDispatcher, Depends(_runner)], +) -> StompConfig: + """Retrieve the stomp configuration for the server.""" + config = runner.run(interface.get_stomp_config) + if config is None: + raise HTTPException(status_code=status.HTTP_204_NO_CONTENT) + return config + + @secure_router.get("/plans", tags=[Tag.PLAN]) @start_as_current_span(TRACER) def get_plans(runner: Annotated[WorkerDispatcher, Depends(_runner)]) -> PlanResponse: