TechEmpower.FrameworkBenchmarks/frameworks/Python/pyramid/frameworkbenchmarks/models.py

"""
Benchmark models.
"""

import json
import os
import psycopg2
from collections import Iterable
from sqlalchemy import create_engine, MetaData, Table, Column, Integer, String
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import QueuePool
from sqlalchemy.ext.declarative import declarative_base, DeclarativeMeta


def get_conn():
    return psycopg2.connect(
        user='benchmarkdbuser',
        password='benchmarkdbpass',
        host='tfb-database',
        port='5432',
        database='hello_world')


conn_pool = QueuePool(get_conn, pool_size=100, max_overflow=25, echo=False)

pg = create_engine('postgresql://', pool=conn_pool)
DBSession = sessionmaker(bind=pg)()
metadata = MetaData()

DatabaseBase = declarative_base()


def sqlalchemy_encoder_factory(system_values):
    return SQLAlchemyEncoder()


class SQLAlchemyEncoder(json.JSONEncoder):
    def __call__(self, obj, system_values):
        if isinstance(obj, Iterable):
            return json.dumps([self.default(x) for x in obj])
        else:
            return json.dumps(self.default(obj))

    def default(self, obj):
        if isinstance(obj.__class__, DeclarativeMeta):
            return obj.__json__()
        return super(SQLAlchemyEncoder, self).default(obj)


class World(DatabaseBase):
    __tablename__ = 'world'

    id = Column('id', Integer, primary_key=True)
    randomNumber = Column(
        'randomnumber', Integer, nullable=False, server_default='0')

    def __json__(self):
        return {'id': self.id, 'randomNumber': self.randomNumber}


class Fortune(DatabaseBase):
    __tablename__ = 'fortune'

    id = Column('id', Integer, primary_key=True)
    message = Column('message', String, nullable=False)

    def __json__(self):
        return {'id': self.id, 'message': self.message}