# 2026-03-19
# stress_writer.py — пишет N строк в terraform_demo_table (по умолчанию 5).
# Проверяет параллельные INSERT'ы и устойчивость соединения с PG при нагрузке.

import os
import psycopg2
import time

_VERSION = "v1"


def run(event):
    n = int(event.get("rows", 5))
    prefix = event.get("prefix", "stress")

    conn = psycopg2.connect(
        host=os.environ["PGHOST"],
        port=int(os.environ.get("PGPORT", "5432")),
        dbname=os.environ["PGDATABASE"],
        user=os.environ["PGUSER"],
        password=os.environ["PGPASSWORD"],
        sslmode=os.environ.get("PGSSLMODE", "require"),
    )
    inserted = []
    try:
        with conn.cursor() as cur:
            for i in range(n):
                title = f"{prefix}-{int(time.time()*1000)}-{i}"
                cur.execute(
                    "INSERT INTO terraform_demo_table (title) VALUES (%s) RETURNING id",
                    (title,),
                )
                row = cur.fetchone()
                inserted.append({"id": row[0], "title": title})
        conn.commit()
    finally:
        conn.close()

    return {"version": _VERSION, "inserted": inserted, "count": len(inserted)}