generate_payments_data.py

# generate_payments_data.py
#
# What this does: creates a SQLite DB with fake-but-realistic B2B payments data
# for a LatAm fintech demo. Three tables: clients, payments, payment_errors.
#
# Run it with:  python generate_payments_data.py
# It will nuke the old DB and rebuild from scratch every time.
#
# ---- LESSONS LEARNED while building this ----
# 1. SQLite doesn't actually enforce foreign keys unless you run
#    PRAGMA foreign_keys = ON — wasted 30 min wondering why bad inserts worked fine
# 2. log-uniform distribution (math.log trick) gives WAY more realistic payment
#    amounts than just random.uniform. uniform gives too many huge amounts.
# 3. Batch inserts (executemany) vs one-by-one: 12k rows went from ~8s to ~0.3s
#    after I switched. Worth it.
# 4. random.choices with weights is perfect for simulating the 85/10/5 status split
# 5. Had to remember to drop tables in the right order (errors → payments → clients)
#    because of FK constraints, even with PRAGMA foreign_keys = OFF
# ---------------------------------------------

import os
import math
import random
import sqlite3
import datetime

# seed so results are reproducible — good for demos
SEED = 42
random.seed(SEED)

DB_PATH = "payments_demo.db"

# country → local currency, keeping it LatAm focused
COUNTRY_CURRENCY = {
    "CO": "COP",
    "MX": "MXN",
    "AR": "ARS",
    "BR": "BRL",
    "CL": "CLP",
    "PE": "PEN",
}
COUNTRIES = list(COUNTRY_CURRENCY.keys())

SEGMENTS = ["SMB", "Mid-Market", "Enterprise"]

# mix and match these to get fake company names
# not super creative but it gets the job done
COMPANY_PREFIXES = [
    "Grupo", "Soluciones", "Inversiones", "Comercial", "Digital",
    "Tech", "Nexo", "Capital", "Plataforma", "Finanzas",
    "Red", "Alianza", "Gestión", "Servi", "Multi",
]
COMPANY_SUFFIXES = [
    "Andina", "Global", "Latina", "Express", "Pro",
    "360", "Hub", "Net", "Plus", "Corp",
    "SAS", "Ltda", "SA", "Partners", "Group",
]

PAYMENT_METHODS    = ["bank_transfer", "card", "wallet"]
PAYMENT_DIRECTIONS = ["incoming", "outgoing"]
PAYMENT_STATUSES   = ["completed", "pending", "failed"]

# realistic split: ~85% completed, ~10% pending, ~5% failed
STATUS_WEIGHTS = [0.85, 0.10, 0.05]

# bank_transfer dominates B2B (65%), then card (20%), wallet (15%)
METHOD_WEIGHTS = [0.65, 0.20, 0.15]

# error codes you'd see in a real payments system
ERROR_CODES = [
    "INSUFFICIENT_FUNDS",
    "NETWORK_ERROR",
    "INVALID_ACCOUNT",
    "LIMIT_EXCEEDED",
    "DUPLICATE_TRANSACTION",
    "AUTHORIZATION_FAILED",
    "TIMEOUT",
    "INVALID_CURRENCY",
]

# map each code to a human-readable message
# FIXME: in production these would come from a shared constants file, not hardcoded here
ERROR_MESSAGES = {
    "INSUFFICIENT_FUNDS":    "Client account has insufficient funds to complete the transaction.",
    "NETWORK_ERROR":         "A network timeout occurred while communicating with the payment processor.",
    "INVALID_ACCOUNT":       "The destination account number is invalid or does not exist.",
    "LIMIT_EXCEEDED":        "Transaction amount exceeds the client's daily or monthly limit.",
    "DUPLICATE_TRANSACTION": "A duplicate transaction was detected within the idempotency window.",
    "AUTHORIZATION_FAILED":  "The transaction was declined by the issuing bank.",
    "TIMEOUT":               "The payment request timed out before a response was received.",
    "INVALID_CURRENCY":      "The specified currency is not supported for this payment corridor.",
}

# payment amount ranges per segment (USD-equivalent)
# enterprise goes way higher, SMB stays small
VOLUME_RANGES = {
    "SMB":        (500.0,     20_000.0),
    "Mid-Market": (5_000.0,  100_000.0),
    "Enterprise": (50_000.0, 500_000.0),
}


# ---- helper functions --------------------------------------------------------

def make_company_name(used_names):
    # keep trying until we get a name we haven't used yet
    # 200 tries is overkill but with 15x15=225 combos and 45 clients it's fine
    for _ in range(200):
        name = f"{random.choice(COMPANY_PREFIXES)} {random.choice(COMPANY_SUFFIXES)}"
        if name not in used_names:
            used_names.add(name)
            return name
    # if we somehow exhaust combos, just append the count — shouldn't happen
    fallback = f"{random.choice(COMPANY_PREFIXES)} {random.choice(COMPANY_SUFFIXES)} {len(used_names)}"
    used_names.add(fallback)
    return fallback


def random_dt(start, end):
    # lol SQLite datetime parsing is annoying so I just store ISO strings
    # pick a random day + random second within that day
    delta    = (end - start).days
    rand_day = random.randint(0, delta)
    rand_sec = random.randint(0, 86_399)
    base = datetime.datetime.combine(start, datetime.time.min)
    return base + datetime.timedelta(days=rand_day, seconds=rand_sec)


def payment_amount(segment):
    # log-uniform so we get a realistic mix of small and large payments
    # instead of everything clustering at the midpoint like uniform would do
    # not sure if this is the best approach but it looks convincing in queries
    lo, hi = VOLUME_RANGES[segment]
    return round(math.exp(random.uniform(math.log(lo), math.log(hi))), 2)


# ---- database setup ----------------------------------------------------------

def drop_and_create_tables(conn):
    # nuke everything and start fresh
    # drop order matters: errors → payments → clients (FK dependencies)
    conn.executescript("""
        PRAGMA foreign_keys = OFF;

        DROP TABLE IF EXISTS payment_errors;
        DROP TABLE IF EXISTS payments;
        DROP TABLE IF EXISTS clients;

        CREATE TABLE clients (
            client_id          INTEGER  PRIMARY KEY AUTOINCREMENT,
            client_name        TEXT     NOT NULL,
            segment            TEXT     NOT NULL
                                        CHECK(segment IN ('SMB','Mid-Market','Enterprise')),
            country            TEXT     NOT NULL,
            created_at         DATETIME NOT NULL,
            monthly_volume_avg NUMERIC  NOT NULL
        );

        CREATE TABLE payments (
            payment_id  INTEGER  PRIMARY KEY AUTOINCREMENT,
            client_id   INTEGER  NOT NULL REFERENCES clients(client_id),
            date        DATE     NOT NULL,
            amount      NUMERIC  NOT NULL,
            currency    TEXT     NOT NULL,
            status      TEXT     NOT NULL
                                  CHECK(status IN ('completed','pending','failed')),
            method      TEXT     NOT NULL,
            direction   TEXT     NOT NULL
                                  CHECK(direction IN ('incoming','outgoing')),
            created_at  DATETIME NOT NULL
        );

        CREATE TABLE payment_errors (
            error_id      INTEGER  PRIMARY KEY AUTOINCREMENT,
            payment_id    INTEGER  NOT NULL REFERENCES payments(payment_id),
            error_code    TEXT     NOT NULL,
            error_message TEXT     NOT NULL,
            timestamp     DATETIME NOT NULL
        );

        PRAGMA foreign_keys = ON;
    """)
    conn.commit()
    print("  done: tables created")


# ---- data generation ---------------------------------------------------------

def build_segment_pool(n):
    # rough distribution: ~44% SMB, ~33% Mid-Market, ~22% Enterprise
    # enterprise is rare in real B2B portfolios but drives most volume
    n_smb        = round(n * 0.44)
    n_mid        = round(n * 0.33)
    n_enterprise = n - n_smb - n_mid
    pool = ["SMB"] * n_smb + ["Mid-Market"] * n_mid + ["Enterprise"] * n_enterprise
    random.shuffle(pool)
    return pool, n_smb, n_mid, n_enterprise


def insert_clients(conn, n=45):
    # generates n fake B2B clients and inserts them
    # returns a list of dicts so generate_payments() knows what to create
    cur = conn.cursor()

    pool, n_smb, n_mid, n_ent = build_segment_pool(n)

    # clients were onboarded in 2022-2023, before the 2024 payment data
    onboard_start = datetime.date(2022, 6, 1)
    onboard_end   = datetime.date(2023, 12, 31)

    clients    = []
    used_names = set()

    for segment in pool:
        country  = random.choice(COUNTRIES)
        currency = COUNTRY_CURRENCY[country]
        lo, hi   = VOLUME_RANGES[segment]

        # monthly avg is somewhere in the upper half of the segment range
        monthly_avg = round(random.uniform(lo * 1.5, hi), 2)
        name        = make_company_name(used_names)
        created_at  = random_dt(onboard_start, onboard_end)

        cur.execute(
            "INSERT INTO clients (client_name, segment, country, created_at, monthly_volume_avg) "
            "VALUES (?, ?, ?, ?, ?)",
            (name, segment, country, created_at.isoformat(), monthly_avg),
        )
        clients.append({
            "client_id": cur.lastrowid,
            "segment":   segment,
            "country":   country,
            "currency":  currency,
        })

    conn.commit()
    print(f"  done: {len(clients)} clients (SMB={n_smb}, Mid-Market={n_mid}, Enterprise={n_ent})")
    return clients


def build_client_weights(clients):
    # enterprise clients get more transactions because they do more volume
    # not sure if 8x is the right multiplier but it looks realistic in queries
    seg_weight = {"SMB": 1, "Mid-Market": 3, "Enterprise": 8}
    return [seg_weight[c["segment"]] for c in clients]


def make_payment_row(client, period_start, period_end):
    # build a single payment tuple ready for executemany
    amount    = payment_amount(client["segment"])
    status    = random.choices(PAYMENT_STATUSES, weights=STATUS_WEIGHTS, k=1)[0]
    method    = random.choices(PAYMENT_METHODS,  weights=METHOD_WEIGHTS,  k=1)[0]
    direction = random.choice(PAYMENT_DIRECTIONS)
    pay_dt    = random_dt(period_start, period_end)

    return (
        client["client_id"],
        pay_dt.date().isoformat(),
        amount,
        client["currency"],
        status,
        method,
        direction,
        pay_dt.isoformat(),
    )


def insert_payments(conn, clients, total=12_000):
    # This is where I got stuck for a bit:
    # first version inserted one row at a time and took forever (~8s for 12k rows)
    # switched to executemany with batches of 500 and now it's basically instant
    cur     = conn.cursor()
    weights = build_client_weights(clients)

    period_start = datetime.date(2024, 1, 1)
    period_end   = datetime.date(2024, 12, 31)

    BATCH = 500
    batch = []

    for _ in range(total):
        client = random.choices(clients, weights=weights, k=1)[0]
        batch.append(make_payment_row(client, period_start, period_end))

        if len(batch) >= BATCH:
            cur.executemany(
                "INSERT INTO payments "
                "(client_id, date, amount, currency, status, method, direction, created_at) "
                "VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
                batch,
            )
            conn.commit()
            batch = []

    # flush whatever's left over
    if batch:
        cur.executemany(
            "INSERT INTO payments "
            "(client_id, date, amount, currency, status, method, direction, created_at) "
            "VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
            batch,
        )
        conn.commit()

    # grab the IDs of failed payments so we can attach error records to them
    cur.execute("SELECT payment_id FROM payments WHERE status = 'failed'")
    failed_ids = [row[0] for row in cur.fetchall()]

    print(f"  done: {total} payments inserted, {len(failed_ids)} failed")
    return failed_ids


def fetch_failed_payment_timestamps(cur, failed_ids):
    # need timestamps so error records happen AFTER the payment, not before
    cur.execute("SELECT payment_id, created_at FROM payments WHERE status = 'failed'")
    # build a dict for fast lookups
    return {
        row[0]: datetime.datetime.fromisoformat(row[1])
        for row in cur.fetchall()
    }


def insert_errors(conn, failed_ids, target=250):
    # attach error records to the failed payments
    # some payments get more than one error (e.g. timeout then auth failure on retry)
    if not failed_ids:
        print("  warning: no failed payments found, skipping errors")
        return

    cur    = conn.cursor()
    ts_map = fetch_failed_payment_timestamps(cur, failed_ids)

    # random.choices with replacement so some payments get multiple errors — realistic
    sampled = random.choices(failed_ids, k=target)

    rows = []
    for pid in sampled:
        code     = random.choice(ERROR_CODES)
        message  = ERROR_MESSAGES[code]
        pay_ts   = ts_map.get(pid, datetime.datetime(2024, 1, 1))
        # error happens within 5 minutes of the payment attempt
        error_ts = pay_ts + datetime.timedelta(seconds=random.randint(1, 300))
        rows.append((pid, code, message, error_ts.isoformat()))

    cur.executemany(
        "INSERT INTO payment_errors (payment_id, error_code, error_message, timestamp) "
        "VALUES (?, ?, ?, ?)",
        rows,
    )
    conn.commit()
    print(f"  done: {len(rows)} payment_errors inserted")


# ---- summary query -----------------------------------------------------------

def print_summary(conn):
    cur = conn.cursor()

    cur.execute("SELECT COUNT(*) FROM clients")
    n_clients = cur.fetchone()[0]

    cur.execute("SELECT COUNT(*) FROM payments")
    n_payments = cur.fetchone()[0]

    cur.execute("SELECT COUNT(*) FROM payment_errors")
    n_errors = cur.fetchone()[0]

    # breakdown by status so we can sanity-check the 85/10/5 split
    cur.execute("SELECT status, COUNT(*) FROM payments GROUP BY status ORDER BY status")
    breakdown = cur.fetchall()

    print("\n" + "=" * 50)
    print("  done! here's what was created:")
    print("=" * 50)
    print(f"  clients        : {n_clients:,}")
    print(f"  payments       : {n_payments:,}")
    for status, count in breakdown:
        pct = count / n_payments * 100
        print(f"    {status:<12} : {count:,}  ({pct:.1f}%)")
    print(f"  payment_errors : {n_errors:,}")
    print(f"  db file        : {DB_PATH}")
    print("=" * 50)


# ---- main --------------------------------------------------------------------

def main():
    # blow away the old DB so we always start clean
    if os.path.exists(DB_PATH):
        os.remove(DB_PATH)
        print(f"deleted old {DB_PATH}")

    print(f"\nbuilding {DB_PATH}...\n")
    conn = sqlite3.connect(DB_PATH)

    print("[1/3] creating tables...")
    drop_and_create_tables(conn)

    print("\n[2/3] generating clients...")
    clients = insert_clients(conn, n=45)

    print("\n[3/3] generating payments + errors...")
    failed_ids = insert_payments(conn, clients, total=12_000)
    insert_errors(conn, failed_ids, target=250)

    print_summary(conn)
    conn.close()


if __name__ == "__main__":
    main()