This repository has been archived on 2020-03-24. You can view files and clone it, but cannot push or open issues or pull requests.
quay/data/read_slave.py
2014-07-02 19:10:24 -04:00

56 lines
1.7 KiB
Python

"""
Adapted from:
https://github.com/coleifer/peewee/blob/master/playhouse/read_slave.py
Support for using a dedicated read-slave. The read database is specified as a
Model.Meta option, and will be used for SELECT statements:
master = PostgresqlDatabase('master')
read_slave = PostgresqlDatabase('read_slave')
class BaseModel(ReadSlaveModel):
class Meta:
database = master
read_slaves = [read_slave] # This database will be used for SELECTs.
# Now define your models as you would normally.
class User(BaseModel):
username = CharField()
# To force a SELECT on the master database, you can instantiate the SelectQuery
# by hand:
master_select = SelectQuery(User).where(...)
"""
from peewee import *
class ReadSlaveModel(Model):
@classmethod
def _get_read_database(cls):
if (not getattr(cls._meta, 'read_slaves', None) or
cls._meta.database.transaction_depth() > 0):
return cls._meta.database
current_idx = getattr(cls, '_read_slave_idx', -1)
cls._read_slave_idx = (current_idx + 1) % len(cls._meta.read_slaves)
selected_read_slave = cls._meta.read_slaves[cls._read_slave_idx]
if isinstance(selected_read_slave, Proxy) and selected_read_slave.obj is None:
# It's possible the read slave was disabled by not initializing it
return cls._meta.database
return selected_read_slave
@classmethod
def select(cls, *args, **kwargs):
query = super(ReadSlaveModel, cls).select(*args, **kwargs)
query.database = cls._get_read_database()
return query
@classmethod
def raw(cls, *args, **kwargs):
query = super(ReadSlaveModel, cls).raw(*args, **kwargs)
if query._sql.lower().startswith('select'):
query.database = cls._get_read_database()
return query