Changed how export works (in context of #34)
This commit is contained in:
@@ -1,39 +0,0 @@
|
||||
from os import sep, makedirs, remove
|
||||
from uuid import uuid1
|
||||
from app import app
|
||||
from pyrogram import filters
|
||||
from pyrogram.types import Message
|
||||
from pyrogram.client import Client
|
||||
from pyrogram.enums.chat_action import ChatAction
|
||||
from modules.logging import logWrite
|
||||
from modules.utils import should_quote, jsonSave
|
||||
from modules.database import col_applications
|
||||
from modules import custom_filters
|
||||
|
||||
|
||||
@app.on_message(
|
||||
custom_filters.enabled_applications
|
||||
& ~filters.scheduled
|
||||
& filters.command(["applications"], prefixes=["/"])
|
||||
& custom_filters.admin
|
||||
)
|
||||
async def cmd_applications(app: Client, msg: Message):
|
||||
logWrite(f"Admin {msg.from_user.id} requested export of a database")
|
||||
await app.send_chat_action(msg.chat.id, ChatAction.UPLOAD_DOCUMENT)
|
||||
filename = uuid1()
|
||||
output = []
|
||||
for entry in col_applications.find():
|
||||
del entry["_id"]
|
||||
entry["date"] = entry["date"].strftime("%d.%m.%Y, %H:%M")
|
||||
entry["application"]["2"] = entry["application"]["2"].strftime(
|
||||
"%d.%m.%Y, %H:%M"
|
||||
)
|
||||
output.append(entry)
|
||||
makedirs("tmp", exist_ok=True)
|
||||
jsonSave(output, f"tmp{sep}{filename}.json")
|
||||
await msg.reply_document(
|
||||
document=f"tmp{sep}{filename}.json",
|
||||
file_name="applications",
|
||||
quote=should_quote(msg),
|
||||
)
|
||||
remove(f"tmp{sep}{filename}.json")
|
||||
182
modules/commands/export.py
Normal file
182
modules/commands/export.py
Normal file
@@ -0,0 +1,182 @@
|
||||
from csv import QUOTE_ALL
|
||||
from os import makedirs, path, remove
|
||||
from uuid import uuid1
|
||||
|
||||
import aiofiles
|
||||
from aiocsv.writers import AsyncDictWriter
|
||||
from pyrogram import filters
|
||||
from pyrogram.client import Client
|
||||
from pyrogram.enums.chat_action import ChatAction
|
||||
from pyrogram.types import Message
|
||||
from ujson import dumps
|
||||
|
||||
from app import app
|
||||
from modules import custom_filters
|
||||
from modules.database import col_applications, col_sponsorships, col_warnings
|
||||
from modules.logging import logWrite
|
||||
from modules.utils import locale, should_quote
|
||||
|
||||
|
||||
@app.on_message(
|
||||
custom_filters.enabled_general
|
||||
& ~filters.scheduled
|
||||
& filters.command(["export"], prefixes=["/"])
|
||||
& custom_filters.admin
|
||||
)
|
||||
async def cmd_export(app: Client, msg: Message):
|
||||
if len(msg.command) <= 1:
|
||||
await msg.reply_text(
|
||||
locale("syntax_export", "message", locale=msg.from_user),
|
||||
quote=should_quote(msg),
|
||||
)
|
||||
return
|
||||
|
||||
selection = msg.command[1].lower()
|
||||
|
||||
if selection not in ["applications", "warnings", "sponsorships", "bans"]:
|
||||
await msg.reply_text(
|
||||
locale("syntax_export", "message", locale=msg.from_user),
|
||||
quote=should_quote(msg),
|
||||
)
|
||||
return
|
||||
|
||||
logWrite(f"Admin {msg.from_user.id} requested export of {selection}")
|
||||
|
||||
makedirs("tmp", exist_ok=True)
|
||||
temp_file = path.join("tmp", str(uuid1()))
|
||||
|
||||
await app.send_chat_action(msg.chat.id, ChatAction.TYPING)
|
||||
|
||||
output_csv = []
|
||||
output_json = []
|
||||
|
||||
if selection == "applications":
|
||||
header_csv = [
|
||||
"user",
|
||||
"date",
|
||||
"admin",
|
||||
"question_1",
|
||||
"question_2",
|
||||
"question_3",
|
||||
"question_4",
|
||||
"question_5",
|
||||
"question_6",
|
||||
"question_7",
|
||||
"question_8",
|
||||
"question_9",
|
||||
"question_10",
|
||||
]
|
||||
|
||||
for entry in list(col_applications.find()):
|
||||
del entry["_id"]
|
||||
entry["date"] = entry["date"].isoformat()
|
||||
entry["application"]["2"] = entry["application"]["2"].isoformat()
|
||||
output_json.append(entry)
|
||||
|
||||
for entry in list(col_applications.find()):
|
||||
del entry["_id"]
|
||||
entry["date"] = entry["date"].isoformat()
|
||||
entry["application"]["2"] = entry["application"]["2"].isoformat()
|
||||
for index, value in enumerate(entry["application"]):
|
||||
entry[f"question_{index+1}"] = entry["application"][value]
|
||||
entry[
|
||||
"question_3"
|
||||
] = f"{entry['application']['3']['name']} ({entry['application']['3']['adminName1']}, {entry['application']['3']['countryName']})"
|
||||
del entry["application"]
|
||||
output_csv.append(entry)
|
||||
|
||||
elif selection == "warnings":
|
||||
header_csv = [
|
||||
"id",
|
||||
"user",
|
||||
"admin",
|
||||
"date",
|
||||
"reason",
|
||||
"active",
|
||||
"revoke_date",
|
||||
]
|
||||
|
||||
for entry in list(col_warnings.find()):
|
||||
for k, v in list(entry.items()):
|
||||
entry[{"_id": "id"}.get(k, k)] = entry.pop(k)
|
||||
entry["id"] = str(entry["id"])
|
||||
entry["date"] = entry["date"].isoformat()
|
||||
if entry["revoke_date"] is not None:
|
||||
entry["revoke_date"] = entry["revoke_date"].isoformat()
|
||||
output_json.append(entry)
|
||||
output_csv.append(entry)
|
||||
|
||||
elif selection == "sponsorships":
|
||||
header_csv = [
|
||||
"user",
|
||||
"date",
|
||||
"admin",
|
||||
"streamer",
|
||||
"expires",
|
||||
"proof",
|
||||
"label",
|
||||
]
|
||||
|
||||
for entry in list(col_sponsorships.find()):
|
||||
del entry["_id"]
|
||||
entry["date"] = entry["date"].isoformat()
|
||||
entry["sponsorship"]["expires"] = entry["sponsorship"][
|
||||
"expires"
|
||||
].isoformat()
|
||||
output_json.append(entry)
|
||||
|
||||
for entry in list(col_sponsorships.find()):
|
||||
del entry["_id"]
|
||||
entry["date"] = entry["date"].isoformat()
|
||||
entry["sponsorship"]["expires"] = entry["sponsorship"][
|
||||
"expires"
|
||||
].isoformat()
|
||||
for index, value in enumerate(entry["sponsorship"]):
|
||||
entry[value] = entry["sponsorship"][value]
|
||||
del entry["sponsorship"]
|
||||
output_csv.append(entry)
|
||||
|
||||
elif selection == "bans":
|
||||
header_csv = ["user", "admin", "date"]
|
||||
|
||||
for entry in list(col_warnings.find()):
|
||||
del entry["id"]
|
||||
entry["date"] = entry["date"].isoformat()
|
||||
output_json.append(entry)
|
||||
output_csv.append(entry)
|
||||
|
||||
# Saving CSV
|
||||
async with aiofiles.open(temp_file + ".csv", mode="w", encoding="utf-8") as file:
|
||||
writer = AsyncDictWriter(file, header_csv, restval="NULL", quoting=QUOTE_ALL)
|
||||
await writer.writeheader()
|
||||
await writer.writerows(output_csv)
|
||||
|
||||
# Saving JSON
|
||||
async with aiofiles.open(temp_file + ".json", mode="w", encoding="utf-8") as file:
|
||||
await file.write(
|
||||
dumps(
|
||||
output_json, ensure_ascii=False, escape_forward_slashes=False, indent=4
|
||||
)
|
||||
)
|
||||
|
||||
# Sending CSV
|
||||
await app.send_chat_action(msg.chat.id, ChatAction.UPLOAD_DOCUMENT)
|
||||
await msg.reply_document(
|
||||
document=temp_file + ".csv",
|
||||
file_name=f"{selection}.csv",
|
||||
quote=should_quote(msg),
|
||||
)
|
||||
|
||||
# Sending JSON
|
||||
await app.send_chat_action(msg.chat.id, ChatAction.UPLOAD_DOCUMENT)
|
||||
await msg.reply_document(
|
||||
document=temp_file + ".json",
|
||||
file_name=f"{selection}.json",
|
||||
quote=should_quote(msg),
|
||||
)
|
||||
|
||||
del output_csv, output_json
|
||||
|
||||
# Removing temp files
|
||||
remove(temp_file + ".csv")
|
||||
remove(temp_file + ".json")
|
||||
Reference in New Issue
Block a user