def ignore(self, model):
"""Ignores the trigger in a single thread of execution"""
with contextlib.ExitStack() as pre_execute_hook:
# Create the table name / trigger name URI to pass down to the
# trigger.
ignore_uri = f'{model._meta.db_table}:{self.get_pgid(model)}'
if not hasattr(_ignore, 'value'):
_ignore.value = set()
if not _ignore.value:
# If this is the first time we are ignoring trigger execution,
# register the pre_execute_hook
pre_execute_hook.enter_context(
pgconnection.pre_execute_hook(_inject_pgtrigger_ignore))
if ignore_uri not in _ignore.value:
try:
_ignore.value.add(ignore_uri)
yield
finally:
_ignore.value.remove(ignore_uri)
else: # The trigger is already being ignored
yield
def __enter__(self):
if not hasattr(_tracker, 'value'):
self._pre_execute_hook = pgconnection.pre_execute_hook(
_inject_history_context)
self._pre_execute_hook.__enter__()
_tracker.value = Context(id=uuid.uuid4(), metadata=self.metadata)
return _tracker.value
def test_pre_executed_modify_sql():
"""Registers a pre_execute hook that modifies the SQL"""
def modify_sql_hook(sql, sql_vars, cursor):
sql = f'--- Hello World\n{sql}'
return sql, sql_vars
with CaptureQueriesContext(connection) as queries:
with pgconnection.pre_execute_hook(modify_sql_hook):
list(User.objects.filter(id=1))
assert queries[0]['sql'].startswith('--- Hello World\n')
def route(destination, using='default'):
"""
Route connections to another database.
If the source database is configured to use pgconnection cursors,
an additional hook is provisioned to ensure that no queries happen
on the source database
Args:
destination (dict): Database configuration dictionary to be routed.
using (str, default='default'): The source database to use when
routing to another database. Defaults to the default database.
"""
pgconnection.check()
if (
not isinstance(destination, dict) or 'ENGINE' not in destination
): # pragma: no cover
raise ValueError(
'Destination database must be a configuration dictionary in the'
' same format as databases from settings.DATABASES.'
)
_check_source_is_not_destination(settings.DATABASES[using], destination)
with pgconnection.pre_execute_hook(
_guard_source_database_access, using=using
):
# Store the original source database and connection so that the
# destination can be put it its place
source_connection = connections[using]
source = connections.databases[using]
# When the connection is deleted, Django will re-establish it
# from the config in connections.databases. Since
# connections.databases is a pgconnection.core._Databases object, it
# is safe for us to patch it
del connections[using]
connections.databases[using] = destination
try:
yield
finally:
# Close out the patched connection and remove it
connections[using].close()
del connections[using]
# Revert the patched connection back to the source
connections.databases[using] = source
connections[using] = source_connection
def test_pre_executed_logging_hook(caplog):
"""Registers a pre_execute hook that logs SQL and verifies it works"""
caplog.set_level(logging.INFO)
def logging_hook(sql, sql_vars, cursor):
logging.info(sql)
# These queries should not be logged
list(User.objects.filter(id=1))
# These queries should be logged
with pgconnection.pre_execute_hook(logging_hook):
list(Group.objects.values_list('id'))
list(User.objects.values_list('id'))
# These queries should not be logged
list(Group.objects.filter(id=1))
assert caplog.record_tuples == [
('root', logging.INFO, 'SELECT "auth_group"."id" FROM "auth_group"'),
('root', logging.INFO, 'SELECT "auth_user"."id" FROM "auth_user"'),
]