oxspeak_server/scripts/generate_messages.py

#!/usr/bin/env python3
"""Generate test messages directly in the project's SQLite database."""

from __future__ import annotations

import argparse
import random
import sqlite3
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path


WORD_POOL = (
    "message", "canal", "serveur", "utilisateur", "test", "donnee", "histoire",
    "discussion", "contenu", "generation", "curseur", "fenetre", "lecture",
    "chargement", "conversation", "exemple", "texte", "systeme", "application",
    "client", "serveur", "base", "requete", "resultat", "information", "session",
    "connexion", "fonction", "version", "contenu", "rapide", "simple", "aleatoire",
    "important", "nouveau", "ancien", "prochain", "precedent", "visible", "local",
    "distant", "stable", "chronologique", "variable", "longueur", "performance",
    "validation", "operation", "transaction", "historique", "position", "defilement",
)
MESSAGE_MARKER_FORMAT = "[{number:04d}]"


def positive_int(value: str) -> int:
    parsed = int(value)
    if parsed <= 0:
        raise argparse.ArgumentTypeError("must be greater than zero")
    return parsed


def parse_uuid(value: str, option_name: str) -> uuid.UUID:
    try:
        return uuid.UUID(value)
    except ValueError as error:
        raise argparse.ArgumentTypeError(f"{option_name} is not a valid UUID: {value}") from error


def next_uuid(previous: uuid.UUID | None) -> uuid.UUID:
    """Return a UUID v7 strictly greater than the previous generated ID."""
    generated = uuid.uuid7()
    if previous is not None and generated.int <= previous.int:
        generated = uuid.UUID(int=previous.int + 1)
    return generated


def random_message(
    rng: random.Random,
    min_words: int,
    max_words: int,
    marker: str,
) -> str:
    # The marker itself counts as one word in the requested range.
    body_count = rng.randint(max(0, min_words - 1), max_words - 1)
    body = " ".join(rng.choices(WORD_POOL, k=body_count))
    return f"{marker} {body}".rstrip()


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--db",
        type=Path,
        default=Path("oxspeak.db"),
        help="SQLite database path (default: oxspeak.db)",
    )
    parser.add_argument("--channel-id", required=True, help="target channel UUID")
    parser.add_argument("--user-id", required=True, help="author user UUID")
    parser.add_argument(
        "--count",
        required=True,
        type=positive_int,
        help="number of messages to insert",
    )
    parser.add_argument(
        "--min-words",
        type=positive_int,
        default=10,
        help="minimum number of words per message (default: 10)",
    )
    parser.add_argument(
        "--max-words",
        type=positive_int,
        default=500,
        help="maximum number of words per message (default: 500)",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=None,
        help="optional seed to reproduce generated contents",
    )
    parser.add_argument(
        "--batch-size",
        type=positive_int,
        default=500,
        help="number of rows inserted per batch (default: 500)",
    )
    return parser


def ensure_target_exists(
    connection: sqlite3.Connection,
    table: str,
    identifier: bytes,
    label: str,
) -> None:
    row = connection.execute(
        f'SELECT 1 FROM "{table}" WHERE id = ? LIMIT 1',
        (identifier,),
    ).fetchone()
    if row is None:
        raise ValueError(f"{label} does not exist in the database")


def generate_messages(
    database: Path,
    channel_id: uuid.UUID,
    user_id: uuid.UUID,
    count: int,
    batch_size: int,
    min_words: int,
    max_words: int,
    seed: int | None,
) -> None:
    started_at = time.monotonic()
    connection = sqlite3.connect(database)
    connection.execute("PRAGMA foreign_keys = ON")
    connection.execute("PRAGMA busy_timeout = 5000")

    try:
        ensure_target_exists(connection, "channel", channel_id.bytes, "channel")
        ensure_target_exists(connection, "user", user_id.bytes, "user")

        previous_id: uuid.UUID | None = None
        inserted = 0
        rng = random.Random(seed)

        connection.execute("BEGIN")
        try:
            while inserted < count:
                current_batch_size = min(batch_size, count - inserted)
                rows = []

                for offset in range(current_batch_size):
                    message_id = next_uuid(previous_id)
                    previous_id = message_id
                    message_number = inserted + offset + 1
                    marker = MESSAGE_MARKER_FORMAT.format(number=message_number)
                    content = random_message(rng, min_words, max_words, marker)
                    created_at = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")
                    rows.append(
                        (
                            message_id.bytes,
                            channel_id.bytes,
                            user_id.bytes,
                            content,
                            created_at,
                            None,
                            None,
                        )
                    )

                connection.executemany(
                    """
                    INSERT INTO message
                        (id, channel_id, user_id, content, created_at, updated_at, reply_to_id)
                    VALUES (?, ?, ?, ?, ?, ?, ?)
                    """,
                    rows,
                )
                inserted += current_batch_size

            connection.commit()
        except Exception:
            connection.rollback()
            raise
    finally:
        connection.close()

    elapsed = time.monotonic() - started_at
    print(f"Inserted {count} messages into {database} in {elapsed:.2f}s")


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()

    if args.max_words < args.min_words:
        parser.error("--max-words must be greater than or equal to --min-words")

    try:
        channel_id = parse_uuid(args.channel_id, "--channel-id")
        user_id = parse_uuid(args.user_id, "--user-id")
        generate_messages(
            database=args.db,
            channel_id=channel_id,
            user_id=user_id,
            count=args.count,
            batch_size=args.batch_size,
            min_words=args.min_words,
            max_words=args.max_words,
            seed=args.seed,
        )
    except (OSError, sqlite3.Error, ValueError) as error:
        parser.error(str(error))

    return 0


if __name__ == "__main__":
    raise SystemExit(main())