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 import logging
  12 import psycopg2
  13
  14 class Connection(object):
  15         """
  16                 A lightweight wrapper around MySQLdb DB-API connections.
  17
  18                 The main value we provide is wrapping rows in a dict/object so that
  19                 columns can be accessed by name. Typical usage::
  20
  21                         db = torndb.Connection("localhost", "mydatabase")
  22                         for article in db.query("SELECT * FROM articles"):
  23                                 print article.title
  24
  25                 Cursors are hidden by the implementation, but other than that, the methods
  26                 are very similar to the DB-API.
  27
  28                 We explicitly set the timezone to UTC and the character encoding to
  29                 UTF-8 on all connections to avoid time zone and encoding errors.
  30         """
  31         def __init__(self, host, database, user=None, password=None):
  32                 self.host = host
  33                 self.database = database
  34
  35                 self._db = None
  36                 self._db_args = {
  37                         "host"     : host,
  38                         "database" : database,
  39                         "user"     : user,
  40                         "password" : password,
  41                 }
  42
  43                 try:
  44                         self.reconnect()
  45                 except Exception:
  46                         logging.error("Cannot connect to database on %s", self.host, exc_info=True)
  47
  48         def __del__(self):
  49                 self.close()
  50
  51         def close(self):
  52                 """
  53                         Closes this database connection.
  54                 """
  55                 if getattr(self, "_db", None) is not None:
  56                         self._db.close()
  57                         self._db = None
  58
  59         def reconnect(self):
  60                 """
  61                         Closes the existing database connection and re-opens it.
  62                 """
  63                 self.close()
  64
  65                 self._db = psycopg2.connect(**self._db_args)
  66                 self._db.autocommit = True
  67
  68                 # Initialize the timezone setting.
  69                 self.execute("SET TIMEZONE TO 'UTC'")
  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(zip(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                         logging.warning("Database connection was lost...")
 161
 162                         self.reconnect()
 163
 164         def _cursor(self):
 165                 self._ensure_connected()
 166                 return self._db.cursor()
 167
 168         def _execute(self, cursor, query, parameters, kwparameters):
 169                 logging.debug("SQL Query: %s" % (query % (kwparameters or parameters)))
 170
 171                 try:
 172                         return cursor.execute(query, kwparameters or parameters)
 173                 except (OperationalError, psycopg2.ProgrammingError):
 174                         logging.error("Error connecting to database on %s", self.host)
 175                         self.close()
 176                         raise
 177
 178         def transaction(self):
 179                 return Transaction(self)
 180
 181
 182 class Row(dict):
 183         """A dict that allows for object-like property access syntax."""
 184         def __getattr__(self, name):
 185                 try:
 186                         return self[name]
 187                 except KeyError:
 188                         raise AttributeError(name)
 189
 190
 191 class Transaction(object):
 192         def __init__(self, db):
 193                 self.db = db
 194
 195                 self.db.execute("START TRANSACTION")
 196
 197         def __enter__(self):
 198                 return self
 199
 200         def __exit__(self, exctype, excvalue, traceback):
 201                 if exctype is not None:
 202                         self.db.execute("ROLLBACK")
 203                 else:
 204                         self.db.execute("COMMIT")
 205
 206
 207 # Alias some common exceptions
 208 IntegrityError = psycopg2.IntegrityError
 209 OperationalError = psycopg2.OperationalError