worker.py

"""
Unified Worker for Kubernetes deployment (Approach 2: Environment-Driven).

This single worker file selects which workflow classes to register based on
the Build ID from environment variables. The Temporal Worker Controller 
injects the Build ID, and Temporal routes tasks to the appropriate worker.

Supports:
- Local/self-hosted Temporal (no authentication)
- Temporal Cloud with API Key authentication (recommended)
- Temporal Cloud with mTLS authentication (alternative)

Benefits:
- Single Docker image for all versions
- Build ID controlled via environment variable
- Workflow selection based on Build ID
- Prometheus metrics for autoscaling
- API Key and mTLS support for Temporal Cloud

Usage:
    # Local (no authentication)
    TEMPORAL_WORKER_BUILD_ID=1.0 python -m worker_versioning.worker
    
    # Temporal Cloud (with API Key - recommended)
    TEMPORAL_ADDRESS=test-100.ihmbh.tmprl.cloud:7233 \
    TEMPORAL_NAMESPACE=test-100.ihmbh \
    TEMPORAL_API_KEY=your-api-key \
    python -m worker_versioning.worker
    
    # Temporal Cloud (with mTLS - alternative)
    TEMPORAL_ADDRESS=test-100.ihmbh.tmprl.cloud:7233 \
    TEMPORAL_NAMESPACE=test-100.ihmbh \
    TEMPORAL_TLS_CERT=/path/to/client.pem \
    TEMPORAL_TLS_KEY=/path/to/client.key \
    python -m worker_versioning.worker

Build ID to Workflow Mapping:
    - 1.0: AutoUpgradingWorkflowV1, PinnedWorkflowV1
    - 1.1: AutoUpgradingWorkflowV1b, PinnedWorkflowV1
    - 2.0: AutoUpgradingWorkflowV1b, PinnedWorkflowV2

Metrics:
    - Prometheus metrics exposed on port 9464 (configurable via METRICS_PORT)
    - Metrics include: temporal_worker_task_slots_available, temporal_worker_task_slots_used
"""

import asyncio
import logging
import os

from temporalio.client import Client, TLSConfig
from temporalio.common import WorkerDeploymentVersion
from temporalio.worker import Worker, WorkerDeploymentConfig
from temporalio.runtime import Runtime, TelemetryConfig, PrometheusConfig

from worker_versioning.activities import slow_activity, some_activity, some_incompatible_activity
from worker_versioning.workflows import (
    AutoUpgradingWorkflowV1,
    AutoUpgradingWorkflowV1b,
    PinnedWorkflowV1,
    PinnedWorkflowV2,
)

logging.basicConfig(level=logging.INFO)

# Configuration from environment (for Kubernetes deployment)
TEMPORAL_ADDRESS = os.getenv("TEMPORAL_ADDRESS", "localhost:7233")
TEMPORAL_NAMESPACE = os.getenv("TEMPORAL_NAMESPACE", "default")
TASK_QUEUE = os.getenv("TEMPORAL_TASK_QUEUE", "worker-versioning")

# Temporal Cloud authentication - API Key (recommended)
# Set TEMPORAL_API_KEY environment variable with your API key
TEMPORAL_API_KEY = os.getenv("TEMPORAL_API_KEY", "")

# mTLS configuration for Temporal Cloud (alternative to API key)
TEMPORAL_TLS_CERT = os.getenv("TEMPORAL_TLS_CERT", "")  # Path to client certificate
TEMPORAL_TLS_KEY = os.getenv("TEMPORAL_TLS_KEY", "")    # Path to client private key

# Worker versioning config - injected by Temporal Worker Controller
# The controller sets this based on the image tag or deployment spec
BUILD_ID = os.getenv("TEMPORAL_WORKER_BUILD_ID", "1.0")
DEPLOYMENT_NAME = os.getenv("TEMPORAL_DEPLOYMENT_NAME", "my-deployment")

# Metrics configuration for GMP-based autoscaling
METRICS_PORT = int(os.getenv("METRICS_PORT", "9464"))
METRICS_ENABLED = os.getenv("METRICS_ENABLED", "true").lower() == "true"

# Workflow mapping based on Build ID
# Each Build ID maps to specific workflow implementations
WORKFLOW_VERSIONS = {
    "1.0": {
        "workflows": [AutoUpgradingWorkflowV1, PinnedWorkflowV1],
        "description": "v1.0 - Original workflows"
    },
    "1.1": {
        "workflows": [AutoUpgradingWorkflowV1b, PinnedWorkflowV1],
        "description": "v1.1 - Compatible changes to AutoUpgrading"
    },
    "2.0": {
        "workflows": [AutoUpgradingWorkflowV1b, PinnedWorkflowV2],
        "description": "v2.0 - Incompatible changes to Pinned"
    },
}

# Default workflows if Build ID not in mapping
DEFAULT_WORKFLOWS = [AutoUpgradingWorkflowV1, PinnedWorkflowV1]


def get_workflows_for_build_id(build_id: str) -> list:
    """Get the appropriate workflow classes for the given Build ID."""
    # Check exact match first
    if build_id in WORKFLOW_VERSIONS:
        return WORKFLOW_VERSIONS[build_id]["workflows"]
    
    # Try to match by major.minor version (ignore suffix like "-f8bf")
    base_version = build_id.split("-")[0] if "-" in build_id else build_id
    if base_version in WORKFLOW_VERSIONS:
        return WORKFLOW_VERSIONS[base_version]["workflows"]
    
    # Default to latest stable version
    logging.warning(f"Unknown Build ID '{build_id}', using default workflows")
    return DEFAULT_WORKFLOWS


def create_runtime_with_metrics() -> Runtime:
    """Create a Temporal runtime with Prometheus metrics enabled."""
    if not METRICS_ENABLED:
        logging.info("Metrics disabled")
        return Runtime.default()
    
    logging.info(f"Enabling Prometheus metrics on port {METRICS_PORT}")
    return Runtime(
        telemetry=TelemetryConfig(
            metrics=PrometheusConfig(bind_address=f"0.0.0.0:{METRICS_PORT}")
        )
    )


def get_tls_config() -> TLSConfig | None:
    """Get TLS configuration for Temporal Cloud mTLS connection."""
    if not TEMPORAL_TLS_CERT or not TEMPORAL_TLS_KEY:
        return None
    
    logging.info(f"Loading mTLS certificates from {TEMPORAL_TLS_CERT}")
    
    # Read certificate and key files
    with open(TEMPORAL_TLS_CERT, "rb") as f:
        client_cert = f.read()
    with open(TEMPORAL_TLS_KEY, "rb") as f:
        client_key = f.read()
    
    return TLSConfig(
        client_cert=client_cert,
        client_private_key=client_key,
    )


async def main() -> None:
    """Run the unified versioned worker."""
    logging.info(f"Connecting to Temporal at {TEMPORAL_ADDRESS}")
    logging.info(f"Namespace: {TEMPORAL_NAMESPACE}")
    logging.info(f"Task Queue: {TASK_QUEUE}")
    logging.info(f"Build ID: {BUILD_ID}")
    logging.info(f"Deployment Name: {DEPLOYMENT_NAME}")

    # Create runtime with metrics
    runtime = create_runtime_with_metrics()
    
    # Determine authentication method
    api_key = TEMPORAL_API_KEY if TEMPORAL_API_KEY else None
    tls_config = get_tls_config()
    
    # Log authentication method
    if api_key:
        logging.info("Using API key authentication for Temporal Cloud")
    elif tls_config:
        logging.info("Using mTLS authentication for Temporal Cloud")
    else:
        logging.info("No authentication configured (local development)")

    # Connect to Temporal
    # When api_key is provided, TLS is automatically enabled
    client = await Client.connect(
        TEMPORAL_ADDRESS, 
        namespace=TEMPORAL_NAMESPACE,
        runtime=runtime,
        api_key=api_key,
        tls=tls_config if tls_config else (True if api_key else False),
    )

    # Select workflows based on Build ID
    workflows = get_workflows_for_build_id(BUILD_ID)
    workflow_names = [w.__name__ for w in workflows]
    
    logging.info(f"Selected workflows for Build ID '{BUILD_ID}': {workflow_names}")

    worker = Worker(
        client,
        task_queue=TASK_QUEUE,
        workflows=workflows,
        activities=[slow_activity, some_activity, some_incompatible_activity],
        deployment_config=WorkerDeploymentConfig(
            version=WorkerDeploymentVersion(
                deployment_name=DEPLOYMENT_NAME,
                build_id=BUILD_ID,
            ),
            use_worker_versioning=True,
        ),
    )

    logging.info(f"Starting unified worker with Build ID: {BUILD_ID}")
    await worker.run()


if __name__ == "__main__":
    asyncio.run(main())