23044656b97d-all-bulk-inserts-examples.py

import contextlib
import time

from sqlalchemy import Column, String, create_engine
from sqlalchemy.dialects.mysql import INTEGER
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import Session

Base = declarative_base()


db_url = "mysql+mysqldb://root:root@127.0.0.1:13306/data"

engine = create_engine(db_url, pool_size=5, pool_recycle=3600)


class Customer(Base):
    __tablename__ = "customers"
    id = Column(INTEGER, primary_key=True)
    name = Column(String(255))


@contextlib.contextmanager
def get_session(cleanup=False):
    session = Session(bind=engine)
    Base.metadata.create_all(engine)

    try:
        yield session
    except Exception:
        session.rollback()
    finally:
        session.close()

    if cleanup:
        Base.metadata.drop_all(engine)


@contextlib.contextmanager
def get_conn(cleanup=False):
    conn = engine.connect()
    Base.metadata.create_all(engine)

    yield conn
    conn.close()

    if cleanup:
        Base.metadata.drop_all(engine)


def test_orm_add_one_by_one_no_pks(cleanup=True, num=20000):
    start_time = time.time()
    with get_session(cleanup) as session:
        for idx in range(num):
            session.add(Customer(name=f"Customer - {idx+1}"))
        session.commit()
    duration = time.time() - start_time
    print(f"ORM add one by one with no PKs takes {duration:.2f} seconds")
    # ~5 seconds


def test_orm_add_all_no_pks(cleanup=True, num=20000):
    start_time = time.time()
    with get_session(cleanup) as session:
        session.add_all(
            [Customer(name=f"Customer - {idx+1}") for idx in range(num)]
        )
        session.commit()
    duration = time.time() - start_time
    print(f"ORM add_all with no PKs takes {duration:.2f} seconds")
    # ~5 seconds


def test_orm_add_all_with_pks(cleanup=True, num=20000):
    start_time = time.time()
    with get_session(cleanup) as session:
        customers = [
            Customer(id=idx + 1, name=f"Customer - {idx+1}")
            for idx in range(num)
        ]
        session.add_all(customers)
        session.commit()

    duration = time.time() - start_time
    print(f"ORM add_all with PKs takes {duration:.2f} seconds")
    # 1 second


def test_orm_bulk_save_objects_return_pks(cleanup=True, num=20000):
    start_time = time.time()
    with get_session(cleanup) as session:
        session.bulk_save_objects(
            [Customer(name=f"Customer - {idx+1}") for idx in range(num)],
            return_defaults=True,
        )
        session.commit()
    duration = time.time() - start_time
    print(f"ORM bulk_save_objects returning PKs: {duration:.2f} seconds")
    # ~5 seconds


def test_orm_bulk_save_objects_no_returns(cleanup=True, num=20000):
    start_time = time.time()
    with get_session(cleanup) as session:
        session.bulk_save_objects(
            [Customer(name=f"Customer - {idx+1}") for idx in range(num)]
        )
        session.commit()
    duration = time.time() - start_time
    print(f"ORM bulk_save_objects returning no PKs: {duration:.2f} seconds")
    # ~1 second


def test_orm_bulk_save_objects_with_pks(cleanup=True, num=20000):
    start_time = time.time()
    with get_session(cleanup) as session:
        session.bulk_save_objects(
            [
                Customer(id=idx + 1, name=f"Customer - {idx+1}")
                for idx in range(num)
            ]
        )
        session.commit()
    duration = time.time() - start_time
    print(f"ORM bulk_save_objects with PKs: {duration:.2f} seconds")
    # ~1 second


def test_orm_bulk_insert_mappings_no_pks(cleanup=True, num=20000):
    start_time = time.time()
    with get_session(cleanup) as session:
        session.bulk_insert_mappings(
            Customer, [{"name": f"Customer - {idx+1}"} for idx in range(num)]
        )
        session.commit()
    duration = time.time() - start_time
    print(f"ORM bulk_insert_mappings no pks: {duration:.2f} seconds.")
    # ~0.5 second


def test_orm_bulk_insert_mappings_with_pks(cleanup=True, num=20000):
    start_time = time.time()
    with get_session(cleanup) as session:
        session.bulk_insert_mappings(
            Customer,
            [
                {"id": idx + 1, "name": f"Customer - {idx+1}"}
                for idx in range(num)
            ],
        )
        session.commit()
    duration = time.time() - start_time
    print(f"ORM bulk_insert_mappings with pks: {duration:.2f} seconds.")
    # ~0.5 second


def test_core_insert_expression_no_pks(cleanup=True, num=20000):
    start_time = time.time()
    with get_conn(cleanup) as conn:
        insert_obj = Customer.__table__.insert()
        conn.execute(
            insert_obj, [{"name": f"Customer - {idx+1}"} for idx in range(num)]
        )
    duration = time.time() - start_time
    print(f"Core insert with expression no pks: {duration:.2f} seconds.")
    # ~0.4 second


def test_core_insert_expression_with_pks(cleanup=True, num=20000):
    start_time = time.time()
    with get_conn(cleanup) as conn:
        insert_obj = Customer.__table__.insert()
        conn.execute(
            insert_obj,
            [
                {"id": idx + 1, "name": f"Customer - {idx+1}"}
                for idx in range(num)
            ],
        )

    duration = time.time() - start_time
    print(f"Core insert with expression: {duration:.2f} seconds.")
    # ~0.4 second


def test_core_insert_plain_sql_no_pks(cleanup=True, num=20000):
    start_time = time.time()
    sql_query = """
        INSERT INTO data.customers
            (name)
        VALUES (%(name)s)
    """

    with get_conn(cleanup) as conn:
        conn.exec_driver_sql(
            sql_query, [{"name": f"Customer - {idx+1}"} for idx in range(num)]
        )
    duration = time.time() - start_time
    print(f"Core insert with plain query no pks: {duration:.2f} seconds.")
    # ~0.4 sec


def test_core_insert_plain_sql_with_pks(cleanup=True, num=20000):
    start_time = time.time()
    sql_query = """
        INSERT INTO data.customers
            (id, name)
        VALUES (%(id)s, %(name)s)
    """

    with get_conn(cleanup) as conn:
        conn.exec_driver_sql(
            sql_query,
            [
                {"id": idx + 1, "name": f"Customer - {idx+1}"}
                for idx in range(num)
            ],
        )
    duration = time.time() - start_time
    print(f"Core insert with plain query with pks: {duration:.2f} seconds.")
    # ~0.4 seconds.


test_orm_add_one_by_one_no_pks()  # ~5 seconds
test_orm_add_all_no_pks()  # ~5 seconds
test_orm_add_all_with_pks()  # ~1 second

test_orm_bulk_save_objects_return_pks()  # ~5 seconds
test_orm_bulk_save_objects_no_returns()  # ~1 second
test_orm_bulk_save_objects_with_pks()  # ~1 second

test_orm_bulk_insert_mappings_no_pks()  # ~0.5 second
test_orm_bulk_insert_mappings_with_pks() # ~0.5 second

test_core_insert_expression_no_pks()  # ~0.4 second
test_core_insert_expression_with_pks() # ~0.4 second

test_core_insert_plain_sql_no_pks()  # ~0.4 second
test_core_insert_plain_sql_with_pks()  # ~0.4 second