def test_shard_retrieve_shard_from_user_when_used_with_model_and_signal(self):
from django.contrib.auth import get_user_model
user = get_user_model().objects.create_user(
username='******', password='******', email='*****@*****.**'
)
sut = RoundRobinBucketingStrategy(shard_group='default', databases=settings.DATABASES)
self.assertEqual(user.shard, sut.get_shard(user))
self.assertIsNotNone(user.shard)
def test_picking_is_cyclic(self):
class FakeUser():
username = None
expected_cycle = ['app_shard_001', 'app_shard_002']
sut = RoundRobinBucketingStrategy(shard_group='default', databases=settings.DATABASES)
resulting_shards = [sut.pick_shard(FakeUser()) for i in xrange(100)]
self.assertEqual(len(set([resulting_shards[i] for i in xrange(0, 100, 2)])), 1)
self.assertEqual(len(set([resulting_shards[i] for i in xrange(1, 100, 2)])), 1)
resulting_cycled_shard = resulting_shards[:2]
resulting_cycled_shard.sort()
self.assertEqual(resulting_cycled_shard, expected_cycle)
def ready(self):
shard_settings = getattr(settings, 'DJANGO_SHARDING_SETTINGS', {})
shard_groups = [
settings.DATABASES[db_settings]['SHARD_GROUP']
for db_settings in settings.DATABASES
]
shard_groups = set(
filter(lambda group: group is not None, shard_groups))
self.bucketers = {}
self.routing_strategies = {}
for shard_group in shard_groups:
group_settings = shard_settings.get(shard_group, {})
self.bucketers[shard_group] = group_settings.get(
'BUCKETING_STRATEGY',
RoundRobinBucketingStrategy(shard_group=shard_group,
databases=settings.DATABASES))
self.routing_strategies[shard_group] = group_settings.get(
'ROUTING_STRATEGY',
PrimaryOnlyRoutingStrategy(databases=settings.DATABASES))
# Unless otherwise instructed, add the signal to save the shard to the model if it has a shard field.
for model in apps.get_models():
if getattr(model,
'django_sharding__stores_shard', False) and getattr(
model, 'django_sharding__shard_group', None):
shard_group = getattr(model, 'django_sharding__shard_group',
None)
shard_field = getattr(model, 'django_sharding__shard_field',
None)
if not shard_field:
raise Exception(
'The model {} must have a `shard_field` attribute'.
format(model))
else:
shard_fields = list(
filter(
lambda field: getattr(
field, 'django_sharding__stores_shard', False),
model._meta.fields))
if not any(shard_fields):
continue
if len(shard_fields) > 1:
raise Exception(
'The model {} has multiple fields for shard storage: {}'
.format(model, shard_fields))
shard_field = shard_fields[0]
shard_group = getattr(shard_field,
'django_sharding__shard_group', None)
if not shard_group:
raise Exception(
'The model {} with the shard field must have a `shard_group` attribute'
.format(model))
if not getattr(shard_field, 'django_sharding__use_signal',
False):
continue
group_settings = shard_settings.get(shard_group, {})
if group_settings.get('SKIP_ADD_SHARDED_SIGNAL', False):
continue
receiver(models.signals.pre_save, sender=model)(save_shard_handler)
