From cc59d354bace9d6c1e54ed2428cea05dca361821 Mon Sep 17 00:00:00 2001 From: Abigail Emery Date: Thu, 8 Jan 2026 14:49:45 +0000 Subject: [PATCH 1/3] 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 4d974d35e2..afbf1e1e04 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 5aa44c533e..46f4e94482 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: From 82976acc128661d4926815734c8ca38d78a13042 Mon Sep 17 00:00:00 2001 From: Abigail Emery Date: Mon, 9 Feb 2026 16:18:00 +0000 Subject: [PATCH 2/3] Use rest client for getting stomp config --- src/blueapi/cli/cli.py | 31 ++++++++++++++----------------- src/blueapi/client/client.py | 3 --- 2 files changed, 14 insertions(+), 20 deletions(-) diff --git a/src/blueapi/cli/cli.py b/src/blueapi/cli/cli.py index afbf1e1e04..256845ff76 100644 --- a/src/blueapi/cli/cli.py +++ b/src/blueapi/cli/cli.py @@ -24,6 +24,7 @@ from blueapi.client import BlueapiClient from blueapi.client.event_bus import AnyEvent, BlueskyStreamingError, EventBusClient from blueapi.client.rest import ( + BlueapiRestClient, BlueskyRemoteControlError, InvalidParametersError, UnauthorisedAccessError, @@ -32,6 +33,7 @@ from blueapi.config import ( ApplicationConfig, ConfigLoader, + RestConfig, ) from blueapi.core import OTLP_EXPORT_ENABLED, DataEvent from blueapi.log import set_up_logging @@ -68,17 +70,6 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> T: 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, ...], @@ -314,7 +305,6 @@ def controller( ctx.ensure_object(dict) ctx.obj["fmt"] = OutputFormat(output) - config: ApplicationConfig = ctx.obj["config"] if url is not None: @@ -324,11 +314,18 @@ def controller( "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 + tmp_client = BlueapiRestClient(RestConfig(url=HttpUrl(url))) + stomp_config = None + try: + stomp_config = tmp_client.get_stomp_config() + config.stomp = stomp_config + config.api.url = url + ctx.obj["config"] = config + except Exception: + LOGGER.error( + "Server does not support --url access for " + "this command. Please use a config file.", + ) set_up_logging(config.logging) ctx.obj["client"] = BlueapiClient.from_config(config) diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py index 3a896c8e6d..0930e240a9 100644 --- a/src/blueapi/client/client.py +++ b/src/blueapi/client/client.py @@ -71,9 +71,6 @@ 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: """ From 4f0b4a42fba59ee25960229b7c9b73bb2029b2ed Mon Sep 17 00:00:00 2001 From: Abigail Emery Date: Mon, 9 Feb 2026 17:02:17 +0000 Subject: [PATCH 3/3] Bump rest api version --- docs/reference/openapi.yaml | 58 ++++++++++++++++++++++++++++++++++++- src/blueapi/service/main.py | 2 +- 2 files changed, 58 insertions(+), 2 deletions(-) diff --git a/docs/reference/openapi.yaml b/docs/reference/openapi.yaml index 499e8cc52d..a190decee5 100644 --- a/docs/reference/openapi.yaml +++ b/docs/reference/openapi.yaml @@ -1,5 +1,22 @@ components: schemas: + BasicAuthentication: + additionalProperties: false + description: User credentials for basic authentication + properties: + password: + description: Password to verify user's identity + title: Password + type: string + username: + description: Unique identifier for user + title: Username + type: string + required: + - username + - password + title: BasicAuthentication + type: object DeviceModel: additionalProperties: false description: Representation of a device @@ -224,6 +241,29 @@ components: - new_state title: StateChangeRequest type: object + StompConfig: + additionalProperties: false + description: Config for connecting to stomp broker + properties: + auth: + anyOf: + - $ref: '#/components/schemas/BasicAuthentication' + - type: 'null' + description: Auth information for communicating with STOMP broker, if required + enabled: + default: false + description: True if blueapi should connect to stomp for asynchronous event + publishing + title: Enabled + type: boolean + url: + default: tcp://localhost:61613 + format: uri + minLength: 1 + title: Url + type: string + title: StompConfig + type: object Task: additionalProperties: false description: Task that will run a plan @@ -377,7 +417,7 @@ info: name: Apache 2.0 url: https://www.apache.org/licenses/LICENSE-2.0.html title: BlueAPI Control - version: 1.1.2 + version: 1.2.0 openapi: 3.1.0 paths: /config/oidc: @@ -396,6 +436,22 @@ paths: summary: Get Oidc Config tags: - Meta + /config/stomp: + get: + description: Retrieve the stomp configuration for the server. + operationId: get_stomp_config_config_stomp_get + responses: + '200': + content: + application/json: + schema: + $ref: '#/components/schemas/StompConfig' + description: Successful Response + '204': + description: No Stomp configured + summary: Get Stomp Config + tags: + - Meta /devices: get: description: Retrieve information about all available devices. diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py index 46f4e94482..74a701c67e 100644 --- a/src/blueapi/service/main.py +++ b/src/blueapi/service/main.py @@ -58,7 +58,7 @@ from .runner import WorkerDispatcher #: API version to publish in OpenAPI schema -REST_API_VERSION = "1.1.2" +REST_API_VERSION = "1.2.0" LICENSE_INFO: dict[str, str] = { "name": "Apache 2.0",