VersionAlchemy benchmark code

benchmark.py

from contextlib import contextmanager
import sys
import time

import sqlalchemy as sa
from sqlalchemy import create_engine
from sqlalchemy.exc import InternalError
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy.schema import UniqueConstraint
from sqlalchemy.sql.expression import insert
import versionalchemy as va
from versionalchemy.models import VAModelMixin, VALogMixin

# --- Constants ---
TRIALS = int(sys.argv[1])
VALUE = 'TEST'

engine = create_engine('<insert mysql url her>')
Session = sessionmaker(bind=engine)
Base = declarative_base(bind=engine)


# --- Test Model ---
class TestTable(Base):
    __tablename__ = 'test_table'

    id = sa.Column(sa.Integer, primary_key=True)
    value = sa.Column(sa.String(128))


# --- Helpers ---
@contextmanager
def record_time():
    start_time = time.time()
    yield
    end_time = time.time()
    print "execution took {}s".format(end_time - start_time)


def new_environment(Base, c):
    for t in ['test_table', 'test_table_archive']:
        try:
            c.execute('drop table {}'.format(t))
        except InternalError as e:
            if e.orig[0] == 1051:
                continue
    Base.metadata.create_all()


# --- Core test ---
print 'Running core test...',
conn = engine.connect()
new_environment(Base, conn)
with record_time():
    try:
        for i in xrange(TRIALS):
            with conn.begin():
                conn.execute(insert(TestTable).values(value=VALUE))
    finally:
        conn.close()


# --- ORM Test ---
print 'Running ORM test...',
session = Session()
new_environment(Base, session)
with record_time():
    for i in xrange(TRIALS):
        session.add(TestTable(value=VALUE))
        session.commit()
session.close()

# --- Versionalchemy Test ---
VABase = declarative_base(bind=engine)


class TestTableArchive(VABase, VALogMixin):
    __tablename__ = 'test_table_archive'
    __table_args__ = (
        UniqueConstraint('id', 'va_version'),
    )
    id = sa.Column(sa.Integer, primary_key=True)
    user_id = sa.Column(sa.Integer)


class TestTableVA(VABase, VAModelMixin):
    __tablename__ = 'test_table'
    va_version_columns = ['id']

    id = sa.Column(sa.Integer, primary_key=True)
    value = sa.Column(sa.String(128))

print 'Running Versionalchemy test...',
session = Session()
new_environment(VABase, session)
va.init()
TestTableVA.register(TestTableArchive, engine)
with record_time():
    for i in xrange(TRIALS):
        session.add(TestTableVA(value=VALUE))
        session.commit()
session.close()