Fix file too long on tmp file create (#2132)

* Splits the extension on a ? query parameter
* Ensure that the filename in URLPath is truncated

python/cog/types.py

@@ -261,6 +261,9 @@ class URLPath(pathlib.PosixPath):  # pylint: disable=abstract-method
     _path: Optional[Path]
 
     def __init__(self, *, source: str, filename: str, fileobj: io.IOBase) -> None:  # pylint: disable=super-init-not-called
+        if len(filename) > FILENAME_MAX_LENGTH:
+            filename = _truncate_filename_bytes(filename, FILENAME_MAX_LENGTH)
+
         self.source = source
         self.filename = filename
         self.fileobj = fileobj
@@ -540,5 +543,6 @@ def _truncate_filename_bytes(s: str, length: int, encoding: str = "utf-8") -> st
     and avoiding text encoding corruption from truncation.
     """
     root, ext = os.path.splitext(s.encode(encoding))
+    ext = ext.decode(encoding).split("?")[0].encode(encoding)
     root = root[: length - len(ext) - 1]
     return root.decode(encoding, "ignore") + "~" + ext.decode(encoding)

python/tests/test_types.py

@@ -1,10 +1,12 @@
 import io
 import pickle
+import random
+import string
 
 import pytest
 import responses
 
-from cog.types import Secret, URLFile, get_filename
+from cog.types import Secret, URLFile, URLPath, get_filename
 
 
 def test_urlfile_protocol_validation():
@@ -146,3 +148,22 @@ def test_secret_type():
 
     assert secret.get_secret_value() == secret_value
     assert str(secret) == "**********"
+
+
+def test_truncate_filename_if_long():
+    # Test that a file too long exception is not raised.
+    random_str = "".join(
+        random.SystemRandom().choice(string.ascii_uppercase + string.digits)
+        for _ in range(350)
+    )
+    big_query = "query=" + random_str
+    filename = "waffwyyg~.zip"
+    url = "https://www.amazon.com/" + filename + "?" + big_query
+    fileobj = io.BytesIO()
+    url_path = URLPath(
+        source=url,
+        filename=filename + "?" + big_query,
+        fileobj=fileobj,
+    )
+    assert url_path.filename == "waffwyyg~~.zip"
+    _ = url_path.convert()