S127 · Design a URL Shortener — the Classic Warm-Up

#!/usr/bin/env python3
# shortener.py — a minimal but real URL shortener.
# pip install flask redis psycopg2-binary
# docker run -d -p 6379:6379 redis:7
# docker run -d -p 5432:5432 -e POSTGRES_PASSWORD=pw postgres:16
import os
import string
 
import psycopg2
import redis
from flask import Flask, abort, jsonify, redirect, request
 
app = Flask(__name__)
r = redis.Redis(host="localhost", port=6379, decode_responses=True)
pg = psycopg2.connect("dbname=postgres user=postgres password=pw host=localhost")
pg.autocommit = True
 
with pg.cursor() as c:
    c.execute(
        """
        CREATE TABLE IF NOT EXISTS urls (
            code TEXT PRIMARY KEY,
            long_url TEXT NOT NULL,
            created_at TIMESTAMPTZ DEFAULT now()
        )
        """
    )
 
ALPHABET = string.ascii_letters + string.digits  # 62 chars
BASE = len(ALPHABET)
CACHE_TTL_SEC = 24 * 3600
 
def encode_base62(n: int) -> str:
    if n == 0:
        return ALPHABET[0]
    out = []
    while n > 0:
        n, rem = divmod(n, BASE)
        out.append(ALPHABET[rem])
    return "".join(reversed(out))
 
def next_id() -> int:
    # Real systems: a range-allocation service (each app gets 10k IDs at a time
    # from ZooKeeper or a DB sequence) so INCR isn't a single hot key.
    return r.incr("url:counter")
 
@app.post("/shorten")
def shorten():
    body = request.get_json(force=True)
    long_url = (body or {}).get("url", "").strip()
    if not long_url.startswith(("http://", "https://")) or len(long_url) > 2048:
        abort(400, description="Invalid URL")
    code = encode_base62(next_id())
    with pg.cursor() as c:
        c.execute(
            "INSERT INTO urls (code, long_url) VALUES (%s, %s) ON CONFLICT DO NOTHING",
            (code, long_url),
        )
    r.setex(f"url:{code}", CACHE_TTL_SEC, long_url)  # warm the cache on write
    return jsonify({"short": request.host_url + code, "code": code})
 
@app.get("/<code>")
def follow(code):
    # 1. Try Redis (hot path — ~1 ms)
    cached = r.get(f"url:{code}")
    if cached:
        r.incr(f"clicks:{code}")  # in prod: push to Kafka instead
        return redirect(cached, code=302)
 
    # 2. Cache miss → Postgres (rare — ~50 ms)
    with pg.cursor() as c:
        c.execute("SELECT long_url FROM urls WHERE code = %s", (code,))
        row = c.fetchone()
    if not row:
        abort(404)
    long_url = row[0]
 
    # 3. Warm the cache for next time
    r.setex(f"url:{code}", CACHE_TTL_SEC, long_url)
    r.incr(f"clicks:{code}")
    return redirect(long_url, code=302)
 
@app.get("/_stats/<code>")
def stats(code):
    return jsonify({"clicks": int(r.get(f"clicks:{code}") or 0)})
 
if __name__ == "__main__":
    app.run(port=5000, debug=True)