nucliadb 4.0.0.post542__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/writer/tus/azure.py

@@ -0,0 +1,111 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program. If not, see <http://www.gnu.org/licenses/>.
+#
+from __future__ import annotations
+
+from typing import Optional
+
+from nucliadb.writer import logger
+from nucliadb.writer.tus.dm import FileDataManager
+from nucliadb.writer.tus.storage import BlobStore, FileStorageManager
+from nucliadb_protos.resources_pb2 import CloudFile
+from nucliadb_utils.storages import CHUNK_SIZE
+from nucliadb_utils.storages.azure import AzureObjectStore
+from nucliadb_utils.storages.exceptions import ObjectNotFoundError
+from nucliadb_utils.storages.utils import ObjectMetadata
+
+
+class AzureBlobStore(BlobStore):
+    async def finalize(self):
+        if self._object_store is None:
+            return
+        try:
+            await self._object_store.finalize()
+        except Exception:
+            logger.exception("Error closing AzureBlobStore")
+        self._object_store = None
+
+    async def initialize(self, account_url: str, connection_string: Optional[str] = None):
+        self.bucket = "nucliadb-{kbid}"
+        self.source = CloudFile.Source.AZURE
+        self._object_store = AzureObjectStore(account_url, connection_string=connection_string)
+        await self._object_store.initialize()
+
+    @property
+    def object_store(self) -> AzureObjectStore:
+        assert self._object_store is not None
+        return self._object_store
+
+    async def check_exists(self, bucket_name: str) -> bool:
+        return await self.object_store.bucket_exists(bucket_name)
+
+    async def create_bucket(self, bucket_name: str) -> bool:
+        created = await self.object_store.bucket_create(bucket_name)
+        return not created
+
+
+class AzureFileStorageManager(FileStorageManager):
+    storage: AzureBlobStore
+    chunk_size = CHUNK_SIZE
+    min_upload_size = None
+
+    @property
+    def object_store(self) -> AzureObjectStore:
+        return self.storage.object_store
+
+    async def start(self, dm: FileDataManager, path: str, kbid: str):
+        bucket = self.storage.get_bucket_name(kbid)
+        if dm.filename == 0:
+            filename = "file"
+        else:
+            filename = dm.filename
+        metadata = ObjectMetadata(
+            filename=filename,
+            content_type=dm.content_type,
+            size=dm.size,
+        )
+        await self.object_store.upload_multipart_start(bucket, path, metadata)
+        await dm.update(path=path, bucket=bucket)
+
+    async def delete_upload(self, uri: str, kbid: str) -> None:
+        bucket = self.storage.get_bucket_name(kbid)
+        try:
+            await self.object_store.delete(bucket, uri)
+        except ObjectNotFoundError:
+            logger.warning(
+                "Attempt to delete an upload but not found",
+                extra={"uri": uri, "kbid": kbid, "bucket": bucket},
+            )
+
+    async def append(self, dm: FileDataManager, iterable, offset: int) -> int:
+        bucket = dm.get("bucket")
+        assert bucket is not None
+        path = dm.get("path")
+        assert path is not None
+        uploaded_bytes = await self.object_store.upload_multipart_append(bucket, path, iterable)
+        await dm.update(offset=offset)
+        return uploaded_bytes
+
+    async def finish(self, dm: FileDataManager):
+        path = dm.get("path")
+        await dm.finish()
+        return path
+
+    def validate_intermediate_chunk(self, uploaded_bytes: int):
+        pass

nucliadb/writer/tus/dm.py

@@ -20,6 +20,7 @@
 import time
 from typing import Any, Optional
 
+import backoff
 import orjson
 from redis import asyncio as aioredis
 from starlette.requests import Request
@@ -33,6 +34,11 @@ class NoRedisConfigured(Exception):
     pass
 
 
+RETRIABLE_REDIS_ERRORS = (
+    aioredis.ConnectionError,
+    aioredis.TimeoutError,
+)
+
 DATA: dict[str, Any] = {}
 
 
@@ -59,10 +65,7 @@ class FileDataManager:
             # someone else
             last_activity: Optional[int] = self._data.get("last_activity")
             if last_activity and (time.time() - last_activity) < self._ttl:
-                if (
-                    request.headers
-                    and request.headers.get("tus-override-upload", "0") != "1"
-                ):
+                if request.headers and request.headers.get("tus-override-upload", "0") != "1":
                     raise HTTPPreconditionFailed(
                         detail="There is already an active tusupload that conflicts with this one."
                     )
@@ -136,7 +139,7 @@ class RedisFileDataManagerFactory:
 
     async def finalize(self):
         try:
-            await self.redis.close(close_connection_pool=True)
+            await self.redis.aclose(close_connection_pool=True)
         except Exception:
             logger.warning("Error closing redis connection", exc_info=True)
             pass
@@ -146,6 +149,9 @@ class RedisFileDataManager(FileDataManager):
     def __init__(self, redis: aioredis.Redis):
         self.redis = redis
 
+    @backoff.on_exception(
+        backoff.expo, RETRIABLE_REDIS_ERRORS, jitter=backoff.random_jitter, max_tries=3
+    )
     async def load(self, key):
         # preload data
         self.key = key
@@ -157,6 +163,9 @@ class RedisFileDataManager(FileDataManager):
                 self._data = orjson.loads(data)
                 self._loaded = True
 
+    @backoff.on_exception(
+        backoff.expo, RETRIABLE_REDIS_ERRORS, jitter=backoff.random_jitter, max_tries=3
+    )
     async def save(self):
         if self.key is None:
             raise Exception("Not initialized")
@@ -164,6 +173,9 @@ class RedisFileDataManager(FileDataManager):
         value = orjson.dumps(self._data)
         await self.redis.set(self.key, value, ex=self._ttl)
 
+    @backoff.on_exception(
+        backoff.expo, RETRIABLE_REDIS_ERRORS, jitter=backoff.random_jitter, max_tries=3
+    )
     async def _delete_key(self):
         if self.key is None:
             raise Exception("Not initialized")

nucliadb/writer/tus/exceptions.py

@@ -31,9 +31,7 @@ class HTTPException(StarletteHTTPException):
 
     def __init__(self, detail: Optional[str] = None):
         if self._status_code:
-            super(HTTPException, self).__init__(
-                status_code=self._status_code, detail=detail
-            )
+            super(HTTPException, self).__init__(status_code=self._status_code, detail=detail)
         else:
             raise AttributeError("Status code not defined")
 

nucliadb/writer/tus/gcs.py

@@ -28,27 +28,28 @@ import tempfile
 import uuid
 from concurrent.futures import ThreadPoolExecutor
 from copy import deepcopy
-from datetime import datetime
-from typing import AsyncIterator, Optional
+from typing import Optional
 from urllib.parse import quote_plus
 
 import aiohttp
 import backoff
-from nucliadb_protos.resources_pb2 import CloudFile
+import google.auth.compute_engine.credentials  # type: ignore
+import google.auth.transport.requests  # type: ignore
+import google.oauth2.credentials  # type: ignore
+from google.auth.exceptions import DefaultCredentialsError  # type: ignore
 from oauth2client.service_account import ServiceAccountCredentials  # type: ignore
 
 from nucliadb.writer import logger
 from nucliadb.writer.tus.dm import FileDataManager
 from nucliadb.writer.tus.exceptions import (
-    CloudFileNotFound,
     HTTPBadRequest,
-    HTTPNotFound,
     HTTPPreconditionFailed,
     ResumableURINotAvailable,
 )
 from nucliadb.writer.tus.storage import BlobStore, FileStorageManager
 from nucliadb.writer.tus.utils import to_str
-from nucliadb_utils.storages.gcs import CHUNK_SIZE, MIN_UPLOAD_SIZE
+from nucliadb_protos.resources_pb2 import CloudFile
+from nucliadb_utils.storages.gcs import CHUNK_SIZE, MIN_UPLOAD_SIZE, TIMEOUT
 
 
 class GoogleCloudException(Exception):
@@ -76,7 +77,7 @@ class GCloudBlobStore(BlobStore):
     loop = None
     upload_url: str
     object_base_url: str
-    json_credentials: str
+    json_credentials: Optional[str]
     bucket: str
     location: str
     project: str
@@ -90,9 +91,18 @@ class GCloudBlobStore(BlobStore):
         return {"AUTHORIZATION": f"Bearer {token}"}
 
     def _get_access_token(self):
-        access_token = self._credentials.get_access_token()
-        self._creation_access_token = datetime.now()
-        return access_token.access_token
+        if isinstance(
+            self._credentials, google.auth.compute_engine.credentials.Credentials
+        ) or isinstance(self._credentials, google.oauth2.credentials.Credentials):
+            # google default auth object
+            if self._credentials.expired or self._credentials.valid is False:
+                request = google.auth.transport.requests.Request()
+                self._credentials.refresh(request)
+
+            return self._credentials.token
+        else:
+            access_token = self._credentials.get_access_token()
+            return access_token.access_token
 
     async def finalize(self):
         if self.session is not None:
@@ -113,32 +123,35 @@ class GCloudBlobStore(BlobStore):
         self.project = project
         self.bucket_labels = bucket_labels
         self.object_base_url = object_base_url + "/storage/v1/b"
-        self.upload_url = (
-            object_base_url + "/upload/storage/v1/b/{bucket}/o?uploadType=resumable"
-        )  # noqa
-
+        self.upload_url = object_base_url + "/upload/storage/v1/b/{bucket}/o?uploadType=resumable"  # noqa
+        self.json_credentials = json_credentials
         self._credentials = None
 
-        if json_credentials is not None:
-            self.json_credentials_file = os.path.join(
-                tempfile.mkdtemp(), "gcs_credentials.json"
-            )
-            open(self.json_credentials_file, "w").write(
-                base64.b64decode(json_credentials).decode("utf-8")
-            )
+        if self.json_credentials is not None and self.json_credentials.strip() != "":
+            self.json_credentials_file = os.path.join(tempfile.mkdtemp(), "gcs_credentials.json")
+            with open(self.json_credentials_file, "w") as file:
+                file.write(base64.b64decode(self.json_credentials).decode("utf-8"))
             self._credentials = ServiceAccountCredentials.from_json_keyfile_name(
                 self.json_credentials_file, SCOPES
             )
+        else:
+            try:
+                self._credentials, self.project = google.auth.default()
+            except DefaultCredentialsError:
+                logger.warning("Setting up without credentials as couldn't find workload identity")
+                self._credentials = None
 
         loop = asyncio.get_event_loop()
-        self.session = aiohttp.ClientSession(loop=loop)
+        self.session = aiohttp.ClientSession(loop=loop, timeout=TIMEOUT)
 
     async def check_exists(self, bucket_name: str):
         if self.session is None:
             raise AttributeError()
 
         headers = await self.get_access_headers()
-        url = f"{self.object_base_url}/{bucket_name}?project={self.project}"
+        # Using object access url instead of bucket access to avoid
+        # giving admin permission to the SA, needed to GET a bucket
+        url = f"{self.object_base_url}/{bucket_name}/o"
         async with self.session.get(
             url,
             headers=headers,
@@ -177,9 +190,7 @@ class GCloudFileStorageManager(FileStorageManager):
     chunk_size = CHUNK_SIZE
     min_upload_size = MIN_UPLOAD_SIZE
 
-    @backoff.on_exception(
-        backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=4
-    )
+    @backoff.on_exception(backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=4)
     async def start(self, dm: FileDataManager, path: str, kbid: str):
         """Init an upload.
 
@@ -187,12 +198,15 @@ class GCloudFileStorageManager(FileStorageManager):
         _resumable_uri : uri to resumable upload
         _uri : finished uploaded image
         """
+
         if self.storage.session is None:
             raise AttributeError()
 
-        upload_file_id = dm.get("upload_file_id", str(uuid.uuid4()))
+        upload_file_id = dm.get("upload_file_id")
         if upload_file_id is not None:
             await self.delete_upload(upload_file_id, kbid)
+        else:
+            upload_file_id = str(uuid.uuid4())
 
         bucket = self.storage.get_bucket_name(kbid)
         init_url = "{}&name={}".format(
@@ -237,13 +251,9 @@ class GCloudFileStorageManager(FileStorageManager):
                 raise GoogleCloudException(text)
             resumable_uri = call.headers["Location"]
 
-        await dm.update(
-            resumable_uri=resumable_uri, upload_file_id=upload_file_id, path=path
-        )
+        await dm.update(resumable_uri=resumable_uri, upload_file_id=upload_file_id, path=path)
 
-    @backoff.on_exception(
-        backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=4
-    )
+    @backoff.on_exception(backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=4)
     async def delete_upload(self, uri, kbid):
         bucket = self.storage.get_bucket_name(kbid)
 
@@ -266,8 +276,7 @@ class GCloudFileStorageManager(FileStorageManager):
                 if resp.status not in (200, 204, 404):
                     if resp.status == 404:
                         logger.error(
-                            f"Attempt to delete not found gcloud: {data}, "
-                            f"status: {resp.status}",
+                            f"Attempt to delete not found gcloud: {data}, " f"status: {resp.status}",
                             exc_info=True,
                         )
                     else:
@@ -275,9 +284,7 @@ class GCloudFileStorageManager(FileStorageManager):
         else:
             raise AttributeError("No valid uri")
 
-    @backoff.on_exception(
-        backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=4
-    )
+    @backoff.on_exception(backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=4)
     async def _append(self, dm: FileDataManager, data, offset):
         if self.storage.session is None:
             raise AttributeError()
@@ -342,9 +349,7 @@ class GCloudFileStorageManager(FileStorageManager):
                 break
         return count
 
-    @backoff.on_exception(
-        backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=4
-    )
+    @backoff.on_exception(backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=4)
     async def finish(self, dm: FileDataManager):
         if dm.size == 0:
             if self.storage.session is None:
@@ -370,46 +375,6 @@ class GCloudFileStorageManager(FileStorageManager):
         await dm.finish()
         return path
 
-    async def iter_data(self, uri, kbid: str, headers: Optional[dict[str, str]] = None):
-        if self.storage.session is None:
-            raise AttributeError()
-        if headers is None:
-            headers = {}
-
-        url = "{}/{}/o/{}".format(
-            self.storage.object_base_url,
-            self.storage.get_bucket_name(kbid),
-            quote_plus(uri),
-        )
-        headers_auth = await self.storage.get_access_headers()
-        headers.update(headers_auth)
-        async with self.storage.session.get(
-            url, headers=headers, params={"alt": "media"}, timeout=-1
-        ) as api_resp:
-            if api_resp.status not in (200, 206):
-                text = await api_resp.text()
-                if api_resp.status == 404:
-                    raise CloudFileNotFound("Google cloud file not found")
-                elif api_resp.status == 401:
-                    logger.warning(f"Invalid google cloud credentials error: {text}")
-                    raise HTTPNotFound(
-                        detail=f"Google cloud invalid credentials: {text}"
-                    )
-                raise GoogleCloudException(f"{api_resp.status}: {text}")
-            while True:
-                chunk = await api_resp.content.read(1024 * 1024)
-                if len(chunk) > 0:
-                    yield chunk
-                else:
-                    break
-
-    async def read_range(
-        self, uri: str, kbid: str, start: int, end: int
-    ) -> AsyncIterator[bytes]:
-        """
-        Iterate through ranges of data
-        """
-        async for chunk in self.iter_data(
-            uri, kbid, headers={"Range": f"bytes={start}-{end - 1}"}
-        ):
-            yield chunk
+    def validate_intermediate_chunk(self, uploaded_bytes: int):
+        if uploaded_bytes < self.min_upload_size:
+            raise ValueError(f"Intermediate chunks cannot be smaller than {self.min_upload_size} bytes")

nucliadb/writer/tus/local.py

@@ -22,14 +22,13 @@ from __future__ import annotations
 import json
 import os
 import uuid
-from typing import AsyncIterator
+from typing import Any
 
 import aiofiles
-from nucliadb_protos.resources_pb2 import CloudFile
 
 from nucliadb.writer.tus.dm import FileDataManager
-from nucliadb.writer.tus.exceptions import CloudFileNotFound
 from nucliadb.writer.tus.storage import BlobStore, FileStorageManager
+from nucliadb_protos.resources_pb2 import CloudFile
 from nucliadb_utils.storages import CHUNK_SIZE
 
 
@@ -50,51 +49,24 @@ class LocalFileStorageManager(FileStorageManager):
         bucket = self.storage.get_bucket_name(kbid)
         upload_file_id = dm.get("upload_file_id", str(uuid.uuid4()))
         init_url = self.get_file_path(bucket, upload_file_id)
-        metadata_init_url = self.metadata_key(init_url)
         metadata = {
             "FILENAME": dm.filename,
             "CONTENT_TYPE": dm.content_type,
             "SIZE": dm.size,
         }
-        async with aiofiles.open(metadata_init_url, "w+") as resp:
-            await resp.write(json.dumps(metadata))
+        await self.set_metadata(kbid, upload_file_id, metadata)
 
         async with aiofiles.open(init_url, "wb+") as aio_fi:
             await aio_fi.write(b"")
 
-        await dm.update(upload_file_id=upload_file_id, path=path, bucket=bucket)
+        await dm.update(upload_file_id=upload_file_id, path=path, bucket=bucket, kbid=kbid)
 
-    async def iter_data(self, uri, kbid: str, headers=None):
-        bucket = self.storage.get_bucket_name(kbid)
-        file_path = self.get_file_path(bucket, uri)
-        async with aiofiles.open(file_path) as resp:
-            data = await resp.read(CHUNK_SIZE)
-            while data is not None:
-                yield data
-                data = await resp.read(CHUNK_SIZE)
-
-    async def read_range(
-        self, uri: str, kbid: str, start: int, end: int
-    ) -> AsyncIterator[bytes]:
-        """
-        Iterate through ranges of data
-        """
+    async def set_metadata(self, kbid: str, upload_file_id: str, metadata: dict[str, Any]):
         bucket = self.storage.get_bucket_name(kbid)
-        file_path = self.get_file_path(bucket, uri)
-        try:
-            async with aiofiles.open(file_path, "rb") as resp:
-                await resp.seek(start)
-                count = 0
-                data = await resp.read(CHUNK_SIZE)
-                while data and count < end:
-                    if count + len(data) > end:
-                        new_end = end - count
-                        data = data[:new_end]
-                    yield data
-                    count += len(data)
-                    data = await resp.read(CHUNK_SIZE)
-        except FileNotFoundError:
-            raise CloudFileNotFound()
+        init_url = self.get_file_path(bucket, upload_file_id)
+        metadata_init_url = self.metadata_key(init_url)
+        async with aiofiles.open(metadata_init_url, "w+") as resp:
+            await resp.write(json.dumps(metadata))
 
     async def append(self, dm: FileDataManager, iterable, offset) -> int:
         count = 0
@@ -118,6 +90,15 @@ class LocalFileStorageManager(FileStorageManager):
         upload_file_id = dm.get("upload_file_id")
         from_url = self.get_file_path(bucket, upload_file_id)
 
+        if dm.size > 0:
+            kbid = dm.get("kbid")
+            metadata = {
+                "FILENAME": dm.filename,
+                "CONTENT_TYPE": dm.content_type,
+                "SIZE": dm.size,
+            }
+            await self.set_metadata(kbid, upload_file_id, metadata)
+
         path = dm.get("path")
         to_url = self.get_file_path(bucket, path)
         to_url_dirs = os.path.dirname(to_url)
@@ -138,6 +119,9 @@ class LocalFileStorageManager(FileStorageManager):
         file_path = self.get_file_path(bucket, uri)
         os.remove(file_path)
 
+    def validate_intermediate_chunk(self, uploaded_bytes: int):
+        pass
+
 
 class LocalBlobStore(BlobStore):
     def __init__(self, local_testing_files: str):