src/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                 # Initialize the timezone setting.
  75                 self.execute("SET TIMEZONE TO 'UTC'")
  76
  77         def query(self, query, *parameters, **kwparameters):
  78                 """
  79                         Returns a row list for the given query and parameters.
  80                 """
  81                 cursor = self._cursor()
  82                 try:
  83                         self._execute(cursor, query, parameters, kwparameters)
  84                         column_names = [d[0] for d in cursor.description]
  85                         return [Row(itertools.izip(column_names, row)) for row in cursor]
  86                 finally:
  87                         cursor.close()
  88
  89         def get(self, query, *parameters, **kwparameters):
  90                 """
  91                         Returns the first row returned for the given query.
  92                 """
  93                 rows = self.query(query, *parameters, **kwparameters)
  94                 if not rows:
  95                         return None
  96                 elif len(rows) > 1:
  97                         raise Exception("Multiple rows returned for Database.get() query")
  98                 else:
  99                         return rows[0]
 100
 101         def execute(self, query, *parameters, **kwparameters):
 102                 """
 103                         Executes the given query, returning the lastrowid from the query.
 104                 """
 105                 return self.execute_lastrowid(query, *parameters, **kwparameters)
 106
 107         def execute_lastrowid(self, query, *parameters, **kwparameters):
 108                 """
 109                         Executes the given query, returning the lastrowid from the query.
 110                 """
 111                 cursor = self._cursor()
 112                 try:
 113                         self._execute(cursor, query, parameters, kwparameters)
 114                         return cursor.lastrowid
 115                 finally:
 116                         cursor.close()
 117
 118         def execute_rowcount(self, query, *parameters, **kwparameters):
 119                 """
 120                         Executes the given query, returning the rowcount from the query.
 121                 """
 122                 cursor = self._cursor()
 123                 try:
 124                         self._execute(cursor, query, parameters, kwparameters)
 125                         return cursor.rowcount
 126                 finally:
 127                         cursor.close()
 128
 129         def executemany(self, query, parameters):
 130                 """
 131                         Executes the given query against all the given param sequences.
 132
 133                         We return the lastrowid from the query.
 134                 """
 135                 return self.executemany_lastrowid(query, parameters)
 136
 137         def executemany_lastrowid(self, query, parameters):
 138                 """
 139                         Executes the given query against all the given param sequences.
 140
 141                         We return the lastrowid from the query.
 142                 """
 143                 cursor = self._cursor()
 144                 try:
 145                         cursor.executemany(query, parameters)
 146                         return cursor.lastrowid
 147                 finally:
 148                         cursor.close()
 149
 150         def executemany_rowcount(self, query, parameters):
 151                 """
 152                         Executes the given query against all the given param sequences.
 153
 154                         We return the rowcount from the query.
 155                 """
 156                 cursor = self._cursor()
 157
 158                 try:
 159                         cursor.executemany(query, parameters)
 160                         return cursor.rowcount
 161                 finally:
 162                         cursor.close()
 163
 164         def _ensure_connected(self):
 165                 if self._db is None:
 166                         logging.warning("Database connection was lost...")
 167
 168                         self.reconnect()
 169
 170         def _cursor(self):
 171                 self._ensure_connected()
 172                 return self._db.cursor()
 173
 174         def _execute(self, cursor, query, parameters, kwparameters):
 175                 logging.debug("SQL Query: %s" % (query % (kwparameters or parameters)))
 176
 177                 try:
 178                         return cursor.execute(query, kwparameters or parameters)
 179                 except (OperationalError, psycopg2.ProgrammingError):
 180                         logging.error("Error connecting to database on %s", self.host)
 181                         self.close()
 182                         raise
 183
 184         def transaction(self):
 185                 return Transaction(self)
 186
 187
 188 class Row(dict):
 189         """A dict that allows for object-like property access syntax."""
 190         def __getattr__(self, name):
 191                 try:
 192                         return self[name]
 193                 except KeyError:
 194                         raise AttributeError(name)
 195
 196
 197 class Transaction(object):
 198         def __init__(self, db):
 199                 self.db = db
 200
 201                 self.db.execute("START TRANSACTION")
 202
 203         def __enter__(self):
 204                 return self
 205
 206         def __exit__(self, exctype, excvalue, traceback):
 207                 if exctype is not None:
 208                         self.db.execute("ROLLBACK")
 209                 else:
 210                         self.db.execute("COMMIT")
 211
 212
 213 # Alias some common exceptions
 214 IntegrityError = psycopg2.IntegrityError
 215 OperationalError = psycopg2.OperationalError