def test_read_load_balancing(self):
reporting_dbs = {
'ucr': {
'WRITE': 'ucr',
'READ': [('ucr', 8), ('other', 1), ('default', 1)]
},
}
with override_settings(REPORTING_DATABASES=reporting_dbs):
manager = ConnectionManager()
self.assertEqual(manager.db_connection_map, {
'default': 'postgresql+psycopg2://:@localhost:5432/default',
'ucr': 'postgresql+psycopg2://:@localhost:5432/ucr',
'other': 'postgresql+psycopg2://:@localhost:5432/other'
})
# test that load balancing works with a 10% margin for randomness
total_requests = 10000
randomness_margin = total_requests * 0.1
total_weighting = sum(db[1] for db in reporting_dbs['ucr']['READ'])
expected = {
alias: weight * total_requests // total_weighting
for alias, weight in reporting_dbs['ucr']['READ']
}
balanced = Counter(manager.get_load_balanced_read_engine_id('ucr') for i in range(total_requests))
for db, requests in balanced.items():
self.assertAlmostEqual(requests, expected[db], delta=randomness_margin)
with override_settings(REPORTING_DATABASES={'default': 'default'}):
manager = ConnectionManager()
self.assertEqual(
['default', 'default', 'default'],
[manager.get_load_balanced_read_engine_id('default') for i in range(3)]
)
def test_load_balanced_read_apps(self, mock):
load_balanced_apps = {
'users': [
('users_db1', 5),
]
}
with override_settings(LOAD_BALANCED_APPS=load_balanced_apps,
DATABASES={
'default': _get_db_config('default'),
'users_db1': _get_db_config('users_db1')
}):
manager = ConnectionManager()
self.assertEqual(
manager.get_load_balanced_read_db_alias(
'users', default="default_option"), 'users_db1')
with override_settings(LOAD_BALANCED_APPS=load_balanced_apps):
# load_balanced_db should be part of settings.DATABASES
with self.assertRaises(AssertionError):
ConnectionManager().get_load_balanced_read_db_alias('users')
# If `LOAD_BALANCED_APPS` is not set for an app, it should point to default kwarg
manager = ConnectionManager()
self.assertEqual(
manager.get_load_balanced_read_db_alias('users',
default='default_option'),
'default_option')
def test_read_load_balancing(self, *args):
reporting_dbs = {
'ucr': {
'WRITE': 'ucr',
'READ': [('ucr', 8), ('other', 1), ('default', 1)]
},
}
with override_settings(REPORTING_DATABASES=reporting_dbs):
manager = ConnectionManager()
self.assertEqual(manager.engine_id_django_db_map, {
'default': 'default',
'ucr': 'ucr',
})
# test that load balancing works with a 10% margin for randomness
total_requests = 10000
randomness_margin = total_requests * 0.1
total_weighting = sum(db[1] for db in reporting_dbs['ucr']['READ'])
expected = {
alias: weight * total_requests // total_weighting
for alias, weight in reporting_dbs['ucr']['READ']
}
balanced = Counter(manager.get_load_balanced_read_db_alias('ucr') for i in range(total_requests))
for db, requests in balanced.items():
self.assertAlmostEqual(requests, expected[db], delta=randomness_margin)
with override_settings(REPORTING_DATABASES={'default': DEFAULT_DB_ALIAS}):
manager = ConnectionManager()
self.assertEqual(
[DEFAULT_DB_ALIAS] * 3,
[manager.get_load_balanced_read_db_alias(DEFAULT_DB_ALIAS) for i in range(3)]
)
def test_new_settings(self):
manager = ConnectionManager()
self.assertEqual(manager.engine_id_django_db_map, {
'default': 'default',
'ucr': 'ucr',
'other': 'other'
})
def test_read_load_balancing_session(self, *args):
reporting_dbs = {
'ucr': {
'WRITE': 'ucr',
'READ': [('ucr', 8), ('other', 1), ('default', 1)]
},
}
with override_settings(REPORTING_DATABASES=reporting_dbs):
manager = ConnectionManager()
self.assertEqual(manager.engine_id_django_db_map, {
'default': 'default',
'ucr': 'ucr',
})
urls = {
manager.get_connection_string(alias)
for alias, _ in reporting_dbs['ucr']['READ']
}
self.assertEqual(len(urls), 3)
# withing 50 iterations we should have seen all 3 databases at least once
for i in range(50):
url = manager.get_session_helper('ucr', readonly=True).url
if url in urls:
urls.remove(url)
if not urls:
break
if urls:
self.fail(f'DBs skipped in load balancing: {urls}')
def test_new_settings(self):
manager = ConnectionManager()
self.assertEqual(manager.db_connection_map, {
'default': 'postgresql+psycopg2://:@localhost:5432/default',
'ucr': 'postgresql+psycopg2://:@localhost:5432/ucr',
'other': 'postgresql+psycopg2://:@localhost:5432/other'
})
def test_standby_filtering(self, *args):
reporting_dbs = {
'ucr_engine': {
'WRITE': 'ucr',
'READ': [('ucr', 8), ('other', 1)]
},
}
with override_settings(REPORTING_DATABASES=reporting_dbs):
# should always return the `ucr` db since `other` has bad replication delay
manager = ConnectionManager()
self.assertEqual(['ucr', 'ucr', 'ucr'], [
manager.get_load_balanced_read_db_alias('ucr_engine')
for i in range(3)
])
def check_for_ucr_tables_without_existing_domain():
all_domain_names = Domain.get_all_names()
connection_name = ConnectionManager().get_django_db_alias(UCR_ENGINE_ID)
table_names = connections[connection_name].introspection.table_names()
ucr_table_names = [name for name in table_names if name.startswith('config_report')]
missing_domains_to_tables = defaultdict(list)
for ucr_table_name in ucr_table_names:
table_domain = ucr_table_name.split('_')[2]
if table_domain not in all_domain_names:
missing_domains_to_tables[table_domain].append(ucr_table_name)
if missing_domains_to_tables:
for missing_domain in missing_domains_to_tables:
mail_admins_async.delay(
'Missing domain "%s" has remaining UCR tables' % missing_domain,
six.text_type(missing_domains_to_tables[missing_domain])
)
elif _is_monday():
mail_admins_async.delay('All UCR tables belong to valid domains', '')