Function ingest

misc/python/materialize/workload_replay/ingest.py:338–500

Ingest data into a source (Postgres, MySQL, Kafka, SQL Server).

(
    c: Composition,
    child: dict[str, Any],
    source: dict[str, Any],
    columns: list[Column],
    num_rows: int,
    rng: random.Random,
)

def ingest(
    c: Composition,
    child: dict[str, Any],
    source: dict[str, Any],
    columns: list[Column],
    num_rows: int,
    rng: random.Random,
) -> None:
    """Ingest data into a source (Postgres, MySQL, Kafka, SQL Server)."""
    if source["type"] == "postgres":
        ref_database, ref_schema, ref_table = get_postgres_reference_db_schema_table(
            child
        )
        conn = psycopg.connect(
            host="127.0.0.1",
            port=c.default_port("postgres"),
            user="postgres",
            password="postgres",
            dbname=ref_database,
        )
        conn.autocommit = True

        col_names = [col.name for col in columns]

        with conn.cursor() as cur:
            copy_stmt = SQL("COPY {}.{} ({}) FROM STDIN").format(
                Identifier(ref_schema),
                Identifier(ref_table),
                SQL(", ").join(map(Identifier, col_names)),
            )

            with cur.copy(copy_stmt) as copy:
                for _ in range(num_rows):
                    row = [col.value(rng, in_query=False) for col in columns]
                    copy.write_row(row)

    elif source["type"] == "mysql":
        ref_database, ref_table = get_mysql_reference_db_table(child)

        conn = pymysql.connect(
            host="127.0.0.1",
            user="root",
            password=MySql.DEFAULT_ROOT_PASSWORD,
            database=ref_database,
            port=c.default_port("mysql"),
            autocommit=False,
        )

        value_funcs = [col.value for col in columns]
        rows_sql = []
        for _ in range(num_rows):
            row = [fn(rng) for fn in value_funcs]
            rows_sql.append("(" + ", ".join(row) + ")")

        stmt = f"INSERT INTO {ref_table} VALUES " + ", ".join(rows_sql)

        with conn.cursor() as cur:
            cur.execute(stmt)