Deploy Lambda Function with Version Publishing & Alias Traffic Shifting

import boto3
import zipfile
import io
import hashlib
import time
import logging

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")


class LambdaDeployer:
    def __init__(self, region: str = "us-east-1"):
        """
        boto3.client("lambda") is the low-level Lambda client.
        We use it directly (not the resource API) because we need
        fine-grained control over publish_version and create_alias.
        """
        self.lmb = boto3.client("lambda", region_name=region)

    # ── Step 1: Create deployment package ────────────────────────
    def create_deployment_package(self, source_file: str) -> bytes:
        """
        Lambda requires the function code as a ZIP archive.
        zipfile.ZipFile + io.BytesIO creates an in-memory ZIP
        so we don't need to write a temp file to disk.

        ZIP_DEFLATED compresses the archive (required for large files).
        zf.write(source_file, arcname="lambda_function.py") writes the
        file into the archive with a specific internal name — Lambda
        uses this name to find the handler module.
        """
        zip_buffer = io.BytesIO()
        with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf:
            zf.write(source_file, arcname="lambda_function.py")
        zip_buffer.seek(0)
        return zip_buffer.read()

    # ── Step 2: Deploy or update ──────────────────────────────────
    def deploy_or_update(
        self,
        function_name: str,
        source_file: str,
        role_arn: str,
        env_vars: dict,
        runtime: str  = "python3.12",
        memory: int   = 256,
        timeout: int  = 30,
    ) -> str:
        """
        Returns the published version number (as a string, e.g. "5").

        We use a SHA-256 hash of the zip bytes as a human-readable
        deployment identifier — shows up in logs and Lambda tags.
        """
        zip_bytes = self.create_deployment_package(source_file)
        code_hash = hashlib.sha256(zip_bytes).hexdigest()[:8]
        logger.info(f"Deploying {function_name} | code hash: {code_hash}")

        try:
            # ── UPDATE existing function ──────────────────────────
            self.lmb.get_function(FunctionName=function_name)
            logger.info(f"Function exists — updating")

            # update_function_code() replaces the function code.
            # Publish=False means we upload the code first and publish
            # a new version separately after updating config too.
            self.lmb.update_function_code(
                FunctionName=function_name,
                ZipFile=zip_bytes,
                Publish=False,
            )
            # Lambda updates are asynchronous — must wait before next call
            self._wait_for_update(function_name)

            # update_function_configuration() changes runtime settings.
            # Must be a separate call from update_function_code.
            self.lmb.update_function_configuration(
                FunctionName=function_name,
                Environment={"Variables": env_vars},
                MemorySize=memory,
                Timeout=timeout,
            )
            self._wait_for_update(function_name)

        except self.lmb.exceptions.ResourceNotFoundException:
            # ── CREATE new function ───────────────────────────────
            logger.info(f"Function does not exist — creating")
            self.lmb.create_function(
                FunctionName=function_name,
                Runtime=runtime,
                Role=role_arn,
                # Handler = "module_name.function_name"
                # Lambda imports lambda_function.py and calls lambda_handler()
                Handler="lambda_function.lambda_handler",
                Code={"ZipFile": zip_bytes},
                Environment={"Variables": env_vars},
                MemorySize=memory,
                Timeout=timeout,
                Tags={"DeployedBy": "DeployScript", "CodeHash": code_hash},
            )
            self._wait_for_active(function_name)

        # ── Step 3: Publish version ───────────────────────────────
        # publish_version() takes a snapshot of the current code + config
        # and assigns it an immutable version number.
        # Versions are permanent — they can only be deleted explicitly.
        version_resp = self.lmb.publish_version(
            FunctionName=function_name,
            Description=f"Deploy {code_hash}",
        )
        version = version_resp["Version"]
        logger.info(f"Published version: {version}")
        return version

    # ── Step 4: Create / update alias with traffic shifting ───────
    def create_or_update_alias(
        self,
        function_name: str,
        alias: str,
        version: str,
        canary_weight: float = None,
        previous_version: str = None,
    ) -> None:
        """
        An alias is a named pointer to one (or two) Lambda versions.
        Used so callers reference "prod" instead of a version number.

        RoutingConfig.AdditionalVersionWeights enables canary deployment:
        - {"5": 0.9} means 90% traffic to version 5, 10% to this alias's
          FunctionVersion (the new one). Weights must sum to < 1.0.

        Set canary_weight=None to route 100% to the new version.
        """
        config: dict = {
            "FunctionName": function_name,
            "Name":         alias,
            "FunctionVersion": version,
            "Description":  f"Points to version {version}",
        }

        if canary_weight is not None and previous_version is not None:
            # Send `canary_weight` fraction of traffic to old version
            # Remaining fraction (1 - canary_weight) goes to `version`
            config["RoutingConfig"] = {
                "AdditionalVersionWeights": {
                    previous_version: canary_weight
                }
            }
            logger.info(
                f"Canary: {(1-canary_weight)*100:.0f}% → v{version}, "
                f"{canary_weight*100:.0f}% → v{previous_version}"
            )

        try:
            self.lmb.get_alias(FunctionName=function_name, Name=alias)
            self.lmb.update_alias(**config)
            logger.info(f"Updated alias '{alias}' → version {version}")
        except self.lmb.exceptions.ResourceNotFoundException:
            self.lmb.create_alias(**config)
            logger.info(f"Created alias '{alias}' → version {version}")

    def promote_alias(self, function_name: str, alias: str, version: str) -> None:
        """After canary looks healthy, promote to 100% traffic."""
        self.lmb.update_alias(
            FunctionName=function_name,
            Name=alias,
            FunctionVersion=version,
            RoutingConfig={"AdditionalVersionWeights": {}},  # Clear canary weights
        )
        logger.info(f"Alias '{alias}' promoted to 100% → version {version}")

    # ── Waiters ───────────────────────────────────────────────────
    def _wait_for_update(self, function_name: str, max_wait: int = 60) -> None:
        """
        Lambda updates are asynchronous. get_function_configuration()
        returns LastUpdateStatus which cycles:
          InProgress → Successful (or Failed)
        Polling with 1-second intervals is fine; updates usually complete in 2-5s.
        """
        for _ in range(max_wait):
            resp = self.lmb.get_function_configuration(FunctionName=function_name)
            status = resp["LastUpdateStatus"]
            if status == "Successful":
                return
            if status == "Failed":
                raise RuntimeError(
                    f"Lambda update failed: {resp.get('LastUpdateStatusReasonCode')}"
                )
            time.sleep(1)
        raise TimeoutError(f"Lambda update timed out for {function_name}")

    def _wait_for_active(self, function_name: str, max_wait: int = 60) -> None:
        """
        After create_function(), the State cycles: Pending → Active.
        Must wait before calling publish_version() or the API will error.
        """
        for _ in range(max_wait):
            resp = self.lmb.get_function_configuration(FunctionName=function_name)
            if resp["State"] == "Active":
                return
            time.sleep(1)
        raise TimeoutError(f"Lambda activation timed out for {function_name}")


# ── Entry point ───────────────────────────────────────────────────
if __name__ == "__main__":
    deployer = LambdaDeployer(region="us-east-1")

    env_vars = {
        "ENVIRONMENT": "production",
        "DB_HOST":     "prod-db.cluster.us-east-1.rds.amazonaws.com",
        "LOG_LEVEL":   "INFO",
    }

    # Deploy and get new version number
    new_version = deployer.deploy_or_update(
        function_name="my-app-processor",
        source_file="lambda_function.py",
        role_arn="arn:aws:iam::123456789012:role/LambdaExecRole",
        env_vars=env_vars,
        memory=512,
        timeout=60,
    )

    # Canary: route 10% to new version, 90% stays on previous version "4"
    deployer.create_or_update_alias(
        function_name="my-app-processor",
        alias="prod",
        version=new_version,        # e.g., "5"
        canary_weight=0.9,          # 90% to old version
        previous_version="4",       # Previous stable version
    )

    # After observing metrics for 15 minutes...
    # deployer.promote_alias("my-app-processor", "prod", new_version)

zip_buffer = io.BytesIO()
with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf:
    zf.write(source_file, arcname="lambda_function.py")
zip_buffer.seek(0)
return zip_buffer.read()

zip_bytes = self.create_deployment_package(source_file)
code_hash = hashlib.sha256(zip_bytes).hexdigest()[:8]

self.lmb.get_function(FunctionName=function_name)

self.lmb.update_function_code(
    FunctionName=function_name,
    ZipFile=zip_bytes,
    Publish=False,
)
self._wait_for_update(function_name)

self.lmb.update_function_configuration(
    FunctionName=function_name,
    Environment={"Variables": env_vars},
    MemorySize=memory,
    Timeout=timeout,
)

version_resp = self.lmb.publish_version(
    FunctionName=function_name,
    Description=f"Deploy {code_hash}",
)
version = version_resp["Version"]

config["RoutingConfig"] = {
    "AdditionalVersionWeights": {previous_version: canary_weight}
}

try:
    self.lmb.get_alias(FunctionName=function_name, Name=alias)
    self.lmb.update_alias(**config)
except self.lmb.exceptions.ResourceNotFoundException:
    self.lmb.create_alias(**config)

for _ in range(max_wait):
    resp = self.lmb.get_function_configuration(FunctionName=function_name)
    status = resp["LastUpdateStatus"]
    if status == "Successful":
        return
    if status == "Failed":
        raise RuntimeError(...)
    time.sleep(1)
raise TimeoutError(...)