Coverage for cli / commands / analytics_cmd.py: 100.00%
331 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +0000
1"""GCO analytics environment command group.
3Provides the ``gco analytics`` sub-commands:
5* ``enable`` / ``disable`` / ``status`` — flip the
6 ``analytics_environment.enabled`` toggle in ``cdk.json``.
7* ``users add`` / ``users list`` / ``users remove`` — manage Cognito
8 users against the auto-discovered pool id from ``gco-analytics``.
9* ``studio login`` — SRP-authenticate against Cognito and fetch a
10 SageMaker Studio presigned URL from ``/studio/login`` on the
11 existing ``gco-api-gateway``.
12* ``doctor`` — pre-flight checks before ``gco stacks deploy
13 gco-analytics``.
15The Click wiring mirrors ``stacks_cmd.py::fsx_cmd`` exactly. Every
16command delegates to helpers in :mod:`cli.analytics_user_mgmt` so the
17command layer stays thin and testable via ``click.testing.CliRunner``.
18"""
20from __future__ import annotations
22import json
23import os
24import sys
25import urllib.error
26from typing import Any
28import click
30from ..config import GCOConfig
31from ..output import confirm, emit_structured_document, get_output_formatter, prompt
33pass_config = click.make_pass_decorator(GCOConfig, ensure=True)
36def _stack_missing_message(project_name: str) -> str:
37 """Error text when the analytics stack isn't deployed (#139 project-scoped)."""
38 stack = f"{project_name}-analytics"
39 return (
40 f"{stack} stack not deployed — run `gco analytics enable` then `gco stacks deploy {stack}`"
41 )
44@click.group()
45@pass_config
46def analytics(config: Any) -> None:
47 """Manage the GCO analytics (SageMaker Studio + EMR) environment."""
50# ---------------------------------------------------------------------------
51# Toggle commands — enable / disable / status
52# ---------------------------------------------------------------------------
55@analytics.command("status")
56@pass_config
57def analytics_status(config: Any) -> None:
58 """Show the current analytics environment toggle state from cdk.json."""
59 from ..stacks import get_analytics_config
61 formatter = get_output_formatter(config)
62 try:
63 current = get_analytics_config()
64 formatter.print_info("Analytics environment config:")
65 formatter.print(current)
66 except Exception as exc: # surface every loader error
67 formatter.print_error(f"Failed to read analytics config: {exc}")
68 sys.exit(1)
71@analytics.command("enable")
72@click.option("--hyperpod", is_flag=True, help="Also enable SageMaker HyperPod job submission.")
73@click.option(
74 "--canvas",
75 is_flag=True,
76 help="Also enable the SageMaker Canvas no-code ML app.",
77)
78@click.option("--yes", "-y", is_flag=True, help="Skip confirmation.")
79@pass_config
80def analytics_enable(config: Any, hyperpod: bool, canvas: bool, yes: bool) -> None:
81 """Enable the analytics environment in cdk.json.
83 Flips ``analytics_environment.enabled`` to ``true``; ``--hyperpod``
84 additionally flips ``analytics_environment.hyperpod.enabled``, and
85 ``--canvas`` flips ``analytics_environment.canvas.enabled`` (which
86 attaches ``AmazonSageMakerCanvasFullAccess`` to the SageMaker
87 execution role). Prints the follow-up ``gco stacks deploy
88 gco-analytics`` command — does not deploy automatically.
89 """
90 from ..stacks import get_analytics_config, update_analytics_config
92 formatter = get_output_formatter(config)
94 if not yes:
95 formatter.print_info("Analytics environment will be enabled in cdk.json.")
96 if hyperpod:
97 formatter.print_info(" Hyperpod sub-toggle will also be enabled.")
98 if canvas:
99 formatter.print_info(" Canvas sub-toggle will also be enabled.")
100 confirm("\nEnable the analytics environment?", abort=True)
102 try:
103 current = get_analytics_config()
104 # Preserve everything the operator has set under ``hyperpod`` /
105 # ``canvas`` — the underlying helper replaces nested blocks
106 # wholesale, so we rebuild each sub-dict with only the field we own.
107 hyperpod_block = dict(current.get("hyperpod") or {})
108 if hyperpod:
109 hyperpod_block["enabled"] = True
110 hyperpod_block.setdefault("enabled", False)
112 canvas_block = dict(current.get("canvas") or {})
113 if canvas:
114 canvas_block["enabled"] = True
115 canvas_block.setdefault("enabled", False)
117 update_analytics_config(
118 {
119 "enabled": True,
120 "hyperpod": hyperpod_block,
121 "canvas": canvas_block,
122 }
123 )
124 formatter.print_success("Analytics environment enabled in cdk.json")
125 formatter.print_info(
126 f"Run `gco stacks deploy {config.project_name}-analytics` to apply changes"
127 )
128 except Exception as exc: # user-facing error from file I/O
129 formatter.print_error(f"Failed to enable analytics environment: {exc}")
130 sys.exit(1)
133@analytics.command("disable")
134@click.option("--yes", "-y", is_flag=True, help="Skip confirmation.")
135@pass_config
136def analytics_disable(config: Any, yes: bool) -> None:
137 """Disable the analytics environment in cdk.json.
139 Only flips ``analytics_environment.enabled`` to ``false``; the
140 ``hyperpod`` / ``canvas`` / ``cognito`` / ``efs`` sub-blocks are
141 left untouched so the operator's existing preferences survive a
142 disable/enable cycle.
143 """
144 from ..stacks import update_analytics_config
146 formatter = get_output_formatter(config)
148 if not yes:
149 formatter.print_warning("This will disable the analytics environment.")
150 formatter.print_warning(
151 "Existing SageMaker Studio / Cognito / EMR resources will be destroyed on next deploy."
152 )
153 confirm("Are you sure?", abort=True)
155 try:
156 update_analytics_config({"enabled": False})
157 formatter.print_success("Analytics environment disabled in cdk.json")
158 formatter.print_info(
159 f"Run `gco stacks destroy {config.project_name}-analytics` to tear down resources"
160 )
161 except Exception as exc: # user-facing error from file I/O
162 formatter.print_error(f"Failed to disable analytics environment: {exc}")
163 sys.exit(1)
166# ---------------------------------------------------------------------------
167# Users subgroup
168# ---------------------------------------------------------------------------
171@analytics.group("users")
172@pass_config
173def users_cmd(config: Any) -> None:
174 """Manage Cognito users who can sign in to SageMaker Studio."""
177def _require_cognito_pool_id(config: Any) -> tuple[str, str]:
178 """Return ``(pool_id, region)`` or exit with the documented error message."""
179 from ..analytics_user_mgmt import discover_cognito_pool_id
181 formatter = get_output_formatter(config)
182 region = config.api_gateway_region
183 pool_id = discover_cognito_pool_id(region, config.project_name)
184 if not pool_id:
185 formatter.print_error(_stack_missing_message(config.project_name))
186 sys.exit(1)
187 return pool_id, region
190@users_cmd.command("add")
191@click.option("--username", required=True, help="Cognito username to create.")
192@click.option("--email", help="Email address for the new user (optional).")
193@click.option(
194 "--no-email",
195 is_flag=True,
196 help="Suppress the Cognito welcome email (MessageAction=SUPPRESS).",
197)
198@click.option(
199 "--password",
200 envvar="GCO_STUDIO_PASSWORD",
201 help=(
202 "Set a permanent password via admin_set_user_password (also read "
203 "from $GCO_STUDIO_PASSWORD). Mutually exclusive with --generate-password."
204 ),
205)
206@click.option(
207 "--generate-password",
208 is_flag=True,
209 help=(
210 "Generate a strong random password, set it as permanent via "
211 "admin_set_user_password, and print it once. Mutually exclusive "
212 "with --password."
213 ),
214)
215@pass_config
216def users_add(
217 config: Any,
218 username: str,
219 email: str | None,
220 no_email: bool,
221 password: str | None,
222 generate_password: bool,
223) -> None:
224 """Create a Cognito user and print the temporary password exactly once.
226 When ``--password`` or ``--generate-password`` is passed, the user is
227 created and then has a permanent password set via
228 ``admin_set_user_password`` — this skips the ``NEW_PASSWORD_REQUIRED``
229 challenge on first login, so the resulting credentials work directly
230 with ``gco analytics studio login``.
231 """
232 from botocore.exceptions import ClientError
234 from ..analytics_user_mgmt import (
235 admin_create_user,
236 admin_set_user_password,
237 generate_strong_password,
238 )
240 formatter = get_output_formatter(config)
242 if password and generate_password:
243 formatter.print_error("--password and --generate-password are mutually exclusive")
244 sys.exit(1)
246 pool_id, region = _require_cognito_pool_id(config)
248 try:
249 _, temporary_password = admin_create_user(
250 pool_id=pool_id,
251 region=region,
252 username=username,
253 email=email,
254 suppress_email=no_email,
255 )
256 except ClientError as exc:
257 error_code = exc.response.get("Error", {}).get("Code", "Unknown")
258 formatter.print_error(f"Failed to create user {username}: {error_code}")
259 sys.exit(1)
261 if config.output_format == "table":
262 formatter.print_success(f"Created Cognito user: {username}")
264 # Password path — explicit or generated — takes precedence over the
265 # temporary-password path so the resulting credentials don't get
266 # blocked by NEW_PASSWORD_REQUIRED on first sign-in.
267 if password or generate_password:
268 final_password = password or generate_strong_password()
269 try:
270 admin_set_user_password(
271 pool_id=pool_id,
272 region=region,
273 username=username,
274 password=final_password,
275 permanent=True,
276 )
277 except ClientError as exc:
278 error_code = exc.response.get("Error", {}).get("Code", "Unknown")
279 formatter.print_error(
280 f"User {username} created, but setting the password "
281 f"failed: {error_code}. Retry with "
282 "`aws cognito-idp admin-set-user-password --permanent`."
283 )
284 sys.exit(1)
286 if config.output_format == "table":
287 if generate_password:
288 formatter.print_info(f"Generated password (printed exactly once): {final_password}")
289 else:
290 formatter.print_info(f"Password set (permanent) for {username}")
291 else:
292 result: dict[str, Any] = {
293 "created": True,
294 "username": username,
295 "email": email,
296 "user_pool_id": pool_id,
297 "region": region,
298 "password_state": "permanent", # nosec B105 # state label, not a credential
299 "password_generated": generate_password,
300 "password_source": "generated" if generate_password else "provided",
301 "password_permanent": True, # nosec B105 # boolean flag, not a credential
302 }
303 if generate_password:
304 result["password"] = final_password
305 formatter.print(result)
306 return
308 if config.output_format == "table":
309 if temporary_password:
310 formatter.print_info(f"Temporary password (printed exactly once): {temporary_password}")
311 else:
312 formatter.print_info(
313 "Cognito did not return a temporary password. "
314 "If --no-email was passed, set one via "
315 "`aws cognito-idp admin-set-user-password` "
316 "or re-run `gco analytics users add` with --password or --generate-password."
317 )
318 else:
319 result = {
320 "created": True,
321 "username": username,
322 "email": email,
323 "user_pool_id": pool_id,
324 "region": region,
325 "password_state": "temporary" if temporary_password else "not_returned",
326 "password_generated": False, # nosec B105 # boolean flag, not a credential
327 "password_source": "cognito" if temporary_password else "unavailable",
328 "password_permanent": False if temporary_password else None,
329 }
330 if temporary_password:
331 result["password"] = temporary_password
332 formatter.print(result)
335@users_cmd.command("list")
336@click.option("--as-json", "as_json", is_flag=True, help="Emit JSON instead of a table.")
337@pass_config
338def users_list(config: Any, as_json: bool) -> None:
339 """List Cognito users in the analytics user pool."""
340 from botocore.exceptions import ClientError
342 from ..analytics_user_mgmt import list_users as _list_users
344 formatter = get_output_formatter(config)
345 pool_id, region = _require_cognito_pool_id(config)
347 try:
348 users = _list_users(pool_id, region)
349 except ClientError as exc:
350 error_code = exc.response.get("Error", {}).get("Code", "Unknown")
351 formatter.print_error(f"Failed to list users: {error_code}")
352 sys.exit(1)
354 if as_json:
355 emit_structured_document(
356 users,
357 output_format="json",
358 rendered=json.dumps(users, indent=2),
359 )
360 return
361 formatter.print(users)
364@users_cmd.command("remove")
365@click.option("--username", required=True, help="Cognito username to remove.")
366@click.option("--yes", is_flag=True, help="Skip the confirmation prompt.")
367@pass_config
368def users_remove(config: Any, username: str, yes: bool) -> None:
369 """Delete a Cognito user from the analytics user pool."""
370 from botocore.exceptions import ClientError
372 from ..analytics_user_mgmt import admin_delete_user
374 formatter = get_output_formatter(config)
375 pool_id, region = _require_cognito_pool_id(config)
377 if not yes:
378 confirm(f"Delete Cognito user '{username}'?", abort=True)
380 try:
381 admin_delete_user(pool_id, region, username)
382 except ClientError as exc:
383 error_code = exc.response.get("Error", {}).get("Code", "Unknown")
384 formatter.print_error(f"Failed to delete user {username}: {error_code}")
385 sys.exit(1)
387 formatter.print_success(f"Deleted Cognito user: {username}")
390@users_cmd.command("set-password")
391@click.option("--username", required=True, help="Cognito username whose password to change.")
392@click.option(
393 "--password",
394 envvar="GCO_STUDIO_PASSWORD",
395 help=(
396 "New password (also read from $GCO_STUDIO_PASSWORD; prompted "
397 "otherwise). Mutually exclusive with --generate-password."
398 ),
399)
400@click.option(
401 "--generate-password",
402 is_flag=True,
403 help=(
404 "Generate a strong random password, set it, and print it once. "
405 "Mutually exclusive with --password."
406 ),
407)
408@click.option(
409 "--temporary",
410 is_flag=True,
411 help=(
412 "Set the password as temporary so the user is forced to change "
413 "it on first login (Permanent=false). Default is permanent."
414 ),
415)
416@click.option("--yes", "-y", is_flag=True, help="Skip the confirmation prompt.")
417@pass_config
418def users_set_password(
419 config: Any,
420 username: str,
421 password: str | None,
422 generate_password: bool,
423 temporary: bool,
424 yes: bool,
425) -> None:
426 """Change a Cognito user's password via AdminSetUserPassword.
428 By default the new password is marked ``Permanent=true`` so the
429 user can sign in directly with ``gco analytics studio login``
430 without the ``NEW_PASSWORD_REQUIRED`` challenge. Pass
431 ``--temporary`` to require the user to choose their own password
432 on first sign-in.
433 """
434 from botocore.exceptions import ClientError
436 from ..analytics_user_mgmt import admin_set_user_password, generate_strong_password
438 formatter = get_output_formatter(config)
440 if password and generate_password:
441 formatter.print_error("--password and --generate-password are mutually exclusive")
442 sys.exit(1)
444 pool_id, region = _require_cognito_pool_id(config)
446 if generate_password:
447 new_password = generate_strong_password()
448 elif password is not None:
449 new_password = password
450 else:
451 prompt_kwargs: dict[str, Any] = {
452 "hide_input": True,
453 "confirmation_prompt": True,
454 }
455 if config.output_format != "table":
456 prompt_kwargs["err"] = True
457 new_password = prompt("New password", **prompt_kwargs)
459 if not yes:
460 qualifier = "temporary" if temporary else "permanent"
461 confirm_kwargs: dict[str, Any] = {"abort": True}
462 if config.output_format != "table":
463 confirm_kwargs["err"] = True
464 confirm(
465 f"Set a new {qualifier} password for Cognito user '{username}'?",
466 **confirm_kwargs,
467 )
469 try:
470 admin_set_user_password(
471 pool_id=pool_id,
472 region=region,
473 username=username,
474 password=new_password,
475 permanent=not temporary,
476 )
477 except ClientError as exc:
478 error_code = exc.response.get("Error", {}).get("Code", "Unknown")
479 formatter.print_error(f"Failed to set password for {username}: {error_code}")
480 sys.exit(1)
482 qualifier = "temporary" if temporary else "permanent"
483 if config.output_format == "table":
484 formatter.print_success(f"Password set ({qualifier}) for {username}")
485 if generate_password:
486 formatter.print_info(f"Generated password (printed exactly once): {new_password}")
487 else:
488 result: dict[str, Any] = {
489 "password_set": True, # nosec B105 # boolean flag, not a credential
490 "username": username,
491 "user_pool_id": pool_id,
492 "region": region,
493 "password_state": qualifier,
494 "password_generated": generate_password,
495 "password_source": "generated" if generate_password else "provided",
496 "password_permanent": not temporary,
497 }
498 if generate_password:
499 result["password"] = new_password
500 formatter.print(result)
503# ---------------------------------------------------------------------------
504# Studio login subgroup
505# ---------------------------------------------------------------------------
508@analytics.group("studio")
509@pass_config
510def studio_cmd(config: Any) -> None:
511 """SageMaker Studio helpers (login, etc.)."""
514@studio_cmd.command("login")
515@click.option("--username", required=True, help="Cognito username to sign in with.")
516@click.option(
517 "--password",
518 envvar="GCO_STUDIO_PASSWORD",
519 help="Password (also read from $GCO_STUDIO_PASSWORD; prompted otherwise).",
520)
521@click.option("--api-url", help="Override the API Gateway base URL (otherwise auto-discovered).")
522@click.option("--open", "open_browser", is_flag=True, help="Open the URL in the default browser.")
523@pass_config
524def studio_login(
525 config: Any,
526 username: str,
527 password: str | None,
528 api_url: str | None,
529 open_browser: bool,
530) -> None:
531 """Sign in to SageMaker Studio via Cognito SRP and print the presigned URL."""
532 from botocore.exceptions import ClientError
534 from ..analytics_user_mgmt import (
535 discover_api_endpoint,
536 discover_cognito_client_id,
537 discover_cognito_pool_id,
538 fetch_studio_url,
539 srp_authenticate,
540 )
542 formatter = get_output_formatter(config)
543 region = config.api_gateway_region
544 project_name = config.project_name
546 pool_id = discover_cognito_pool_id(region, project_name)
547 client_id = discover_cognito_client_id(region, project_name)
548 if not pool_id or not client_id:
549 formatter.print_error(_stack_missing_message(config.project_name))
550 sys.exit(1)
552 api_base = (
553 api_url
554 or discover_api_endpoint(region, project_name)
555 or os.environ.get("GCO_API_GATEWAY_URL")
556 )
557 if not api_base:
558 formatter.print_error(
559 "Could not resolve API Gateway endpoint — pass --api-url or deploy gco-api-gateway."
560 )
561 sys.exit(1)
563 if password is None:
564 password = prompt("Password", hide_input=True)
566 try:
567 tokens = srp_authenticate(
568 pool_id=pool_id,
569 client_id=client_id,
570 username=username,
571 password=password,
572 region=region,
573 )
574 except ClientError as exc:
575 error_code = exc.response.get("Error", {}).get("Code", "Unknown")
576 formatter.print_error(f"Cognito authentication failed: {error_code}")
577 sys.exit(1)
579 id_token = tokens.get("IdToken")
580 if not id_token:
581 formatter.print_error("Cognito authentication failed: no IdToken returned")
582 sys.exit(1)
584 try:
585 # Poll until the Lambda returns HTTP 200 with the presigned URL.
586 # First-time logins trigger user-profile provisioning (30-60s);
587 # the Lambda returns HTTP 202 while the profile is pending.
588 import time as _time
590 max_wait = 120 # seconds
591 poll_interval = 5 # seconds
592 elapsed = 0
593 url = ""
595 while elapsed < max_wait:
596 # Only the URL matters here; the presigned expiry is reported by
597 # `analytics studio-url`, not by this open-in-browser path.
598 url, _expires_in, _ = fetch_studio_url(api_base, id_token)
599 if url:
600 break
601 # 202 -- profile still provisioning.
602 if elapsed == 0:
603 click.echo(" Waiting for user profile to provision...", nl=False)
604 click.echo(".", nl=False)
605 _time.sleep(poll_interval)
606 elapsed += poll_interval
608 if elapsed > 0 and url:
609 click.echo(" ready")
610 elif not url:
611 click.echo("")
612 formatter.print_error(
613 f"User profile did not become ready within {max_wait}s. Try again in a minute."
614 )
615 sys.exit(2)
616 except urllib.error.HTTPError as exc:
617 correlation_id = exc.headers.get("x-amzn-RequestId") if exc.headers else "N/A"
618 formatter.print_error(
619 f"login failed: HTTP {exc.code}, correlation_id={correlation_id or 'N/A'}"
620 )
621 sys.exit(2)
622 except urllib.error.URLError as exc:
623 formatter.print_error(f"login failed: network error: {exc.reason!r}")
624 sys.exit(2)
625 except ValueError as exc:
626 formatter.print_error(f"login failed: {exc}")
627 sys.exit(2)
629 # Print the URL on its own line for pipe-friendliness.
630 click.echo(url)
631 if open_browser:
632 click.launch(url)
635# ---------------------------------------------------------------------------
636# Doctor subcommand
637# ---------------------------------------------------------------------------
640@analytics.command("doctor")
641@pass_config
642def analytics_doctor(config: Any) -> None:
643 """Run pre-flight checks before `gco stacks deploy gco-analytics`.
645 Exits non-zero on any failing check. Each check prints ``✓``/``✗``
646 plus a short remediation line so the operator knows exactly what
647 to fix.
648 """
649 from ..analytics_user_mgmt import (
650 check_ssm_parameter,
651 check_stack_complete,
652 scan_orphan_analytics_resources,
653 )
654 from ..config import _load_cdk_json
655 from ..stacks import _find_cdk_json
657 formatter = get_output_formatter(config)
658 any_failed = False
660 def _emit(name: str, ok: bool, remediation: str) -> None:
661 nonlocal any_failed
662 if ok:
663 click.echo(f" ✓ {name}")
664 else:
665 any_failed = True
666 click.echo(f" ✗ {name}")
667 if remediation:
668 click.echo(f" → {remediation}")
670 # 1. cdk.json parses
671 cdk_json_path = _find_cdk_json()
672 if cdk_json_path is None:
673 _emit(
674 "cdk.json present",
675 False,
676 "run `gco analytics doctor` from the project root (cdk.json not found).",
677 )
678 else:
679 try:
680 with open(cdk_json_path, encoding="utf-8") as fh:
681 json.load(fh)
682 _emit("cdk.json parses as JSON", True, "")
683 except json.JSONDecodeError as exc:
684 _emit(
685 "cdk.json parses as JSON",
686 False,
687 f"fix malformed JSON at {cdk_json_path}: {exc.msg} (line {exc.lineno})",
688 )
690 # 2. Prerequisite stacks healthy
691 for region, stack_name in (
692 (config.global_region, f"{config.project_name}-global"),
693 (config.api_gateway_region, f"{config.project_name}-api-gateway"),
694 ):
695 ok, remediation = check_stack_complete(region, stack_name)
696 _emit(
697 f"{stack_name} is CREATE_COMPLETE",
698 ok,
699 remediation or f"deploy with `gco stacks deploy {stack_name}`",
700 )
702 cdk_regions = _load_cdk_json()
703 regional_regions = cdk_regions.get("regional", []) if isinstance(cdk_regions, dict) else []
704 for region in regional_regions:
705 stack_name = f"{config.project_name}-{region}"
706 ok, remediation = check_stack_complete(region, stack_name)
707 _emit(
708 f"{stack_name} is CREATE_COMPLETE",
709 ok,
710 remediation or f"deploy with `gco stacks deploy {stack_name}`",
711 )
713 # 3. SSM cluster-shared-bucket parameters exist
714 from gco.stacks.constants import cluster_shared_ssm_parameter_prefix
716 ssm_prefix = cluster_shared_ssm_parameter_prefix(config.project_name)
717 for suffix in ("name", "arn", "region"):
718 param = f"{ssm_prefix}/{suffix}"
719 ok, remediation = check_ssm_parameter(config.global_region, param)
720 _emit(
721 f"SSM parameter {param} exists",
722 ok,
723 remediation and f"deploy {config.project_name}-global first ({remediation})",
724 )
726 # 4. No orphaned retained analytics resources
727 orphan_cmds = scan_orphan_analytics_resources(config.api_gateway_region)
728 _emit(
729 "no orphaned retained analytics resources",
730 not orphan_cmds,
731 "; ".join(orphan_cmds) if orphan_cmds else "",
732 )
734 if any_failed:
735 formatter.print_error("Doctor checks failed — see remediation lines above.")
736 sys.exit(1)
737 formatter.print_success("All pre-flight checks passed.")