added db read-only mode implementation

This commit is contained in:
Abhimanyu Saharan 2023-01-14 17:15:12 +05:30
parent d7c37d9dd6
commit 6346e4dd60
4 changed files with 178 additions and 3 deletions

115
netbox/netbox/cursor.py Normal file
View File

@ -0,0 +1,115 @@
import logging
import time
from django.conf import settings
from django.db.backends.utils import CursorWrapper as _CursorWrapper
from netbox.exceptions import DatabaseWriteDenied
logger = logging.getLogger('netbox.db')
class ReadOnlyCursorWrapper:
"""
A read-only wrapper around a database cursor.
This wrapper prevents write operations from being performed on the database. It is used to prevent changes to the
database during a read-only request. It is not intended to be used directly; rather, it is applied automatically by
the ReadOnlyMiddleware. See the documentation for that class for more information.
"""
SQL_BLACKLIST = (
# Data definition
'CREATE',
'ALTER',
'DROP',
'TRUNCATE',
'RENAME',
# Data manipulation
'INSERT',
'UPDATE',
'DELETE',
'MERGE',
'REPLACE',
)
def __init__(self, cursor, db, *args, **kwargs):
self.cursor = cursor
self.db = db
self.read_only = settings.MAINTENANCE_MODE
def execute(self, sql, params=()):
# Check the SQL
if self.read_only and self._write_sql(sql):
raise DatabaseWriteDenied
return self.cursor.execute(sql, params)
def executemany(self, sql, param_list):
# Check the SQL
if self.read_only and self._write_sql(sql):
raise DatabaseWriteDenied
return self.cursor.executemany(sql, param_list)
def __getattr__(self, item):
return getattr(self.cursor, item)
def __iter__(self):
return iter(self.cursor)
def _write_sql(self, sql):
"""
Check the SQL to determine if it is a write operation.
"""
return any(
s.strip().upper().startswith(self.SQL_BLACKLIST) for s in sql.split(';')
)
@property
def _last_executed(self):
return getattr(self.cursor, '_last_executed', '')
class CursorWrapper(_CursorWrapper):
def __init__(self, cursor, db):
self.cursor = ReadOnlyCursorWrapper(cursor, db)
self.db = db
class CursorDebugWrapper(CursorWrapper):
def execute(self, sql, params=None):
start = time.time()
try:
return self.cursor.execute(sql, params)
finally:
stop = time.time()
duration = stop - start
sql = self.db.ops.last_executed_query(self.cursor, sql, params)
self.db.queries_log.append({
'sql': sql,
'time': '%.3f' % duration,
})
logger.debug(
"(%.3f) %s; args=%s",
duration,
sql,
params,
extra={"duration": duration, "sql": sql, "params": params},
)
def executemany(self, sql, param_list):
start = time.time()
try:
return self.cursor.executemany(sql, param_list)
finally:
stop = time.time()
duration = stop - start
self.db.queries.append({
"sql": "%s times: %s" % (len(param_list), sql),
"time": "%.3f" % duration,
})
logger.debug(
"(%.3f) %s; args=%s",
duration,
sql,
param_list,
extra={"duration": duration, "sql": sql, "params": param_list},
)

View File

@ -0,0 +1,8 @@
from django.db.utils import DatabaseError
class DatabaseWriteDenied(DatabaseError):
"""
Custom exception raised when a write operation is attempted in maintenance mode.
"""
pass

View File

@ -3,16 +3,19 @@ import uuid
from urllib import parse from urllib import parse
from django.conf import settings from django.conf import settings
from django.contrib import auth from django.contrib import auth, messages
from django.contrib.auth.middleware import RemoteUserMiddleware as RemoteUserMiddleware_ from django.contrib.auth.middleware import RemoteUserMiddleware as RemoteUserMiddleware_
from django.core.exceptions import ImproperlyConfigured from django.core.exceptions import ImproperlyConfigured
from django.db import ProgrammingError from django.db import ProgrammingError
from django.http import Http404, HttpResponseRedirect from django.http import Http404, HttpResponse, HttpResponseRedirect, JsonResponse
from django.utils.deprecation import MiddlewareMixin
from django.utils.encoding import iri_to_uri
from extras.context_managers import change_logging from extras.context_managers import change_logging
from netbox.config import clear_config from netbox.config import clear_config
from netbox.exceptions import DatabaseWriteDenied
from netbox.views import handler_500 from netbox.views import handler_500
from utilities.api import is_api_request, rest_api_server_error from utilities.api import is_api_request
class LoginRequiredMiddleware: class LoginRequiredMiddleware:
@ -202,3 +205,44 @@ class ExceptionHandlingMiddleware:
# Return a custom error message, or fall back to Django's default 500 error handling # Return a custom error message, or fall back to Django's default 500 error handling
if custom_template: if custom_template:
return handler_500(request, template_name=custom_template) return handler_500(request, template_name=custom_template)
class HttpResponseReload(HttpResponse):
"""
Reload page and stay on the same page from where request was made.
"""
status_code = 302
def __init__(self, request, *args, **kwargs):
super().__init__(*args, **kwargs)
referrer = request.META.get('HTTP_REFERER', '/')
self['Location'] = iri_to_uri(referrer)
class DatabaseReadOnlyMiddleware(MiddlewareMixin):
"""
Process exceptions raised by the database when a write operation is attempted while the database is in read-only
"""
def process_exception(self, request, exception):
# Only process DatabaseWriteDenied exceptions
if not isinstance(exception, DatabaseWriteDenied):
return None
not_allowed_methods = ['POST', 'PUT', 'PATCH', 'DELETE']
error_message = 'The database is currently in read-only mode. Please try again later.'
status_code = 503
# If the request is an API request, return a 503 Service Unavailable response
if is_api_request(request) and request.method in not_allowed_methods:
return JsonResponse({'detail': error_message, }, status=status_code)
else:
# Handle exceptions
if request.method in not_allowed_methods:
# Display a message to the user
messages.error(request, error_message)
# Redirect back to the referring page
return HttpResponseReload(request)
else:
return HttpResponse(error_message, status=status_code)

View File

@ -357,6 +357,7 @@ MIDDLEWARE = [
'django.contrib.messages.middleware.MessageMiddleware', 'django.contrib.messages.middleware.MessageMiddleware',
'django.middleware.clickjacking.XFrameOptionsMiddleware', 'django.middleware.clickjacking.XFrameOptionsMiddleware',
'django.middleware.security.SecurityMiddleware', 'django.middleware.security.SecurityMiddleware',
'netbox.middleware.DatabaseReadOnlyMiddleware',
'netbox.middleware.ExceptionHandlingMiddleware', 'netbox.middleware.ExceptionHandlingMiddleware',
'netbox.middleware.RemoteUserMiddleware', 'netbox.middleware.RemoteUserMiddleware',
'netbox.middleware.LoginRequiredMiddleware', 'netbox.middleware.LoginRequiredMiddleware',
@ -751,3 +752,10 @@ for plugin_name in PLUGINS:
RQ_QUEUES.update({ RQ_QUEUES.update({
f"{plugin_name}.{queue}": RQ_PARAMS for queue in plugin_config.queues f"{plugin_name}.{queue}": RQ_PARAMS for queue in plugin_config.queues
}) })
# Monkey patch CursorWrapper and CursorDebugWrapper to enable read-only mode
if getattr(configuration, 'MAINTENANCE_MODE', False):
from django.db.backends import utils
from netbox.cursor import CursorWrapper, CursorDebugWrapper
utils.CursorWrapper = CursorWrapper
utils.CursorDebugWrapper = CursorDebugWrapper