Coverage for backend/django/core/auxiliary/services/object_storage/s3.py: 81%
150 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-07-22 05:22 +0000
1import hashlib
2import json
3import re
4from datetime import UTC, datetime
5from functools import lru_cache
7import boto3
8from botocore.client import BaseClient
9from botocore.config import Config
10from botocore.exceptions import ClientError
11from botocore.response import StreamingBody
12from django.conf import settings
13from opentelemetry import trace
15from core.auxiliary.services.upload_types import (
16 CompletedMultipartPart,
17 CompletedMultipartUploadStorageResult,
18 UploadedPart,
19)
22DEFAULT_REGION = "us-east-1"
23tracer = trace.get_tracer(settings.OPEN_TELEMETRY_TRACER_NAME)
26def _set_span_attributes(span, attributes: dict[str, object | None]) -> None:
27 """Attach non-empty diagnostic attributes without leaking object keys."""
28 for key, value in attributes.items():
29 if value is not None: 29 ↛ 28line 29 didn't jump to line 28 because the condition on line 29 was always true
30 span.set_attribute(key, value)
33def get_bucket_name() -> str:
34 """Return the configured default object-storage bucket name."""
35 bucket = settings.SEAWEED_S3_BUCKET
36 if not bucket: 36 ↛ 37line 36 didn't jump to line 37 because the condition on line 36 was never true
37 raise RuntimeError("SEAWEED_S3_BUCKET is not configured.")
38 return bucket
41def _client_config() -> Config:
42 """Build a botocore config for the configured S3-compatible object store."""
43 config_kwargs = {"signature_version": "s3v4"}
44 if settings.SEAWEED_S3_FORCE_PATH_STYLE: 44 ↛ 46line 44 didn't jump to line 46 because the condition on line 44 was always true
45 config_kwargs["s3"] = {"addressing_style": "path"}
46 return Config(**config_kwargs)
49def _build_client(endpoint_url: str) -> BaseClient:
50 """Create an S3 client for the provided internal or public endpoint URL."""
51 return boto3.client(
52 "s3",
53 endpoint_url=endpoint_url,
54 aws_access_key_id=settings.SEAWEED_S3_ACCESS_KEY,
55 aws_secret_access_key=settings.SEAWEED_S3_SECRET_KEY,
56 region_name=settings.SEAWEED_S3_REGION or DEFAULT_REGION,
57 config=_client_config(),
58 )
61@lru_cache(maxsize=1)
62def get_s3_client() -> BaseClient:
63 """Return the cached internal S3 client used for storage operations."""
64 endpoint_url = settings.SEAWEED_S3_ENDPOINT
65 if not endpoint_url: 65 ↛ 66line 65 didn't jump to line 66 because the condition on line 65 was never true
66 raise RuntimeError("SEAWEED_S3_ENDPOINT is not configured.")
67 return _build_client(endpoint_url)
70@lru_cache(maxsize=1)
71def get_presign_client() -> BaseClient:
72 """Return the cached S3 client used to generate browser-facing presigned URLs."""
73 endpoint_url = settings.SEAWEED_S3_PUBLIC_ENDPOINT or settings.SEAWEED_S3_ENDPOINT
74 if not endpoint_url: 74 ↛ 75line 74 didn't jump to line 75 because the condition on line 74 was never true
75 raise RuntimeError("SEAWEED_S3_PUBLIC_ENDPOINT or SEAWEED_S3_ENDPOINT must be configured.")
76 return _build_client(endpoint_url)
79def _lifecycle_rule_id(prefix: str) -> str:
80 """Build a deterministic lifecycle rule identifier for a prefix."""
81 return f"expire-{hashlib.sha1(prefix.encode('utf-8')).hexdigest()[:16]}"
84@tracer.start_as_current_span("csv_upload.s3.get_bucket_lifecycle")
85def _get_lifecycle_rules(bucket: str) -> list[dict]:
86 """Fetch lifecycle rules for a bucket, tolerating an empty configuration."""
87 span = trace.get_current_span()
88 _set_span_attributes(span, {"csv_upload.bucket": bucket})
89 try:
90 response = get_s3_client().get_bucket_lifecycle_configuration(Bucket=bucket)
91 except ClientError as exc:
92 error_code = exc.response.get("Error", {}).get("Code")
93 span.set_attribute("csv_upload.s3.error_code", error_code or "")
94 if error_code == "NoSuchLifecycleConfiguration": 94 ↛ 97line 94 didn't jump to line 97 because the condition on line 94 was always true
95 span.set_attribute("csv_upload.lifecycle.rule_count", 0)
96 return []
97 raise
99 rules = list(response.get("Rules", []))
100 span.set_attribute("csv_upload.lifecycle.rule_count", len(rules))
101 return rules
104@tracer.start_as_current_span("csv_upload.s3.put_bucket_lifecycle")
105def _put_lifecycle_rule(bucket: str, rule: dict) -> None:
106 """Upsert a lifecycle rule on the bucket."""
107 rules = [existing for existing in _get_lifecycle_rules(bucket) if existing.get("ID") != rule["ID"]]
108 rules.append(rule)
109 _set_span_attributes(
110 trace.get_current_span(),
111 {
112 "csv_upload.bucket": bucket,
113 "csv_upload.lifecycle.rule_count": len(rules),
114 "csv_upload.lifecycle.mode": "per_object",
115 },
116 )
117 get_s3_client().put_bucket_lifecycle_configuration(
118 Bucket=bucket,
119 LifecycleConfiguration={"Rules": rules},
120 )
123def _canonical_lifecycle_rules(rules: list[dict]) -> str:
124 """Serialize lifecycle rules for process-local idempotency caching."""
125 return json.dumps(rules, sort_keys=True, separators=(",", ":"))
128def _is_obsolete_prefix_expiration_rule(
129 rule: dict,
130 *,
131 obsolete_expiration_prefixes: tuple[str, ...],
132 required_rule_ids: set[str],
133) -> bool:
134 """Return whether a legacy exact-date expiration rule should be pruned."""
135 if rule.get("ID") in required_rule_ids or "Expiration" not in rule: 135 ↛ 136line 135 didn't jump to line 136 because the condition on line 135 was never true
136 return False
137 if "Date" not in rule["Expiration"]:
138 return False
140 prefix = rule.get("Filter", {}).get("Prefix")
141 return isinstance(prefix, str) and prefix.startswith(obsolete_expiration_prefixes)
144@lru_cache(maxsize=32)
145@tracer.start_as_current_span("csv_upload.s3.ensure_bucket_lifecycle")
146def _ensure_lifecycle_rules_cached(
147 bucket: str,
148 canonical_rules: str,
149 obsolete_expiration_prefixes: tuple[str, ...],
150 obsolete_rule_ids: tuple[str, ...],
151) -> None:
152 """Install stable lifecycle rules once per process and bucket."""
153 required_rules = json.loads(canonical_rules)
154 required_rule_ids = {rule["ID"] for rule in required_rules}
155 existing_rules = _get_lifecycle_rules(bucket)
156 existing_by_id = {rule.get("ID"): rule for rule in existing_rules}
158 next_rules = []
159 pruned_rule_count = 0
160 for rule in existing_rules:
161 if rule.get("ID") in required_rule_ids:
162 continue
163 if rule.get("ID") in obsolete_rule_ids: 163 ↛ 164line 163 didn't jump to line 164 because the condition on line 163 was never true
164 pruned_rule_count += 1
165 continue
166 if _is_obsolete_prefix_expiration_rule(
167 rule,
168 obsolete_expiration_prefixes=obsolete_expiration_prefixes,
169 required_rule_ids=required_rule_ids,
170 ):
171 pruned_rule_count += 1
172 continue
173 next_rules.append(rule)
174 next_rules.extend(required_rules)
175 changed = pruned_rule_count > 0 or not all(existing_by_id.get(rule["ID"]) == rule for rule in required_rules)
176 _set_span_attributes(
177 trace.get_current_span(),
178 {
179 "csv_upload.bucket": bucket,
180 "csv_upload.lifecycle.mode": "stable_prefix",
181 "csv_upload.lifecycle.existing_rule_count": len(existing_rules),
182 "csv_upload.lifecycle.required_rule_count": len(required_rules),
183 "csv_upload.lifecycle.next_rule_count": len(next_rules),
184 "csv_upload.lifecycle.pruned_rule_count": pruned_rule_count,
185 "csv_upload.lifecycle.changed": changed,
186 },
187 )
188 if not changed: 188 ↛ 189line 188 didn't jump to line 189 because the condition on line 188 was never true
189 return
191 get_s3_client().put_bucket_lifecycle_configuration(
192 Bucket=bucket,
193 LifecycleConfiguration={"Rules": next_rules},
194 )
197def ensure_lifecycle_rules(
198 bucket: str,
199 rules: list[dict],
200 *,
201 obsolete_expiration_prefixes: tuple[str, ...] = (),
202 obsolete_rule_ids: tuple[str, ...] = (),
203) -> None:
204 """Ensure stable lifecycle rules and remove explicitly retired rules."""
205 _ensure_lifecycle_rules_cached(
206 bucket,
207 _canonical_lifecycle_rules(rules),
208 obsolete_expiration_prefixes,
209 obsolete_rule_ids,
210 )
213def _sanitize_content_disposition_filename(filename: str) -> str:
214 sanitized = re.sub(r'[\r\n"\\;]', "_", filename).strip()
215 return sanitized or "download"
218def schedule_object_expiration(bucket: str, key: str, expires_at: datetime) -> None:
219 """Schedule expiration for a single object using an exact-key prefix rule."""
220 # S3-compatible lifecycle prefix filters are starts-with matches rather than strict
221 # equality checks. We use the full object key here; because upload keys include a
222 # UUID segment, the effective match remains unique for our generated objects.
223 _put_lifecycle_rule(
224 bucket,
225 {
226 "ID": _lifecycle_rule_id(key),
227 "Status": "Enabled",
228 "Filter": {"Prefix": key},
229 "Expiration": {"Date": expires_at.astimezone(UTC).replace(microsecond=0)},
230 },
231 )
234def create_multipart_upload(bucket: str, key: str, content_type: str) -> str:
235 """Start a multipart upload and return the storage upload id."""
236 response = get_s3_client().create_multipart_upload(
237 Bucket=bucket,
238 Key=key,
239 ContentType=content_type or "text/csv",
240 )
241 return response["UploadId"]
244def presign_upload_part(
245 bucket: str,
246 key: str,
247 upload_id: str,
248 part_number: int,
249 expires_seconds: int,
250) -> str:
251 """Generate a presigned URL for one multipart upload part."""
252 return get_presign_client().generate_presigned_url(
253 "upload_part",
254 Params={
255 "Bucket": bucket,
256 "Key": key,
257 "UploadId": upload_id,
258 "PartNumber": part_number,
259 },
260 ExpiresIn=expires_seconds,
261 )
264def presign_download_url(
265 bucket: str,
266 key: str,
267 filename: str,
268 expires_seconds: int = 3600,
269) -> str:
270 """Generate a presigned URL for downloading an object from S3-compatible storage."""
271 safe_filename = _sanitize_content_disposition_filename(filename)
272 return get_presign_client().generate_presigned_url(
273 "get_object",
274 Params={
275 "Bucket": bucket,
276 "Key": key,
277 "ResponseContentDisposition": f'attachment; filename="{safe_filename}"',
278 },
279 ExpiresIn=expires_seconds,
280 )
283def list_uploaded_parts(bucket: str, key: str, upload_id: str) -> list[UploadedPart]:
284 """Return all uploaded multipart parts currently stored for an upload id."""
285 client = get_s3_client()
286 parts: list[UploadedPart] = []
287 part_number_marker = 0
289 while True:
290 response = client.list_parts(
291 Bucket=bucket,
292 Key=key,
293 UploadId=upload_id,
294 PartNumberMarker=part_number_marker,
295 )
296 parts.extend(
297 UploadedPart(
298 part_number=part["PartNumber"],
299 etag=part["ETag"],
300 size_bytes=part["Size"],
301 )
302 for part in response.get("Parts", [])
303 )
304 if not response.get("IsTruncated"): 304 ↛ 306line 304 didn't jump to line 306 because the condition on line 304 was always true
305 break
306 part_number_marker = response.get("NextPartNumberMarker", 0)
308 return parts
311def complete_multipart_upload(
312 bucket: str,
313 key: str,
314 upload_id: str,
315 parts: list[CompletedMultipartPart],
316) -> CompletedMultipartUploadStorageResult:
317 """Finalize a multipart upload and return the storage completion metadata."""
318 sorted_parts = sorted(parts, key=lambda part: part.part_number)
319 response = get_s3_client().complete_multipart_upload(
320 Bucket=bucket,
321 Key=key,
322 UploadId=upload_id,
323 MultipartUpload={
324 "Parts": [
325 {"PartNumber": part.part_number, "ETag": part.etag}
326 for part in sorted_parts
327 ]
328 },
329 )
330 return CompletedMultipartUploadStorageResult(
331 bucket=bucket,
332 object_key=key,
333 etag=response.get("ETag"),
334 location=response.get("Location"),
335 )
338def abort_multipart_upload(bucket: str, key: str, upload_id: str) -> None:
339 """Abort a multipart upload that should no longer accept browser parts."""
340 get_s3_client().abort_multipart_upload(Bucket=bucket, Key=key, UploadId=upload_id)
343def get_range(bucket: str, key: str, byte_range: str) -> bytes:
344 """Read a byte range from an object for lightweight CSV inspection."""
345 response = get_s3_client().get_object(Bucket=bucket, Key=key, Range=f"bytes={byte_range}")
346 return response["Body"].read()
349def stream_object(bucket: str, key: str) -> StreamingBody:
350 """Open a streaming reader for the full object body."""
351 response = get_s3_client().get_object(Bucket=bucket, Key=key)
352 return response["Body"]
355def object_exists(bucket: str, key: str) -> bool:
356 """Return whether the configured object is still available in object storage."""
357 try:
358 get_s3_client().head_object(Bucket=bucket, Key=key)
359 except ClientError as exc:
360 error_code = exc.response.get("Error", {}).get("Code")
361 if error_code in {"404", "NoSuchKey", "NotFound"}:
362 return False
363 raise
364 return True
367def delete_object(bucket: str, key: str) -> None:
368 """Delete an object from S3-compatible object storage."""
369 get_s3_client().delete_object(Bucket=bucket, Key=key)