Select the random row function based on DB driver.
This commit is contained in:
parent
511ee12a58
commit
efb66f7c1e
3 changed files with 22 additions and 2 deletions
|
@ -21,8 +21,24 @@ SCHEME_DRIVERS = {
|
|||
'postgresql+psycopg2': PostgresqlDatabase,
|
||||
}
|
||||
|
||||
SCHEME_RANDOM_FUNCTION = {
|
||||
'mysql': fn.Rand,
|
||||
'mysql+pymysql': fn.Rand,
|
||||
'sqlite': fn.Random,
|
||||
'postgresql': fn.Random,
|
||||
'postgresql+psycopg2': fn.Random,
|
||||
}
|
||||
|
||||
class CallableProxy(Proxy):
|
||||
def __call__(self, *args, **kwargs):
|
||||
if self.obj is None:
|
||||
raise AttributeError('Cannot use uninitialized Proxy.')
|
||||
return self.obj(*args, **kwargs)
|
||||
|
||||
db = Proxy()
|
||||
read_slave = Proxy()
|
||||
db_random_func = CallableProxy()
|
||||
|
||||
|
||||
def _db_from_url(url, db_kwargs):
|
||||
parsed_url = make_url(url)
|
||||
|
@ -38,11 +54,15 @@ def _db_from_url(url, db_kwargs):
|
|||
|
||||
return SCHEME_DRIVERS[parsed_url.drivername](parsed_url.database, **db_kwargs)
|
||||
|
||||
|
||||
def configure(config_object):
|
||||
db_kwargs = dict(config_object['DB_CONNECTION_ARGS'])
|
||||
write_db_uri = config_object['DB_URI']
|
||||
db.initialize(_db_from_url(write_db_uri, db_kwargs))
|
||||
|
||||
parsed_write_uri = make_url(write_db_uri)
|
||||
db_random_func.initialize(SCHEME_RANDOM_FUNCTION[parsed_write_uri.drivername])
|
||||
|
||||
read_slave_uri = config_object.get('DB_READ_SLAVE_URI', None)
|
||||
if read_slave_uri is not None:
|
||||
read_slave.initialize(_db_from_url(read_slave_uri, db_kwargs))
|
||||
|
|
Reference in a new issue