def next_id_like(cls, entity_id, pool_id = None):
if not isinstance(entity_id, Id):
entity_id = Id(entity_id)
return MySQL.next_id(
entity_id.get_shard_id(),
pool_id or entity_id.get_pool_id()
)
def build_like_id(cls, length, id):
if isinstance(id, int):
id = Id(id)
return cls.build(
length,
id.get_shard_id(),
id.get_pool_id())
def get(cls, id):
if not isinstance(id,Id):
id = Id(id)
shard = id.get_shard_id()
pool_id = id.get_pool_id()
pool = cls.get_pool(pool_id)
return pool.get_shard(shard)
def get(cls, id):
if not isinstance(id, Id):
id = Id(id)
shard = id.get_shard_id()
pool_id = id.get_pool_id()
pool = cls.get_pool(pool_id)
return pool.get_shard(shard)
def multi_shard_in_list_query(self, in_list, query, other_params = None):
"""
Looks for a %l in the query string and replaces it with (%s, %s ...)
It will parse the in_list parameter and run the correct list on each shard
This is an optimized way to do a multi-shard lookup on a list of primary keys
Limitations:
- only 1 %l param can be in the query
- %l must be the last param in the query
"""
if other_params is None:
other_params = ()
shard_id_to_in_list = {}
for id in in_list:
shard_id = Id(id).get_shard_id()
if not shard_id in shard_id_to_in_list:
shard_id_to_in_list[shard_id] = [id]
else:
shard_id_to_in_list[shard_id].append(id)
qr = InListQueryRunner(self._pool, query, other_params, shard_id_to_in_list, self._use_multi, self._catch_errors)
if self._num_threads <= 1:
res = [qr.run_one_query(i) for i in shard_id_to_in_list.keys()]
else:
with Pool(self._num_threads) as p:
res = p.map(qr.run_one_query, shard_id_to_in_list.keys())
return self._prepare_result(res)
def multi_shard_insert(self, table_name, shard_by_col_name, dicts_to_insert, cols_to_update = None):
count = 0
#construct mapping of inserted objects to shard that they go to
shard_to_dicts = {}
for d in dicts_to_insert:
try:
primary_id = d[shard_by_col_name]
shard_id = Id(primary_id).get_shard_id()
if shard_id not in shard_to_dicts:
shard_to_dicts[shard_id] = []
shard_to_dicts[shard_id].append(d)
except Exception as e:
#skip objects that don't have the shard_by_col, or in wrong format
if not self._catch_errors:
raise e
for shard_id, dict_list in shard_to_dicts.items():
#make sure we only have a separate call for each column set
dict_lists_by_cols = {}
for d in dict_list:
cols = list(d.keys())
cols.sort()
col_str = ",".join(cols)
if not col_str in dict_lists_by_cols:
dict_lists_by_cols[col_str] = []
dict_lists_by_cols[col_str].append(d)
for col_str, dict_list in dict_lists_by_cols.items():
#get vals array
vals = [ ["%s" for k in d.keys()] for d in dict_list ]
#create parameter placeholders
params = [ v for d in dict_list for v in d.values()]
cols = dict_list[0].keys()
qb = SQLQueryBuilder.insert(table_name).columns(cols).values(vals)
if cols_to_update:
update_list = [ ("`" + c + "`", "VALUES(`" + c+ "`)") for c in cols_to_update ]
qb.on_duplicate_key_update(update_list)
#do insert
shard = MySQL.get_by_shard_id(shard_id, self._pool.get_id())
try:
count = count + shard.query(qb.build(), params, self._use_multi)
except Exception as e:
if not self._catch_errors:
raise e
count = 0
return count
def next_id(cls,shard_id = None,pool_id = None, number_ids = 1):
if pool_id is None:
pool_id = MySQLPool.MAIN
pool = cls.get_pool(pool_id)
#pick a random shard id if one wasn't provided
if shard_id is None:
shard_id = random.randrange(0, pool.get_num_shards())
return Id.next(shard_id, pool_id, number_ids)
def next_id(cls, shard_id=None, pool_id=None, number_ids=1):
if pool_id is None:
pool_id = MySQLPool.MAIN
pool = cls.get_pool(pool_id)
#pick a random shard id if one wasn't provided
if shard_id is None:
shard_id = random.randrange(0, pool.get_num_shards())
return Id.next(shard_id, pool_id, number_ids)
def multi_shard_query(self, id_list, query, params = None):
"""
Runs a query on just the shards used for the provided id_list
"""
seen_shard_ids = {}
for id in id_list:
seen_shard_ids[Id(id).get_shard_id()] = True
return self.multi_shard_query_by_shard_id(
seen_shard_ids.keys(), query, params
)
def new_session(self,
auth,
user_agent_string,
log_out_ts=None,
session_flags=None):
user = auth.user
new_session_id = MySQL.next_id(Id(user.id).get_shard_id())
try:
user_agent = self._get_user_agent_by_string(user_agent_string)
except UserAgentNotFoundException:
shard_id = MySQL.get_shard_id_for_string(user_agent_string)
id = MySQL.next_id(shard_id)
user_agent = UserAgent(id, user_agent_string)
self._save_user_agent(user_agent)
return Session(new_session_id,
user,
user_agent,
auth,
log_out_ts=log_out_ts,
flags=session_flags)
def get_temp_id_from_token(cls, token):
shard_id = cls.get_shard_from_token(token)
pool_id = cls.get_pool_from_token(token)
return Id.dummy(shard_id, pool_id)
def build_like_id(cls, length, id):
if isinstance(id, int):
id = Id(id)
return cls.build(length, id.get_shard_id(), id.get_pool_id())
def get_temp_id_from_token(cls, token):
shard_id = cls.get_shard_from_token(token)
pool_id = cls.get_pool_from_token(token)
return Id.dummy(shard_id, pool_id)
def next_id_like(cls, entity_id, pool_id=None):
if not isinstance(entity_id, Id):
entity_id = Id(entity_id)
return MySQL.next_id(entity_id.get_shard_id(), pool_id
or entity_id.get_pool_id())
