1 import psycopg2

     2 import psycopg2.extensions

     3 import psycopg2.extras

     4 

     5 

     6 class DatabaseError(Exception):

     7     def __init__(self, msg, query=None):

     8         self.query = query

     9         Exception.__init__(self, msg)

    10 

    11 

    12 class BadConnectionError(Exception):

    13     pass

    14 

    15 

    16 class Row(dict):

    17     def __getattr__(self, key):

    18         return self[key]

    19 

    20 

    21 class Database:

    22     def __init__(self):

    23         # pool of database connections

    24         # indexed by conninfo, items are lists of connections

    25         self.pool = {}

    26         # number of unused connections per conninfo to keep open

    27         self.pool_keep_open = 1

    28 

    29 

    30     def __del__(self):

    31         for conninfo in self.pool.keys():

    32             for conn in self.pool[conninfo]:

    33                 conn.close()

    34 

    35 

    36     def connect(self, conninfo):

    37         try:

    38             conn = psycopg2.connect(conninfo, async=1)

    39             psycopg2.extras.wait_select(conn)

    40         except psycopg2.DatabaseError, e:

    41             raise DatabaseError(str(e))

    42         return conn

    43 

    44 

    45     def get_conn(self, conninfo):

    46         if not conninfo in self.pool:

    47             self.pool[conninfo] = []

    48             return self.connect(conninfo)

    49         else:

    50             conn = None

    51             while len(self.pool[conninfo]) and conn is None:

    52                 conn = self.pool[conninfo].pop()

    53                 if conn.closed:

    54                     conn = None

    55             if conn is None:

    56                 return self.connect(conninfo)

    57         return conn

    58 

    59 

    60     def put_conn(self, conninfo, conn):

    61         if len(self.pool[conninfo]) >= self.pool_keep_open:

    62             conn.close()

    63         else:

    64             self.pool[conninfo].append(conn)

    65 

    66 

    67     def execute(self, q, args=[]):

    68         conn = self.get_conn()

    69         try:

    70             curs = conn.cursor()

    71             curs.execute(q, args)

    72             psycopg2.extras.wait_select(curs.connection)

    73 #            conn.commit()

    74         except psycopg2.OperationalError, e:

    75             # disconnected?

    76 #            conn.rollback()

    77             conn.close()

    78             raise BadConnectionError(str(e))

    79         except psycopg2.DatabaseError, e:

    80 #            conn.rollback()

    81             raise DatabaseError(str(e), curs.query)

    82         return curs

    83 

    84 

    85     def finish(self, curs):

    86         self.put_conn(curs.connection)

    87 

    88 

    89     def row(self, curs, row):

    90         return Row(zip([x[0] for x in curs.description], row))

    91 

    92 

    93     def fetchone(self, curs):

    94         return self.row(curs, curs.fetchone())

    95 

    96 

    97     def fetchall(self, curs):

    98         rows = curs.fetchall()

    99         return [self.row(curs, row) for row in rows]

   100 

   101