python(feat): add data import api to sift_client (#515)

Author: wei-qlu

python/lib/sift_client/_internal/low_level_wrappers/data_imports.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, cast
+
+from sift.data_imports.v2.data_imports_pb2 import (
+    CreateDataImportFromUploadRequest,
+    CreateDataImportFromUploadResponse,
+    DetectConfigRequest,
+    DetectConfigResponse,
+    GetDataImportRequest,
+    GetDataImportResponse,
+)
+from sift.data_imports.v2.data_imports_pb2_grpc import DataImportServiceStub
+
+from sift_client._internal.low_level_wrappers.base import LowLevelClientBase
+from sift_client.sift_types.data_import import (
+    CsvImportConfig,
+    Hdf5ImportConfig,
+    ImportConfig,
+    ParquetFlatDatasetImportConfig,
+    ParquetSingleChannelPerRowImportConfig,
+    TdmsImportConfig,
+)
+from sift_client.transport import WithGrpcClient
+
+if TYPE_CHECKING:
+    from sift.data_imports.v2.data_imports_pb2 import DataTypeKey
+
+    from sift_client.transport.grpc_transport import GrpcClient
+
+
+def _set_config_on_request(
+    request: CreateDataImportFromUploadRequest,
+    config: ImportConfig,
+) -> None:
+    """Set the appropriate config field on a proto request based on the config type."""
+    if isinstance(config, CsvImportConfig):
+        request.csv_config.CopyFrom(config._to_proto())
+    elif isinstance(
+        config, (ParquetFlatDatasetImportConfig, ParquetSingleChannelPerRowImportConfig)
+    ):
+        request.parquet_config.CopyFrom(config._to_proto())
+    elif isinstance(config, TdmsImportConfig):
+        request.tdms_config.CopyFrom(config._to_proto())
+    elif isinstance(config, Hdf5ImportConfig):
+        request.hdf5_config.CopyFrom(config._to_proto())
+    else:
+        raise TypeError(f"Unsupported import config type: {type(config).__name__}")
+
+
+class DataImportsLowLevelClient(LowLevelClientBase, WithGrpcClient):
+    """Low-level client for the DataImportService.
+
+    This class provides a thin wrapper around the autogenerated bindings for the DataImportsAPI.
+    """
+
+    def __init__(self, grpc_client: GrpcClient):
+        WithGrpcClient.__init__(self, grpc_client=grpc_client)
+
+    async def create_from_upload(self, config: ImportConfig) -> tuple[str, str]:
+        """Create a data import and get back a presigned upload URL.
+
+        Args:
+            config: The import configuration.
+
+        Returns:
+            A tuple of (data_import_id, upload_url).
+        """
+        request = CreateDataImportFromUploadRequest()
+        _set_config_on_request(request, config)
+        response = await self._grpc_client.get_stub(
+            DataImportServiceStub
+        ).CreateDataImportFromUpload(request)
+        response = cast("CreateDataImportFromUploadResponse", response)
+        return response.data_import_id, response.upload_url
+
+    async def get(self, data_import_id: str) -> GetDataImportResponse:
+        """Get a data import by ID.
+
+        Args:
+            data_import_id: The ID of the data import.
+
+        Returns:
+            The GetDataImportResponse proto.
+        """
+        request = GetDataImportRequest(data_import_id=data_import_id)
+        response = await self._grpc_client.get_stub(DataImportServiceStub).GetDataImport(request)
+        return cast("GetDataImportResponse", response)
+
+    async def detect_config(
+        self, data: bytes, data_type_key: DataTypeKey.ValueType
+    ) -> DetectConfigResponse:
+        """Call the DetectConfig RPC to auto-detect import configuration.
+
+        Args:
+            data: A sample of the file content (e.g. the first 64 KiB).
+            data_type_key: The file type hint.
+
+        Returns:
+            The raw DetectConfigResponse proto.
+        """
+        request = DetectConfigRequest(data=data, type=data_type_key)
+        response = await self._grpc_client.get_stub(DataImportServiceStub).DetectConfig(request)
+        return cast("DetectConfigResponse", response)

python/lib/sift_client/_internal/util/file.py

@@ -1,11 +1,14 @@
 from __future__ import annotations
 
+import os
+import struct
 import warnings
 import zipfile
 from typing import TYPE_CHECKING
 
 from alive_progress import alive_bar  # type: ignore[import-untyped]
 
+import sift_client as _sift_client_module
 from sift_client.errors import SiftWarning
 
 if TYPE_CHECKING:
@@ -14,6 +17,78 @@
     from sift_client.transport.rest_transport import RestClient
 
 
+class _ProgressReader:
+    """Wraps a file object to report read progress to an alive_bar callback."""
+
+    def __init__(self, file_object, progress_bar):
+        self._file_object = file_object
+        self._progress_bar = progress_bar
+
+    def read(self, size=-1):
+        chunk = self._file_object.read(size)
+        if chunk:
+            self._progress_bar(len(chunk))
+        return chunk
+
+    def __getattr__(self, name):
+        return getattr(self._file_object, name)
+
+
+def resolve_show_progress(*, is_sync: bool) -> bool:
+    """Resolve the show_progress setting from the global config.
+
+    Returns the global ``sift_client.config.show_progress`` value when set,
+    otherwise defaults to ``is_sync``.
+    """
+    global_setting = _sift_client_module.config.show_progress
+    if global_setting is not None:
+        return global_setting
+    return is_sync
+
+
+def upload_file(
+    signed_url: str,
+    file_path: Path,
+    *,
+    rest_client: RestClient,
+    show_progress: bool = False,
+) -> dict:
+    """Upload a file to a presigned URL.
+
+    Args:
+        signed_url: The presigned URL to upload to.
+        file_path: Path to the file to upload.
+        rest_client: The SDK rest client to use for the upload.
+        show_progress: If True, display a progress spinner during upload.
+
+    Returns:
+        The parsed JSON response from the server.
+
+    Raises:
+        ValueError: If the upload request fails.
+    """
+    file_size = file_path.stat().st_size
+
+    with alive_bar(
+        file_size,
+        title=f"Upload [{file_path.name}]",
+        spinner="dots_waves",
+        spinner_length=7,
+        unit="B",
+        scale="SI",
+        disable=not show_progress,
+    ) as bar:
+        with open(file_path, "rb") as file:
+            wrapped = _ProgressReader(file, bar)
+            response = rest_client.post(
+                signed_url,
+                data=wrapped,
+                headers={"Content-Disposition": f'attachment; filename="{file_path.name}"'},
+            )
+            response.raise_for_status()
+            return response.json()
+
+
 def download_file(
     signed_url: str,
     output_path: Path,
@@ -82,3 +157,27 @@ def extract_zip(zip_path: Path, output_dir: Path, *, delete_zip: bool = True) ->
         except OSError:
             warnings.warn(f"Failed to delete zip file '{zip_path}'", SiftWarning, stacklevel=2)
     return [output_dir / name for name in names if not name.endswith("/")]
+
+
+def extract_parquet_footer(path: Path) -> tuple[bytes, int]:
+    """Extract the Parquet footer bytes and compute the footer offset.
+
+    Args:
+        path: Path to the Parquet file.
+
+    Returns:
+        A tuple of (footer_bytes, footer_offset).
+
+    Raises:
+        ValueError: If the file is not a valid Parquet file.
+    """
+    with open(path, "rb") as f:
+        f.seek(-8, 2)
+        footer_tail = f.read(8)
+        footer_len = struct.unpack("<I", footer_tail[:4])[0]
+        magic = footer_tail[4:]
+        if magic != b"PAR1":
+            raise ValueError(f"Invalid Parquet file: missing magic bytes in {path}")
+        f.seek(-(footer_len + 8), 2)
+        footer_bytes = f.read(footer_len)
+    return footer_bytes, os.path.getsize(path) - len(footer_bytes) - 8