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 glob
17import json
18import logging
19import os
20import shutil
21import tempfile
22from collections.abc import Callable, Iterator
23from datetime import datetime, timezone
24from enum import Enum
25from io import BytesIO, StringIO
26from typing import IO, Any, Optional, TypeVar, Union
27
28import xattr
29
30from ..telemetry import Telemetry
31from ..types import AWARE_DATETIME_MIN, ObjectMetadata, Range, SymlinkHandling
32from ..utils import (
33 create_attribute_filter_evaluator,
34 matches_attribute_filter_expression,
35 safe_makedirs,
36 validate_attributes,
37)
38from .base import BaseStorageProvider
39
40_T = TypeVar("_T")
41
42PROVIDER = "file"
43READ_CHUNK_SIZE = 8192
44
45logger = logging.getLogger(__name__)
46
47
48class _EntryType(Enum):
49 """
50 An enum representing the type of an entry in a directory.
51 """
52
53 FILE = 1
54 DIRECTORY = 2
55 DIRECTORY_TO_EXPLORE = 3
56 SYMLINK = 4
57
58
[docs]
59def atomic_write(source: Union[str, IO], destination: str, attributes: Optional[dict[str, str]] = None):
60 """
61 Writes the contents of a file to the specified destination path.
62
63 This function ensures that the file write operation is atomic, meaning the output file is either fully written or not modified at all.
64 This is achieved by writing to a temporary file first and then renaming it to the destination path.
65
66 :param source: The input file to read from. It can be a string representing the path to a file, or an open file-like object (IO).
67 :param destination: The path to the destination file where the contents should be written.
68 :param attributes: The attributes to set on the file.
69 """
70
71 with tempfile.NamedTemporaryFile(mode="wb", delete=False, dir=os.path.dirname(destination), prefix=".") as fp:
72 temp_file_path = fp.name
73 if isinstance(source, str):
74 with open(source, mode="rb") as src:
75 while chunk := src.read(READ_CHUNK_SIZE):
76 fp.write(chunk)
77 else:
78 while chunk := source.read(READ_CHUNK_SIZE):
79 fp.write(chunk)
80
81 # Set attributes on temp file if provided
82 validated_attributes = validate_attributes(attributes)
83 if validated_attributes:
84 try:
85 xattr.setxattr(temp_file_path, "user.json", json.dumps(validated_attributes).encode("utf-8"))
86 except OSError as e:
87 logger.debug(f"Failed to set extended attributes on temp file {temp_file_path}: {e}")
88
89 os.rename(src=temp_file_path, dst=destination)
90
91
[docs]
92class PosixFileStorageProvider(BaseStorageProvider):
93 """
94 A concrete implementation of the :py:class:`multistorageclient.types.StorageProvider` for interacting with POSIX file systems.
95 """
96
97 def __init__(
98 self,
99 base_path: str,
100 config_dict: Optional[dict[str, Any]] = None,
101 telemetry_provider: Optional[Callable[[], Telemetry]] = None,
102 **kwargs: Any,
103 ) -> None:
104 """
105 :param base_path: The root prefix path within the POSIX file system where all operations will be scoped.
106 :param config_dict: Resolved MSC config.
107 :param telemetry_provider: A function that provides a telemetry instance.
108 """
109
110 # Validate POSIX path
111 if base_path == "":
112 base_path = "/"
113
114 if not base_path.startswith("/"):
115 raise ValueError(f"The base_path {base_path} must be an absolute path.")
116
117 super().__init__(
118 base_path=base_path,
119 provider_name=PROVIDER,
120 config_dict=config_dict,
121 telemetry_provider=telemetry_provider,
122 )
123
124 def _translate_errors(
125 self,
126 func: Callable[[], _T],
127 operation: str,
128 path: str,
129 ) -> _T:
130 """
131 Translates errors like timeouts and client errors.
132
133 :param func: The function that performs the actual file operation.
134 :param operation: The type of operation being performed (e.g., "PUT", "GET", "DELETE").
135 :param path: The path to the object.
136
137 :return: The result of the file operation, typically the return value of the `func` callable.
138 """
139 try:
140 return func()
141 except FileNotFoundError:
142 raise
143 except Exception as error:
144 raise RuntimeError(f"Failed to {operation} object(s) at {path}, error: {error}") from error
145
146 def _put_object(
147 self,
148 path: str,
149 body: bytes,
150 if_match: Optional[str] = None,
151 if_none_match: Optional[str] = None,
152 attributes: Optional[dict[str, str]] = None,
153 ) -> int:
154 def _invoke_api() -> int:
155 safe_makedirs(os.path.dirname(path))
156 atomic_write(source=BytesIO(body), destination=path, attributes=attributes)
157 return len(body)
158
159 return self._translate_errors(_invoke_api, operation="PUT", path=path)
160
161 def _get_object(self, path: str, byte_range: Optional[Range] = None) -> bytes:
162 def _invoke_api() -> bytes:
163 if byte_range:
164 with open(path, "rb") as f:
165 f.seek(byte_range.offset)
166 return f.read(byte_range.size)
167 else:
168 with open(path, "rb") as f:
169 return f.read()
170
171 return self._translate_errors(_invoke_api, operation="GET", path=path)
172
173 def _copy_object(self, src_path: str, dest_path: str) -> int:
174 src_object = self._get_object_metadata(src_path)
175
176 def _invoke_api() -> int:
177 safe_makedirs(os.path.dirname(dest_path))
178 atomic_write(source=src_path, destination=dest_path, attributes=src_object.metadata)
179
180 return src_object.content_length
181
182 return self._translate_errors(_invoke_api, operation="COPY", path=src_path)
183
184 def _delete_object(self, path: str, if_match: Optional[str] = None) -> None:
185 def _invoke_api() -> None:
186 if os.path.exists(path) and os.path.isfile(path):
187 os.remove(path)
188
189 return self._translate_errors(_invoke_api, operation="DELETE", path=path)
190
191 def _make_symlink(self, path: str, target: str) -> None:
192 def _invoke_api() -> None:
193 safe_makedirs(os.path.dirname(path))
194 relative_target = ObjectMetadata.encode_symlink_target(path, target)
195 if os.path.lexists(path):
196 os.remove(path)
197 os.symlink(relative_target, path)
198
199 self._translate_errors(_invoke_api, operation="SYMLINK", path=path)
200
201 def _get_object_metadata(self, path: str, strict: bool = True) -> ObjectMetadata:
202 is_dir = os.path.isdir(path)
203 if is_dir:
204 path = self._append_delimiter(path)
205
206 def _invoke_api() -> ObjectMetadata:
207 metadata_dict = {}
208 try:
209 json_bytes = xattr.getxattr(path, "user.json")
210 metadata_dict = json.loads(json_bytes.decode("utf-8"))
211 except (OSError, IOError, KeyError, json.JSONDecodeError, AttributeError) as e:
212 logger.debug(f"Failed to read extended attributes from {path}: {e}")
213 pass
214
215 # ``os.readlink`` may return an absolute path; normalise to the
216 # parent-relative form used by every backend.
217 symlink_target: Optional[str] = None
218 if os.path.islink(path):
219 raw = os.readlink(path)
220 symlink_target = ObjectMetadata.encode_symlink_target(path, raw) if os.path.isabs(raw) else raw
221
222 return ObjectMetadata(
223 key=path,
224 type="directory" if is_dir else "file",
225 content_length=0 if is_dir else os.path.getsize(path),
226 last_modified=datetime.fromtimestamp(os.path.getmtime(path), tz=timezone.utc),
227 metadata=metadata_dict if metadata_dict else None,
228 symlink_target=symlink_target,
229 )
230
231 return self._translate_errors(_invoke_api, operation="HEAD", path=path)
232
233 def _list_objects(
234 self,
235 path: str,
236 start_after: Optional[str] = None,
237 end_at: Optional[str] = None,
238 include_directories: bool = False,
239 symlink_handling: SymlinkHandling = SymlinkHandling.FOLLOW,
240 ) -> Iterator[ObjectMetadata]:
241 start_after = os.path.relpath(start_after, self._base_path) if start_after else None
242 end_at = os.path.relpath(end_at, self._base_path) if end_at else None
243
244 def _invoke_api() -> Iterator[ObjectMetadata]:
245 if os.path.isfile(path):
246 yield ObjectMetadata(
247 key=os.path.relpath(path, self._base_path),
248 content_length=os.path.getsize(path),
249 last_modified=datetime.fromtimestamp(os.path.getmtime(path), tz=timezone.utc),
250 )
251 dir_path = path.rstrip("/") + "/"
252 if not os.path.isdir(dir_path):
253 return
254
255 yield from self._explore_directory(dir_path, start_after, end_at, include_directories, symlink_handling)
256
257 return self._translate_errors(_invoke_api, operation="LIST", path=path)
258
259 @property
260 def supports_parallel_listing(self) -> bool:
261 """
262 Whether this provider supports heap-based parallel recursive listing.
263
264 :return: ``True`` for POSIX file storage.
265 """
266 return True
267
268 def _shallow_list(self, path: str, symlink_handling: SymlinkHandling) -> tuple[list[str], list[ObjectMetadata]]:
269 """
270 Adapt POSIX relative listing keys to the full-key contract expected by the recursive listing heap.
271 """
272 prefixes: list[str] = []
273 objects: list[ObjectMetadata] = []
274
275 for item in self._list_objects(path, include_directories=True, symlink_handling=symlink_handling):
276 full_key = self._prepend_base_path(item.key)
277 if item.type == "directory" and item.symlink_target is None:
278 child_prefix = full_key + "/"
279 if child_prefix != path.rstrip("/") + "/":
280 prefixes.append(child_prefix)
281 else:
282 item.key = full_key
283 objects.append(item)
284
285 return prefixes, objects
286
287 def _explore_directory(
288 self,
289 dir_path: str,
290 start_after: Optional[str],
291 end_at: Optional[str],
292 include_directories: bool,
293 symlink_handling: SymlinkHandling = SymlinkHandling.FOLLOW,
294 ) -> Iterator[ObjectMetadata]:
295 """
296 Recursively explore a directory and yield objects in lexicographical order.
297
298 :param dir_path: The directory path to explore
299 :param start_after: The key to start after
300 :param end_at: The key to end at
301 :param include_directories: Whether to include directories in the result
302 :param symlink_handling: How to handle symbolic links during listing.
303 """
304 try:
305 dir_entries = os.listdir(dir_path)
306 dir_entries.sort()
307
308 entries: list[tuple[str, str, _EntryType]] = []
309 symlink_info: dict[str, tuple[str, str]] = {}
310
311 for entry in dir_entries:
312 full_path = os.path.join(dir_path, entry)
313 is_link = os.path.islink(full_path)
314
315 if is_link and symlink_handling == SymlinkHandling.SKIP:
316 continue
317
318 if is_link and symlink_handling == SymlinkHandling.PRESERVE:
319 relative_path = os.path.relpath(full_path, self._base_path)
320
321 real_target = os.path.realpath(full_path)
322 if not os.path.exists(real_target):
323 raise FileNotFoundError(
324 f"Broken symlink '{relative_path}' points to a missing target "
325 f"({full_path} -> {real_target}). Use symlink_handling=SKIP to ignore."
326 )
327
328 try:
329 target_key = os.path.relpath(real_target, self._base_path)
330 except ValueError:
331 target_key = None
332
333 if target_key is None or target_key.startswith(".."):
334 raise ValueError(
335 f"Symlink '{relative_path}' points outside the base directory "
336 f"({full_path} -> {real_target}). Use symlink_handling=FOLLOW "
337 f"to dereference or symlink_handling=SKIP to ignore."
338 )
339
340 raw_link = os.readlink(full_path)
341 immediate_target = (
342 raw_link
343 if os.path.isabs(raw_link)
344 else os.path.normpath(os.path.join(os.path.dirname(full_path), raw_link))
345 )
346 relative_target = ObjectMetadata.encode_symlink_target(full_path, immediate_target)
347
348 target_type = "directory" if os.path.isdir(full_path) else "file"
349
350 if (start_after is None or start_after < relative_path) and (
351 end_at is None or relative_path <= end_at
352 ):
353 entries.append((relative_path, full_path, _EntryType.SYMLINK))
354 symlink_info[relative_path] = (relative_target, target_type)
355 continue
356
357 if is_link and symlink_handling == SymlinkHandling.FOLLOW:
358 # Broken symlinks have no target to dereference: isfile/isdir both return
359 # False, so the entry would be silently dropped. Fail fast instead.
360 real_target = os.path.realpath(full_path)
361 if not os.path.exists(real_target):
362 relative_path = os.path.relpath(full_path, self._base_path)
363 raise FileNotFoundError(
364 f"Broken symlink '{relative_path}' points to a missing target "
365 f"({full_path} -> {real_target}). Use symlink_handling=SKIP to ignore."
366 )
367
368 relative_path = os.path.relpath(full_path, self._base_path)
369
370 if (start_after is None or start_after < relative_path) and (end_at is None or relative_path <= end_at):
371 if os.path.isfile(full_path):
372 entries.append((relative_path, full_path, _EntryType.FILE))
373 elif os.path.isdir(full_path):
374 if include_directories:
375 entries.append((relative_path, full_path, _EntryType.DIRECTORY))
376 else:
377 entries.append((relative_path, full_path, _EntryType.DIRECTORY_TO_EXPLORE))
378
379 # Sort keys must mirror the keys S3 ``list_objects_v2`` would return so POSIX
380 # listings stay in the same raw-UTF-8-byte order. For directories (expanded or
381 # returned as-is) the emitted keys live under ``<name>/``, so the trailing
382 # delimiter must be part of the sort key. Otherwise the bare name ``a`` sorts
383 # before sibling file ``a.txt`` (since ``""`` < ``".txt"``), but S3 orders the
384 # nested key ``a/b.txt`` *after* ``a.txt`` because ``.`` (0x2E) < ``/`` (0x2F).
385 def _sort_key(entry: tuple[str, str, _EntryType]) -> str:
386 relative, _, entry_type = entry
387 if entry_type in (_EntryType.DIRECTORY, _EntryType.DIRECTORY_TO_EXPLORE):
388 return relative + "/"
389 return relative
390
391 entries.sort(key=_sort_key)
392
393 for relative_path, full_path, entry_type in entries:
394 if entry_type == _EntryType.FILE:
395 yield ObjectMetadata(
396 key=relative_path,
397 content_length=os.path.getsize(full_path),
398 last_modified=datetime.fromtimestamp(os.path.getmtime(full_path), tz=timezone.utc),
399 )
400 elif entry_type == _EntryType.DIRECTORY:
401 yield ObjectMetadata(
402 key=relative_path,
403 content_length=0,
404 type="directory",
405 last_modified=AWARE_DATETIME_MIN,
406 )
407 elif entry_type == _EntryType.DIRECTORY_TO_EXPLORE:
408 yield from self._explore_directory(
409 full_path, start_after, end_at, include_directories, symlink_handling
410 )
411 elif entry_type == _EntryType.SYMLINK:
412 relative_target, target_type = symlink_info[relative_path]
413 yield ObjectMetadata(
414 key=relative_path,
415 content_length=0,
416 last_modified=datetime.fromtimestamp(os.path.getmtime(full_path), tz=timezone.utc),
417 type=target_type,
418 symlink_target=relative_target,
419 )
420
421 except FileNotFoundError:
422 raise
423 except (OSError, PermissionError) as e:
424 logger.warning(f"Failed to list contents of {dir_path}, caused by: {e}")
425 return
426
427 def _upload_file(self, remote_path: str, f: Union[str, IO], attributes: Optional[dict[str, str]] = None) -> int:
428 safe_makedirs(os.path.dirname(remote_path))
429
430 filesize: int = 0
431 if isinstance(f, str):
432 filesize = os.path.getsize(f)
433 elif isinstance(f, StringIO):
434 filesize = len(f.getvalue().encode("utf-8"))
435 else:
436 filesize = len(f.getvalue()) # type: ignore
437
438 def _invoke_api() -> int:
439 atomic_write(source=f, destination=remote_path, attributes=attributes)
440
441 return filesize
442
443 return self._translate_errors(_invoke_api, operation="PUT", path=remote_path)
444
445 def _download_file(self, remote_path: str, f: Union[str, IO], metadata: Optional[ObjectMetadata] = None) -> int:
446 filesize = metadata.content_length if metadata else os.path.getsize(remote_path)
447
448 if isinstance(f, str):
449
450 def _invoke_api() -> int:
451 if os.path.dirname(f):
452 safe_makedirs(os.path.dirname(f))
453 atomic_write(source=remote_path, destination=f)
454
455 return filesize
456
457 return self._translate_errors(_invoke_api, operation="GET", path=remote_path)
458 elif isinstance(f, StringIO):
459
460 def _invoke_api() -> int:
461 with open(remote_path, "r", encoding="utf-8") as src:
462 while chunk := src.read(READ_CHUNK_SIZE):
463 f.write(chunk)
464
465 return filesize
466
467 return self._translate_errors(_invoke_api, operation="GET", path=remote_path)
468 else:
469
470 def _invoke_api() -> int:
471 with open(remote_path, "rb") as src:
472 while chunk := src.read(READ_CHUNK_SIZE):
473 f.write(chunk)
474
475 return filesize
476
477 return self._translate_errors(_invoke_api, operation="GET", path=remote_path)
478
[docs]
479 def glob(self, pattern: str, attribute_filter_expression: Optional[str] = None) -> list[str]:
480 pattern = self._prepend_base_path(pattern)
481 keys = list(glob.glob(pattern, recursive=True))
482 if attribute_filter_expression:
483 filtered_keys = []
484 evaluator = create_attribute_filter_evaluator(attribute_filter_expression)
485 for key in keys:
486 obj_metadata = self._get_object_metadata(key)
487 if matches_attribute_filter_expression(obj_metadata, evaluator):
488 filtered_keys.append(key)
489 keys = filtered_keys
490 if self._base_path == "/":
491 return keys
492 else:
493 # NOTE: PosixStorageProvider does not have the concept of bucket and prefix.
494 # So we drop the base_path from it.
495 return [key.replace(self._base_path, "", 1).lstrip("/") for key in keys]
496
[docs]
497 def is_file(self, path: str) -> bool:
498 path = self._prepend_base_path(path)
499 return os.path.isfile(path)
500
[docs]
501 def rmtree(self, path: str) -> None:
502 path = self._prepend_base_path(path)
503 shutil.rmtree(path)