Coverage for cli/commands/analytics_cmd.py: 89.74%
310 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-30 21:22 +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 get_output_formatter
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: # noqa: BLE001 — 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: 96 ↛ 98line 96 didn't jump to line 98 because the condition on line 96 was always true
97 formatter.print_info(" Hyperpod sub-toggle will also be enabled.")
98 if canvas: 98 ↛ 99line 98 didn't jump to line 99 because the condition on line 98 was never true
99 formatter.print_info(" Canvas sub-toggle will also be enabled.")
100 click.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: # noqa: BLE001 — 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 click.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: # noqa: BLE001 — 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 formatter.print_success(f"Created Cognito user: {username}")
263 # Password path — explicit or generated — takes precedence over the
264 # temporary-password path so the resulting credentials don't get
265 # blocked by NEW_PASSWORD_REQUIRED on first sign-in.
266 if password or generate_password:
267 final_password = password or generate_strong_password()
268 try:
269 admin_set_user_password(
270 pool_id=pool_id,
271 region=region,
272 username=username,
273 password=final_password,
274 permanent=True,
275 )
276 except ClientError as exc:
277 error_code = exc.response.get("Error", {}).get("Code", "Unknown")
278 formatter.print_error(
279 f"User {username} created, but setting the password "
280 f"failed: {error_code}. Retry with "
281 "`aws cognito-idp admin-set-user-password --permanent`."
282 )
283 sys.exit(1)
285 if generate_password:
286 formatter.print_info(f"Generated password (printed exactly once): {final_password}")
287 else:
288 formatter.print_info(f"Password set (permanent) for {username}")
289 return
291 if temporary_password: 291 ↛ 292line 291 didn't jump to line 292 because the condition on line 291 was never true
292 formatter.print_info(f"Temporary password (printed exactly once): {temporary_password}")
293 else:
294 formatter.print_info(
295 "Cognito did not return a temporary password. "
296 "If --no-email was passed, set one via "
297 "`aws cognito-idp admin-set-user-password` "
298 "or re-run `gco analytics users add` with --password or --generate-password."
299 )
302@users_cmd.command("list")
303@click.option("--as-json", "as_json", is_flag=True, help="Emit JSON instead of a table.")
304@pass_config
305def users_list(config: Any, as_json: bool) -> None:
306 """List Cognito users in the analytics user pool."""
307 from botocore.exceptions import ClientError
309 from ..analytics_user_mgmt import list_users as _list_users
311 formatter = get_output_formatter(config)
312 pool_id, region = _require_cognito_pool_id(config)
314 try:
315 users = _list_users(pool_id, region)
316 except ClientError as exc:
317 error_code = exc.response.get("Error", {}).get("Code", "Unknown")
318 formatter.print_error(f"Failed to list users: {error_code}")
319 sys.exit(1)
321 if as_json:
322 print(json.dumps(users, indent=2))
323 return
324 formatter.print(users)
327@users_cmd.command("remove")
328@click.option("--username", required=True, help="Cognito username to remove.")
329@click.option("--yes", is_flag=True, help="Skip the confirmation prompt.")
330@pass_config
331def users_remove(config: Any, username: str, yes: bool) -> None:
332 """Delete a Cognito user from the analytics user pool."""
333 from botocore.exceptions import ClientError
335 from ..analytics_user_mgmt import admin_delete_user
337 formatter = get_output_formatter(config)
338 pool_id, region = _require_cognito_pool_id(config)
340 if not yes:
341 click.confirm(f"Delete Cognito user '{username}'?", abort=True)
343 try:
344 admin_delete_user(pool_id, region, username)
345 except ClientError as exc:
346 error_code = exc.response.get("Error", {}).get("Code", "Unknown")
347 formatter.print_error(f"Failed to delete user {username}: {error_code}")
348 sys.exit(1)
350 formatter.print_success(f"Deleted Cognito user: {username}")
353@users_cmd.command("set-password")
354@click.option("--username", required=True, help="Cognito username whose password to change.")
355@click.option(
356 "--password",
357 envvar="GCO_STUDIO_PASSWORD",
358 help=(
359 "New password (also read from $GCO_STUDIO_PASSWORD; prompted "
360 "otherwise). Mutually exclusive with --generate-password."
361 ),
362)
363@click.option(
364 "--generate-password",
365 is_flag=True,
366 help=(
367 "Generate a strong random password, set it, and print it once. "
368 "Mutually exclusive with --password."
369 ),
370)
371@click.option(
372 "--temporary",
373 is_flag=True,
374 help=(
375 "Set the password as temporary so the user is forced to change "
376 "it on first login (Permanent=false). Default is permanent."
377 ),
378)
379@click.option("--yes", "-y", is_flag=True, help="Skip the confirmation prompt.")
380@pass_config
381def users_set_password(
382 config: Any,
383 username: str,
384 password: str | None,
385 generate_password: bool,
386 temporary: bool,
387 yes: bool,
388) -> None:
389 """Change a Cognito user's password via AdminSetUserPassword.
391 By default the new password is marked ``Permanent=true`` so the
392 user can sign in directly with ``gco analytics studio login``
393 without the ``NEW_PASSWORD_REQUIRED`` challenge. Pass
394 ``--temporary`` to require the user to choose their own password
395 on first sign-in.
396 """
397 from botocore.exceptions import ClientError
399 from ..analytics_user_mgmt import admin_set_user_password, generate_strong_password
401 formatter = get_output_formatter(config)
403 if password and generate_password:
404 formatter.print_error("--password and --generate-password are mutually exclusive")
405 sys.exit(1)
407 pool_id, region = _require_cognito_pool_id(config)
409 if generate_password:
410 new_password = generate_strong_password()
411 elif password is not None: 411 ↛ 414line 411 didn't jump to line 414 because the condition on line 411 was always true
412 new_password = password
413 else:
414 new_password = click.prompt(
415 "New password",
416 hide_input=True,
417 confirmation_prompt=True,
418 )
420 if not yes: 420 ↛ 421line 420 didn't jump to line 421 because the condition on line 420 was never true
421 qualifier = "temporary" if temporary else "permanent"
422 click.confirm(
423 f"Set a new {qualifier} password for Cognito user '{username}'?",
424 abort=True,
425 )
427 try:
428 admin_set_user_password(
429 pool_id=pool_id,
430 region=region,
431 username=username,
432 password=new_password,
433 permanent=not temporary,
434 )
435 except ClientError as exc:
436 error_code = exc.response.get("Error", {}).get("Code", "Unknown")
437 formatter.print_error(f"Failed to set password for {username}: {error_code}")
438 sys.exit(1)
440 qualifier = "temporary" if temporary else "permanent"
441 formatter.print_success(f"Password set ({qualifier}) for {username}")
442 if generate_password:
443 formatter.print_info(f"Generated password (printed exactly once): {new_password}")
446# ---------------------------------------------------------------------------
447# Studio login subgroup
448# ---------------------------------------------------------------------------
451@analytics.group("studio")
452@pass_config
453def studio_cmd(config: Any) -> None:
454 """SageMaker Studio helpers (login, etc.)."""
457@studio_cmd.command("login")
458@click.option("--username", required=True, help="Cognito username to sign in with.")
459@click.option(
460 "--password",
461 envvar="GCO_STUDIO_PASSWORD",
462 help="Password (also read from $GCO_STUDIO_PASSWORD; prompted otherwise).",
463)
464@click.option("--api-url", help="Override the API Gateway base URL (otherwise auto-discovered).")
465@click.option("--open", "open_browser", is_flag=True, help="Open the URL in the default browser.")
466@pass_config
467def studio_login(
468 config: Any,
469 username: str,
470 password: str | None,
471 api_url: str | None,
472 open_browser: bool,
473) -> None:
474 """Sign in to SageMaker Studio via Cognito SRP and print the presigned URL."""
475 from botocore.exceptions import ClientError
477 from ..analytics_user_mgmt import (
478 discover_api_endpoint,
479 discover_cognito_client_id,
480 discover_cognito_pool_id,
481 fetch_studio_url,
482 srp_authenticate,
483 )
485 formatter = get_output_formatter(config)
486 region = config.api_gateway_region
487 project_name = config.project_name
489 pool_id = discover_cognito_pool_id(region, project_name)
490 client_id = discover_cognito_client_id(region, project_name)
491 if not pool_id or not client_id: 491 ↛ 492line 491 didn't jump to line 492 because the condition on line 491 was never true
492 formatter.print_error(_stack_missing_message(config.project_name))
493 sys.exit(1)
495 api_base = (
496 api_url
497 or discover_api_endpoint(region, project_name)
498 or os.environ.get("GCO_API_GATEWAY_URL")
499 )
500 if not api_base:
501 formatter.print_error(
502 "Could not resolve API Gateway endpoint — pass --api-url or deploy gco-api-gateway."
503 )
504 sys.exit(1)
506 if password is None:
507 password = click.prompt("Password", hide_input=True)
509 try:
510 tokens = srp_authenticate(
511 pool_id=pool_id,
512 client_id=client_id,
513 username=username,
514 password=password,
515 region=region,
516 )
517 except ClientError as exc:
518 error_code = exc.response.get("Error", {}).get("Code", "Unknown")
519 formatter.print_error(f"Cognito authentication failed: {error_code}")
520 sys.exit(1)
522 id_token = tokens.get("IdToken")
523 if not id_token:
524 formatter.print_error("Cognito authentication failed: no IdToken returned")
525 sys.exit(1)
527 try:
528 # Poll until the Lambda returns HTTP 200 with the presigned URL.
529 # First-time logins trigger user-profile provisioning (30-60s);
530 # the Lambda returns HTTP 202 while the profile is pending.
531 import time as _time
533 max_wait = 120 # seconds
534 poll_interval = 5 # seconds
535 elapsed = 0
536 url = ""
537 expires_in = 0
539 while elapsed < max_wait: 539 ↛ 550line 539 didn't jump to line 550 because the condition on line 539 was always true
540 url, expires_in, _ = fetch_studio_url(api_base, id_token)
541 if url: 541 ↛ 544line 541 didn't jump to line 544 because the condition on line 541 was always true
542 break
543 # 202 -- profile still provisioning.
544 if elapsed == 0:
545 click.echo(" Waiting for user profile to provision...", nl=False)
546 click.echo(".", nl=False)
547 _time.sleep(poll_interval)
548 elapsed += poll_interval
550 if elapsed > 0 and url: 550 ↛ 551line 550 didn't jump to line 551 because the condition on line 550 was never true
551 click.echo(" ready")
552 elif not url: 552 ↛ 553line 552 didn't jump to line 553 because the condition on line 552 was never true
553 click.echo("")
554 formatter.print_error(
555 f"User profile did not become ready within {max_wait}s. Try again in a minute."
556 )
557 sys.exit(2)
558 except urllib.error.HTTPError as exc:
559 correlation_id = exc.headers.get("x-amzn-RequestId") if exc.headers else "N/A"
560 formatter.print_error(
561 f"login failed: HTTP {exc.code}, correlation_id={correlation_id or 'N/A'}"
562 )
563 sys.exit(2)
564 except urllib.error.URLError as exc:
565 formatter.print_error(f"login failed: network error: {exc.reason!r}")
566 sys.exit(2)
567 except ValueError as exc:
568 formatter.print_error(f"login failed: {exc}")
569 sys.exit(2)
571 # Print the URL on its own line for pipe-friendliness.
572 click.echo(url)
573 if open_browser: 573 ↛ 574line 573 didn't jump to line 574 because the condition on line 573 was never true
574 click.launch(url)
577# ---------------------------------------------------------------------------
578# Doctor subcommand
579# ---------------------------------------------------------------------------
582@analytics.command("doctor")
583@pass_config
584def analytics_doctor(config: Any) -> None:
585 """Run pre-flight checks before `gco stacks deploy gco-analytics`.
587 Exits non-zero on any failing check. Each check prints ``✓``/``✗``
588 plus a short remediation line so the operator knows exactly what
589 to fix.
590 """
591 from ..analytics_user_mgmt import (
592 check_ssm_parameter,
593 check_stack_complete,
594 scan_orphan_analytics_resources,
595 )
596 from ..config import _load_cdk_json
597 from ..stacks import _find_cdk_json
599 formatter = get_output_formatter(config)
600 any_failed = False
602 def _emit(name: str, ok: bool, remediation: str) -> None:
603 nonlocal any_failed
604 if ok:
605 click.echo(f" ✓ {name}")
606 else:
607 any_failed = True
608 click.echo(f" ✗ {name}")
609 if remediation: 609 ↛ exitline 609 didn't return from function '_emit' because the condition on line 609 was always true
610 click.echo(f" → {remediation}")
612 # 1. cdk.json parses
613 cdk_json_path = _find_cdk_json()
614 if cdk_json_path is None:
615 _emit(
616 "cdk.json present",
617 False,
618 "run `gco analytics doctor` from the project root (cdk.json not found).",
619 )
620 else:
621 try:
622 with open(cdk_json_path, encoding="utf-8") as fh:
623 json.load(fh)
624 _emit("cdk.json parses as JSON", True, "")
625 except json.JSONDecodeError as exc:
626 _emit(
627 "cdk.json parses as JSON",
628 False,
629 f"fix malformed JSON at {cdk_json_path}: {exc.msg} (line {exc.lineno})",
630 )
632 # 2. Prerequisite stacks healthy
633 for region, stack_name in (
634 (config.global_region, f"{config.project_name}-global"),
635 (config.api_gateway_region, f"{config.project_name}-api-gateway"),
636 ):
637 ok, remediation = check_stack_complete(region, stack_name)
638 _emit(
639 f"{stack_name} is CREATE_COMPLETE",
640 ok,
641 remediation or f"deploy with `gco stacks deploy {stack_name}`",
642 )
644 cdk_regions = _load_cdk_json()
645 regional_regions = cdk_regions.get("regional", []) if isinstance(cdk_regions, dict) else []
646 for region in regional_regions:
647 stack_name = f"{config.project_name}-{region}"
648 ok, remediation = check_stack_complete(region, stack_name)
649 _emit(
650 f"{stack_name} is CREATE_COMPLETE",
651 ok,
652 remediation or f"deploy with `gco stacks deploy {stack_name}`",
653 )
655 # 3. SSM cluster-shared-bucket parameters exist
656 from gco.stacks.constants import cluster_shared_ssm_parameter_prefix
658 ssm_prefix = cluster_shared_ssm_parameter_prefix(config.project_name)
659 for suffix in ("name", "arn", "region"):
660 param = f"{ssm_prefix}/{suffix}"
661 ok, remediation = check_ssm_parameter(config.global_region, param)
662 _emit(
663 f"SSM parameter {param} exists",
664 ok,
665 remediation and f"deploy {config.project_name}-global first ({remediation})",
666 )
668 # 4. No orphaned retained analytics resources
669 orphan_cmds = scan_orphan_analytics_resources(config.api_gateway_region)
670 _emit(
671 "no orphaned retained analytics resources",
672 not orphan_cmds,
673 "; ".join(orphan_cmds) if orphan_cmds else "",
674 )
676 if any_failed:
677 formatter.print_error("Doctor checks failed — see remediation lines above.")
678 sys.exit(1)
679 formatter.print_success("All pre-flight checks passed.")