From cd8d593191da260c3e1f21f6d495401a06057297 Mon Sep 17 00:00:00 2001 From: DavidZheyuYang Date: Tue, 7 May 2024 03:03:29 -0700 Subject: [PATCH 1/3] feat (add support for async operations) --- dynamic_db_router/router.py | 58 +++++++++++++++++-------------------- 1 file changed, 26 insertions(+), 32 deletions(-) diff --git a/dynamic_db_router/router.py b/dynamic_db_router/router.py index f768832..fcadb93 100644 --- a/dynamic_db_router/router.py +++ b/dynamic_db_router/router.py @@ -1,29 +1,35 @@ import threading from functools import wraps from uuid import uuid4 +from contextvars import ContextVar from django.db import connections THREAD_LOCAL = threading.local() +DB_FOR_READ_OVERRIDE = ContextVar('DB_FOR_READ_OVERRIDE', default='default') +DB_FOR_WRITE_OVERRIDE = ContextVar('DB_FOR_WRITE_OVERRIDE', default='default') + class DynamicDbRouter(object): """A router that decides what db to read from based on a variable local to the current thread. """ - + def db_for_read(self, model, **hints): - return getattr(THREAD_LOCAL, 'DB_FOR_READ_OVERRIDE', ['default'])[-1] - + return DB_FOR_READ_OVERRIDE.get() + # return getattr(THREAD_LOCAL, 'DB_FOR_READ_OVERRIDE', ['default'])[-1] + def db_for_write(self, model, **hints): - return getattr(THREAD_LOCAL, 'DB_FOR_WRITE_OVERRIDE', ['default'])[-1] - + return DB_FOR_WRITE_OVERRIDE.get() + # return getattr(THREAD_LOCAL, 'DB_FOR_WRITE_OVERRIDE', ['default'])[-1] + def allow_relation(self, *args, **kwargs): return True - + def allow_syncdb(self, *args, **kwargs): return None - + def allow_migrate(self, *args, **kwargs): return None @@ -83,45 +89,33 @@ def lowest_id_account(): def __init__(self, database, read=True, write=False): self.read = read self.write = write + self.database = database self.created_db_config = False - if isinstance(database, str): - self.database = database - elif isinstance(database, dict): - # Note: this invalidates the docs above. Update them - # eventually. + if isinstance(database, dict): self.created_db_config = True self.unique_db_id = str(uuid4()) connections.databases[self.unique_db_id] = database self.database = self.unique_db_id - else: - msg = ("database must be an identifier for an existing db, " - "or a complete configuration.") - raise ValueError(msg) - + def __enter__(self): - if not hasattr(THREAD_LOCAL, 'DB_FOR_READ_OVERRIDE'): - THREAD_LOCAL.DB_FOR_READ_OVERRIDE = ['default'] - if not hasattr(THREAD_LOCAL, 'DB_FOR_WRITE_OVERRIDE'): - THREAD_LOCAL.DB_FOR_WRITE_OVERRIDE = ['default'] - read_db = (self.database if self.read - else THREAD_LOCAL.DB_FOR_READ_OVERRIDE[-1]) - write_db = (self.database if self.write - else THREAD_LOCAL.DB_FOR_WRITE_OVERRIDE[-1]) - THREAD_LOCAL.DB_FOR_READ_OVERRIDE.append(read_db) - THREAD_LOCAL.DB_FOR_WRITE_OVERRIDE.append(write_db) + self.original_read_db = DB_FOR_READ_OVERRIDE.get() + self.original_write_db = DB_FOR_WRITE_OVERRIDE.get() + if self.read: + DB_FOR_READ_OVERRIDE.set(self.database) + if self.write: + DB_FOR_WRITE_OVERRIDE.set(self.database) return self - + def __exit__(self, exc_type, exc_value, traceback): - THREAD_LOCAL.DB_FOR_READ_OVERRIDE.pop() - THREAD_LOCAL.DB_FOR_WRITE_OVERRIDE.pop() + DB_FOR_READ_OVERRIDE.set(self.original_read_db) + DB_FOR_WRITE_OVERRIDE.set(self.original_write_db) if self.created_db_config: connections[self.unique_db_id].close() del connections.databases[self.unique_db_id] - + def __call__(self, querying_func): @wraps(querying_func) def inner(*args, **kwargs): - # Call the function in our context manager with self: return querying_func(*args, **kwargs) return inner From 5043f157fa626d667b7ec981415dc6486e1b421d Mon Sep 17 00:00:00 2001 From: davidbm04 Date: Tue, 7 May 2024 16:11:10 -0700 Subject: [PATCH 2/3] feat(refactor): add if condition to check wether database parameter is str along with initial else condition that throws an error specifying what the database parameter should be of type. add comments and clean up --- dynamic_db_router/router.py | 107 ++++++++++++------------------------ 1 file changed, 35 insertions(+), 72 deletions(-) diff --git a/dynamic_db_router/router.py b/dynamic_db_router/router.py index fcadb93..14fda85 100644 --- a/dynamic_db_router/router.py +++ b/dynamic_db_router/router.py @@ -1,119 +1,82 @@ -import threading +from contextvars import ContextVar from functools import wraps from uuid import uuid4 -from contextvars import ContextVar - from django.db import connections -THREAD_LOCAL = threading.local() - +# Define context variables for read and write database settings +# These variables will maintain database preferences per context DB_FOR_READ_OVERRIDE = ContextVar('DB_FOR_READ_OVERRIDE', default='default') DB_FOR_WRITE_OVERRIDE = ContextVar('DB_FOR_WRITE_OVERRIDE', default='default') - -class DynamicDbRouter(object): - """A router that decides what db to read from based on a variable - local to the current thread. + +class DynamicDbRouter: + """ + A router that dynamically determines which database to perform read and write operations + on based on the current execution context. It supports both synchronous and asynchronous code. """ - + def db_for_read(self, model, **hints): return DB_FOR_READ_OVERRIDE.get() - # return getattr(THREAD_LOCAL, 'DB_FOR_READ_OVERRIDE', ['default'])[-1] - + def db_for_write(self, model, **hints): return DB_FOR_WRITE_OVERRIDE.get() - # return getattr(THREAD_LOCAL, 'DB_FOR_WRITE_OVERRIDE', ['default'])[-1] - + def allow_relation(self, *args, **kwargs): return True - + def allow_syncdb(self, *args, **kwargs): return None - + def allow_migrate(self, *args, **kwargs): return None - -class in_database(object): - """A decorator and context manager to do queries on a given database. - - :type database: str or dict - :param database: The database to run queries on. A string - will route through the matching database in - ``django.conf.settings.DATABASES``. A dictionary will set up a - connection with the given configuration and route queries to it. - - :type read: bool, optional - :param read: Controls whether database reads will route through - the provided database. If ``False``, reads will route through - the ``'default'`` database. Defaults to ``True``. - - :type write: bool, optional - :param write: Controls whether database writes will route to - the provided database. If ``False``, writes will route to - the ``'default'`` database. Defaults to ``False``. - - When used as eithe a decorator or a context manager, `in_database` - requires a single argument, which is the name of the database to - route queries to, or a configuration dictionary for a database to - route to. - - Usage as a context manager: - - .. code-block:: python - - from my_django_app.utils import tricky_query - - with in_database('Database_A'): - results = tricky_query() - - Usage as a decorator: - - .. code-block:: python - - from my_django_app.models import Account - - @in_database('Database_B') - def lowest_id_account(): - Account.objects.order_by('-id')[0] - - Used with a configuration dictionary: - - .. code-block:: python - - db_config = {'ENGINE': 'django.db.backends.sqlite3', - 'NAME': 'path/to/mydatabase.db'} - with in_database(db_config): - # Run queries +class in_database: + """ + A context manager and decorator for setting a specific database for the duration of a block of code. """ def __init__(self, database, read=True, write=False): self.read = read self.write = write self.database = database self.created_db_config = False - if isinstance(database, dict): + + # Handle database parameter either as a string (alias) or as a dict (configuration) + if isinstance(database, str): + self.database = database + elif isinstance(database, dict): + # If it's a dict, create a unique database configuration self.created_db_config = True self.unique_db_id = str(uuid4()) connections.databases[self.unique_db_id] = database self.database = self.unique_db_id - + else: + raise ValueError("database must be an identifier (str) for an existing db, " + "or a complete configuration (dict).") + def __enter__(self): + # Capture the current database settings self.original_read_db = DB_FOR_READ_OVERRIDE.get() self.original_write_db = DB_FOR_WRITE_OVERRIDE.get() + + # Override the database settings for the duration of the context if self.read: DB_FOR_READ_OVERRIDE.set(self.database) if self.write: DB_FOR_WRITE_OVERRIDE.set(self.database) return self - + def __exit__(self, exc_type, exc_value, traceback): + # Restore the original database settings after the context. DB_FOR_READ_OVERRIDE.set(self.original_read_db) DB_FOR_WRITE_OVERRIDE.set(self.original_write_db) + + # Close and delete created database configuration if self.created_db_config: connections[self.unique_db_id].close() del connections.databases[self.unique_db_id] - + def __call__(self, querying_func): + # Allow the object to be used as a decorator @wraps(querying_func) def inner(*args, **kwargs): with self: From 8ed828a783377b48f6c134499806e71d89330807 Mon Sep 17 00:00:00 2001 From: davidbm04 Date: Tue, 7 May 2024 17:37:09 -0700 Subject: [PATCH 3/3] refactor: add initial comments back, add type hint to database parameter --- dynamic_db_router/router.py | 38 +++++++++++++++++++++++++++++++++++-- 1 file changed, 36 insertions(+), 2 deletions(-) diff --git a/dynamic_db_router/router.py b/dynamic_db_router/router.py index 14fda85..f500070 100644 --- a/dynamic_db_router/router.py +++ b/dynamic_db_router/router.py @@ -32,9 +32,43 @@ def allow_migrate(self, *args, **kwargs): class in_database: """ - A context manager and decorator for setting a specific database for the duration of a block of code. + A decorator and context manager to do queries on a given database. + :type database: str or dict + :param database: The database to run queries on. A string + will route through the matching database in + ``django.conf.settings.DATABASES``. A dictionary will set up a + connection with the given configuration and route queries to it. + :type read: bool, optional + :param read: Controls whether database reads will route through + the provided database. If ``False``, reads will route through + the ``'default'`` database. Defaults to ``True``. + :type write: bool, optional + :param write: Controls whether database writes will route to + the provided database. If ``False``, writes will route to + the ``'default'`` database. Defaults to ``False``. + When used as eithe a decorator or a context manager, `in_database` + requires a single argument, which is the name of the database to + route queries to, or a configuration dictionary for a database to + route to. + Usage as a context manager: + .. code-block:: python + from my_django_app.utils import tricky_query + with in_database('Database_A'): + results = tricky_query() + Usage as a decorator: + .. code-block:: python + from my_django_app.models import Account + @in_database('Database_B') + def lowest_id_account(): + Account.objects.order_by('-id')[0] + Used with a configuration dictionary: + .. code-block:: python + db_config = {'ENGINE': 'django.db.backends.sqlite3', + 'NAME': 'path/to/mydatabase.db'} + with in_database(db_config): + # Run queries """ - def __init__(self, database, read=True, write=False): + def __init__(self, database: str | dict, read=True, write=False): self.read = read self.write = write self.database = database