src/python/location/database.py

   1 """
   2         A lightweight wrapper around psycopg2.
   3
   4         Originally part of the Tornado framework.  The tornado.database module
   5         is slated for removal in Tornado 3.0, and it is now available separately
   6         as torndb.
   7 """
   8
   9 import logging
  10 import psycopg2
  11 import time
  12
  13 log = logging.getLogger("location.database")
  14 log.propagate = 1
  15
  16 class Connection(object):
  17         """
  18                 A lightweight wrapper around MySQLdb DB-API connections.
  19
  20                 The main value we provide is wrapping rows in a dict/object so that
  21                 columns can be accessed by name. Typical usage::
  22
  23                         db = torndb.Connection("localhost", "mydatabase")
  24                         for article in db.query("SELECT * FROM articles"):
  25                                 print article.title
  26
  27                 Cursors are hidden by the implementation, but other than that, the methods
  28                 are very similar to the DB-API.
  29
  30                 We explicitly set the timezone to UTC and the character encoding to
  31                 UTF-8 on all connections to avoid time zone and encoding errors.
  32         """
  33         def __init__(self, host, database, user=None, password=None):
  34                 self.host = host
  35                 self.database = database
  36
  37                 self._db = None
  38                 self._db_args = {
  39                         "host"     : host,
  40                         "database" : database,
  41                         "user"     : user,
  42                         "password" : password,
  43                         "sslmode"  : "require",
  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(
 173                                 "Executing query: %s" % \
 174                                                 cursor.mogrify(query, kwparameters or parameters).decode(),
 175                 )
 176
 177                 # Store the time when the query started
 178                 t = time.monotonic()
 179
 180                 try:
 181                         return cursor.execute(query, kwparameters or parameters)
 182
 183                 # Catch any errors
 184                 except OperationalError:
 185                         log.error("Error connecting to database on %s", self.host)
 186                         self.close()
 187                         raise
 188
 189                 # Log how long the query took
 190                 finally:
 191                         # Determine duration the query took
 192                         d = time.monotonic() - t
 193
 194                         log.debug("Query took %.2fms" % (d * 1000.0))
 195
 196         def transaction(self):
 197                 return Transaction(self)
 198
 199
 200 class Row(dict):
 201         """A dict that allows for object-like property access syntax."""
 202         def __getattr__(self, name):
 203                 try:
 204                         return self[name]
 205                 except KeyError:
 206                         raise AttributeError(name)
 207
 208
 209 class Transaction(object):
 210         def __init__(self, db):
 211                 self.db = db
 212
 213                 self.db.execute("START TRANSACTION")
 214
 215         def __enter__(self):
 216                 return self
 217
 218         def __exit__(self, exctype, excvalue, traceback):
 219                 if exctype is not None:
 220                         self.db.execute("ROLLBACK")
 221                 else:
 222                         self.db.execute("COMMIT")
 223
 224
 225 # Alias some common exceptions
 226 IntegrityError = psycopg2.IntegrityError
 227 OperationalError = psycopg2.OperationalError