mirror of
https://github.com/Cog-Creators/Red-DiscordBot.git
synced 2025-11-22 02:37:57 -05:00
PostgreSQL driver, tests against DB backends, and general drivers cleanup (#2723)
* PostgreSQL driver and general drivers cleanup Signed-off-by: Toby Harradine <tobyharradine@gmail.com> * Make tests pass Signed-off-by: Toby Harradine <tobyharradine@gmail.com> * Add black --target-version flag in make.bat Signed-off-by: Toby Harradine <tobyharradine@gmail.com> * Rewrite postgres driver Most of the logic is now in PL/pgSQL. This completely avoids the use of Python f-strings to format identifiers into queries. Although an SQL-injection attack would have been impossible anyway (only the owner would have ever had the ability to do that), using PostgreSQL's format() is more reliable for unusual identifiers. Performance-wise, I'm not sure whether this is an improvement, but I highly doubt that it's worse. Signed-off-by: Toby Harradine <tobyharradine@gmail.com> * Reformat Signed-off-by: Toby Harradine <tobyharradine@gmail.com> * Fix PostgresDriver.delete_all_data() Signed-off-by: Toby Harradine <tobyharradine@gmail.com> * Clean up PL/pgSQL code Signed-off-by: Toby Harradine <tobyharradine@gmail.com> * More PL/pgSQL cleanup Signed-off-by: Toby Harradine <tobyharradine@gmail.com> * PL/pgSQL function optimisations Signed-off-by: Toby Harradine <tobyharradine@gmail.com> * Ensure compatibility with PostgreSQL 10 and below Signed-off-by: Toby Harradine <tobyharradine@gmail.com> * More/better docstrings for PG functions Signed-off-by: Toby Harradine <tobyharradine@gmail.com> * Fix typo in docstring Signed-off-by: Toby Harradine <tobyharradine@gmail.com> * Return correct value on toggle() Signed-off-by: Toby Harradine <tobyharradine@gmail.com> * Use composite type for PG function parameters Signed-off-by: Toby Harradine <tobyharradine@gmail.com> * Fix JSON driver's Config.clear_all() Signed-off-by: Toby Harradine <tobyharradine@gmail.com> * Correct description for Mongo tox recipe Signed-off-by: Toby Harradine <tobyharradine@gmail.com> * Fix linting errors Signed-off-by: Toby Harradine <tobyharradine@gmail.com> * Update dep specification after merging bumpdeps Signed-off-by: Toby Harradine <tobyharradine@gmail.com> * Add towncrier entries Signed-off-by: Toby Harradine <tobyharradine@gmail.com> * Update from merge Signed-off-by: Toby Harradine <tobyharradine@gmail.com> * Mention [postgres] extra in install docs Signed-off-by: Toby Harradine <tobyharradine@gmail.com> * Support more connection options and use better defaults Signed-off-by: Toby Harradine <tobyharradine@gmail.com> * Actually pass PG env vars in tox Signed-off-by: Toby Harradine <tobyharradine@gmail.com> * Replace event trigger with manual DELETE queries Signed-off-by: Toby Harradine <tobyharradine@gmail.com>
This commit is contained in:
Toby Harradine
committed by
Michael H
Michael H
parent
57fa29dd64
commit
d1a46acc9a
@@ -3,15 +3,12 @@ import contextlib
|
||||
import datetime
|
||||
import importlib
|
||||
import itertools
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import pathlib
|
||||
import sys
|
||||
import platform
|
||||
import getpass
|
||||
import pip
|
||||
import tarfile
|
||||
import traceback
|
||||
from collections import namedtuple
|
||||
from pathlib import Path
|
||||
@@ -23,15 +20,18 @@ import aiohttp
|
||||
import discord
|
||||
import pkg_resources
|
||||
|
||||
from redbot.core import (
|
||||
from . import (
|
||||
__version__,
|
||||
version_info as red_version_info,
|
||||
VersionInfo,
|
||||
checks,
|
||||
commands,
|
||||
drivers,
|
||||
errors,
|
||||
i18n,
|
||||
config,
|
||||
)
|
||||
from .utils import create_backup
|
||||
from .utils.predicates import MessagePredicate
|
||||
from .utils.chat_formatting import humanize_timedelta, pagify, box, inline, humanize_list
|
||||
from .commands.requires import PrivilegeLevel
|
||||
@@ -1307,105 +1307,71 @@ class Core(commands.Cog, CoreLogic):
|
||||
|
||||
@commands.command()
|
||||
@checks.is_owner()
|
||||
async def backup(self, ctx: commands.Context, *, backup_path: str = None):
|
||||
"""Creates a backup of all data for the instance."""
|
||||
if backup_path:
|
||||
path = pathlib.Path(backup_path)
|
||||
if not (path.exists() and path.is_dir()):
|
||||
return await ctx.send(
|
||||
_("That path doesn't seem to exist. Please provide a valid path.")
|
||||
)
|
||||
from redbot.core.data_manager import basic_config, instance_name
|
||||
from redbot.core.drivers.red_json import JSON
|
||||
async def backup(self, ctx: commands.Context, *, backup_dir: str = None):
|
||||
"""Creates a backup of all data for the instance.
|
||||
|
||||
data_dir = Path(basic_config["DATA_PATH"])
|
||||
if basic_config["STORAGE_TYPE"] == "MongoDB":
|
||||
from redbot.core.drivers.red_mongo import Mongo
|
||||
|
||||
m = Mongo("Core", "0", **basic_config["STORAGE_DETAILS"])
|
||||
db = m.db
|
||||
collection_names = await db.list_collection_names()
|
||||
for c_name in collection_names:
|
||||
if c_name == "Core":
|
||||
c_data_path = data_dir / basic_config["CORE_PATH_APPEND"]
|
||||
else:
|
||||
c_data_path = data_dir / basic_config["COG_PATH_APPEND"] / c_name
|
||||
docs = await db[c_name].find().to_list(None)
|
||||
for item in docs:
|
||||
item_id = str(item.pop("_id"))
|
||||
target = JSON(c_name, item_id, data_path_override=c_data_path)
|
||||
target.data = item
|
||||
await target._save()
|
||||
backup_filename = "redv3-{}-{}.tar.gz".format(
|
||||
instance_name, ctx.message.created_at.strftime("%Y-%m-%d %H-%M-%S")
|
||||
)
|
||||
if data_dir.exists():
|
||||
if not backup_path:
|
||||
backup_pth = data_dir.home()
|
||||
else:
|
||||
backup_pth = Path(backup_path)
|
||||
backup_file = backup_pth / backup_filename
|
||||
|
||||
to_backup = []
|
||||
exclusions = [
|
||||
"__pycache__",
|
||||
"Lavalink.jar",
|
||||
os.path.join("Downloader", "lib"),
|
||||
os.path.join("CogManager", "cogs"),
|
||||
os.path.join("RepoManager", "repos"),
|
||||
]
|
||||
downloader_cog = ctx.bot.get_cog("Downloader")
|
||||
if downloader_cog and hasattr(downloader_cog, "_repo_manager"):
|
||||
repo_output = []
|
||||
repo_mgr = downloader_cog._repo_manager
|
||||
for repo in repo_mgr._repos.values():
|
||||
repo_output.append({"url": repo.url, "name": repo.name, "branch": repo.branch})
|
||||
repo_filename = data_dir / "cogs" / "RepoManager" / "repos.json"
|
||||
with open(str(repo_filename), "w") as f:
|
||||
f.write(json.dumps(repo_output, indent=4))
|
||||
instance_data = {instance_name: basic_config}
|
||||
instance_file = data_dir / "instance.json"
|
||||
with open(str(instance_file), "w") as instance_out:
|
||||
instance_out.write(json.dumps(instance_data, indent=4))
|
||||
for f in data_dir.glob("**/*"):
|
||||
if not any(ex in str(f) for ex in exclusions):
|
||||
to_backup.append(f)
|
||||
with tarfile.open(str(backup_file), "w:gz") as tar:
|
||||
for f in to_backup:
|
||||
tar.add(str(f), recursive=False)
|
||||
print(str(backup_file))
|
||||
await ctx.send(
|
||||
_("A backup has been made of this instance. It is at {}.").format(backup_file)
|
||||
)
|
||||
if backup_file.stat().st_size > 8_000_000:
|
||||
await ctx.send(_("This backup is too large to send via DM."))
|
||||
return
|
||||
await ctx.send(_("Would you like to receive a copy via DM? (y/n)"))
|
||||
|
||||
pred = MessagePredicate.yes_or_no(ctx)
|
||||
try:
|
||||
await ctx.bot.wait_for("message", check=pred, timeout=60)
|
||||
except asyncio.TimeoutError:
|
||||
await ctx.send(_("Response timed out."))
|
||||
else:
|
||||
if pred.result is True:
|
||||
await ctx.send(_("OK, it's on its way!"))
|
||||
try:
|
||||
async with ctx.author.typing():
|
||||
await ctx.author.send(
|
||||
_("Here's a copy of the backup"),
|
||||
file=discord.File(str(backup_file)),
|
||||
)
|
||||
except discord.Forbidden:
|
||||
await ctx.send(
|
||||
_("I don't seem to be able to DM you. Do you have closed DMs?")
|
||||
)
|
||||
except discord.HTTPException:
|
||||
await ctx.send(_("I could not send the backup file."))
|
||||
else:
|
||||
await ctx.send(_("OK then."))
|
||||
You may provide a path to a directory for the backup archive to
|
||||
be placed in. If the directory does not exist, the bot will
|
||||
attempt to create it.
|
||||
"""
|
||||
if backup_dir is None:
|
||||
dest = Path.home()
|
||||
else:
|
||||
await ctx.send(_("That directory doesn't seem to exist..."))
|
||||
dest = Path(backup_dir)
|
||||
|
||||
driver_cls = drivers.get_driver_class()
|
||||
if driver_cls != drivers.JsonDriver:
|
||||
await ctx.send(_("Converting data to JSON for backup..."))
|
||||
async with ctx.typing():
|
||||
await config.migrate(driver_cls, drivers.JsonDriver)
|
||||
|
||||
log.info("Creating backup for this instance...")
|
||||
try:
|
||||
backup_fpath = await create_backup(dest)
|
||||
except OSError as exc:
|
||||
await ctx.send(
|
||||
_(
|
||||
"Creating the backup archive failed! Please check your console or logs for "
|
||||
"details."
|
||||
)
|
||||
)
|
||||
log.exception("Failed to create backup archive", exc_info=exc)
|
||||
return
|
||||
|
||||
if backup_fpath is None:
|
||||
await ctx.send(_("Your datapath appears to be empty."))
|
||||
return
|
||||
|
||||
log.info("Backup archive created successfully at '%s'", backup_fpath)
|
||||
await ctx.send(
|
||||
_("A backup has been made of this instance. It is located at `{path}`.").format(
|
||||
path=backup_fpath
|
||||
)
|
||||
)
|
||||
if backup_fpath.stat().st_size > 8_000_000:
|
||||
await ctx.send(_("This backup is too large to send via DM."))
|
||||
return
|
||||
await ctx.send(_("Would you like to receive a copy via DM? (y/n)"))
|
||||
|
||||
pred = MessagePredicate.yes_or_no(ctx)
|
||||
try:
|
||||
await ctx.bot.wait_for("message", check=pred, timeout=60)
|
||||
except asyncio.TimeoutError:
|
||||
await ctx.send(_("Response timed out."))
|
||||
else:
|
||||
if pred.result is True:
|
||||
await ctx.send(_("OK, it's on its way!"))
|
||||
try:
|
||||
async with ctx.author.typing():
|
||||
await ctx.author.send(
|
||||
_("Here's a copy of the backup"), file=discord.File(str(backup_fpath))
|
||||
)
|
||||
except discord.Forbidden:
|
||||
await ctx.send(_("I don't seem to be able to DM you. Do you have closed DMs?"))
|
||||
except discord.HTTPException:
|
||||
await ctx.send(_("I could not send the backup file."))
|
||||
else:
|
||||
await ctx.send(_("OK then."))
|
||||
|
||||
@commands.command()
|
||||
@commands.cooldown(1, 60, commands.BucketType.user)
|
||||
|
||||
Reference in New Issue
Block a user