webapp/backend/database.py

   1 #!/usr/bin/env python
   2
   3 """
   4         A lightweight wrapper around psycopg2.
   5
   6         Originally part of the Tornado framework.  The tornado.database module
   7         is slated for removal in Tornado 3.0, and it is now available separately
   8         as torndb.
   9 """
  10
  11 from __future__ import absolute_import, division, with_statement
  12
  13 import copy
  14 import itertools
  15 import logging
  16 import os
  17 import psycopg2
  18 import time
  19
  20 class Connection(object):
  21         """
  22                 A lightweight wrapper around MySQLdb DB-API connections.
  23
  24                 The main value we provide is wrapping rows in a dict/object so that
  25                 columns can be accessed by name. Typical usage::
  26
  27                         db = torndb.Connection("localhost", "mydatabase")
  28                         for article in db.query("SELECT * FROM articles"):
  29                                 print article.title
  30
  31                 Cursors are hidden by the implementation, but other than that, the methods
  32                 are very similar to the DB-API.
  33
  34                 We explicitly set the timezone to UTC and the character encoding to
  35                 UTF-8 on all connections to avoid time zone and encoding errors.
  36         """
  37         def __init__(self, host, database, user=None, password=None):
  38                 self.host = host
  39                 self.database = database
  40
  41                 self._db = None
  42                 self._db_args = {
  43                         "host"     : host,
  44                         "database" : database,
  45                         "user"     : user,
  46                         "password" : password,
  47                 }
  48
  49                 try:
  50                         self.reconnect()
  51                 except Exception:
  52                         logging.error("Cannot connect to database on %s", self.host, exc_info=True)
  53
  54         def __del__(self):
  55                 self.close()
  56
  57         def close(self):
  58                 """
  59                         Closes this database connection.
  60                 """
  61                 if getattr(self, "_db", None) is not None:
  62                         self._db.close()
  63                         self._db = None
  64
  65         def reconnect(self):
  66                 """
  67                         Closes the existing database connection and re-opens it.
  68                 """
  69                 self.close()
  70
  71                 self._db = psycopg2.connect(**self._db_args)
  72                 self._db.autocommit = True
  73
  74         def query(self, query, *parameters, **kwparameters):
  75                 """
  76                         Returns a row list for the given query and parameters.
  77                 """
  78                 cursor = self._cursor()
  79                 try:
  80                         self._execute(cursor, query, parameters, kwparameters)
  81                         column_names = [d[0] for d in cursor.description]
  82                         return [Row(itertools.izip(column_names, row)) for row in cursor]
  83                 finally:
  84                         cursor.close()
  85
  86         def get(self, query, *parameters, **kwparameters):
  87                 """
  88                         Returns the first row returned for the given query.
  89                 """
  90                 rows = self.query(query, *parameters, **kwparameters)
  91                 if not rows:
  92                         return None
  93                 elif len(rows) > 1:
  94                         raise Exception("Multiple rows returned for Database.get() query")
  95                 else:
  96                         return rows[0]
  97
  98         def execute(self, query, *parameters, **kwparameters):
  99                 """
 100                         Executes the given query, returning the lastrowid from the query.
 101                 """
 102                 return self.execute_lastrowid(query, *parameters, **kwparameters)
 103
 104         def execute_lastrowid(self, query, *parameters, **kwparameters):
 105                 """
 106                         Executes the given query, returning the lastrowid from the query.
 107                 """
 108                 cursor = self._cursor()
 109                 try:
 110                         self._execute(cursor, query, parameters, kwparameters)
 111                         return cursor.lastrowid
 112                 finally:
 113                         cursor.close()
 114
 115         def execute_rowcount(self, query, *parameters, **kwparameters):
 116                 """
 117                         Executes the given query, returning the rowcount from the query.
 118                 """
 119                 cursor = self._cursor()
 120                 try:
 121                         self._execute(cursor, query, parameters, kwparameters)
 122                         return cursor.rowcount
 123                 finally:
 124                         cursor.close()
 125
 126         def executemany(self, query, parameters):
 127                 """
 128                         Executes the given query against all the given param sequences.
 129
 130                         We return the lastrowid from the query.
 131                 """
 132                 return self.executemany_lastrowid(query, parameters)
 133
 134         def executemany_lastrowid(self, query, parameters):
 135                 """
 136                         Executes the given query against all the given param sequences.
 137
 138                         We return the lastrowid from the query.
 139                 """
 140                 cursor = self._cursor()
 141                 try:
 142                         cursor.executemany(query, parameters)
 143                         return cursor.lastrowid
 144                 finally:
 145                         cursor.close()
 146
 147         def executemany_rowcount(self, query, parameters):
 148                 """
 149                         Executes the given query against all the given param sequences.
 150
 151                         We return the rowcount from the query.
 152                 """
 153                 cursor = self._cursor()
 154
 155                 try:
 156                         cursor.executemany(query, parameters)
 157                         return cursor.rowcount
 158                 finally:
 159                         cursor.close()
 160
 161         def _ensure_connected(self):
 162                 if self._db is None:
 163                         self.reconnect()
 164
 165         def _cursor(self):
 166                 self._ensure_connected()
 167                 return self._db.cursor()
 168
 169         def _execute(self, cursor, query, parameters, kwparameters):
 170                 try:
 171                         return cursor.execute(query, kwparameters or parameters)
 172                 except OperationalError:
 173                         logging.error("Error connecting to database on %s", self.host)
 174                         self.close()
 175                         raise
 176
 177
 178 class Row(dict):
 179         """A dict that allows for object-like property access syntax."""
 180         def __getattr__(self, name):
 181                 try:
 182                         return self[name]
 183                 except KeyError:
 184                         raise AttributeError(name)
 185
 186
 187 # Alias some common exceptions
 188 IntegrityError = psycopg2.IntegrityError
 189 OperationalError = psycopg2.OperationalError