Coverage for lambda / analytics-presigned-url / handler.py: 100.00%
130 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +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-18T02:11:36Z
57# Generated from Git commit: b8faa9689385cea16155a285a7f70cf6d488e512
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: object) -> 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`.
115 ``event`` is typed ``object`` because every layer of the path is
116 shape-checked here, the event itself included; the helper makes no
117 assumption about what API Gateway (or a test) hands it.
118 """
119 if not isinstance(event, dict):
120 return {}
121 request_context = event.get("requestContext")
122 if not isinstance(request_context, dict):
123 return {}
124 authorizer = request_context.get("authorizer")
125 if not isinstance(authorizer, dict):
126 return {}
127 claims = authorizer.get("claims")
128 if not isinstance(claims, dict):
129 return {}
130 return claims
133def _derive_posix_ids(username: str) -> tuple[int, int]:
134 """Derive a deterministic POSIX ``(uid, gid)`` pair from a username.
136 Uses SHA-256 over the UTF-8 encoded username; the first four bytes
137 are interpreted as a big-endian unsigned int, reduced modulo
138 ``_POSIX_ID_MODULUS``, and shifted by ``_POSIX_ID_OFFSET`` so the
139 result is always ``>= 100000`` and fits within the 32-bit signed
140 range EFS expects.
142 The gid is always equal to the uid -- per-user home directories own
143 a single-user group, matching the ``0700`` permissions on
144 ``/home/<username>`` access points.
145 """
146 digest = hashlib.sha256(username.encode("utf-8")).digest()
147 raw = int.from_bytes(digest[:4], byteorder="big", signed=False)
148 uid = (raw % _POSIX_ID_MODULUS) + _POSIX_ID_OFFSET
149 return uid, uid
152def _format_success(url: str, expires: int) -> dict[str, Any]:
153 """Format an HTTP 200 API Gateway proxy response with the presigned URL."""
154 return {
155 "statusCode": 200,
156 "headers": {"Content-Type": "application/json"},
157 "body": json.dumps({"url": url, "expires_in": expires}),
158 }
161def _format_provisioning() -> dict[str, Any]:
162 """Format an HTTP 202 response indicating the profile is still provisioning.
164 The CLI polls on this status code until the profile reaches InService
165 and the Lambda returns HTTP 200 with the presigned URL.
166 """
167 return {
168 "statusCode": 202,
169 "headers": {"Content-Type": "application/json"},
170 "body": json.dumps({"status": "provisioning"}),
171 }
174def _format_error(status: int, token: str) -> dict[str, Any]:
175 """Format an HTTP error API Gateway proxy response.
177 ``token`` is one of the module-level ``_ERR_*`` constants; ``status``
178 is the HTTP status code. The body never contains an exception
179 message -- callers log the underlying exception via :data:`logger`
180 before returning.
181 """
182 return {
183 "statusCode": status,
184 "headers": {"Content-Type": "application/json"},
185 "body": json.dumps({"error": token}),
186 }
189# ==========================================================================
190# Effectful helpers (wrap boto3 calls)
191# ==========================================================================
194def _resolve_domain_id(domain_name: str) -> str | None:
195 """Return the Studio ``DomainId`` for ``domain_name`` or ``None``.
197 Paginates ``sagemaker:ListDomains`` by following ``NextToken``. A
198 ``None`` return signals the caller to emit HTTP 404 ``SagemakerDomainNotFound``.
199 """
200 next_token: str | None = None
201 while True:
202 kwargs: dict[str, Any] = {}
203 if next_token is not None:
204 kwargs["NextToken"] = next_token
205 response = sagemaker.list_domains(**kwargs)
206 for domain in response.get("Domains", []):
207 if domain.get("DomainName") == domain_name:
208 domain_id = domain.get("DomainId")
209 return str(domain_id) if domain_id is not None else None
210 next_token = response.get("NextToken")
211 if not next_token:
212 return None
215def _ensure_user_profile(domain_id: str, username: str, efs_id: str) -> str:
216 """Ensure a ``sagemaker:UserProfile`` exists for ``username``.
218 Returns the profile status:
220 * ``"InService"`` -- profile is ready; caller can mint a presigned URL.
221 * ``"Provisioning"`` -- profile was just created or is still starting;
222 caller should return HTTP 202 so the CLI can poll.
224 If the profile is in ``Failed`` state, it is deleted and recreated
225 so the next poll attempt finds a fresh ``Pending`` profile.
226 """
227 try:
228 resp = sagemaker.describe_user_profile(
229 DomainId=domain_id,
230 UserProfileName=username,
231 )
232 status = resp.get("Status", "")
234 if status == "InService":
235 return "InService"
237 if status == "Failed":
238 logger.warning(
239 "User profile %s is Failed (%s) -- deleting and recreating",
240 username,
241 resp.get("FailureReason", "unknown"),
242 )
243 try:
244 sagemaker.delete_user_profile(
245 DomainId=domain_id,
246 UserProfileName=username,
247 )
248 except ClientError as del_exc:
249 logger.warning("Could not delete failed profile %s: %s", username, del_exc)
250 # Create a fresh profile (may race with the delete; ResourceInUse
251 # is caught below).
252 _create_user_profile(domain_id, username, efs_id)
253 return "Provisioning"
255 # Pending / Updating / any other transient state.
256 logger.info("User profile %s status=%s, still provisioning", username, status)
257 return "Provisioning"
259 except ClientError as exc:
260 code = exc.response.get("Error", {}).get("Code", "")
261 if code not in {"ValidationException", "ResourceNotFound"}:
262 raise
264 # Profile doesn't exist -- create it.
265 _create_user_profile(domain_id, username, efs_id)
266 return "Provisioning"
269def _create_user_profile(domain_id: str, username: str, efs_id: str) -> None:
270 """Create a SageMaker user profile. Silently ignores ResourceInUse."""
271 try:
272 sagemaker.create_user_profile(
273 DomainId=domain_id,
274 UserProfileName=username,
275 UserSettings={
276 "ExecutionRole": SAGEMAKER_EXECUTION_ROLE_ARN,
277 "CustomFileSystemConfigs": [
278 {
279 "EFSFileSystemConfig": {
280 "FileSystemId": efs_id,
281 "FileSystemPath": f"/home/{username}",
282 }
283 }
284 ],
285 },
286 )
287 logger.info("Created user profile %s in domain %s", username, domain_id)
288 except ClientError as exc:
289 code = exc.response.get("Error", {}).get("Code", "")
290 if code != "ResourceInUse":
291 raise
292 logger.info("User profile %s already exists (ResourceInUse)", username)
295def _ensure_access_point(username: str, efs_id: str) -> str:
296 """Ensure a per-user EFS access point at ``/home/<username>`` exists.
298 Searches existing access points on ``efs_id`` for one whose
299 ``RootDirectory.Path`` equals ``/home/<username>``; creates one if
300 absent with a POSIX ``(uid, gid)`` derived from
301 :func:`_derive_posix_ids` and ``0700`` permissions.
303 Returns the ``AccessPointArn`` so the caller can associate it with
304 the user profile if needed.
305 """
306 uid, gid = _derive_posix_ids(username)
307 target_path = f"/home/{username}"
309 # Paginate describe_access_points. EFS returns at most 100 APs per
310 # page by default; we explicitly cap at 1000 to stay inside a single
311 # request for typical deployments.
312 next_token: str | None = None
313 while True:
314 kwargs: dict[str, Any] = {"FileSystemId": efs_id, "MaxResults": 1000}
315 if next_token is not None:
316 kwargs["NextToken"] = next_token
317 response = efs.describe_access_points(**kwargs)
318 for ap in response.get("AccessPoints", []):
319 root_dir = ap.get("RootDirectory", {})
320 if root_dir.get("Path") == target_path:
321 return str(ap.get("AccessPointArn", ""))
322 next_token = response.get("NextToken")
323 if not next_token:
324 break
326 created = efs.create_access_point(
327 FileSystemId=efs_id,
328 PosixUser={"Uid": uid, "Gid": gid},
329 RootDirectory={
330 "Path": target_path,
331 "CreationInfo": {
332 "OwnerUid": uid,
333 "OwnerGid": gid,
334 "Permissions": "0700",
335 },
336 },
337 Tags=[
338 {"Key": "gco:analytics:user", "Value": username},
339 {"Key": "gco:analytics:managed", "Value": "true"},
340 ],
341 )
342 return str(created.get("AccessPointArn", ""))
345# ==========================================================================
346# Entry point
347# ==========================================================================
350def lambda_handler(event: dict[str, Any], context: Any) -> dict[str, Any]:
351 """Exchange a Cognito-authorized event for a presigned Studio URL.
353 Returns:
355 * **HTTP 200** with ``{"url": "...", "expires_in": N}`` when the
356 user profile is ``InService`` and the presigned URL is ready.
357 * **HTTP 202** with ``{"status": "provisioning"}`` when the user
358 profile is still being created. The CLI should retry after a
359 few seconds.
360 * **HTTP 4xx / 5xx** with ``{"error": "<token>"}`` on failure.
362 The Lambda never blocks waiting for profile provisioning. The CLI
363 is responsible for polling until it receives HTTP 200.
364 """
365 try:
366 # Step 1: extract + validate the Cognito username claim.
367 claims = _parse_claims(event)
368 username = claims.get("cognito:username") or claims.get("username")
369 if not isinstance(username, str) or not username:
370 logger.warning("Request missing Cognito username claim")
371 return _format_error(401, _ERR_MISSING_CLAIM)
373 # Step 2: resolve the Studio DomainId.
374 domain_id = STUDIO_DOMAIN_ID if STUDIO_DOMAIN_ID else _resolve_domain_id("")
375 if not domain_id:
376 logger.error(
377 "SageMaker domain not found (STUDIO_DOMAIN_ID=%r)",
378 STUDIO_DOMAIN_ID,
379 )
380 return _format_error(404, _ERR_DOMAIN_NOT_FOUND)
382 # Step 3: describe-or-create the user profile.
383 profile_status = _ensure_user_profile(domain_id, username, STUDIO_EFS_ID)
384 if profile_status != "InService":
385 return _format_provisioning()
387 # Step 4: lazy per-user EFS access point.
388 _ensure_access_point(username, STUDIO_EFS_ID)
390 # Step 5: mint the presigned URL.
391 response = sagemaker.create_presigned_domain_url(
392 DomainId=domain_id,
393 UserProfileName=username,
394 SessionExpirationDurationInSeconds=SESSION_EXPIRES_SECONDS,
395 ExpiresInSeconds=URL_EXPIRES_SECONDS,
396 )
397 url = response.get("AuthorizedUrl", "")
398 return _format_success(url, URL_EXPIRES_SECONDS)
400 except Exception as exc: # outer catch-all so every failure returns an opaque error token
401 # Log with exception info so CloudWatch captures the stack trace,
402 # but never leak the message to the HTTP response body.
403 logger.error("Presigned URL generation failed: %s", exc, exc_info=True)
404 return _format_error(500, _ERR_GENERIC)