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 kwargs["request_checksum_calculation"] = "when_required"
62 kwargs["response_checksum_validation"] = "when_required"
63
64 # "legacy" retry mode is required for SwiftStack (retry on HTTP 429 errors)
65 kwargs["retries"] = kwargs.get("retries", {}) | {"mode": "legacy"}
66
67 super().__init__(*args, **kwargs)
68
69 # override the provider name from "s3"
70 self._provider_name = PROVIDER
71
72 def _guess_content_type(self, file_path: str) -> Optional[str]:
73 """
74 Guess the content type based on the file extension using Python's mimetypes module.
75
76 This method mimics the behavior of python-swiftclient, which automatically infers
77 MIME types from file extensions (e.g., .wav → audio/x-wav).
78
79 :param file_path: The path or key of the file (can be local path or remote key).
80 :return: The guessed MIME type, or None if inference is disabled or type cannot be determined.
81 """
82 if not self._infer_content_type:
83 return None
84
85 # Use mimetypes to guess the content type based on file extension
86 content_type, _ = mimetypes.guess_type(file_path)
87 return content_type
88
89 def _put_object(
90 self,
91 path: str,
92 body: bytes,
93 if_match: Optional[str] = None,
94 if_none_match: Optional[str] = None,
95 attributes: Optional[dict[str, str]] = None,
96 content_type: Optional[str] = None,
97 ) -> int:
98 """
99 Uploads an object with optional content type inference.
100
101 Infers the content type from the file extension if enabled and not explicitly provided.
102
103 :param path: The S3 path where the object will be uploaded.
104 :param body: The content of the object as bytes.
105 :param if_match: Optional If-Match header value.
106 :param if_none_match: Optional If-None-Match header value.
107 :param attributes: Optional attributes to attach to the object.
108 :param content_type: Optional explicit Content-Type. If not provided, will be inferred if enabled.
109 """
110 # Infer content type from the path if not explicitly provided
111 if content_type is None:
112 content_type = self._guess_content_type(path)
113
114 # Delegate to parent with inferred or explicit content_type
115 return super()._put_object(path, body, if_match, if_none_match, attributes, content_type)
116
117 def _upload_file(
118 self,
119 remote_path: str,
120 f: Union[str, IO],
121 attributes: Optional[dict[str, str]] = None,
122 content_type: Optional[str] = None,
123 ) -> int:
124 """
125 Uploads a file with optional content type inference.
126
127 Infers the content type from the file extension if enabled and not explicitly provided.
128
129 :param remote_path: The remote path where the file will be uploaded.
130 :param f: The source file to upload (file path or file object).
131 :param attributes: Optional attributes to attach to the file.
132 :param content_type: Optional explicit Content-Type. If not provided, will be inferred if enabled.
133 """
134 # Infer content type if not explicitly provided
135 if content_type is None:
136 # For file paths, infer from the local file path
137 # For file objects, infer from the remote path (destination key)
138 if isinstance(f, str):
139 content_type = self._guess_content_type(f)
140 else:
141 content_type = self._guess_content_type(remote_path)
142
143 # Delegate to parent with inferred or explicit content_type
144 return super()._upload_file(remote_path, f, attributes, content_type)