Building a Serverless URL Shortener with Amazon API Gateway, AWS Lambda, and DynamoDB

import json
import os
import string
import random
import boto3

dynamodb = boto3.resource("dynamodb")
table_name = os.environ.get("TABLE_NAME", "UrlShortenerLinks")
table = dynamodb.Table(table_name)

def generate_code(length=6):
chars = string.ascii_letters + string.digits
return "".join(random.choices(chars, k=length))

def lambda_handler(event, context):
route_key = event.get("routeKey", "")

# 1. Handle POST /shorten to create a new short link
if route_key == "POST /shorten":
try:
body = json.loads(event.get("body") or "{}")
long_url = body.get("url")

if not long_url or not long_url.startswith(("http://", "https://")):
return {
"statusCode": 400,
"body": json.dumps({"error": "Please provide a valid http:// or https:// URL"})
}

short_code = body.get("custom_code") or generate_code()

table.put_item(
Item={
"short_code": short_code,
"long_url": long_url
}
)

domain = event.get("requestContext", {}).get("domainName", "localhost")
return {
"statusCode": 201,
"headers": {"Content-Type": "application/json"},
"body": json.dumps({
"short_code": short_code,
"short_url": f"https://{domain}/{short_code}",
"long_url": long_url
})
}
except Exception as e:
return {"statusCode": 500, "body": json.dumps({"error": str(e)})}

# 2. Handle GET /{code} to redirect to the destination URL
elif route_key == "GET /{code}":
short_code = event.get("pathParameters", {}).get("code")

response = table.get_item(Key={"short_code": short_code})
item = response.get("Item")

if not item:
return {
"statusCode": 404,
"headers": {"Content-Type": "application/json"},
"body": json.dumps({"error": "Short URL not found"})
}

return {
"statusCode": 301,
"headers": {
"Location": item["long_url"],
"Cache-Control": "no-cache"
},
"body": ""
}

return {"statusCode": 400, "body": json.dumps({"error": "Unsupported route"})}

# Create IAM Trust Policy
cat <<EOF> trust-policy.json
{
"Version": "2012-10-17",
"Statement": [
{
"Effect": "Allow",
"Principal": { "Service": "lambda.amazonaws.com" },
"Action": "sts:AssumeRole"
}
]
}
EOF

ROLE_ARN=$(aws iam create-role \
--role-name UrlShortenerLambdaRole \
--assume-role-policy-document file://trust-policy.json \
--query "Role.Arn" \
--output text)

# Attach CloudWatch Logs policy
aws iam attach-role-policy \
--role-name UrlShortenerLambdaRole \
--policy-arn arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole

# Attach least-privilege DynamoDB access policy
TABLE_ARN=$(aws dynamodb describe-table \
--table-name "$TABLE_NAME" \
--query "Table.Arn" \
--output text \
--region "$AWS_REGION")

cat <<EOF> dynamodb-policy.json
{
"Version": "2012-10-17",
"Statement": [
{
"Effect": "Allow",
"Action": [
"dynamodb:GetItem",
"dynamodb:PutItem"
],
"Resource": "$TABLE_ARN"
}
]
}
EOF

aws iam put-role-policy \
--role-name UrlShortenerLambdaRole \
--policy-name DynamoDBShortenerAccess \
--policy-document file://dynamodb-policy.json

# 1. Delete Lambda Function
aws lambda delete-function --function-name UrlShortenerFunction --region "$AWS_REGION"

# 2. Delete DynamoDB Table
aws dynamodb delete-table --table-name "$TABLE_NAME" --region "$AWS_REGION"

# 3. Delete IAM Role & Policies
aws iam delete-role-policy --role-name UrlShortenerLambdaRole --policy-name DynamoDBShortenerAccess
aws iam detach-role-policy --role-name UrlShortenerLambdaRole --policy-arn arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
aws iam delete-role --role-name UrlShortenerLambdaRole