Claim Check Pattern with Temporal

import uuid
import redis.asyncio as redis
from typing import Iterable, List

from temporalio.api.common.v1 import Payload
from temporalio.converter import PayloadCodec

class ClaimCheckCodec(PayloadCodec):
    def __init__(self, redis_host: str = "localhost", redis_port: int = 6379):
        self.redis_client = redis.Redis(host=redis_host, port=redis_port)

    async def encode(self, payloads: Iterable[Payload]) -> List[Payload]:
        """Replace large payloads with keys and store original data in Redis."""
        out: List[Payload] = []
        for payload in payloads:
            encoded = await self.encode_payload(payload)
            out.append(encoded)
        return out

    async def decode(self, payloads: Iterable[Payload]) -> List[Payload]:
        """Retrieve original payloads from Redis using stored keys."""
        out: List[Payload] = []
        for payload in payloads:
            if payload.metadata.get("temporal.io/claim-check-codec", b"").decode() != "v1":
                # Not a claim-checked payload, pass through unchanged
                out.append(payload)
                continue

            redis_key = payload.data.decode("utf-8")
            stored_data = await self.redis_client.get(redis_key)
            if stored_data is None:
                raise ValueError(f"Claim check key not found in Redis: {redis_key}")
            
            original_payload = Payload.FromString(stored_data)
            out.append(original_payload)
        return out

    async def encode_payload(self, payload: Payload) -> Payload:
        """Store payload in Redis and return a key-based payload."""
        key = str(uuid.uuid4())
        serialized_data = payload.SerializeToString()
        
        # Store the original payload data in Redis
        await self.redis_client.set(key, serialized_data)
        
        # Return a lightweight payload containing only the key
        return Payload(
            metadata={
                "encoding": b"claim-checked",
                "temporal.io/claim-check-codec": b"v1",
            },
            data=key.encode("utf-8"),
        )

import os
from temporalio.client import Plugin, ClientConfig
from temporalio.converter import DataConverter

from claim_check_codec import ClaimCheckCodec

class ClaimCheckPlugin(Plugin):
    def __init__(self):
        self.redis_host = os.getenv("REDIS_HOST", "localhost")
        self.redis_port = int(os.getenv("REDIS_PORT", "6379"))

    def get_data_converter(self, config: ClientConfig) -> DataConverter:
        """Configure the data converter with claim check codec."""
        default_converter_class = config["data_converter"].payload_converter_class
        claim_check_codec = ClaimCheckCodec(self.redis_host, self.redis_port)
        
        return DataConverter(
            payload_converter_class=default_converter_class,
            payload_codec=claim_check_codec
        )

    def configure_client(self, config: ClientConfig) -> ClientConfig:
        """Apply the claim check configuration to the client."""
        config["data_converter"] = self.get_data_converter(config)
        return super().configure_client(config)

from temporalio import workflow, activity
from dataclasses import dataclass
from typing import List
import json

@dataclass
class LargeDataset:
    """Represents a large dataset that would benefit from claim check pattern."""
    data: List[dict]
    metadata: dict

@dataclass
class ProcessingResult:
    """Result of processing the large dataset."""
    processed_count: int
    summary: dict
    errors: List[str]

@activity.defn
async def process_large_dataset(dataset: LargeDataset) -> ProcessingResult:
    """Process a large dataset - this would normally cause payload size issues.
    
    This activity demonstrates how the claim check pattern allows processing
    of large datasets without hitting Temporal's payload size limits.
    
    Note: For production workloads processing very large datasets, consider adding
    activity heartbeats to prevent timeout issues during long-running processing.
    
    Args:
        dataset: Large dataset to process
        
    Returns:
        ProcessingResult with processing statistics and any errors
    """
    processed_count = 0
    errors = []
    summary = {"total_items": len(dataset.data)}
    
    for item in dataset.data:
        try:
            # Simulate processing each item
            if "value" in item:
                item["processed_value"] = item["value"] * 2
                processed_count += 1
            elif "text" in item:
                # Simulate text processing
                item["word_count"] = len(item["text"].split())
                processed_count += 1
        except Exception as e:
            errors.append(f"Error processing item {item.get('id', 'unknown')}: {str(e)}")
    
    summary["processed_items"] = processed_count
    summary["error_count"] = len(errors)
    
    return ProcessingResult(
        processed_count=processed_count,
        summary=summary,
        errors=errors
    )

@workflow.defn
class LargeDataProcessingWorkflow:
    @workflow.run
    async def run(self, dataset: LargeDataset) -> ProcessingResult:
        """Process large dataset using claim check pattern."""
        result = await workflow.execute_activity(
            process_large_dataset,
            dataset,
            start_to_close_timeout=timedelta(minutes=10),
            summary="Process large dataset"
        )
        return result

@activity.defn
async def process_large_dataset(dataset: LargeDataset) -> ProcessingResult:
    total_items = len(dataset.data)
    
    for i, item in enumerate(dataset.data):
        # Send heartbeat every 100 items to prevent timeout
        if i % 100 == 0:
            await activity.heartbeat(f"Processed {i}/{total_items} items")
        
        # Process item...