adding signed ingestion

Author: bgmello

alephclient/api.py

@@ -1,5 +1,6 @@
 import importlib.metadata
 import json
+import mimetypes
 import uuid
 import logging
 from itertools import count
@@ -471,6 +472,7 @@ def ingest_upload(
         metadata: Optional[Dict] = None,
         sync: bool = False,
         index: bool = True,
+        signed_url: bool = False,
     ) -> Dict:
         """
         Create an empty folder in a collection or upload a document to it
@@ -483,6 +485,9 @@ def ingest_upload(
         files, metadata contains foreign_id of the parent. Metadata for a
         directory contains foreign_id for itself as well as its parent and the
         name of the directory.
+        signed_url: use the signed URL workflow for file uploads. When True,
+        files are uploaded via a signed URL instead of multipart ingest.
+        Directories always use the standard ingest endpoint.
         """
         url_path = "collections/{0}/ingest".format(collection_id)
         params = {"sync": sync, "index": index}
@@ -491,6 +496,9 @@ def ingest_upload(
             data = {"meta": json.dumps(metadata)}
             return self._request("POST", url, data=data)
 
+        if signed_url:
+            return self._signed_url_upload(collection_id, file_path, metadata, index)
+
         for attempt in count(1):
             try:
                 with file_path.open("rb") as fh:
@@ -509,6 +517,54 @@ def ingest_upload(
                 backoff(ae, attempt)
         return {}
 
+    def _signed_url_upload(
+        self,
+        collection_id: str,
+        file_path: Path,
+        metadata: Optional[Dict],
+        index: bool,
+    ) -> Dict:
+        mime_type = mimetypes.guess_type(file_path.name)[0] or MIME
+        meta = dict(metadata or {})
+        meta["file_name"] = file_path.name
+        meta["mime_type"] = mime_type
+
+        for attempt in count(1):
+            try:
+                # Request a signed upload URL
+                upload_url = self._make_url("file/uploadUrl")
+                result = self._request("POST", upload_url)
+                signed_url = result["url"]
+                upload_id = result["id"]
+
+                # PUT file content to the signed URL
+                try:
+                    with file_path.open("rb") as fh:
+                        response = self.session.put(
+                            signed_url,
+                            data=fh,
+                            headers={"Content-Type": "application/octet-stream"},
+                        )
+                        response.raise_for_status()
+                except (RequestException, HTTPError) as exc:
+                    raise AlephException(exc) from exc
+
+                # Create document record.
+                # The server returns an empty 200 when a document with
+                # the same foreign_id already exists in the collection.
+                doc_url_path = f"collections/{collection_id}/document"
+                doc_url = self._make_url(doc_url_path, params={"index": index})
+                payload = {"upload_id": upload_id, "meta": meta}
+                result = self._request("POST", doc_url, json=payload)
+                if not result:
+                    return {"id": upload_id, "status": "ok"}
+                return result
+            except AlephException as ae:
+                if not ae.transient or attempt > self.retries:
+                    raise ae from ae
+                backoff(ae, attempt)
+        return {}
+
     def create_entityset(
         self, collection_id: str, type: str, label: str, summary: Optional[str]
     ) -> Dict:

alephclient/cli.py

@@ -87,6 +87,12 @@ def cli(ctx, host, api_key, retries):
     help="maximum number of parallel uploads",
 )
 @click.option("-f", "--foreign-id", required=True, help="foreign-id of the collection")
+@click.option(
+    "--signed-url",
+    is_flag=True,
+    default=False,
+    help="use signed URL workflow for file uploads",
+)
 @click.argument("path", type=click.Path(exists=True))
 @click.pass_context
 def crawldir(
@@ -98,6 +104,7 @@ def crawldir(
     noindex=False,
     nojunk=False,
     parallel=1,
+    signed_url=False,
 ):
     """Crawl a directory recursively and upload the documents in it to a
     collection."""
@@ -112,6 +119,7 @@ def crawldir(
             index=not noindex,
             nojunk=nojunk,
             parallel=parallel,
+            signed_url=signed_url,
         )
     except AlephException as exc:
         raise click.ClickException(str(exc))

alephclient/crawldir.py

@@ -22,9 +22,11 @@ def __init__(
         path: Path,
         index: bool = True,
         nojunk: bool = False,
+        signed_url: bool = False,
     ):
         self.api = api
         self.index = index
+        self.signed_url = signed_url
         self.exclude = (
             {
                 "f": re.compile(r"\..*|thumbs\.db|desktop\.ini", re.I),
@@ -128,11 +130,15 @@ def ingest_upload(self, path: Path, parent_id: str, foreign_id: str) -> str:
         log.info("Upload [%s->%s]: %s", self.collection_id, parent_id, foreign_id)
         if parent_id is not None:
             metadata["parent_id"] = parent_id
+        kwargs = {}
+        if self.signed_url:
+            kwargs["signed_url"] = True
         result = self.api.ingest_upload(
             self.collection_id,
             path,
             metadata=metadata,
             index=self.index,
+            **kwargs,
         )
         if "id" not in result and not hasattr(result, "id"):
             raise AlephException("Upload failed")
@@ -147,6 +153,7 @@ def crawl_dir(
     index: bool = True,
     nojunk: bool = False,
     parallel: int = 1,
+    signed_url: bool = False,
 ):
     """Crawl a directory and upload its content to a collection
 
@@ -158,7 +165,9 @@ def crawl_dir(
     """
     root = Path(path).resolve()
     collection = api.load_collection_by_foreign_id(foreign_id, config)
-    crawler = CrawlDirectory(api, collection, root, index=index, nojunk=nojunk)
+    crawler = CrawlDirectory(
+        api, collection, root, index=index, nojunk=nojunk, signed_url=signed_url
+    )
     consumers = []
 
     # Use one thread to produce using scandir and at least one to consume

alephclient/tests/test_crawldir.py

@@ -48,3 +48,13 @@ def test_is_excluded_exclude_dir(self):
         crawldir.exclude["d"] = re.compile(r"week1\/*", re.I)
         is_excluded = crawldir.is_excluded(path)
         assert is_excluded
+
+    def test_signed_url_default_false(self):
+        path = Path(os.path.join(self.base_path, "jan/week1"))
+        crawldir = CrawlDirectory(AlephAPI, {}, path)
+        assert crawldir.signed_url is False
+
+    def test_signed_url_true(self):
+        path = Path(os.path.join(self.base_path, "jan/week1"))
+        crawldir = CrawlDirectory(AlephAPI, {}, path, signed_url=True)
+        assert crawldir.signed_url is True

alephclient/tests/test_tasks.py

@@ -120,3 +120,22 @@ def test_ingest(self, mocker):
         ]
         for call in expected_calls:
             assert call in self.api.ingest_upload.mock_calls
+
+    def test_ingest_signed_url(self, mocker):
+        mocker.patch.object(self.api, "ingest_upload", return_value={"id": 42})
+        mocker.patch.object(
+            self.api, "load_collection_by_foreign_id", return_value={"id": 2}
+        )
+        mocker.patch.object(self.api, "update_collection")
+        crawl_dir(
+            self.api,
+            "alephclient/tests/testdata",
+            "test153",
+            {},
+            True,
+            True,
+            signed_url=True,
+        )
+        assert self.api.ingest_upload.call_count == 6
+        for call in self.api.ingest_upload.call_args_list:
+            assert call.kwargs.get("signed_url") is True