channel.messages: add foreign key deletions before message delete

- images: export try_unlink, and allow it to handle pathlib.Path
This commit is contained in:
Luna 2019-03-23 21:19:30 -03:00
parent 306444061d
commit 15fa2f3231
2 changed files with 45 additions and 5 deletions

View File

@ -18,6 +18,7 @@ along with this program. If not, see <http://www.gnu.org/licenses/>.
""" """
import json import json
from pathlib import Path
from PIL import Image from PIL import Image
from quart import Blueprint, request, current_app as app, jsonify from quart import Blueprint, request, current_app as app, jsonify
@ -35,6 +36,7 @@ from litecord.utils import pg_set_json
from litecord.embed.sanitizer import fill_embed from litecord.embed.sanitizer import fill_embed
from litecord.embed.messages import process_url_embed from litecord.embed.messages import process_url_embed
from litecord.blueprints.channel.dm_checks import dm_pre_check from litecord.blueprints.channel.dm_checks import dm_pre_check
from litecord.images import try_unlink
log = Logger(__name__) log = Logger(__name__)
@ -515,6 +517,37 @@ async def edit_message(channel_id, message_id):
return jsonify(message) return jsonify(message)
async def _del_msg_fkeys(message_id: int):
attachs = await app.db.fetch("""
SELECT id FROM attachments
WHERE message_id = $1
""", message_id)
attachs = [r['id'] for r in attachs]
attachments = Path('./attachments')
for attach_id in attachs:
# anything starting with the given attachment shall be
# deleted, because there may be resizes of the original
# attachment laying around.
for filepath in attachments.glob(f'{attach_id}*'):
try_unlink(filepath)
# after trying to delete all available attachments, delete
# them from the database.
# take the chance and delete all the data from the other tables too!
tables = ['attachments', 'message_webhook_info',
'message_reactions', 'channel_pins']
for table in tables:
await app.db.execute(f"""
DELETE FROM {table}
WHERE message_id = $1
""", message_id)
@bp.route('/<int:channel_id>/messages/<int:message_id>', methods=['DELETE']) @bp.route('/<int:channel_id>/messages/<int:message_id>', methods=['DELETE'])
async def delete_message(channel_id, message_id): async def delete_message(channel_id, message_id):
user_id = await token_check() user_id = await token_check()
@ -535,6 +568,8 @@ async def delete_message(channel_id, message_id):
if not can_delete: if not can_delete:
raise Forbidden('You can not delete this message') raise Forbidden('You can not delete this message')
await _del_msg_fkeys(message_id)
await app.db.execute(""" await app.db.execute("""
DELETE FROM messages DELETE FROM messages
WHERE messages.id = $1 WHERE messages.id = $1

View File

@ -22,7 +22,7 @@ import mimetypes
import asyncio import asyncio
import base64 import base64
import tempfile import tempfile
from typing import Optional from typing import Optional, Union
from dataclasses import dataclass from dataclasses import dataclass
from hashlib import sha256 from hashlib import sha256
@ -205,9 +205,14 @@ def _invalid(kwargs: dict) -> Optional[Icon]:
return Icon(None, None, '') return Icon(None, None, '')
def _try_unlink(path: str): def try_unlink(path: Union[Path, str]):
"""Try unlinking a file. Does not do anything if the file
does not exist."""
try: try:
os.remove(path) if isinstance(path, Path):
path.unlink()
else:
os.remove(path)
except FileNotFoundError: except FileNotFoundError:
pass pass
@ -258,8 +263,8 @@ async def resize_gif(raw_data: bytes, target: tuple) -> tuple:
output_handler.close() output_handler.close()
# delete the files we created with mkstemp # delete the files we created with mkstemp
_try_unlink(input_path) try_unlink(input_path)
_try_unlink(output_path) try_unlink(output_path)
# reseek, save to raw_data, reseek again. # reseek, save to raw_data, reseek again.
# TODO: remove raw_data altogether as its inefficient # TODO: remove raw_data altogether as its inefficient