src/python/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 log = logging.getLogger("location.database")
  15 log.propagate = 1
  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                         log.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                 # Initialize the timezone setting.
  72                 self.execute("SET TIMEZONE TO 'UTC'")
  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(zip(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                         log.warning("Database connection was lost...")
 164
 165                         self.reconnect()
 166
 167         def _cursor(self):
 168                 self._ensure_connected()
 169                 return self._db.cursor()
 170
 171         def _execute(self, cursor, query, parameters, kwparameters):
 172                 log.debug("SQL Query: %s" % (query % (kwparameters or parameters)))
 173
 174                 try:
 175                         return cursor.execute(query, kwparameters or parameters)
 176                 except (OperationalError, psycopg2.ProgrammingError):
 177                         log.error("Error connecting to database on %s", self.host)
 178                         self.close()
 179                         raise
 180
 181         def transaction(self):
 182                 return Transaction(self)
 183
 184
 185 class Row(dict):
 186         """A dict that allows for object-like property access syntax."""
 187         def __getattr__(self, name):
 188                 try:
 189                         return self[name]
 190                 except KeyError:
 191                         raise AttributeError(name)
 192
 193
 194 class Transaction(object):
 195         def __init__(self, db):
 196                 self.db = db
 197
 198                 self.db.execute("START TRANSACTION")
 199
 200         def __enter__(self):
 201                 return self
 202
 203         def __exit__(self, exctype, excvalue, traceback):
 204                 if exctype is not None:
 205                         self.db.execute("ROLLBACK")
 206                 else:
 207                         self.db.execute("COMMIT")
 208
 209
 210 # Alias some common exceptions
 211 IntegrityError = psycopg2.IntegrityError
 212 OperationalError = psycopg2.OperationalError