Implementing maintenance mode

.env_template

@@ -1,4 +1,4 @@
 HOST_ADDRESS = "http://localhost:6752"
 TIMEOUT=3
 ADMIN_SECRET_KEY="test"
-MEDIUM_AUTH_COOKIES="Get your premium subscription account coockies here, uid and sid properties is required"
+# MEDIUM_AUTH_COOKIES="Get your premium subscription account coockies here, uid and sid properties is required"

database-lib/database_lib/cache_db.py

@@ -4,6 +4,7 @@ import json
 from loguru import logger
 from warnings import warn
 import threading
+
 try:
     import sqlite_zstd
 except ImportError:
@@ -11,6 +12,7 @@ except ImportError:
     warn("Can't use zstd compression. Please install 'sqlite_zstd' package")
     sqlite_zstd = None
 
+
 class CacheResponse:
     __slots__ = ('data',)
     def __init__(self, data: str):
@@ -25,6 +27,7 @@ class CacheResponse:
     def __str__(self):
         return self.data
 
+
 class SQLiteCacheBackend:
     __slots__ = ('connection', 'cursor', 'database', 'lock')
     def __init__(self, database: str):
@@ -51,18 +54,17 @@ class SQLiteCacheBackend:
     def random(self, size: int):
         with self.connection:
             return self.cursor.execute("SELECT * FROM cache ORDER BY RANDOM() LIMIT ?", (size,)).fetchall()
-    
+
     def enable_zstd(self):
         if sqlite_zstd is None:
             raise ValueError("Can't use zstd compression. Please install 'sqlite_zstd' package")
-        
+
         with self.connection:
             try:
                 self.cursor.execute("SELECT zstd_enable_transparent('{\"table\": \"cache\", \"column\": \"value\", \"compression_level\": 9, \"dict_chooser\": \"''a''\"}')")
             except Exception as error:
                 print(error)
             self.connection.execute("PRAGMA auto_vacuum=full")
-            self.maintenance_thread()
 
     def init_db(self):
         with self.connection:
@@ -87,11 +89,11 @@ class SQLiteCacheBackend:
                 raise ValueError(f"Unable to serialize value to JSON: {e}")
         elif not isinstance(value, str):
             raise ValueError(f"value argument should be a string or dict, not {type(value).__name__}")
-        
+
         with self.lock:
             with self.connection:
                 self.cursor.execute("INSERT OR REPLACE INTO cache VALUES (:0, :1)", {'0': key, '1': value})
-            
+
     def delete(self, key: str) -> None:
         with self.connection:
             result = self.cursor.execute("SELECT 1 FROM cache WHERE key = :0", {'0': key}).fetchone()
@@ -100,7 +102,7 @@ class SQLiteCacheBackend:
                 logger.debug(f"Deleted key: {key}")
             else:
                 logger.debug(f"Attempted to delete non-existing key: {key}")
-    
+
     def maintenance(self, time: int = None, blocking_time: int = 0.5):
         connection = sqlite3.connect(self.database)
         cursor = connection.cursor()
@@ -108,10 +110,10 @@ class SQLiteCacheBackend:
         connection.execute("PRAGMA foreign_keys = ON;")  # Need for working with foreign keys in db
         connection.execute("PRAGMA journal_mode=WAL;") # Need to properly work with ZSTD compression
         connection.execute("PRAGMA auto_vacuum=full;") # Same as above thing
-        
+
         if sqlite_zstd is not None:
             sqlite_zstd.load(connection)
-        
+
         with connection:
             if time is not None:
                 cursor.execute("SELECT zstd_incremental_maintenance(?, ?);", (time, blocking_time))
@@ -119,7 +121,7 @@ class SQLiteCacheBackend:
                 cursor.execute("SELECT zstd_incremental_maintenance(null, ?);", (blocking_time,))
             cursor.execute("VACUUM")
             cursor.execute("ANALYZE")
-        
+
         cursor.close()
         connection.close()
 
@@ -127,18 +129,6 @@ class SQLiteCacheBackend:
         maintenance_thread = threading.Thread(target=self.maintenance, daemon=True)
         maintenance_thread.start()
 
-    # Doesn't works with sqlite_zstd
-    # def migrate_add_index_to_key(self):
-    #     with self.connection:
-    #         # Check if the index already exists
-    #         index_exists = self.cursor.execute("SELECT name FROM sqlite_master WHERE type='index' AND name='idx_key'").fetchone()
-    #         if not index_exists:
-    #             # Create the index if it doesn't exist
-    #             self.cursor.execute("CREATE INDEX idx_key ON cache (key)")
-    #             logger.info("Index 'idx_key' on column 'key' created successfully.")
-    #         else:
-    #             logger.info("Index 'idx_key' on column 'key' already exists.")
-
     def show_schema_info(self):
         with self.connection:
             return self.connection.execute("SELECT sql FROM sqlite_master").fetchall()

requirements.txt

@@ -10,3 +10,4 @@ starlette==0.24.0
 gunicorn==21.2.0
 redis[hiredis]==4.6.0
 xkcdpass==1.19.3
+apscheduler==3.10.4

server/__init__.py

@@ -3,6 +3,7 @@ import pickledb
 import logging
 from contextvars import ContextVar
 from typing import Optional
+from apscheduler.schedulers.asyncio import AsyncIOScheduler
 
 from jinja2 import Environment, DebugUndefined, FileSystemLoader
 import redis.asyncio as redis
@@ -12,6 +13,8 @@ from server.utils.loguru_handler import InterceptHandler
 
 redis_storage = redis.Redis(host="dragonfly", port=6379, db=0)
 
+scheduler = AsyncIOScheduler()
+
 jinja_env = Environment(enable_async=True)
 jinja_safe_env = Environment(undefined=DebugUndefined)
 template_env = Environment(loader=FileSystemLoader("./server/templates"), enable_async=True)
@@ -41,4 +44,6 @@ ban_db = pickledb.load('ban_post_list.db', True)
 START_TIME = dt.datetime.now().strftime("%H-%M-%S")
 WORDS_LIST_FILE = "xkcdpass/static/legac"
 
-xkcd_passwd = xp.generate_wordlist(wordfile=WORDS_LIST_FILE, min_length=5, max_length=8)
+xkcd_passwd = xp.generate_wordlist(wordfile=WORDS_LIST_FILE, min_length=5, max_length=8)
+
+maintenance_mode = False

server/cli.py

@@ -24,5 +24,8 @@ def server_cmd(cmd, opts):
         cmd.error(f"Port {opts.port} is in use or permission denied")
 
     from server.worker import execute_server_worker
+    from server.utils.maintenance_scheduler import enable_maintenance_mode
+
+    enable_maintenance_mode()
 
     execute_server_worker(host="0.0.0.0", port=opts.port)

server/main.py

@@ -1,7 +1,5 @@
 from math import ceil
 import sentry_sdk
-import asyncio
-from contextlib import suppress
 from fastapi.exceptions import HTTPException
 from fastapi import FastAPI, Depends, APIRouter
 from loguru import logger

server/middlewares/__init__.py

@@ -1,5 +1,6 @@
 from starlette.middleware.cors import CORSMiddleware
 from server.middlewares.logger import LoggerMiddleware
+from server.middlewares.maintenance_mode import MaintenanceModeMiddleware
 
 origins = ["*"]
 
@@ -12,4 +13,4 @@ def register_middlewares(app):
         allow_methods=["*"],
         allow_headers=["*"],
     )
-
+    app.add_middleware(MaintenanceModeMiddleware)

server/middlewares/maintenance_mode.py

@@ -0,0 +1,35 @@
+from fastapi import Request
+from fastapi.responses import HTMLResponse
+from starlette.middleware.base import BaseHTTPMiddleware
+from server import maintenance_mode
+
+HTML_RESPONSE_BODY = """
+<!doctype html>
+<title>Site Maintenance</title>
+<style>
+  body { text-align: center; padding: 150px; }
+  h1 { font-size: 50px; }
+  body { font: 20px Helvetica, sans-serif; color: #333; }
+  article { display: block; text-align: left; width: 650px; margin: 0 auto; }
+  a { color: #dc8100; text-decoration: none; }
+  a:hover { color: #333; text-decoration: none; }
+</style>
+
+<article>
+    <h1>We&rsquo;ll be back soon!</h1>
+    <div>
+        <p>Sorry for the inconvenience but we&rsquo;re performing some maintenance at the moment. If you need to you can always <a href="mailto:#admin@freedium.cfd">contact us</a>, otherwise we&rsquo;ll be back online shortly!</p>
+        <p>&mdash; Freedium developers</p>
+    </div>
+</article>
+"""
+
+class MaintenanceModeMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request: Request, call_next):
+        if maintenance_mode:
+            return HTMLResponse(
+                content=HTML_RESPONSE_BODY,
+                status_code=503
+            )
+        response = await call_next(request)
+        return response

server/utils/maintenance_scheduler.py

@@ -0,0 +1,31 @@
+from contextlib import suppress
+from datetime import datetime
+
+from server import maintenance_mode, scheduler
+from server.utils.notify import send_message
+from medium_parser import cache as medium_cache
+
+from loguru import logger
+
+
+def enable_maintenance_mode():
+    global maintenance_mode
+    maintenance_mode = True
+    logger.debug("Maintenance mode enabled")
+
+    send_message("Maintenance mode enabled")
+
+    now = datetime.now()
+    logger.debug(f"Current time: {now}")
+    send_message(f"Current time: {now}")
+
+    with suppress(Exception):
+        medium_cache.maintenance()
+
+    maintenance_mode = False
+    logger.debug("Maintenance mode disabled")
+    send_message("Maintenance mode disabled")
+
+scheduler.add_job(enable_maintenance_mode, 'cron', hour=1)
+
+scheduler.start()