def get_db_alias_for_partitioned_doc(partition_value):
if settings.USE_PARTITIONED_DATABASE:
from corehq.form_processor.backends.sql.dbaccessors import ShardAccessor
db_name = ShardAccessor.get_database_for_doc(partition_value)
else:
db_name = 'default'
return db_name
def get_db_alias_for_partitioned_doc(partition_value):
if settings.USE_PARTITIONED_DATABASE:
from corehq.form_processor.backends.sql.dbaccessors import ShardAccessor
db_name = ShardAccessor.get_database_for_doc(partition_value)
else:
db_name = 'default'
return db_name
def test_settings(self):
"""
The tests in this class assume a certain partitioned setup to ensure the
partitioning is working properly, so this test makes sure those assumptions
are valid.
"""
self.assertEqual(len(settings.PARTITION_DATABASE_CONFIG['shards']), 2)
self.assertIn(self.db1, settings.PARTITION_DATABASE_CONFIG['shards'])
self.assertIn(self.db2, settings.PARTITION_DATABASE_CONFIG['shards'])
self.assertEqual(
settings.PARTITION_DATABASE_CONFIG['shards'][self.db1], [0, 1])
self.assertEqual(
settings.PARTITION_DATABASE_CONFIG['shards'][self.db2], [2, 3])
self.assertEqual(set(partition_config.get_form_processing_dbs()),
set([self.db1, self.db2]))
self.assertEqual(ShardAccessor.get_database_for_doc(self.p1_uuid),
self.db1)
self.assertEqual(ShardAccessor.get_database_for_doc(self.p2_uuid),
self.db2)
def test_uuids_used(self):
self.assertEqual(ShardAccessor.get_database_for_doc(self.p1_uuid1), self.db1)
self.assertEqual(ShardAccessor.get_database_for_doc(self.p1_uuid2), self.db1)
self.assertEqual(ShardAccessor.get_database_for_doc(self.p1_uuid3), self.db1)
self.assertEqual(ShardAccessor.get_database_for_doc(self.p2_uuid1), self.db2)
self.assertEqual(ShardAccessor.get_database_for_doc(self.p2_uuid2), self.db2)
self.assertEqual(ShardAccessor.get_database_for_doc(self.p2_uuid3), self.db2)
def test_uuids_used(self):
self.assertEqual(ShardAccessor.get_database_for_doc(self.p1_uuid1), self.db1)
self.assertEqual(ShardAccessor.get_database_for_doc(self.p1_uuid2), self.db1)
self.assertEqual(ShardAccessor.get_database_for_doc(self.p1_uuid3), self.db1)
self.assertEqual(ShardAccessor.get_database_for_doc(self.p2_uuid1), self.db2)
self.assertEqual(ShardAccessor.get_database_for_doc(self.p2_uuid2), self.db2)
self.assertEqual(ShardAccessor.get_database_for_doc(self.p2_uuid3), self.db2)
def _group_objects_by_db(objects):
"""
:param objects: Deserialized object dictionaries
:return: List of tuples of (db_alias, [object,...])
"""
objects_by_db = defaultdict(list)
for obj in objects:
app_label = obj['model']
model = apps.get_model(app_label)
db_alias = router.db_for_write(model)
if settings.USE_PARTITIONED_DATABASE and db_alias == partition_config.get_proxy_db():
doc_id = _get_doc_id(app_label, obj)
db_alias = ShardAccessor.get_database_for_doc(doc_id)
objects_by_db[db_alias].append(obj)
return list(objects_by_db.items())
def _group_objects_by_db(objects):
"""
:param objects: Deserialized object dictionaries
:return: List of tuples of (db_alias, [object,...])
"""
objects_by_db = defaultdict(list)
for obj in objects:
app_label = obj['model']
model = apps.get_model(app_label)
db_alias = router.db_for_write(model)
if settings.USE_PARTITIONED_DATABASE and db_alias == partition_config.proxy_db:
doc_id = _get_doc_id(app_label, obj)
db_alias = ShardAccessor.get_database_for_doc(doc_id)
objects_by_db[db_alias].append(obj)
return list(objects_by_db.items())
def delete_object_from_partitioned_database(obj, partition_value):
"""
Determines from which database to delete a partitioned model object and
deletes it there.
:param obj: A Django model object
:param parition_value: The value that is used to partition the model; this
value will be used to select the database
"""
if settings.USE_PARTITIONED_DATABASE:
db_name = ShardAccessor.get_database_for_doc(partition_value)
else:
db_name = 'default'
obj.delete(using=db_name)
def get_object_from_partitioned_database(model_class, partition_value, partitioned_field_name):
"""
Determines from which database to retrieve a paritioned model object and
retrieves it.
:param model_class: A Django model class
:param parition_value: The value that is used to partition the model; this
value will be used to select the database
:param partitioned_field_name: The model field on which the object is partitioned; the
object whose partitioned_field_name attribute equals partition_value is returned
:return: The model object
"""
if settings.USE_PARTITIONED_DATABASE:
db_name = ShardAccessor.get_database_for_doc(partition_value)
else:
db_name = 'default'
kwargs = {
partitioned_field_name: partition_value,
}
return model_class.objects.using(db_name).get(**kwargs)
def get_db_alias_for_partitioned_doc(partition_value):
if settings.USE_PARTITIONED_DATABASE:
db_name = ShardAccessor.get_database_for_doc(partition_value)
else:
db_name = 'default'
return db_name
