From 6346e4dd6036751f58559a2e72a950b83dcc013e Mon Sep 17 00:00:00 2001 From: Abhimanyu Saharan Date: Sat, 14 Jan 2023 17:15:12 +0530 Subject: [PATCH] added db read-only mode implementation --- netbox/netbox/cursor.py | 115 ++++++++++++++++++++++++++++++++++++ netbox/netbox/exceptions.py | 8 +++ netbox/netbox/middleware.py | 50 +++++++++++++++- netbox/netbox/settings.py | 8 +++ 4 files changed, 178 insertions(+), 3 deletions(-) create mode 100644 netbox/netbox/cursor.py create mode 100644 netbox/netbox/exceptions.py diff --git a/netbox/netbox/cursor.py b/netbox/netbox/cursor.py new file mode 100644 index 000000000..257f6e61c --- /dev/null +++ b/netbox/netbox/cursor.py @@ -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}, + ) diff --git a/netbox/netbox/exceptions.py b/netbox/netbox/exceptions.py new file mode 100644 index 000000000..56821b34b --- /dev/null +++ b/netbox/netbox/exceptions.py @@ -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 diff --git a/netbox/netbox/middleware.py b/netbox/netbox/middleware.py index edf88a234..fa5a920f0 100644 --- a/netbox/netbox/middleware.py +++ b/netbox/netbox/middleware.py @@ -3,16 +3,19 @@ import uuid from urllib import parse 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.core.exceptions import ImproperlyConfigured 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 netbox.config import clear_config +from netbox.exceptions import DatabaseWriteDenied 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: @@ -202,3 +205,44 @@ class ExceptionHandlingMiddleware: # Return a custom error message, or fall back to Django's default 500 error handling if 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) diff --git a/netbox/netbox/settings.py b/netbox/netbox/settings.py index 1b47520be..02c3c2aeb 100644 --- a/netbox/netbox/settings.py +++ b/netbox/netbox/settings.py @@ -357,6 +357,7 @@ MIDDLEWARE = [ 'django.contrib.messages.middleware.MessageMiddleware', 'django.middleware.clickjacking.XFrameOptionsMiddleware', 'django.middleware.security.SecurityMiddleware', + 'netbox.middleware.DatabaseReadOnlyMiddleware', 'netbox.middleware.ExceptionHandlingMiddleware', 'netbox.middleware.RemoteUserMiddleware', 'netbox.middleware.LoginRequiredMiddleware', @@ -751,3 +752,10 @@ for plugin_name in PLUGINS: RQ_QUEUES.update({ 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