sift_py.file_attachment.service

  1from pathlib import Path
  2from typing import List, Optional, Union, cast
  3
  4from sift.remote_files.v1.remote_files_pb2 import (
  5    BatchDeleteRemoteFilesRequest,
  6    GetRemoteFileDownloadUrlRequest,
  7    GetRemoteFileDownloadUrlResponse,
  8    GetRemoteFileRequest,
  9    GetRemoteFileResponse,
 10    ListRemoteFilesRequest,
 11    ListRemoteFilesResponse,
 12    RemoteFile,
 13)
 14from sift.remote_files.v1.remote_files_pb2_grpc import RemoteFileServiceStub
 15
 16from sift_py.file_attachment._internal.download import download_remote_file
 17from sift_py.file_attachment._internal.upload import UploadService
 18from sift_py.file_attachment.entity import Entity
 19from sift_py.file_attachment.metadata import Metadata
 20from sift_py.grpc.transport import SiftChannel
 21from sift_py.rest import SiftRestConfig
 22
 23
 24class FileAttachmentService:
 25    """
 26    Service used to retrieve, upload, download, and delete file attachments. Seee `sift_py.file_attachment`
 27    for more information and examples on how to use this service.
 28    """
 29
 30    _remote_file_service_stub: RemoteFileServiceStub
 31    _upload_service: UploadService
 32
 33    def __init__(self, channel: SiftChannel, restconf: SiftRestConfig):
 34        self._remote_file_service_stub = RemoteFileServiceStub(channel)
 35        self._upload_service = UploadService(restconf)
 36
 37    def retrieve_attachments(self, entity: Entity) -> List[RemoteFile]:
 38        """
 39        Retrieves all file attachments for the provided `entity`.
 40        """
 41
 42        filter = f'entity_id=="{entity.entity_id}" && entity_type=="{entity.entity_type.value}"'
 43        page_size = 1_000
 44        next_page_token = ""
 45
 46        remote_files: List[RemoteFile] = []
 47
 48        while True:
 49            req = ListRemoteFilesRequest(
 50                filter=filter,
 51                page_size=page_size,
 52                page_token=next_page_token,
 53            )
 54            res = cast(ListRemoteFilesResponse, self._remote_file_service_stub.ListRemoteFiles(req))
 55            remote_files.extend(res.remote_files)
 56            next_page_token = res.next_page_token
 57
 58            if not next_page_token:
 59                break
 60
 61        return remote_files
 62
 63    def upload_attachment(
 64        self,
 65        path: Union[str, Path],
 66        entity: Entity,
 67        metadata: Optional[Metadata],
 68        description: Optional[str] = None,
 69        organization_id: Optional[str] = None,
 70    ) -> RemoteFile:
 71        """
 72        Uploads a file pointed to by `path` and attaches it to the provided `entity`.
 73
 74        - `path`: A path to the file to upload to Sift as a file attachment.
 75        - `entity`: The entity to attach the file to.
 76        - `metadata`: Optional metadata to include with the specific file.
 77        - `description`: An optional description to provide for the file attachment.
 78        - `organization_id`: Only required if your user belongs to multiple organizations.
 79        """
 80        remote_file_id = self._upload_service.upload_attachment(
 81            path,
 82            entity,
 83            metadata,
 84            description,
 85            organization_id,
 86        )
 87        req = GetRemoteFileRequest(remote_file_id=remote_file_id)
 88        res = cast(GetRemoteFileResponse, self._remote_file_service_stub.GetRemoteFile(req))
 89        return res.remote_file
 90
 91    def download_attachment(
 92        self,
 93        file: Union[RemoteFile, str],
 94        out: Optional[Union[str, Path]] = None,
 95    ) -> Path:
 96        """
 97        Downloads a file attachment and saves it locally.
 98
 99        - `remote_file`: Could either be an instance of `RemoteFile` or the ID of the remote file to download.
100        - `out`: If unspecified, then the file will be downloaded to the current working directory with the original name.
101        """
102
103        if isinstance(file, RemoteFile):
104            remote_file = file
105        else:
106            req = GetRemoteFileRequest(remote_file_id=file)
107            res = cast(GetRemoteFileResponse, self._remote_file_service_stub.GetRemoteFile(req))
108            remote_file = res.remote_file
109
110        output_file_path = (
111            Path(out) if isinstance(out, str) else Path(remote_file.file_name).resolve()
112        )
113
114        download_url_req = GetRemoteFileDownloadUrlRequest(
115            remote_file_id=remote_file.remote_file_id
116        )
117        download_url_res = cast(
118            GetRemoteFileDownloadUrlResponse,
119            self._remote_file_service_stub.GetRemoteFileDownloadUrl(download_url_req),
120        )
121        url = download_url_res.download_url
122
123        download_remote_file(url, output_file_path)
124
125        return output_file_path
126
127    def delete_file_attachments(self, *to_delete: Union[str, RemoteFile]):
128        """
129        Deletes remote files given a set of arguments that could either be instances of `RemoteFile` or the ID
130        of remote files to delete
131        """
132        remote_file_ids = [
133            remote_file.remote_file_id if isinstance(remote_file, RemoteFile) else remote_file
134            for remote_file in to_delete
135        ]
136
137        batch_size = 1_000
138        for i in range(0, len(remote_file_ids), batch_size):
139            batch = remote_file_ids[i : i + batch_size]
140            self._remote_file_service_stub.BatchDeleteRemoteFiles(
141                BatchDeleteRemoteFilesRequest(remote_file_ids=batch)
142            )