1# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2# SPDX-License-Identifier: Apache-2.0
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8# http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16import mimetypes
17from typing import IO, Optional, Union
18
19from .s3 import S3StorageProvider
20
21PROVIDER = "s8k"
22
23
[docs]
24class S8KStorageProvider(S3StorageProvider):
25 """
26 A concrete implementation of the :py:class:`multistorageclient.types.StorageProvider` for interacting with SwiftStack.
27
28 This provider extends S3StorageProvider with SwiftStack-specific features:
29
30 - Legacy retry mode for handling HTTP 429 errors
31 - Optional content type inference from file extensions
32
33 :param infer_content_type: When True, automatically infers MIME types from file extensions during upload operations.
34 For example, ``.wav`` files are uploaded with ``Content-Type: audio/x-wav``, enabling browsers to play
35 media files inline rather than downloading them. Uses Python's built-in ``mimetypes`` module for inference.
36 Default is False. Only affects write operations (``upload_file``, ``write``, ``put_object``).
37
38 .. note::
39 This provider inherits all parameters from :py:class:`multistorageclient.providers.s3.S3StorageProvider`.
40 See the S3 provider documentation for additional configuration options like ``region_name``, ``endpoint_url``,
41 ``base_path``, ``credentials_provider``, ``multipart_threshold``, ``max_concurrency``, etc.
42
43 Example:
44 .. code-block:: yaml
45
46 profiles:
47 my-s8k-profile:
48 storage_provider:
49 type: s8k
50 options:
51 base_path: my-bucket
52 region_name: us-east-1
53 endpoint_url: https://s8k.example.com
54 infer_content_type: true # Enable MIME type inference
55 """
56
57 def __init__(self, *args, **kwargs):
58 # Extract the infer_content_type option before passing to parent
59 self._infer_content_type = kwargs.pop("infer_content_type", False)
60
61 # "legacy" retry mode is required for SwiftStack (retry on HTTP 429 errors)
62 kwargs["retries"] = kwargs.get("retries", {}) | {"mode": "legacy"}
63
64 super().__init__(*args, **kwargs)
65
66 # override the provider name from "s3"
67 self._provider_name = PROVIDER
68
69 def _guess_content_type(self, file_path: str) -> Optional[str]:
70 """
71 Guess the content type based on the file extension using Python's mimetypes module.
72
73 This method mimics the behavior of python-swiftclient, which automatically infers
74 MIME types from file extensions (e.g., .wav → audio/x-wav).
75
76 :param file_path: The path or key of the file (can be local path or remote key).
77 :return: The guessed MIME type, or None if inference is disabled or type cannot be determined.
78 """
79 if not self._infer_content_type:
80 return None
81
82 # Use mimetypes to guess the content type based on file extension
83 content_type, _ = mimetypes.guess_type(file_path)
84 return content_type
85
86 def _put_object(
87 self,
88 path: str,
89 body: bytes,
90 if_match: Optional[str] = None,
91 if_none_match: Optional[str] = None,
92 attributes: Optional[dict[str, str]] = None,
93 content_type: Optional[str] = None,
94 ) -> int:
95 """
96 Uploads an object with optional content type inference.
97
98 Infers the content type from the file extension if enabled and not explicitly provided.
99
100 :param path: The S3 path where the object will be uploaded.
101 :param body: The content of the object as bytes.
102 :param if_match: Optional If-Match header value.
103 :param if_none_match: Optional If-None-Match header value.
104 :param attributes: Optional attributes to attach to the object.
105 :param content_type: Optional explicit Content-Type. If not provided, will be inferred if enabled.
106 """
107 # Infer content type from the path if not explicitly provided
108 if content_type is None:
109 content_type = self._guess_content_type(path)
110
111 # Delegate to parent with inferred or explicit content_type
112 return super()._put_object(path, body, if_match, if_none_match, attributes, content_type)
113
114 def _upload_file(
115 self,
116 remote_path: str,
117 f: Union[str, IO],
118 attributes: Optional[dict[str, str]] = None,
119 content_type: Optional[str] = None,
120 ) -> int:
121 """
122 Uploads a file with optional content type inference.
123
124 Infers the content type from the file extension if enabled and not explicitly provided.
125
126 :param remote_path: The remote path where the file will be uploaded.
127 :param f: The source file to upload (file path or file object).
128 :param attributes: Optional attributes to attach to the file.
129 :param content_type: Optional explicit Content-Type. If not provided, will be inferred if enabled.
130 """
131 # Infer content type if not explicitly provided
132 if content_type is None:
133 # For file paths, infer from the local file path
134 # For file objects, infer from the remote path (destination key)
135 if isinstance(f, str):
136 content_type = self._guess_content_type(f)
137 else:
138 content_type = self._guess_content_type(remote_path)
139
140 # Delegate to parent with inferred or explicit content_type
141 return super()._upload_file(remote_path, f, attributes, content_type)