Coverage for lambda / analytics-presigned-url / handler.py: 100.00%
130 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""Presigned-URL Lambda for SageMaker Studio (analytics environment).
3Exchanges a Cognito-authorized API Gateway event for a time-limited
4``sagemaker:CreatePresignedDomainUrl`` link.
6Flow (happy path):
81. Extract ``claims = event["requestContext"]["authorizer"]["claims"]`` and
9 read ``cognito:username`` (falling back to ``username`` if the token
10 shape doesn't namespace Cognito claims).
112. Resolve the Studio ``DomainId`` from the ``STUDIO_DOMAIN_ID`` env var
12 (preferred) or by calling ``sagemaker:ListDomains`` as a fallback.
133. ``sagemaker:DescribeUserProfile`` -- if the profile doesn't exist yet
14 (``ValidationException`` / ``ResourceNotFound``), create it via
15 ``sagemaker:CreateUserProfile``. If the profile is still provisioning
16 (``Pending`` / ``Updating``), return HTTP 202 so the CLI can poll.
17 If the profile is ``Failed``, delete it and recreate.
184. Once the profile is ``InService``, ensure the per-user EFS access
19 point exists (lazy creation).
205. ``sagemaker:CreatePresignedDomainUrl`` and return the URL (HTTP 200).
22The Lambda never blocks waiting for profile provisioning. Instead it
23returns HTTP 202 ``{"status": "provisioning"}`` and the CLI retries
24every few seconds until it receives HTTP 200 with the presigned URL.
25This avoids hitting the API Gateway 29-second integration timeout.
27All failures funnel through the outer ``try/except`` in
28:func:`lambda_handler`; the response body is always JSON and never
29leaks an exception string.
31Environment variables (set by ``GCOAnalyticsStack._create_presigned_url_lambda``):
33- ``STUDIO_DOMAIN_ID`` -- the Studio domain ID (e.g. ``d-abc123xyz``).
34- ``SAGEMAKER_EXECUTION_ROLE_ARN`` -- passed on ``CreateUserProfile``.
35- ``STUDIO_EFS_ID`` -- used by ``_ensure_access_point``.
36- ``URL_EXPIRES_SECONDS`` -- default ``300`` (5 minutes).
37- ``SESSION_EXPIRES_SECONDS`` -- default ``43200`` (12 hours).
39The module-level boto3 clients (``sagemaker`` and ``efs``) are created
40once at cold start so repeat invocations inside a warm container reuse
41the same HTTP connection pools.
42"""
44from __future__ import annotations
46import hashlib
47import json
48import logging
49import os
50from typing import Any
52import boto3
53from botocore.exceptions import ClientError
55# <pyflowchart-code-diagram> BEGIN - auto-inserted, do not edit
56# Generated at (UTC): 2026-09-01T14:42:56Z
57# Generated from Git commit: 89b000378ed5a912a38c06f4feab2b029936ebcc
58# Flowchart(s) generated from this file:
59# * ``lambda_handler`` -> ``diagrams/code_diagrams/lambda/analytics-presigned-url/handler.lambda_handler.html``
60# (PNG: ``diagrams/code_diagrams/lambda/analytics-presigned-url/handler.lambda_handler.png``)
61# Regenerate with ``SOURCE_DATE_EPOCH=<unix-seconds> GCO_DIAGRAM_SOURCE_COMMIT=<40-char-sha> python diagrams/generate.py --code-only``.
62# <pyflowchart-code-diagram> END
65# ---------------------------------------------------------------------------
66# Module-level logger + boto3 clients
67# ---------------------------------------------------------------------------
68# Created once per cold start. boto3 clients are thread-safe for the
69# method calls this Lambda makes (list/describe/create presigned URL).
71logger = logging.getLogger(__name__)
72logger.setLevel(logging.INFO)
74sagemaker = boto3.client("sagemaker")
75efs = boto3.client("efs")
77# ---------------------------------------------------------------------------
78# Environment variables (read at module import, i.e. cold start)
79# ---------------------------------------------------------------------------
81STUDIO_DOMAIN_ID = os.environ.get("STUDIO_DOMAIN_ID", "")
82SAGEMAKER_EXECUTION_ROLE_ARN = os.environ.get("SAGEMAKER_EXECUTION_ROLE_ARN", "")
83STUDIO_EFS_ID = os.environ.get("STUDIO_EFS_ID", "")
84URL_EXPIRES_SECONDS = int(os.environ.get("URL_EXPIRES_SECONDS", "300"))
85SESSION_EXPIRES_SECONDS = int(os.environ.get("SESSION_EXPIRES_SECONDS", "43200"))
87# ---------------------------------------------------------------------------
88# Error tokens -- opaque strings returned in the ``error`` body key.
89# ---------------------------------------------------------------------------
90# Keep these short, stable, and free of implementation details so clients
91# can switch on them without parsing exception messages.
93_ERR_MISSING_CLAIM = "MissingCognitoClaim"
94_ERR_DOMAIN_NOT_FOUND = "SagemakerDomainNotFound"
95_ERR_GENERIC = "PresignedUrlGenerationFailed"
97# POSIX id derivation constants. 2**31 - 19 keeps the result comfortably
98# within the 32-bit signed-int range EFS accepts; the 100000 offset pushes
99# the uid/gid out of the system-user range reserved for the base image.
100_POSIX_ID_MODULUS = 2147483629
101_POSIX_ID_OFFSET = 100000
103# ==========================================================================
104# Pure helpers (unit-testable without mocking)
105# ==========================================================================
108def _parse_claims(event: dict[str, Any]) -> dict[str, Any]:
109 """Extract the Cognito claims dict from an API Gateway proxy event.
111 Returns an empty dict if ``event["requestContext"]["authorizer"]["claims"]``
112 is not present or not a dict. The caller decides whether an empty
113 result warrants a 401 -- see :func:`lambda_handler`.
114 """
115 if not isinstance(event, dict):
116 return {}
117 request_context = event.get("requestContext")
118 if not isinstance(request_context, dict):
119 return {}
120 authorizer = request_context.get("authorizer")
121 if not isinstance(authorizer, dict):
122 return {}
123 claims = authorizer.get("claims")
124 if not isinstance(claims, dict):
125 return {}
126 return claims
129def _derive_posix_ids(username: str) -> tuple[int, int]:
130 """Derive a deterministic POSIX ``(uid, gid)`` pair from a username.
132 Uses SHA-256 over the UTF-8 encoded username; the first four bytes
133 are interpreted as a big-endian unsigned int, reduced modulo
134 ``_POSIX_ID_MODULUS``, and shifted by ``_POSIX_ID_OFFSET`` so the
135 result is always ``>= 100000`` and fits within the 32-bit signed
136 range EFS expects.
138 The gid is always equal to the uid -- per-user home directories own
139 a single-user group, matching the ``0700`` permissions on
140 ``/home/<username>`` access points.
141 """
142 digest = hashlib.sha256(username.encode("utf-8")).digest()
143 raw = int.from_bytes(digest[:4], byteorder="big", signed=False)
144 uid = (raw % _POSIX_ID_MODULUS) + _POSIX_ID_OFFSET
145 return uid, uid
148def _format_success(url: str, expires: int) -> dict[str, Any]:
149 """Format an HTTP 200 API Gateway proxy response with the presigned URL."""
150 return {
151 "statusCode": 200,
152 "headers": {"Content-Type": "application/json"},
153 "body": json.dumps({"url": url, "expires_in": expires}),
154 }
157def _format_provisioning() -> dict[str, Any]:
158 """Format an HTTP 202 response indicating the profile is still provisioning.
160 The CLI polls on this status code until the profile reaches InService
161 and the Lambda returns HTTP 200 with the presigned URL.
162 """
163 return {
164 "statusCode": 202,
165 "headers": {"Content-Type": "application/json"},
166 "body": json.dumps({"status": "provisioning"}),
167 }
170def _format_error(status: int, token: str) -> dict[str, Any]:
171 """Format an HTTP error API Gateway proxy response.
173 ``token`` is one of the module-level ``_ERR_*`` constants; ``status``
174 is the HTTP status code. The body never contains an exception
175 message -- callers log the underlying exception via :data:`logger`
176 before returning.
177 """
178 return {
179 "statusCode": status,
180 "headers": {"Content-Type": "application/json"},
181 "body": json.dumps({"error": token}),
182 }
185# ==========================================================================
186# Effectful helpers (wrap boto3 calls)
187# ==========================================================================
190def _resolve_domain_id(domain_name: str) -> str | None:
191 """Return the Studio ``DomainId`` for ``domain_name`` or ``None``.
193 Paginates ``sagemaker:ListDomains`` by following ``NextToken``. A
194 ``None`` return signals the caller to emit HTTP 404 ``SagemakerDomainNotFound``.
195 """
196 next_token: str | None = None
197 while True:
198 kwargs: dict[str, Any] = {}
199 if next_token is not None:
200 kwargs["NextToken"] = next_token
201 response = sagemaker.list_domains(**kwargs)
202 for domain in response.get("Domains", []):
203 if domain.get("DomainName") == domain_name:
204 domain_id = domain.get("DomainId")
205 return str(domain_id) if domain_id is not None else None
206 next_token = response.get("NextToken")
207 if not next_token:
208 return None
211def _ensure_user_profile(domain_id: str, username: str, efs_id: str) -> str:
212 """Ensure a ``sagemaker:UserProfile`` exists for ``username``.
214 Returns the profile status:
216 * ``"InService"`` -- profile is ready; caller can mint a presigned URL.
217 * ``"Provisioning"`` -- profile was just created or is still starting;
218 caller should return HTTP 202 so the CLI can poll.
220 If the profile is in ``Failed`` state, it is deleted and recreated
221 so the next poll attempt finds a fresh ``Pending`` profile.
222 """
223 try:
224 resp = sagemaker.describe_user_profile(
225 DomainId=domain_id,
226 UserProfileName=username,
227 )
228 status = resp.get("Status", "")
230 if status == "InService":
231 return "InService"
233 if status == "Failed":
234 logger.warning(
235 "User profile %s is Failed (%s) -- deleting and recreating",
236 username,
237 resp.get("FailureReason", "unknown"),
238 )
239 try:
240 sagemaker.delete_user_profile(
241 DomainId=domain_id,
242 UserProfileName=username,
243 )
244 except ClientError as del_exc:
245 logger.warning("Could not delete failed profile %s: %s", username, del_exc)
246 # Create a fresh profile (may race with the delete; ResourceInUse
247 # is caught below).
248 _create_user_profile(domain_id, username, efs_id)
249 return "Provisioning"
251 # Pending / Updating / any other transient state.
252 logger.info("User profile %s status=%s, still provisioning", username, status)
253 return "Provisioning"
255 except ClientError as exc:
256 code = exc.response.get("Error", {}).get("Code", "")
257 if code not in {"ValidationException", "ResourceNotFound"}:
258 raise
260 # Profile doesn't exist -- create it.
261 _create_user_profile(domain_id, username, efs_id)
262 return "Provisioning"
265def _create_user_profile(domain_id: str, username: str, efs_id: str) -> None:
266 """Create a SageMaker user profile. Silently ignores ResourceInUse."""
267 try:
268 sagemaker.create_user_profile(
269 DomainId=domain_id,
270 UserProfileName=username,
271 UserSettings={
272 "ExecutionRole": SAGEMAKER_EXECUTION_ROLE_ARN,
273 "CustomFileSystemConfigs": [
274 {
275 "EFSFileSystemConfig": {
276 "FileSystemId": efs_id,
277 "FileSystemPath": f"/home/{username}",
278 }
279 }
280 ],
281 },
282 )
283 logger.info("Created user profile %s in domain %s", username, domain_id)
284 except ClientError as exc:
285 code = exc.response.get("Error", {}).get("Code", "")
286 if code != "ResourceInUse":
287 raise
288 logger.info("User profile %s already exists (ResourceInUse)", username)
291def _ensure_access_point(username: str, efs_id: str) -> str:
292 """Ensure a per-user EFS access point at ``/home/<username>`` exists.
294 Searches existing access points on ``efs_id`` for one whose
295 ``RootDirectory.Path`` equals ``/home/<username>``; creates one if
296 absent with a POSIX ``(uid, gid)`` derived from
297 :func:`_derive_posix_ids` and ``0700`` permissions.
299 Returns the ``AccessPointArn`` so the caller can associate it with
300 the user profile if needed.
301 """
302 uid, gid = _derive_posix_ids(username)
303 target_path = f"/home/{username}"
305 # Paginate describe_access_points. EFS returns at most 100 APs per
306 # page by default; we explicitly cap at 1000 to stay inside a single
307 # request for typical deployments.
308 next_token: str | None = None
309 while True:
310 kwargs: dict[str, Any] = {"FileSystemId": efs_id, "MaxResults": 1000}
311 if next_token is not None:
312 kwargs["NextToken"] = next_token
313 response = efs.describe_access_points(**kwargs)
314 for ap in response.get("AccessPoints", []):
315 root_dir = ap.get("RootDirectory", {})
316 if root_dir.get("Path") == target_path:
317 return str(ap.get("AccessPointArn", ""))
318 next_token = response.get("NextToken")
319 if not next_token:
320 break
322 created = efs.create_access_point(
323 FileSystemId=efs_id,
324 PosixUser={"Uid": uid, "Gid": gid},
325 RootDirectory={
326 "Path": target_path,
327 "CreationInfo": {
328 "OwnerUid": uid,
329 "OwnerGid": gid,
330 "Permissions": "0700",
331 },
332 },
333 Tags=[
334 {"Key": "gco:analytics:user", "Value": username},
335 {"Key": "gco:analytics:managed", "Value": "true"},
336 ],
337 )
338 return str(created.get("AccessPointArn", ""))
341# ==========================================================================
342# Entry point
343# ==========================================================================
346def lambda_handler(event: dict[str, Any], context: Any) -> dict[str, Any]:
347 """Exchange a Cognito-authorized event for a presigned Studio URL.
349 Returns:
351 * **HTTP 200** with ``{"url": "...", "expires_in": N}`` when the
352 user profile is ``InService`` and the presigned URL is ready.
353 * **HTTP 202** with ``{"status": "provisioning"}`` when the user
354 profile is still being created. The CLI should retry after a
355 few seconds.
356 * **HTTP 4xx / 5xx** with ``{"error": "<token>"}`` on failure.
358 The Lambda never blocks waiting for profile provisioning. The CLI
359 is responsible for polling until it receives HTTP 200.
360 """
361 try:
362 # Step 1: extract + validate the Cognito username claim.
363 claims = _parse_claims(event)
364 username = claims.get("cognito:username") or claims.get("username")
365 if not isinstance(username, str) or not username:
366 logger.warning("Request missing Cognito username claim")
367 return _format_error(401, _ERR_MISSING_CLAIM)
369 # Step 2: resolve the Studio DomainId.
370 domain_id = STUDIO_DOMAIN_ID if STUDIO_DOMAIN_ID else _resolve_domain_id("")
371 if not domain_id:
372 logger.error(
373 "SageMaker domain not found (STUDIO_DOMAIN_ID=%r)",
374 STUDIO_DOMAIN_ID,
375 )
376 return _format_error(404, _ERR_DOMAIN_NOT_FOUND)
378 # Step 3: describe-or-create the user profile.
379 profile_status = _ensure_user_profile(domain_id, username, STUDIO_EFS_ID)
380 if profile_status != "InService":
381 return _format_provisioning()
383 # Step 4: lazy per-user EFS access point.
384 _ensure_access_point(username, STUDIO_EFS_ID)
386 # Step 5: mint the presigned URL.
387 response = sagemaker.create_presigned_domain_url(
388 DomainId=domain_id,
389 UserProfileName=username,
390 SessionExpirationDurationInSeconds=SESSION_EXPIRES_SECONDS,
391 ExpiresInSeconds=URL_EXPIRES_SECONDS,
392 )
393 url = response.get("AuthorizedUrl", "")
394 return _format_success(url, URL_EXPIRES_SECONDS)
396 except Exception as exc: # noqa: BLE001 -- outer catch-all so every failure returns an opaque error token
397 # Log with exception info so CloudWatch captures the stack trace,
398 # but never leak the message to the HTTP response body.
399 logger.error("Presigned URL generation failed: %s", exc, exc_info=True)
400 return _format_error(500, _ERR_GENERIC)