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