1 #

     2 # See LICENCE for detail

     3 # Copyright (c) 2014 IRI

     4 #

     5 

     6 from sqlalchemy import create_engine, MetaData

     7 from sqlalchemy.orm import scoped_session, sessionmaker

     8 from sqlalchemy.ext.declarative import declarative_base

     9 

    10 import psycopg2.extras

    11 

    12 from twisted.python import log

    13 

    14 from txpostgres.txpostgres import Connection, ConnectionPool

    15 

    16 import config

    17 

    18 

    19 engine = create_engine(config.CONN_STR, convert_unicode=True)

    20 db_session = scoped_session(sessionmaker(autocommit=False,

    21                                          autoflush=False,

    22                                          bind=engine))

    23 

    24 Base = declarative_base()

    25 Base.query = db_session.query_property()

    26 

    27 def init_db():

    28     # import all modules here that might define models so that

    29     # they will be registered properly on the metadata.  Otherwise

    30     # you will have to import them first before calling init_db()

    31     import models

    32     Base.metadata.create_all(bind=engine)

    33 

    34 

    35 class DictConnection(Connection):

    36 

    37     @staticmethod

    38     def connectionFactory(*args, **kwargs):

    39         kwargs['connection_factory'] = psycopg2.extras.DictConnection

    40         return psycopg2.connect(*args, **kwargs)

    41 

    42 

    43 class DictConnectionPool(ConnectionPool):

    44     connectionFactory = DictConnection

    45 

    46 

    47 def create_connection_pool(conn_string):

    48 

    49     created_connection_pool = DictConnectionPool(None, conn_string)

    50 

    51     d = created_connection_pool.start()

    52     d.addErrback(log.err)

    53 

    54     return created_connection_pool, d