diff --git a/python/lib/sift_client/_internal/low_level_wrappers/upload.py b/python/lib/sift_client/_internal/low_level_wrappers/upload.py index ba7e595f3..e1c9c6777 100644 --- a/python/lib/sift_client/_internal/low_level_wrappers/upload.py +++ b/python/lib/sift_client/_internal/low_level_wrappers/upload.py @@ -2,6 +2,7 @@ import asyncio import logging +import mimetypes from pathlib import Path from typing import TYPE_CHECKING, Any @@ -16,6 +17,10 @@ # Configure logging logger = logging.getLogger(__name__) +# Parquet MIME types are not included in mimetypes by default +mimetypes.add_type("application/vnd.apache.parquet", ".pqt") +mimetypes.add_type("application/vnd.apache.parquet", ".parquet") + class UploadLowLevelClient(LowLevelClientBase, WithRestClient): """Low-level client for file upload operations. @@ -182,7 +187,6 @@ def _mime_and_content_type_from_path(path: Path) -> tuple[str, str | None, str | Returns: A tuple of (file_name, mime_type, content_encoding). """ - import mimetypes file_name = path.name mime, encoding = mimetypes.guess_type(path) diff --git a/python/lib/sift_client/_tests/_internal/low_level_wrappers/test_upload.py b/python/lib/sift_client/_tests/_internal/low_level_wrappers/test_upload.py new file mode 100644 index 000000000..d40fd9aec --- /dev/null +++ b/python/lib/sift_client/_tests/_internal/low_level_wrappers/test_upload.py @@ -0,0 +1,16 @@ +"""Tests for UploadLowLevelClient functionality.""" + +from pathlib import Path + +from sift_client._internal.low_level_wrappers.upload import UploadLowLevelClient + + +class TestUploadLowLevelClient: + class TestMimeAndContentTypeFromPath: + def test_parquet_file_extension(self): + _, mime, _ = UploadLowLevelClient._mime_and_content_type_from_path(Path("data.parquet")) + assert mime == "application/vnd.apache.parquet" + + def test_pqt_file_extension(self): + _, mime, _ = UploadLowLevelClient._mime_and_content_type_from_path(Path("data.pqt")) + assert mime == "application/vnd.apache.parquet"