def set_shard_at_rest(realm, shard_key, location, force=False):
"""Marks a shard as being at rest in the given location. This is used for
initiating shards in preparation for migration. Unless force is True this
will raise an exception if a shard is already at rest in a specific
location.
:param str realm: The name of the realm for the shard
:param shard_key: The key of the shard
:param str location: The location that the data is at (or should be in the
case of a brand new shard)
:param bool force: Force a shard to be placed at rest in a specific location
even if it has already been placed somewhere.
:return: None
"""
_assert_valid_location(location)
shards_coll = _get_shards_coll()
query = {'realm': realm, 'shard_key': shard_key}
if shards_coll.find(query).count() and not force:
raise Exception(
'Shard with key %s has already been placed. Use force=true if '
'you really want to do this' % shard_key)
shards_coll.update(query,
{
'$set': {
'location': location,
'status': ShardStatus.AT_REST,
},
'$unset': {
'new_location': 1,
},
},
upsert=True)
def set_shard_at_rest(realm, shard_key, location):
"""Marks a shard as being at rest in the given location. This is used for
initiating shards in preparation for migration.
:param str realm: The name of the realm for the shard
:param shard_key: The key of the shard
:param str location: The location that the data is at (or should be in the
case of a brand new shard)
:return: None
"""
shards_coll = _get_shards_coll()
shards_coll.update({
'realm': realm,
'shard_key': shard_key,
},
{
'$set': {
'location': location,
'status': ShardStatus.AT_REST,
},
'$unset': {
'new_location': 1,
},
},
upsert=True)
def set_shard_at_rest(realm, shard_key, location, force=False):
"""Marks a shard as being at rest in the given location. This is used for
initiating shards in preparation for migration. Unless force is True this
will raise an exception if a shard is already at rest in a specific
location.
:param str realm: The name of the realm for the shard
:param shard_key: The key of the shard
:param str location: The location that the data is at (or should be in the
case of a brand new shard)
:param bool force: Force a shard to be placed at rest in a specific location
even if it has already been placed somewhere.
:return: None
"""
_assert_valid_location(location)
shards_coll = _get_shards_coll()
query = {'realm': realm, 'shard_key': shard_key}
if shards_coll.find(query).count() and not force:
raise Exception(
'Shard with key %s has already been placed. Use force=true if '
'you really want to do this' % shard_key)
shards_coll.update(query, {
'$set': {
'location': location,
'status': ShardStatus.AT_REST,
},
'$unset': {
'new_location': 1,
},
},
upsert=True)
realm_changed(realm)
def set_shard_to_migration_status(realm, shard_key, status):
"""Marks a shard as being at a specific migration status.
"""
shards_coll = _get_shards_coll()
shards_coll.update(
{'realm': realm, 'shard_key': shard_key},
{'$set': {'status': status}}
)
def create_indices():
realm_coll = _get_realm_coll()
realm_coll.ensure_index([('name', 1)], unique=True)
realm_coll.ensure_index([('collection', 1)], unique=True)
shards_coll = _get_shards_coll()
shards_coll.ensure_index(
[('realm', 1), ('shard_key', 1)], unique=True)
cluster_coll = _get_cluster_coll()
cluster_coll.ensure_index([('name', 1)], unique=True)
def _should_pause_write(collection_name, query):
realm = _get_realm_for_collection(collection_name)
shard_key = _get_query_target(collection_name, query)
if shard_key:
meta = _get_metadata_for_shard(realm, shard_key)
return \
meta['status'] == ShardStatus.POST_MIGRATION_PAUSED_AT_DESTINATION
else:
paused_query = {
'realm': realm['name'],
'status': ShardStatus.POST_MIGRATION_PAUSED_AT_DESTINATION
}
shards_coll = _get_shards_coll()
return shards_coll.find(paused_query).count() > 0
def _should_pause_write(collection_name, query):
realm = _get_realm_for_collection(collection_name)
shard_key = _get_query_target(collection_name, query)
if shard_key:
meta = _get_metadata_for_shard(realm, shard_key)
return \
meta['status'] == ShardStatus.POST_MIGRATION_PAUSED_AT_DESTINATION
else:
paused_query = {
'realm': realm['name'],
'status': ShardStatus.POST_MIGRATION_PAUSED_AT_DESTINATION
}
shards_coll = _get_shards_coll()
return shards_coll.find(paused_query).count() > 0
def start_migration(realm_name, shard_key, new_location):
"""Marks a shard as being in the process of being migrated.
"""
shards_coll = _get_shards_coll()
realm = _get_realm_by_name(realm_name)
existing_location = _get_location_for_shard(realm, shard_key)
if existing_location.location == new_location:
raise Exception('Shard is already at %s' % new_location)
shards_coll.update(
{'realm': realm_name, 'shard_key': shard_key},
{'$set': {
'status': ShardStatus.MIGRATING_COPY,
'new_location': new_location,
}},
)
def start_migration(realm_name, shard_key, new_location):
"""Marks a shard as being in the process of being migrated.
"""
shards_coll = _get_shards_coll()
realm = _get_realm_by_name(realm_name)
existing_location = _get_location_for_shard(realm, shard_key)
if existing_location.location == new_location:
raise Exception('Shard is already at %s' % new_location)
shards_coll.update(
{
'realm': realm_name,
'shard_key': shard_key
},
{
'$set': {
'status': ShardStatus.MIGRATING_COPY,
'new_location': new_location,
}
},
)
def start_migration(realm, shard_key, new_location):
"""Marks a shard as being in the process of being migrated.
"""
shards_coll = _get_shards_coll()
shard_metas = list(
shards_coll.find({'realm': realm, 'shard_key': shard_key},))
if not shard_metas:
raise Exception(
'Could not find shard metadata - use set_shard_at_rest first')
shard_meta, = shard_metas
if shard_meta['location'] == new_location:
raise Exception('Shard is already at %s' % (new_location,))
shards_coll.update(
{'realm': realm, 'shard_key': shard_key},
{'$set': {
'status': ShardStatus.MIGRATING_COPY,
'new_location': new_location,
}},
)
def _reset_sharding_info():
"""Wipes all shard info. For internal test use only.
"""
_get_realm_coll().remove({})
_get_shards_coll().remove({})