75 lines
2.4 KiB
Python
75 lines
2.4 KiB
Python
|
"""S3 destination implementation for Focus export."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
from datetime import timezone
|
||
|
|
from typing import Any, Optional
|
||
|
|
|
||
|
|
import boto3
|
||
|
|
|
||
|
|
from .base import FocusDestination, FocusTimeWindow
|
||
|
|
|
||
|
|
|
||
|
|
class FocusS3Destination(FocusDestination):
|
||
|
|
"""Handles uploading serialized exports to S3 buckets."""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
prefix: str,
|
||
|
|
config: Optional[dict[str, Any]] = None,
|
||
|
|
) -> None:
|
||
|
|
config = config or {}
|
||
|
|
bucket_name = config.get("bucket_name")
|
||
|
|
if not bucket_name:
|
||
|
|
raise ValueError("bucket_name must be provided for S3 destination")
|
||
|
|
self.bucket_name = bucket_name
|
||
|
|
self.prefix = prefix.rstrip("/")
|
||
|
|
self.config = config
|
||
|
|
|
||
|
|
async def deliver(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
content: bytes,
|
||
|
|
time_window: FocusTimeWindow,
|
||
|
|
filename: str,
|
||
|
|
) -> None:
|
||
|
|
object_key = self._build_object_key(time_window=time_window, filename=filename)
|
||
|
|
await asyncio.to_thread(self._upload, content, object_key)
|
||
|
|
|
||
|
|
def _build_object_key(self, *, time_window: FocusTimeWindow, filename: str) -> str:
|
||
|
|
start_utc = time_window.start_time.astimezone(timezone.utc)
|
||
|
|
date_component = f"date={start_utc.strftime('%Y-%m-%d')}"
|
||
|
|
parts = [self.prefix, date_component]
|
||
|
|
if time_window.frequency == "hourly":
|
||
|
|
parts.append(f"hour={start_utc.strftime('%H')}")
|
||
|
|
key_prefix = "/".join(filter(None, parts))
|
||
|
|
return f"{key_prefix}/{filename}" if key_prefix else filename
|
||
|
|
|
||
|
|
def _upload(self, content: bytes, object_key: str) -> None:
|
||
|
|
client_kwargs: dict[str, Any] = {}
|
||
|
|
region_name = self.config.get("region_name")
|
||
|
|
if region_name:
|
||
|
|
client_kwargs["region_name"] = region_name
|
||
|
|
endpoint_url = self.config.get("endpoint_url")
|
||
|
|
if endpoint_url:
|
||
|
|
client_kwargs["endpoint_url"] = endpoint_url
|
||
|
|
|
||
|
|
session_kwargs: dict[str, Any] = {}
|
||
|
|
for key in (
|
||
|
|
"aws_access_key_id",
|
||
|
|
"aws_secret_access_key",
|
||
|
|
"aws_session_token",
|
||
|
|
):
|
||
|
|
if self.config.get(key):
|
||
|
|
session_kwargs[key] = self.config[key]
|
||
|
|
|
||
|
|
s3_client = boto3.client("s3", **client_kwargs, **session_kwargs)
|
||
|
|
s3_client.put_object(
|
||
|
|
Bucket=self.bucket_name,
|
||
|
|
Key=object_key,
|
||
|
|
Body=content,
|
||
|
|
ContentType="application/octet-stream",
|
||
|
|
)
|