Merge pull request #2165 from python-discord/improve-pastebin-error-handling

9 files changed, 102 insertions, 44 deletions

diff --git a/bot/exts/events/code_jams/_cog.py b/bot/exts/events/code_jams/_cog.py
index 452199f5f..86c357863 100644
--- a/bot/exts/events/code_jams/_cog.py
+++ b/bot/exts/events/code_jams/_cog.py
@@ -12,7 +12,7 @@ from bot.constants import Emojis, Roles
 from bot.exts.events.code_jams import _channels
 from bot.log import get_logger
 from bot.utils.members import get_or_fetch_member
-from bot.utils.services import send_to_paste_service
+from bot.utils.services import PasteTooLongError, PasteUploadError, send_to_paste_service
 
 log = get_logger(__name__)
 
@@ -139,11 +139,14 @@ class CodeJams(commands.Cog):
             format_category_info(category, channels) for category, channels in categories.items()
         )
 
-        url = await send_to_paste_service(deletion_details)
-        if url is None:
-            url = "**Unable to send deletion details to the pasting service.**"
+        try:
+            message = await send_to_paste_service(deletion_details)
+        except PasteTooLongError:
+            message = "**Too long to upload to paste service.**"
+        except PasteUploadError:
+            message = "**Failed to upload to paste service.**"
 
-        return f"Are you sure you want to delete all code jam channels?\n\nThe channels to be deleted: {url}"
+        return f"Are you sure you want to delete all code jam channels?\n\nThe channels to be deleted: {message}"
 
     @codejam.command()
     @commands.has_any_role(Roles.admins, Roles.code_jam_event_team)
diff --git a/bot/exts/moderation/dm_relay.py b/bot/exts/moderation/dm_relay.py
index a86c9e409..bf0b96a58 100644
--- a/bot/exts/moderation/dm_relay.py
+++ b/bot/exts/moderation/dm_relay.py
@@ -5,7 +5,7 @@ from bot.bot import Bot
 from bot.constants import Emojis, MODERATION_ROLES
 from bot.log import get_logger
 from bot.utils.channel import is_mod_channel
-from bot.utils.services import send_to_paste_service
+from bot.utils.services import PasteTooLongError, PasteUploadError, send_to_paste_service
 
 log = get_logger(__name__)
 
@@ -53,14 +53,14 @@ class DMRelay(Cog):
             f"User: {user} ({user.id})\n"
             f"Channel ID: {user.dm_channel.id}\n\n"
         )
-
-        paste_link = await send_to_paste_service(metadata + output, extension="txt")
-
-        if paste_link is None:
-            await ctx.send(f"{Emojis.cross_mark} Failed to upload output to hastebin.")
-            return
-
-        await ctx.send(paste_link)
+        try:
+            message = await send_to_paste_service(metadata + output, extension="txt")
+        except PasteTooLongError:
+            message = f"{Emojis.cross_mark} Too long to upload to paste service."
+        except PasteUploadError:
+            message = f"{Emojis.cross_mark} Failed to upload to paste service."
+
+        await ctx.send(message)
 
     async def cog_check(self, ctx: Context) -> bool:
         """Only allow moderators to invoke the commands in this cog in mod channels."""
diff --git a/bot/exts/moderation/metabase.py b/bot/exts/moderation/metabase.py
index 0ca9fd4a5..1b4b0cc71 100644
--- a/bot/exts/moderation/metabase.py
+++ b/bot/exts/moderation/metabase.py
@@ -17,6 +17,7 @@ from bot.converters import allowed_strings
 from bot.log import get_logger
 from bot.utils import send_to_paste_service
 from bot.utils.channel import is_mod_channel
+from bot.utils.services import PasteTooLongError, PasteUploadError
 
 log = get_logger(__name__)
 
@@ -141,11 +142,15 @@ class Metabase(Cog):
                 # Format it nicely for human eyes
                 out = json.dumps(out, indent=4, sort_keys=True)
 
-        paste_link = await send_to_paste_service(out, extension=extension)
-        if paste_link:
-            message = f":+1: {ctx.author.mention} Here's your link: {paste_link}"
+        try:
+            paste_link = await send_to_paste_service(out, extension=extension)
+        except PasteTooLongError:
+            message = f":x: {ctx.author.mention} Too long to upload to paste service."
+        except PasteUploadError:
+            message = f":x: {ctx.author.mention} Failed to upload to paste service."
         else:
-            message = f":x: {ctx.author.mention} Link service is unavailible."
+            message = f":+1: {ctx.author.mention} Here's your link: {paste_link}"
+
         await ctx.send(
             f"{message}\nYou can also access this data within internal eval by doing: "
             f"`bot.get_cog('Metabase').exports[{question_id}]`"
diff --git a/bot/exts/utils/internal.py b/bot/exts/utils/internal.py
index 2148a3676..3125cee75 100644
--- a/bot/exts/utils/internal.py
+++ b/bot/exts/utils/internal.py
@@ -16,6 +16,7 @@ from bot.bot import Bot
 from bot.constants import DEBUG_MODE, Roles
 from bot.log import get_logger
 from bot.utils import find_nth_occurrence, send_to_paste_service
+from bot.utils.services import PasteTooLongError, PasteUploadError
 
 log = get_logger(__name__)
 
@@ -194,11 +195,14 @@ async def func():  # (None,) -> Any
             truncate_index = newline_truncate_index
 
         if len(out) > truncate_index:
-            paste_link = await send_to_paste_service(out, extension="py")
-            if paste_link is not None:
-                paste_text = f"full contents at {paste_link}"
-            else:
+            try:
+                paste_link = await send_to_paste_service(out, extension="py")
+            except PasteTooLongError:
+                paste_text = "too long to upload to paste service."
+            except PasteUploadError:
                 paste_text = "failed to upload contents to paste service."
+            else:
+                paste_text = f"full contents at {paste_link}"
 
             await ctx.send(
                 f"```py\n{out[:truncate_index]}\n```"
diff --git a/bot/exts/utils/snekbox.py b/bot/exts/utils/snekbox.py
index 32b7c8761..98dfd2efa 100644
--- a/bot/exts/utils/snekbox.py
+++ b/bot/exts/utils/snekbox.py
@@ -18,6 +18,7 @@ from bot.decorators import redirect_output
 from bot.log import get_logger
 from bot.utils import send_to_paste_service
 from bot.utils.messages import wait_for_deletion
+from bot.utils.services import PasteTooLongError, PasteUploadError
 
 log = get_logger(__name__)
 
@@ -65,7 +66,7 @@ if not hasattr(sys, "_setup_finished"):
 {setup}
 """
 
-MAX_PASTE_LEN = 10000
+MAX_PASTE_LENGTH = 10_000
 
 # The Snekbox commands' whitelists and blacklists.
 NO_SNEKBOX_CHANNELS = (Channels.python_general,)
@@ -137,10 +138,12 @@ class Snekbox(Cog):
         """Upload the job's output to a paste service and return a URL to it if successful."""
         log.trace("Uploading full output to paste service...")
 
-        if len(output) > MAX_PASTE_LEN:
-            log.info("Full output is too long to upload")
+        try:
+            return await send_to_paste_service(output, extension="txt", max_length=MAX_PASTE_LENGTH)
+        except PasteTooLongError:
             return "too long to upload"
-        return await send_to_paste_service(output, extension="txt")
+        except PasteUploadError:
+            return "unable to upload"
 
     @staticmethod
     def prepare_timeit_input(codeblocks: list[str]) -> tuple[str, list[str]]:
diff --git a/bot/utils/__init__.py b/bot/utils/__init__.py
index 13533a467..567821126 100644
--- a/bot/utils/__init__.py
+++ b/bot/utils/__init__.py
@@ -1,4 +1,12 @@
 from bot.utils.helpers import CogABCMeta, find_nth_occurrence, has_lines, pad_base64
-from bot.utils.services import send_to_paste_service
+from bot.utils.services import PasteTooLongError, PasteUploadError, send_to_paste_service
 
-__all__ = ['CogABCMeta', 'find_nth_occurrence', 'has_lines', 'pad_base64', 'send_to_paste_service']
+__all__ = [
+    'CogABCMeta',
+    'find_nth_occurrence',
+    'has_lines',
+    'pad_base64',
+    'send_to_paste_service',
+    'PasteUploadError',
+    'PasteTooLongError',
+]
diff --git a/bot/utils/services.py b/bot/utils/services.py
index 439c8d500..a752ac0ec 100644
--- a/bot/utils/services.py
+++ b/bot/utils/services.py
@@ -1,5 +1,3 @@
-from typing import Optional
-
 from aiohttp import ClientConnectorError
 
 import bot
@@ -9,18 +7,40 @@ from bot.log import get_logger
 log = get_logger(__name__)
 
 FAILED_REQUEST_ATTEMPTS = 3
+MAX_PASTE_LENGTH = 100_000
+
+
+class PasteUploadError(Exception):
+    """Raised when an error is encountered uploading to the paste service."""
+
+
+class PasteTooLongError(Exception):
+    """Raised when content is too large to upload to the paste service."""
 
 
-async def send_to_paste_service(contents: str, *, extension: str = "") -> Optional[str]:
+async def send_to_paste_service(contents: str, *, extension: str = "", max_length: int = MAX_PASTE_LENGTH) -> str:
     """
     Upload `contents` to the paste service.
 
-    `extension` is added to the output URL
+    Add `extension` to the output URL. Use `max_length` to limit the allowed contents length
+    to lower than the maximum allowed by the paste service.
 
-    When an error occurs, `None` is returned, otherwise the generated URL with the suffix.
+    Raise `ValueError` if `max_length` is greater than the maximum allowed by the paste service.
+    Raise `PasteTooLongError` if `contents` is too long to upload, and `PasteUploadError` if uploading fails.
+
+    Return the generated URL with the extension.
     """
+    if max_length > MAX_PASTE_LENGTH:
+        raise ValueError(f"`max_length` must not be greater than {MAX_PASTE_LENGTH}")
+
     extension = extension and f".{extension}"
-    log.debug(f"Sending contents of size {len(contents.encode())} bytes to paste service.")
+
+    contents_size = len(contents.encode())
+    if contents_size > max_length:
+        log.info("Contents too large to send to paste service.")
+        raise PasteTooLongError(f"Contents of size {contents_size} greater than maximum size {max_length}")
+
+    log.debug(f"Sending contents of size {contents_size} bytes to paste service.")
     paste_url = URLs.paste_service.format(key="documents")
     for attempt in range(1, FAILED_REQUEST_ATTEMPTS + 1):
         try:
@@ -59,3 +79,5 @@ async def send_to_paste_service(contents: str, *, extension: str = "") -> Option
             f"Got unexpected JSON response from paste service: {response_json}\n"
             f"trying again ({attempt}/{FAILED_REQUEST_ATTEMPTS})."
         )
+
+    raise PasteUploadError("Failed to upload contents to paste service")
diff --git a/tests/bot/exts/utils/test_snekbox.py b/tests/bot/exts/utils/test_snekbox.py
index 3c555c051..b870a9945 100644
--- a/tests/bot/exts/utils/test_snekbox.py
+++ b/tests/bot/exts/utils/test_snekbox.py
@@ -35,15 +35,15 @@ class SnekboxTests(unittest.IsolatedAsyncioTestCase):
         resp.json.assert_awaited_once()
 
     async def test_upload_output_reject_too_long(self):
-        """Reject output longer than MAX_PASTE_LEN."""
-        result = await self.cog.upload_output("-" * (snekbox.MAX_PASTE_LEN + 1))
+        """Reject output longer than MAX_PASTE_LENGTH."""
+        result = await self.cog.upload_output("-" * (snekbox.MAX_PASTE_LENGTH + 1))
         self.assertEqual(result, "too long to upload")
 
     @patch("bot.exts.utils.snekbox.send_to_paste_service")
     async def test_upload_output(self, mock_paste_util):
         """Upload the eval output to the URLs.paste_service.format(key="documents") endpoint."""
         await self.cog.upload_output("Test output.")
-        mock_paste_util.assert_called_once_with("Test output.", extension="txt")
+        mock_paste_util.assert_called_once_with("Test output.", extension="txt", max_length=snekbox.MAX_PASTE_LENGTH)
 
     async def test_codeblock_converter(self):
         ctx = MockContext()
diff --git a/tests/bot/utils/test_services.py b/tests/bot/utils/test_services.py
index 3b71022db..d0e801299 100644
--- a/tests/bot/utils/test_services.py
+++ b/tests/bot/utils/test_services.py
@@ -4,7 +4,9 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
 from aiohttp import ClientConnectorError
 
-from bot.utils.services import FAILED_REQUEST_ATTEMPTS, send_to_paste_service
+from bot.utils.services import (
+    FAILED_REQUEST_ATTEMPTS, MAX_PASTE_LENGTH, PasteTooLongError, PasteUploadError, send_to_paste_service
+)
 from tests.helpers import MockBot
 
 
@@ -55,23 +57,34 @@ class PasteTests(unittest.IsolatedAsyncioTestCase):
         for error_json in test_cases:
             with self.subTest(error_json=error_json):
                 response.json = AsyncMock(return_value=error_json)
-                result = await send_to_paste_service("")
+                with self.assertRaises(PasteUploadError):
+                    await send_to_paste_service("")
                 self.assertEqual(self.bot.http_session.post.call_count, FAILED_REQUEST_ATTEMPTS)
-                self.assertIsNone(result)
 
             self.bot.http_session.post.reset_mock()
 
     async def test_request_repeated_on_connection_errors(self):
         """Requests are repeated in the case of connection errors."""
         self.bot.http_session.post = MagicMock(side_effect=ClientConnectorError(Mock(), Mock()))
-        result = await send_to_paste_service("")
+        with self.assertRaises(PasteUploadError):
+            await send_to_paste_service("")
         self.assertEqual(self.bot.http_session.post.call_count, FAILED_REQUEST_ATTEMPTS)
-        self.assertIsNone(result)
 
     async def test_general_error_handled_and_request_repeated(self):
         """All `Exception`s are handled, logged and request repeated."""
         self.bot.http_session.post = MagicMock(side_effect=Exception)
-        result = await send_to_paste_service("")
+        with self.assertRaises(PasteUploadError):
+            await send_to_paste_service("")
         self.assertEqual(self.bot.http_session.post.call_count, FAILED_REQUEST_ATTEMPTS)
         self.assertLogs("bot.utils", logging.ERROR)
-        self.assertIsNone(result)
+
+    async def test_raises_error_on_too_long_input(self):
+        """Ensure PasteTooLongError is raised if `contents` is longer than `MAX_PASTE_LENGTH`."""
+        contents = "a" * (MAX_PASTE_LENGTH + 1)
+        with self.assertRaises(PasteTooLongError):
+            await send_to_paste_service(contents)
+
+    async def test_raises_on_too_large_max_length(self):
+        """Ensure ValueError is raised if `max_length` passed is greater than `MAX_PASTE_LENGTH`."""
+        with self.assertRaises(ValueError):
+            await send_to_paste_service("Hello World!", max_length=MAX_PASTE_LENGTH + 1)