Advertisement
LyndenSylvester

Untitled

Nov 19th, 2024
52
0
22 hours
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 5.58 KB | None | 0 0
  1. import asyncio
  2. from datetime import datetime, timedelta
  3. import os
  4. from ncogs.Utils import Utils as utils
  5. import yaml
  6. import threading
  7. import sqlite3
  8. from ncogs.Utils import Utils
  9. from discord import app_commands
  10. from discord.ext import commands
  11. from typing import NoReturn
  12.  
  13. class Scheduler:
  14.  
  15.     def __init__(self, guild_id: int) -> NoReturn:
  16.         self.tasks = {}
  17.         self.running = False
  18.         self.server_name = Utils.get_server_name(guild_id)
  19.         self.conn = sqlite3.connect(f'{self.server_name}/{Utils.shorten_server_name(guild_id)}.db')
  20.         self.cursor = self.conn.cursor()
  21.         self.cursor.execute("""CREATE TABLE IF NOT EXISTS tasks (name TEXT PRIMARY KEY, func_name TEXT, interval INTEGER, last_run TEXT, next_run TEXT)""")
  22.         self.guild_id = guild_id
  23.         self.schedulers = {}
  24.  
  25.     def get_scheduler(self, guild_id):
  26.         return self.schedulers.get(guild_id)
  27.  
  28.     def set_scheduler(self, guild_id, scheduler):
  29.         self.schedulers[guild_id] = scheduler
  30.  
  31.     async def load_tasks(self):
  32.         self.cursor.execute("SELECT * FROM tasks")
  33.         rows = self.cursor.fetchall()
  34.         for row in rows:
  35.             name, func_name, interval, last_run, next_run = row
  36.             self.tasks[name] = {
  37.                 'func': getattr(self.bot, func_name),
  38.                 'interval': interval,
  39.                 'last_run': datetime.fromisoformat(last_run),
  40.                 'next_run': datetime.fromisoformat(next_run)
  41.             }
  42.  
  43.     async def save_task(self, name, func_name, interval, last_run, next_run):
  44.         self.cursor.execute("INSERT OR REPLACE INTO tasks VALUES (?, ?, ?, ?, ?)",
  45.                             (name, func_name, interval, last_run.isoformat(), next_run.isoformat()))
  46.         self.conn.commit()
  47.  
  48.     async def start(self):
  49.         self.running = True
  50.         await self.check_tasks()
  51.  
  52.     async def add_task(self, name, func, interval):
  53.         self.tasks[name] = {
  54.             'func': func,
  55.             'interval': interval,
  56.             'last_run': datetime.now(),
  57.             'next_run': datetime.now() + timedelta(seconds=interval)
  58.         }
  59.         await self.save_task(name, func.__name__, interval, datetime.now(), self.tasks[name]['next_run'])
  60.         if not self.running:
  61.             await self.start()
  62.  
  63.     async def backup_database(self, backup_type):
  64.         ServerName = await asyncio.to_thread(Utils.get_server_name, self.guild_id)
  65.         BasePath = await asyncio.to_thread(Utils.server_folder, ServerName)
  66.         SourceDatabasePath = await asyncio.to_thread(Utils.primary_database_folder, BasePath, ServerName)
  67.         BackupsRoute = await asyncio.to_thread(Utils.backup_folder, BasePath, ServerName)
  68.         if backup_type == "full":
  69.             await Utils.full_backup(SourceDatabasePath, BackupsRoute)
  70.  
  71.         elif backup_type == 'differential':
  72.             await Utils.differential_backup(SourceDatabasePath, BackupsRoute)
  73.  
  74.    
  75.     async def schedule_full_backup(self, db_name):
  76.         current_time = datetime.now()
  77.         midnight = current_time.replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=1)
  78.         self.tasks[f"full_backup_{db_name}"] = {
  79.             'func': lambda: self.backup_database(db_name, 'full'),
  80.             'interval': 86400,
  81.             'last_run': current_time,
  82.             'next_run': midnight
  83.         }
  84.         await self.save_task(f"full_backup_{db_name}", 'backup_database', 86400, current_time, midnight)
  85.  
  86.  
  87.     async def schedule_differential_backup(self, db_name):
  88.         current_time = datetime.now()
  89.         midnight = current_time.replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=1)
  90.         self.tasks[f"differential_backup_{db_name}"] = {
  91.             'func': lambda: self.backup_database(db_name, 'differential'),
  92.             'interval': 86400,
  93.             'last_run': current_time,
  94.             'next_run': midnight
  95.         }
  96.         await self.save_task(f"differential_backup_{db_name}", 'backup_database', 86400, current_time, midnight)
  97.  
  98.     async def check_tasks(self):
  99.         while self.running:
  100.             current_time = datetime.now()
  101.             for name, task in self.tasks.items():
  102.                 if current_time >= task['next_run']:
  103.                     if name.startswith('full_backup_'):
  104.                         await task['func']()
  105.                         task['last_run'] = current_time
  106.                         task['next_run'] = current_time + timedelta(seconds=task['interval'])
  107.                         await self.save_task(name, 'backup_database', task['interval'], task['last_run'], task['next_run'])
  108.                     elif name.startswith('differential_backup_'):
  109.                         db_name = name.split('_')[2]
  110.                         full_backup_task = self.tasks.get(f"full_backup_{db_name}")
  111.                         if full_backup_task and full_backup_task['next_run'] - current_time < timedelta(hours=1):
  112.                             task['next_run'] = current_time + timedelta(hours=1)
  113.                             await self.save_task(name, 'backup_database', task['interval'], task['last_run'], task['next_run'])
  114.                         else:
  115.                             await task['func']()
  116.                             task['last_run'] = current_time
  117.                             task['next_run'] = current_time + timedelta(seconds=task['interval'])
  118.                             await self.save_task(name, 'backup_database', task['interval'], task['last_run'], task['next_run'])
  119.             await asyncio.sleep(3600)
  120.  
  121.     def stop(self):
  122.         self.running = False
  123.         self.conn.close()
  124.  
  125.  
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement