57 lines
1.7 KiB
Python
57 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
|