Function _mysql_chunk

misc/python/materialize/workload_replay/data.py:178–208

Generate random data and INSERT into MySQL.

(
    conn_params: dict[str, Any],
    table: str,
    column_dicts: list[dict[str, Any]],
    num_rows: int,
    rng_seed: int,
)

def _mysql_chunk(
    conn_params: dict[str, Any],
    table: str,
    column_dicts: list[dict[str, Any]],
    num_rows: int,
    rng_seed: int,
) -> int:
    """Generate random data and INSERT into MySQL."""
    import pymysql

    rng = random.Random(rng_seed)
    columns = [
        Column(c["name"], c["type"], c["nullable"], c["default"], c.get("data_shape"))
        for c in column_dicts
    ]

    conn = pymysql.connect(**conn_params)

    batch_size = 10000
    for start in range(0, num_rows, batch_size):
        batch_rows = min(batch_size, num_rows - start)
        rows_sql = []
        for _ in range(batch_rows):
            row = [col.value(rng) for col in columns]
            rows_sql.append("(" + ", ".join(row) + ")")
        stmt = f"INSERT INTO {table} VALUES " + ", ".join(rows_sql)
        with conn.cursor() as cur:
            cur.execute(stmt)

    conn.close()
    return num_rows


def _submit_chunks(