1515# specific language governing permissions and limitations
1616# under the License.
1717
18+ import base64
19+ import hashlib
20+ import json
1821import logging
1922import os
2023import re
2124import subprocess
25+ import tempfile
2226import threading
27+ import time
2328from datetime import datetime , timezone
29+ from itertools import islice
2430from pathlib import PurePosixPath
25- from typing import Any , Dict , List , Optional
31+ from typing import Any , Dict , Iterable , List , Optional
2632from urllib .parse import splitport , urlparse
2733
2834import pyarrow
@@ -46,6 +52,9 @@ def _pyarrow_lt_7():
4652 return parse (pyarrow .__version__ ) < parse ("7.0.0" )
4753
4854
55+ _S3_DELETE_TIMEOUT_SECONDS = 3600
56+
57+
4958class LegacyOssDirectoryListingError (RuntimeError ):
5059 """Raised when legacy PyArrow OSS cannot enumerate a directory."""
5160
@@ -63,6 +72,7 @@ def __init__(self, path: str, catalog_options: Options):
6372 self ._get_s3_property ("endpoint" , S3Options .S3_ENDPOINT .key ())
6473 if self ._is_s3 else None
6574 )
75+ self ._s3_delete_client = None
6676 self ._oss_bucket = None
6777 _oss_impl = self .properties .get (OssOptions .OSS_IMPL )
6878 self ._use_jindo = False
@@ -96,6 +106,7 @@ def __init__(self, path: str, catalog_options: Options):
96106 raise ValueError (f"Unrecognized filesystem type in URI: { scheme } " )
97107
98108 def _set_pyarrow_version (self ):
109+ self ._pyarrow_gte_22 = parse (pyarrow .__version__ ) >= parse ("22.0.0" )
99110 self ._pyarrow_gte_8 = parse (pyarrow .__version__ ) >= parse ("8.0.0" )
100111 # force_virtual_addressing landed in PyArrow 16; below it the OSS bucket
101112 # goes into endpoint_override, so keys must omit it (init + path share
@@ -107,6 +118,9 @@ def _uses_s3_compatibility(self) -> bool:
107118 return (not self ._use_jindo
108119 and (self ._is_oss or bool (self ._s3_endpoint )))
109120
121+ def _uses_s3_delete_fallback (self ) -> bool :
122+ return self ._uses_s3_compatibility () and self ._pyarrow_gte_22
123+
110124 def _configure_s3_checksums (self ):
111125 if parse (pyarrow .__version__ ) < parse ("22.0.0" ):
112126 return
@@ -119,13 +133,15 @@ def __getstate__(self):
119133 # threading.Lock cannot be pickled; recreated in __setstate__.
120134 state .pop ("_legacy_bucket_lock" , None )
121135 state .pop ("logger" , None )
136+ state .pop ("_s3_delete_client" , None )
122137 # Recreate S3-compatible clients with the worker's AWS SDK settings.
123138 if self ._uses_s3_compatibility ():
124139 state .pop ("filesystem" , None )
125140 return state
126141
127142 def __setstate__ (self , state ):
128143 self .__dict__ .update (state )
144+ self ._s3_delete_client = None
129145 self .logger = logging .getLogger (__name__ )
130146 self ._set_pyarrow_version ()
131147 if "_is_s3" not in state :
@@ -300,6 +316,25 @@ def _initialize_s3_fs(self) -> FileSystem:
300316
301317 return pafs .S3FileSystem (** client_kwargs )
302318
319+ def _s3_connection_options (self ):
320+ return {
321+ "access_key" : self ._get_property (
322+ S3Options .S3_ACCESS_KEY_ID .key (),
323+ * self ._s3_key_variants ("access-key" , "access.key" )),
324+ "secret_key" : self ._get_property (
325+ S3Options .S3_ACCESS_KEY_SECRET .key (),
326+ * self ._s3_key_variants ("secret-key" , "secret.key" )),
327+ "session_token" : self ._get_property (
328+ S3Options .S3_SECURITY_TOKEN .key (),
329+ * self ._s3_key_variants (
330+ "session-token" , "session.token" ,
331+ "security-token" , "security.token" )),
332+ "region" : self ._get_s3_property ("region" , S3Options .S3_REGION .key ()),
333+ "path_style" : (self ._pyarrow_gte_16 and (
334+ self ._get_s3_boolean_property ("path-style-access" ) or
335+ self ._get_s3_boolean_property ("path.style.access" ))),
336+ }
337+
303338 def _initialize_hdfs_fs (self , scheme : str , netloc : Optional [str ]) -> FileSystem :
304339 if 'HADOOP_HOME' not in os .environ :
305340 raise RuntimeError ("HADOOP_HOME environment variable is not set." )
@@ -486,13 +521,38 @@ def exists_batch(self, paths: List[str]) -> Dict[str, bool]:
486521 }
487522
488523 def delete (self , path : str , recursive : bool = False ) -> bool :
524+ use_fallback = self ._uses_s3_delete_fallback ()
525+ if use_fallback and self ._is_oss :
526+ parsed = urlparse (path )
527+ if parsed .scheme == "oss" and parsed .netloc :
528+ bucket = self ._extract_oss_bucket (path )
529+ key = parsed .path .lstrip ("/" )
530+ if "@" in parsed .netloc :
531+ _ , _ , key = key .partition ("/" )
532+ path = "oss://" + bucket + "/" + key
489533 path_str = self .to_filesystem_path (path )
534+ if use_fallback :
535+ bucket , key = self ._split_s3_path (path_str )
536+ if not bucket or not key .strip ("/" ):
537+ raise OSError (f"Refusing to delete bucket root: { path } " )
490538 file_info = self ._get_file_info (path_str )
491539
492540 if file_info .type == pafs .FileType .NotFound :
493541 return False
494542
495543 if file_info .type == pafs .FileType .Directory :
544+ if use_fallback :
545+ if recursive :
546+ return self ._delete_s3_compatible_directory (path_str )
547+ selector = pafs .FileSelector (
548+ path_str , recursive = False , allow_not_found = True )
549+ if self .filesystem .get_file_info (selector ):
550+ raise OSError (f"Directory { path } is not empty" )
551+ bucket , key = self ._split_s3_path (path_str )
552+ if key :
553+ self ._get_s3_delete_client ().delete_object (
554+ Bucket = bucket , Key = key .rstrip ("/" ) + "/" )
555+ return True
496556 if not recursive :
497557 selector = pafs .FileSelector (path_str , recursive = False , allow_not_found = True )
498558 dir_contents = self .filesystem .get_file_info (selector )
@@ -507,6 +567,163 @@ def delete(self, path: str, recursive: bool = False) -> bool:
507567 self .filesystem .delete_file (path_str )
508568 return True
509569
570+ def _delete_s3_compatible_directory (self , path_str : str ) -> bool :
571+ client = self ._get_s3_delete_client ()
572+ bucket , key = self ._split_s3_path (path_str )
573+ prefix = key .rstrip ("/" )
574+ if prefix :
575+ prefix += "/"
576+ deadline = time .monotonic () + _S3_DELETE_TIMEOUT_SECONDS
577+ # Finish listing before deleting: mutating a paginated listing may skip keys.
578+ with tempfile .TemporaryFile (mode = "w+t" , encoding = "utf-8" ) as listed :
579+ for name in self ._list_s3_keys (client , bucket , prefix , deadline , path_str ):
580+ listed .write (json .dumps (name ) + "\n " )
581+ listed .seek (0 )
582+ self ._delete_s3_objects (
583+ client , bucket , (json .loads (line ) for line in listed ),
584+ deadline , path_str )
585+ if any (self ._list_s3_keys (client , bucket , prefix , deadline , path_str )):
586+ raise OSError (f"S3 directory { path_str } changed during deletion" )
587+ return True
588+
589+ def _list_s3_keys (self , client , bucket : str , prefix : str ,
590+ deadline : float , path_str : str ):
591+ token = None
592+ while True :
593+ self ._check_s3_delete_deadline (deadline , path_str )
594+ params = {"Bucket" : bucket , "Prefix" : prefix , "MaxKeys" : 1000 }
595+ if token is not None :
596+ params ["ContinuationToken" ] = token
597+ response = client .list_objects_v2 (** params )
598+ yield from self ._listed_s3_keys (response , bucket , prefix )
599+ if not response .get ("IsTruncated" ):
600+ return
601+ next_token = response .get ("NextContinuationToken" )
602+ if not next_token or next_token == token :
603+ raise OSError ("S3 listing did not advance" )
604+ token = next_token
605+
606+ @staticmethod
607+ def _listed_s3_keys (response , bucket : str , prefix : str ):
608+ if (response .get ("Name" , bucket ) != bucket
609+ or response .get ("Prefix" , prefix ) != prefix ):
610+ raise OSError ("S3 listing returned a different bucket or prefix" )
611+ keys = [item ["Key" ] for item in response .get ("Contents" , ())]
612+ if any (not key .startswith (prefix ) for key in keys ):
613+ raise OSError (f"S3 listing returned a key outside prefix { prefix } " )
614+ return keys
615+
616+ @staticmethod
617+ def _check_s3_delete_deadline (deadline : float , path_str : str ):
618+ if time .monotonic () >= deadline :
619+ raise TimeoutError (f"Timed out deleting S3 directory { path_str } " )
620+
621+ @staticmethod
622+ def _delete_s3_objects (
623+ client , bucket : str , keys : Iterable [str ],
624+ deadline : float , path_str : str ):
625+ keys = iter (keys )
626+ while True :
627+ batch = list (islice (keys , 1000 ))
628+ if not batch :
629+ return
630+ PyArrowFileIO ._check_s3_delete_deadline (deadline , path_str )
631+ ordinary = []
632+ for key in batch :
633+ if key .isprintable ():
634+ ordinary .append (key )
635+ else :
636+ # Non-printable keys may not round-trip through DeleteObjects XML.
637+ PyArrowFileIO ._check_s3_delete_deadline (deadline , path_str )
638+ client .delete_object (Bucket = bucket , Key = key )
639+ if not ordinary :
640+ continue
641+ PyArrowFileIO ._check_s3_delete_deadline (deadline , path_str )
642+ response = client .delete_objects (
643+ Bucket = bucket ,
644+ Delete = {"Objects" : [{"Key" : key } for key in ordinary ],
645+ "Quiet" : False })
646+ deleted = [item ["Key" ] for item in response .get ("Deleted" , ())]
647+ errors = response .get ("Errors" , ())
648+ if errors :
649+ examples = ", " .join (
650+ "{}: {!r}" .format (item .get ("Code" ), item .get ("Key" ))
651+ for item in errors [:3 ])
652+ raise OSError (
653+ f"S3 batch delete incomplete for { path_str } : { examples } " )
654+ if set (deleted ) != set (ordinary ) or len (deleted ) != len (ordinary ):
655+ raise OSError (f"S3 batch delete incomplete for { path_str } " )
656+
657+ @staticmethod
658+ def _split_s3_path (path_str : str ):
659+ bucket , _ , key = path_str .partition ("/" )
660+ return bucket , key
661+
662+ def _get_s3_delete_client (self ):
663+ if self ._s3_delete_client is not None :
664+ return self ._s3_delete_client
665+
666+ import boto3
667+ from botocore .config import Config
668+
669+ if self ._is_oss :
670+ endpoint = self .properties .get (OssOptions .OSS_ENDPOINT )
671+ access_key = self .properties .get (OssOptions .OSS_ACCESS_KEY_ID )
672+ secret_key = self .properties .get (OssOptions .OSS_ACCESS_KEY_SECRET )
673+ session_token = self .properties .get (OssOptions .OSS_SECURITY_TOKEN )
674+ region = self .properties .get (OssOptions .OSS_REGION )
675+ addressing_style = "virtual"
676+ else :
677+ endpoint = self ._s3_endpoint
678+ connection = self ._s3_connection_options ()
679+ access_key = connection ["access_key" ]
680+ secret_key = connection ["secret_key" ]
681+ session_token = connection ["session_token" ]
682+ region = connection ["region" ]
683+ addressing_style = "path" if connection ["path_style" ] else "virtual"
684+
685+ region = region or self .filesystem .region
686+ if endpoint and "://" not in endpoint :
687+ endpoint = "https://" + endpoint
688+ config_args = {
689+ "retries" : {"max_attempts" : 10 , "mode" : "standard" },
690+ "s3" : {"addressing_style" : addressing_style },
691+ "connect_timeout" : 60 ,
692+ "read_timeout" : 60 ,
693+ }
694+ try :
695+ config = Config (request_checksum_calculation = "when_required" ,
696+ ** config_args )
697+ uses_new_checksums = True
698+ except TypeError :
699+ config = Config (** config_args )
700+ uses_new_checksums = False
701+
702+ client = boto3 .session .Session ().client (
703+ "s3" ,
704+ endpoint_url = endpoint ,
705+ aws_access_key_id = access_key ,
706+ aws_secret_access_key = secret_key ,
707+ aws_session_token = session_token ,
708+ region_name = region ,
709+ config = config ,
710+ )
711+ if uses_new_checksums :
712+ # OSS requires Content-MD5; newer Botocore defaults to CRC32.
713+ def use_content_md5 (request , ** kwargs ):
714+ request .headers ["Content-MD5" ] = base64 .b64encode (
715+ hashlib .md5 (request .body , usedforsecurity = False ).digest ()
716+ ).decode ("ascii" )
717+ for name in list (request .headers ):
718+ if (name .lower ().startswith ("x-amz-checksum-" ) or
719+ name .lower () == "x-amz-sdk-checksum-algorithm" ):
720+ del request .headers [name ]
721+
722+ client .meta .events .register (
723+ "before-sign.s3.DeleteObjects" , use_content_md5 )
724+ self ._s3_delete_client = client
725+ return client
726+
510727 def mkdirs (self , path : str ) -> bool :
511728 path_str = self .to_filesystem_path (path )
512729 file_info = self ._get_file_info (path_str )
0 commit comments