pg_fake.py

#! /usr/bin/env python

from __future__ import print_function

import datetime
import uuid
import random

from faker import Faker

from sqlalchemy import create_engine
from sqlalchemy.orm import scoped_session, sessionmaker, relationship
from sqlalchemy.ext.declarative import declarative_base

from sqlalchemy import (Boolean, BigInteger, Column, DateTime, ForeignKey, Index, Integer,
                        Numeric, String, Table)
from sqlalchemy.dialects.postgresql import JSONB, UUID


Base = declarative_base()

SQLALCHEMY_DATABASE_URI = 'postgres://localhost/example'


class Customer(Base):
    __tablename__ = 'customers'

    id = Column(UUID, primary_key=True)
    name = Column(String(100))
    comment = Column(String(500))

    orders = relationship(
        'Order',
        foreign_keys='[Order.customer_id]',
        back_populates='customer'
    )

class Order(Base):
    __tablename__ = 'orders'

    id = Column(UUID, primary_key=True)
    address_unit = Column(String(50))
    address_line_1 = Column(String(500))
    address_city = Column(String(100))
    address_zip = Column(String(50))
    address_country = Column(String(100))

    customer_id = Column(UUID, ForeignKey('customers.id', ondelete='CASCADE'))
    customer = relationship('Customer', foreign_keys=[customer_id], back_populates='orders')

def setup_db_session(expire_on_commit=True):
    # DB connection added to request globals via Flask.before_request()
    engine = create_engine(
        SQLALCHEMY_DATABASE_URI,
        convert_unicode=True,
    )
    db_session = scoped_session(sessionmaker(
        autocommit=False,
        autoflush=False,
        expire_on_commit=expire_on_commit,
        bind=engine
    ))
    Base.query = db_session.query_property()

    return db_session


def create_tables(db_session):
    # create the DB tables
    engine = create_engine(SQLALCHEMY_DATABASE_URI, convert_unicode=True)
    Base.metadata.create_all(bind=engine)
    db_session.commit()


if __name__ == '__main__':
    db_session = setup_db_session()
    create_tables(db_session)

    faker = Faker()

    for x in xrange(1500000):
        c = Customer(
            id=str(uuid.uuid4()),
            name=faker.name(),
            comment=faker.sentence(),
        )
        db_session.add(c)

        # random number of orders
        for y in xrange(random.randint(1,3)):
            o = Order(
                id=str(uuid.uuid4()),
                customer=c,
                address_unit=faker.building_number(),
                address_line_1=faker.street_address(),
                address_city=faker.city(),
                address_zip=faker.zipcode(),
                address_country=faker.country(),
            )
            db_session.add(o)

        # commit every 10000 inserts
        if x > 0 and x % 10000 == 0:
            print('{} {}'.format(datetime.datetime.now().isoformat(), x))
            db_session.commit()

    print('{} {}'.format(datetime.datetime.now().isoformat(), x))
    db_session.commit()