backend/database.py

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