quay/data/readreplica.py

import random

from collections import namedtuple

from peewee import Model, SENTINEL, OperationalError, Proxy

ReadOnlyConfig = namedtuple('ReadOnlyConfig', ['is_readonly', 'read_replicas'])

class ReadOnlyModeException(Exception):
  """ Exception raised if a write operation was attempted when in read only mode.
  """


class AutomaticFailoverWrapper(object):
  """ Class which wraps a peewee database driver and (optionally) a second driver.
      When executing SQL, if an OperationalError occurs, if a second driver is given,
      the query is attempted again on the fallback DB. Otherwise, the exception is raised.
  """
  def __init__(self, primary_db, fallback_db=None):
    self._primary_db = primary_db
    self._fallback_db = fallback_db

  def __getattr__(self, attribute):
    if attribute != 'execute_sql' and hasattr(self._primary_db, attribute):
      return getattr(self._primary_db, attribute)

    return getattr(self, attribute)

  def execute(self, query, commit=SENTINEL, **context_options):
    ctx = self.get_sql_context(**context_options)
    sql, params = ctx.sql(query).query()
    return self.execute_sql(sql, params, commit=commit)

  def execute_sql(self, sql, params=None, commit=SENTINEL):
    try:
      return self._primary_db.execute_sql(sql, params, commit)
    except OperationalError:
      if self._fallback_db is not None:
        try:
          return self._fallback_db.execute_sql(sql, params, commit)
        except OperationalError:
          raise


class ReadReplicaSupportedModel(Model):
  """ Base model for peewee data models that support using a read replica for SELECT
      requests not under transactions, and automatic failover to the master if the
      read replica fails.

      Read-only queries are initially attempted on one of the read replica databases
      being used; if an OperationalError occurs when attempting to invoke the query,
      then the failure is logged and the query is retried on the database master.

      Queries that are non-SELECTs (or under transactions) are always tried on the
      master.

      If the system is configured into read only mode, then all non-read-only queries
      will raise a ReadOnlyModeException.
  """
  @classmethod
  def _read_only_config(cls):
    read_only_config = getattr(cls._meta, 'read_only_config', None)
    if read_only_config is None:
      return ReadOnlyConfig(False, [])

    if isinstance(read_only_config, Proxy) and read_only_config.obj is None:
      return ReadOnlyConfig(False, [])

    return read_only_config.obj or ReadOnlyConfig(False, [])

  @classmethod
  def _in_readonly_mode(cls):
    return cls._read_only_config().is_readonly

  @classmethod
  def _select_database(cls):
    """ Selects a read replica database if we're configured to support read replicas.
        Otherwise, selects the master database.
    """
    # Select the master DB if read replica support is not enabled.
    read_only_config = cls._read_only_config()
    if not read_only_config.read_replicas:
      return cls._meta.database

    # Select the master DB if we're ever under a transaction.
    if cls._meta.database.transaction_depth() > 0:
      return cls._meta.database

    # Otherwise, return a read replica database with auto-retry onto the main database.
    replicas = read_only_config.read_replicas
    selected_read_replica = replicas[random.randrange(len(replicas))]
    return AutomaticFailoverWrapper(selected_read_replica, cls._meta.database)

  @classmethod
  def select(cls, *args, **kwargs):
    query = super(ReadReplicaSupportedModel, cls).select(*args, **kwargs)
    query._database = cls._select_database()
    return query

  @classmethod
  def insert(cls, *args, **kwargs):
    query = super(ReadReplicaSupportedModel, cls).insert(*args, **kwargs)
    if cls._in_readonly_mode():
      raise ReadOnlyModeException()
    return query

  @classmethod
  def update(cls, *args, **kwargs):
    query = super(ReadReplicaSupportedModel, cls).update(*args, **kwargs)
    if cls._in_readonly_mode():
      raise ReadOnlyModeException()
    return query

  @classmethod
  def delete(cls, *args, **kwargs):
    query = super(ReadReplicaSupportedModel, cls).delete(*args, **kwargs)
    if cls._in_readonly_mode():
      raise ReadOnlyModeException()
    return query

  @classmethod
  def raw(cls, *args, **kwargs):
    query = super(ReadReplicaSupportedModel, cls).raw(*args, **kwargs)
    if query._sql.lower().startswith('select '):
      query._database = cls._select_database()
    elif cls._in_readonly_mode():
      raise ReadOnlyModeException()

    return query