Add experimental streaming remote import from source

Author: ao-picterra

src/picterra/client.py

@@ -67,6 +67,7 @@ def _upload_file_to_blobstore(upload_url: str, filename: str):
         logger.error("Error when uploading to blobstore %s" % upload_url)
         raise APIError(resp.text)
 
+
 class Feature(TypedDict):
     type: Literal["Feature"]
     properties: dict[str, Any]
@@ -1033,3 +1034,51 @@ def create_marker(
         if not resp.ok:
             raise APIError(resp.text)
         return resp.json()
+
+    def import_raster_from_remote_source(
+        self,
+        raster_name: str,
+        folder_id: str,
+        source_id: str,
+        aoi_filename: str,
+        method: Literal["streaming"] = "streaming",
+    ) -> str:
+        """
+        This is an experimental feature
+
+        Import a raster from a remote imagery source given a GeoJSON file for the AOI
+
+        Args:
+            raster_name: Name of the new raster
+            folder_id: The id of the folder / project the raster will live in
+            source_id: The id of the remote imagery source to import from
+            filename: The filename of a GeoJSON file. This should contain a FeatureCollection of
+                Polygon/MultiPolygon representing the AOI of the new raster
+
+        Raises:
+            APIError: There was an error during import
+        """
+        # Get upload URL
+        resp = self.sess.post(self._api_url("rasters/import/"))
+        if not resp.ok:
+            raise APIError(resp.text)
+        data = resp.json()
+        upload_url = data["upload_url"]
+        upload_id = data["upload_id"]
+        # Upload to blobstore
+        _upload_file_to_blobstore(upload_url, aoi_filename)
+        # Commit upload
+        resp = self.sess.post(
+            self._api_url(f"rasters/import/{upload_id}/commit/"),
+            json={
+                "method": method,
+                "source_id": source_id,
+                "folder_id": folder_id,
+                "name": raster_name,
+            },
+        )
+        if not resp.ok:
+            raise APIError(resp.text)
+        # Poll operation and get raster identifier
+        operation = self._wait_until_operation_completes(resp.json())
+        return operation["metadata"]["raster_id"]

tests/test_client.py

@@ -312,6 +312,23 @@ def add_mock_detection_areas_upload_responses(raster_id):
     )
 
 
+def add_mock_remote_import_responses(upload_id, post_body):
+    match = responses.json_params_matcher(post_body)
+    # Upload initiation
+    data = {"upload_url": "http://storage.example.com", "upload_id": upload_id}
+    _add_api_response("rasters/import/", responses.POST, data)
+    # Storage PUT
+    responses.add(responses.PUT, "http://storage.example.com", status=200)
+    # Commit
+    _add_api_response(
+        f"rasters/import/{upload_id}/commit/",
+        responses.POST,
+        OP_RESP,
+        match=match,
+        status=200,
+    )
+
+
 def add_mock_detector_run_responses(detector_id):
     op_id = 43
     _add_api_response("detectors/%s/run/" % detector_id, responses.POST, OP_RESP)
@@ -1051,3 +1068,26 @@ def test_run_advanced_tool():
         == "mock_operation_type"
     )
     assert len(responses.calls) == 2
+
+
+@responses.activate
+def test_import_raster_from_remote_source():
+    body = {
+        "method": "streaming",
+        "source_id": "source",
+        "folder_id": "project",
+        "name": "image",
+    }
+    add_mock_remote_import_responses("upload_id", body)
+    add_mock_operations_responses("success")
+
+    client = _client()
+    # This just tests that this doesn't raise
+    with tempfile.NamedTemporaryFile() as f:
+        assert (
+            client.import_raster_from_remote_source(
+                "image", "project", "source", f.name
+            )
+            == "foo"
+        )
+    assert len(responses.calls) == 4