added decorator for functions

django_replicated/decorators.py

@@ -21,9 +21,29 @@ def my_view(request, ...):
 
 from django.utils.decorators import decorator_from_middleware_with_args
 
+from .utils import routers
 from .middleware import ReplicationMiddleware
 
 
+def use_state_in_function(forced_state='master'):
+
+    def _make_decorator(func):
+
+        def wrapper(*args, **kwargs):
+            prev_state = routers.state()  # Store previous state
+            routers.init(forced_state)  # Set state to the forced_state
+            func(*args, **kwargs)
+            # If previous state was master, set it back
+            # Handling function calls from other functions or classes
+            if prev_state == 'master':
+                routers.init('master')
+        return wrapper
+
+    return _make_decorator
+
+
 use_state = decorator_from_middleware_with_args(ReplicationMiddleware)
 use_master = use_state(forced_state='master')
 use_slave = use_state(forced_state='slave')
+use_master_in_function = use_state_in_function()
+use_slave_in_function = use_state_in_function(forced_state='slave')

django_replicated/middleware.py

@@ -9,7 +9,8 @@
 
 from django import db
 from django.conf import settings
-from django.utils import six, functional
+from django.utils import functional
+import six
 
 try:  # django 1.10+
     from django import urls
@@ -105,7 +106,7 @@ def check_state_override(self, request, state):
         Used to check if a web request should use a master or slave
         database besides default choice.
         '''
-        if request.COOKIES.get(settings.REPLICATED_FORCE_MASTER_COOKIE_NAME) == 'true':
+        if getattr(request, 'COOKIES', {}).get(settings.REPLICATED_FORCE_MASTER_COOKIE_NAME) == 'true':
             return 'master'
 
         override_state = self.get_state_override(request)
@@ -150,7 +151,7 @@ def handle_redirect_after_write(self, request, response):
             log.debug('set force master cookie for %s', request.path)
             self.set_force_master_cookie(response)
         else:
-            if settings.REPLICATED_FORCE_MASTER_COOKIE_NAME in request.COOKIES:
+            if settings.REPLICATED_FORCE_MASTER_COOKIE_NAME in getattr(request, 'COOKIES', {}):
                 response.delete_cookie(settings.REPLICATED_FORCE_MASTER_COOKIE_NAME)
 
     def set_force_master_cookie(self, response):

django_replicated/router.py

@@ -77,6 +77,9 @@ def db_for_write(self, *args, **kwargs):
         if self.CHECK_STATE_ON_WRITE and self.state() != 'master':
             raise RuntimeError('Trying to access master database in slave state')
 
+        if self.state() != 'master':
+            self.init('master')
+
         self.context.chosen['master'] = self.DEFAULT_DB_ALIAS
 
         log.debug('db_for_write: %s', self.DEFAULT_DB_ALIAS)