sqlalchemy_mysql_insert_performance.py

# CREATE USER 'example'@'localhost' IDENTIFIED BY 'example_password';
# GRANT ALL PRIVILEGES ON *.* TO 'example'@'localhost' IDENTIFIED BY 'example_password';
# CREATE DATABASE example;
import sys
import os
import time
from datetime import date, timedelta
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy import Column, Integer, String,  create_engine
from sqlalchemy.dialects.mysql import DATETIME, DATE
from sqlalchemy.orm import scoped_session, sessionmaker

Base = declarative_base()
DBSession = scoped_session(sessionmaker())
engine = None
dbname = "mysql+mysqldb://example:example_password@localhost:3306/example"
date_list = []
date_str_list = []
today = date.today() - timedelta(days = 1000)
j = 1
for i in range(1000):
    j += i
    j = j % 40000
    date_list.append(today - timedelta(days = j))
    date_str_list.append((today - timedelta(days = j)).strftime('%Y-%m-%d'))

class Customer(Base):
    __tablename__ = "customer"
    id              = Column(Integer, primary_key=True, autoincrement=True)
    name            = Column(String(50), nullable=False)
    date_of_birth   = Column(DATE, nullable=False)


def init_sqlalchemy():
    global engine
    engine = create_engine(dbname, echo=False)
    DBSession.remove()
    DBSession.configure(bind=engine, autoflush=False, expire_on_commit=False)
    Base.metadata.drop_all(engine)
    Base.metadata.create_all(engine)


def test_sqlalchemy_orm(n=100000):
    init_sqlalchemy()
    t0 = time.time()
    for i in xrange(n):
        customer = Customer()
        customer.name = 'NAME ' + str(i)
        customer.date_of_birth = date_list[i%len(date_list)]
        DBSession.add(customer)
        if i % 1000 == 0:
            DBSession.flush()
    DBSession.commit()
    print(
        "SQLAlchemy ORM: Total time for " + str(n) +
        " records " + str(time.time() - t0) + " secs")


def test_sqlalchemy_orm_pk_given(n=100000):
    init_sqlalchemy()
    t0 = time.time()
    for i in xrange(n):
        customer = Customer(
            name="NAME " + str(i),
            date_of_birth = date_list[i%len(date_list)]
        )
        DBSession.add(customer)
        if i % 1000 == 0:
            DBSession.flush()
    DBSession.commit()
    print(
        "SQLAlchemy ORM pk given: Total time for " + str(n) +
        " records " + str(time.time() - t0) + " secs")


def test_sqlalchemy_orm_bulk_save_objects(n=100000):
    init_sqlalchemy()
    t0 = time.time()
    n1 = n
    while n1 > 0:
        n1 = n1 - 10000
        DBSession.bulk_save_objects(
            [
                Customer(
                    name="NAME " + str(i),
                    date_of_birth = date_list[i%len(date_list)]
                )
                for i in xrange(min(10000, n1))
            ]
        )
    DBSession.commit()
    print(
        "SQLAlchemy ORM bulk_save_objects(): Total time for " + str(n) +
        " records " + str(time.time() - t0) + " secs")

def test_sqlalchemy_orm_bulk_insert(n=100000):
    init_sqlalchemy()
    t0 = time.time()
    n1 = n
    while n1 > 0:
        n1 = n1 - 10000
        DBSession.bulk_insert_mappings(
            Customer,
            [
                dict(
                    name="NAME " + str(i),
                    date_of_birth = date_list[i%len(date_list)]
                )
                for i in xrange(min(10000, n1))
            ]
        )
    DBSession.commit()
    print(
        "SQLAlchemy ORM bulk_insert_mappings(): Total time for " + str(n) +
        " records " + str(time.time() - t0) + " secs")

def test_sqlalchemy_core(n=100000):
    init_sqlalchemy()
    t0 = time.time()
    engine.execute(
        Customer.__table__.insert(),
        [
            {"name": 'NAME ' + str(i),
            'date_of_birth': date_list[i%len(date_list)]}
            for i in xrange(n)
        ]
    )
    print(
        "SQLAlchemy Core: Total time for " + str(n) +
        " records " + str(time.time() - t0) + " secs")


def test_mysql(n=100000, dbname='sqlite3.db'):
    init_sqlalchemy()
    t0 = time.time()
    for i in xrange(n):
        DBSession.execute(
            "INSERT INTO customer (name, date_of_birth) VALUES ('{0}', '{1}')".format(
                'NAME ' + str(i), date_str_list[i%len(date_str_list)]
            )
        )
    DBSession.commit()
    print(
        "mysql: Total time for " + str(n) +
        " records " + str(time.time() - t0) + " sec")

def test_raw_insert(n=100000):
    init_sqlalchemy()
    t0 = time.time()
    f = open('customers_to_insert', 'w')
    for i in xrange(n):
        f.write(
            "{0},{1}\n".format(
                'NAME ' + str(i), date_str_list[i%len(date_str_list)]
            )
        )
    f.close()
    os.system(
        './raw_data_import_script.sh \
        {file_name} {table_name} {db_name} {username} {password}'.format(
            file_name  = 'customers_to_insert',
            table_name = 'customer',
            db_name    = 'example',
            username   = 'example',
            password   = 'example_password'
        )
    )
    os.remove('./customers_to_insert')
    print(
        "mysql raw insert: Total time for " + str(n) +
        " records " + str(time.time() - t0) + " sec")


if __name__ == '__main__':
    try:
        n = int(sys.argv[1])
    except (ValueError, IndexError) as e:
        n = 500000

    test_sqlalchemy_orm(n)
    test_sqlalchemy_orm_pk_given(n)
    test_sqlalchemy_orm_bulk_save_objects(n)
    test_sqlalchemy_orm_bulk_insert(n)
    test_sqlalchemy_core(n)
    test_mysql(n)
    test_raw_insert(n)