Add icon uploading to objectstore

Author: rbro112

requirements.txt

@@ -19,3 +19,4 @@ sentry-kafka-schemas==2.1.2
 sentry-sdk>=2.36.0
 sortedcontainers>=2.4.0
 typing-extensions>=4.15.0
+zstandard>=0.18.0

src/launchpad/api/update_api_models.py

@@ -57,6 +57,7 @@ class UpdateData(BaseModel):
     artifact_type: int
     apple_app_info: Optional[AppleAppInfo] = None
     dequeued_at: Optional[datetime] = Field(None, description="Timestamp when message was dequeued from Kafka")
+    app_icon_id: Optional[str] = None
 
     @field_serializer("dequeued_at")
     def serialize_datetime(self, dt: datetime | None) -> str | None:

src/launchpad/artifact_processor.py

@@ -13,7 +13,9 @@
 
 import sentry_sdk
 
-from sentry_kafka_schemas.schema_types.preprod_artifact_events_v1 import PreprodArtifactEvents
+from sentry_kafka_schemas.schema_types.preprod_artifact_events_v1 import (
+    PreprodArtifactEvents,
+)
 
 from launchpad.api.update_api_models import AppleAppInfo as AppleAppInfoModel
 from launchpad.api.update_api_models import PutSizeFailed, UpdateData
@@ -41,6 +43,12 @@
 from launchpad.size.models.common import BaseAppInfo
 from launchpad.tracing import request_context
 from launchpad.utils.logging import get_logger
+from launchpad.utils.objectstore.service import (
+    Client as ObjectstoreClient,
+)
+from launchpad.utils.objectstore.service import (
+    ClientBuilder as ObjectstoreClientBuilder,
+)
 from launchpad.utils.statsd import StatsdInterface, get_statsd
 
 logger = get_logger(__name__)
@@ -49,12 +57,23 @@
 class ArtifactProcessor:
     """Handles the processing of artifacts including download, analysis, and upload."""
 
-    def __init__(self, sentry_client: SentryClient, statsd: StatsdInterface) -> None:
+    def __init__(
+        self,
+        sentry_client: SentryClient,
+        statsd: StatsdInterface,
+        objectstore_client: ObjectstoreClient,
+    ) -> None:
         self._sentry_client = sentry_client
         self._statsd = statsd
+        self._objectstore_client = objectstore_client
 
     @staticmethod
-    def process_message(payload: PreprodArtifactEvents, service_config=None, artifact_processor=None, statsd=None):
+    def process_message(
+        payload: PreprodArtifactEvents,
+        service_config=None,
+        artifact_processor=None,
+        statsd=None,
+    ):
         """Process an artifact message with proper context and metrics.
         This is used by the Kafka workers and so has to set up the context from scratch.
         If components are not provided, they will be created.
@@ -66,15 +85,18 @@ def process_message(payload: PreprodArtifactEvents, service_config=None, artifac
 
         initialize_sentry_sdk()
 
+        organization_id = payload["organization_id"]
+        project_id = payload["project_id"]
+        artifact_id = payload["artifact_id"]
+
         if statsd is None:
             statsd = get_statsd()
         if artifact_processor is None:
             sentry_client = SentryClient(base_url=service_config.sentry_base_url)
-            artifact_processor = ArtifactProcessor(sentry_client, statsd)
-
-        organization_id = payload["organization_id"]
-        project_id = payload["project_id"]
-        artifact_id = payload["artifact_id"]
+            objectstore_client = ObjectstoreClientBuilder(
+                usecase="app-icons", options={"base_url": "http://localhost:8888/"}
+            ).for_project(organization_id, project_id)
+            artifact_processor = ArtifactProcessor(sentry_client, statsd, objectstore_client)
 
         requested_features = []
         for feature in payload.get("requested_features", []):
@@ -92,7 +114,10 @@ def process_message(payload: PreprodArtifactEvents, service_config=None, artifac
             stack.enter_context(
                 statsd.timed(
                     "artifact.processing.duration",
-                    tags=[f"project_id:{project_id}", f"organization_id:{organization_id}"],
+                    tags=[
+                        f"project_id:{project_id}",
+                        f"organization_id:{organization_id}",
+                    ],
                 )
             )
             scope = stack.enter_context(sentry_sdk.new_scope())
@@ -116,7 +141,11 @@ def process_message(payload: PreprodArtifactEvents, service_config=None, artifac
                 )
 
     def process_artifact(
-        self, organization_id: str, project_id: str, artifact_id: str, requested_features: list[PreprodFeature]
+        self,
+        organization_id: str,
+        project_id: str,
+        artifact_id: str,
+        requested_features: list[PreprodFeature],
     ) -> None:
         """Process an artifact with the requested features."""
         dequeued_at = datetime.now()
@@ -125,8 +154,16 @@ def process_artifact(
             path = stack.enter_context(self._download_artifact(organization_id, project_id, artifact_id))
             artifact = self._parse_artifact(organization_id, project_id, artifact_id, path)
             analyzer = self._create_analyzer(artifact)
-
-            info = self._preprocess_artifact(organization_id, project_id, artifact_id, artifact, analyzer, dequeued_at)
+            app_icon_object_id = self._process_app_icon(organization_id, project_id, artifact_id, artifact)
+            info = self._preprocess_artifact(
+                organization_id,
+                project_id,
+                artifact_id,
+                artifact,
+                analyzer,
+                dequeued_at,
+                app_icon_object_id,
+            )
 
             if PreprodFeature.SIZE_ANALYSIS in requested_features:
                 self._do_size(organization_id, project_id, artifact_id, artifact, analyzer)
@@ -145,7 +182,8 @@ def _download_artifact(
 
         with tempfile.NamedTemporaryFile(suffix=".zip") as tf:
             with self._statsd.timed(
-                "artifact.download.duration", tags=[f"project_id:{project_id}", f"organization_id:{organization_id}"]
+                "artifact.download.duration",
+                tags=[f"project_id:{project_id}", f"organization_id:{organization_id}"],
             ):
                 size = self._sentry_client.download_artifact(
                     org=organization_id,
@@ -189,14 +227,15 @@ def _preprocess_artifact(
         artifact: Artifact,
         analyzer: AndroidAnalyzer | AppleAppAnalyzer,
         dequeued_at: datetime,
+        app_icon_id: str | None,
     ) -> AppleAppInfo | BaseAppInfo:
         logger.info(f"Preprocessing for {artifact_id} (project: {project_id}, org: {organization_id})")
         try:
             info = self._retry_operation(
                 lambda: analyzer.preprocess(cast(Any, artifact)),
                 OperationName.PREPROCESSING,
             )
-            update_data = self._prepare_update_data(info, artifact, dequeued_at)
+            update_data = self._prepare_update_data(info, artifact, dequeued_at, app_icon_id)
             self._sentry_client.update_artifact(
                 org=organization_id,
                 project=project_id,
@@ -217,6 +256,22 @@ def _preprocess_artifact(
         else:
             return info
 
+    def _process_app_icon(
+        self,
+        organization_id: str,
+        project_id: str,
+        artifact_id: str,
+        artifact: Artifact,
+    ) -> str | None:
+        logger.info(f"Processing app icon for {artifact_id} (project: {project_id}, org: {organization_id})")
+        app_icon = artifact.get_app_icon()
+        if app_icon is None:
+            logger.info(f"No app icon found for {artifact_id} (project: {project_id}, org: {organization_id})")
+            return None
+
+        app_icon_id = self._objectstore_client.put(app_icon, compression="none")
+        return app_icon_id
+
     def _do_distribution(
         self,
         organization_id: str,
@@ -376,7 +431,13 @@ def _update_size_error_from_exception(
         if error_message == ProcessingErrorMessage.UNKNOWN_ERROR:
             error_message = _guess_message(error_code, e)
         self._update_size_error(
-            organization_id, project_id, artifact_id, error_code, error_message, str(e), identifier=identifier
+            organization_id,
+            project_id,
+            artifact_id,
+            error_code,
+            error_message,
+            str(e),
+            identifier=identifier,
         )
 
     def _update_size_error(
@@ -413,7 +474,11 @@ def _update_size_error(
             logger.exception(f"Failed to update artifact with error {message}")
 
     def _prepare_update_data(
-        self, app_info: AppleAppInfo | BaseAppInfo, artifact: Artifact, dequeued_at: datetime
+        self,
+        app_info: AppleAppInfo | BaseAppInfo,
+        artifact: Artifact,
+        dequeued_at: datetime,
+        app_icon_id: str | None,
     ) -> Dict[str, Any]:
         def _get_artifact_type(artifact: Artifact) -> ArtifactType:
             if isinstance(artifact, ZippedXCArchive):
@@ -450,6 +515,7 @@ def _get_artifact_type(artifact: Artifact) -> ArtifactType:
             artifact_type=_get_artifact_type(artifact).value,
             apple_app_info=apple_app_info,
             dequeued_at=dequeued_at,
+            app_icon_id=app_icon_id,
         )
 
         return update_data.model_dump(exclude_none=True)

src/launchpad/service.py

@@ -127,20 +127,23 @@ class ServiceConfig:
 
     sentry_base_url: str
     projects_to_skip: list[str]
+    objectstore_base_url: str
 
 
 def get_service_config() -> ServiceConfig:
     """Get service configuration from environment."""
     sentry_base_url = os.getenv("SENTRY_BASE_URL")
     projects_to_skip_str = os.getenv("PROJECT_IDS_TO_SKIP")
     projects_to_skip = projects_to_skip_str.split(",") if projects_to_skip_str else []
+    objectstore_base_url = os.getenv("OBJECTSTORE_BASE_URL")
 
     if sentry_base_url is None:
         sentry_base_url = "http://getsentry.default"
 
     return ServiceConfig(
         sentry_base_url=sentry_base_url,
         projects_to_skip=projects_to_skip,
+        objectstore_base_url=objectstore_base_url,
     )
 
 

src/launchpad/utils/objectstore/metadata.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+import itertools
+import re
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import Literal, cast
+
+Compression = Literal["zstd"]
+
+HEADER_EXPIRATION = "x-sn-expiration"
+HEADER_META_PREFIX = "x-snme-"
+
+
+@dataclass
+class TimeToIdle:
+    delta: timedelta
+
+
+@dataclass
+class TimeToLive:
+    delta: timedelta
+
+
+ExpirationPolicy = TimeToIdle | TimeToLive
+
+
+@dataclass
+class Metadata:
+    compression: Compression | None
+    expiration_policy: ExpirationPolicy | None
+    custom: dict[str, str]
+
+    @classmethod
+    def from_headers(cls, headers: Mapping[str, str]) -> Metadata:
+        compression = None
+        expiration_policy = None
+        custom_metadata = {}
+        for k, v in headers.items():
+            if k == "content-encoding":
+                compression = cast(Compression | None, v)
+            elif k == HEADER_EXPIRATION:
+                expiration_policy = parse_expiration(v)
+            elif k.startswith(HEADER_META_PREFIX):
+                custom_metadata[k[len(HEADER_META_PREFIX) :]] = v
+        return Metadata(compression, expiration_policy, custom_metadata)
+
+
+def format_expiration(expiration_policy: ExpirationPolicy) -> str:
+    if isinstance(expiration_policy, TimeToIdle):
+        return f"tti:{format_timedelta(expiration_policy.delta)}"
+    elif isinstance(expiration_policy, TimeToLive):
+        return f"ttl:{format_timedelta(expiration_policy.delta)}"
+
+
+def parse_expiration(value: str) -> ExpirationPolicy | None:
+    if value.startswith("tti:"):
+        return TimeToIdle(parse_timedelta(value[4:]))
+    elif value.startswith("ttl:"):
+        return TimeToLive(parse_timedelta(value[4:]))
+
+    return None
+
+
+def format_timedelta(delta: timedelta) -> str:
+    days = delta.days
+    output = f"{days} days" if days else ""
+    if seconds := delta.seconds:
+        if output:
+            output += " "
+        output += f"{seconds} seconds"
+
+    return output
+
+
+TIME_SPLIT = re.compile(r"[^\W\d_]+|\d+")
+
+
+def parse_timedelta(delta: str) -> timedelta:
+    words = TIME_SPLIT.findall(delta)
+    seconds = 0
+
+    for num, unit in itertools.batched(words, n=2, strict=True):
+        num = int(num)
+        multiplier = 0
+
+        if unit.startswith("w"):
+            multiplier = 86400 * 7
+        elif unit.startswith("d"):
+            multiplier = 86400
+        elif unit.startswith("h"):
+            multiplier = 3600
+        elif unit.startswith("m") and not unit.startswith("ms"):
+            multiplier = 60
+        elif unit.startswith("s"):
+            multiplier = 1
+
+        seconds += num * multiplier
+
+    return timedelta(seconds=seconds)