Select the random row function based on DB driver.

This commit is contained in:
Jake Moshenko 2014-09-15 15:58:56 -04:00
parent 511ee12a58
commit efb66f7c1e
3 changed files with 22 additions and 2 deletions

View file

@ -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))