o
    RDi2                     @  s|  d dl mZ d dlZd dlZd dlZd dlZd dlZd dlZd dlZd dl	m
Z
 d dlmZ d dlmZ d dlmZ d dlmZ d dlmZmZmZmZmZ d	d
lmZ d	dlmZ d	dlmZ d	dlmZm Z m!Z!m"Z"m#Z#m$Z$m%Z%m&Z&m'Z'm(Z(m)Z)m*Z*m+Z+ d	dl,m-Z- d	dl.m/Z/m0Z0m1Z1m2Z2m3Z3m4Z4m5Z5m6Z6 d	dl7m8Z8m9Z9m:Z:m;Z;m<Z< d	dl=m>Z>m?Z? d	dl@mAZA d	dlBmCZC d	dlDmEZE d	dlFmGZGmHZH erd	dlImJZJ d	dlKmLZL d	dl=mMZM e$e&e e#gZNd ZOeePZQdEddZRdEdd ZSeG d!d" d"ZTejUd#fdFd/d0ZVdGd3d4ZWdHd7d8ZXG d9d: d:ZYG d;d< d<eYZZG d=d> d>eYZ[G d?d@ d@Z\eG dAdB dBZ]G dCdD dDZ^dS )I    )annotationsN)ThreadPoolExecutor)	dataclass)partial)	getLogger)time)IOTYPE_CHECKINGAnyCallableTypeVar   )/_DEFAULT_VALUE_SERVER_DOP_CAP_FOR_FILE_TRANSFER)SnowflakeAzureRestClient)
IS_WINDOWS)AZURE_CHUNK_SIZEAZURE_FSCMD_TYPE_DOWNLOADCMD_TYPE_UPLOADGCS_FSLOCAL_FSS3_DEFAULT_CHUNK_SIZES3_FSS3_MAX_OBJECT_SIZES3_MAX_PARTSS3_MIN_PART_SIZEResultStatusmegabyte)SnowflakeConverterSnowSQL)ER_COMPRESSION_NOT_SUPPORTED ER_FAILED_TO_DOWNLOAD_FROM_STAGEER_FAILED_TO_UPLOAD_TO_STAGEER_FILE_NOT_EXISTS&ER_INTERNAL_NOT_MATCH_ENCRYPT_MATERIALER_INVALID_STAGE_FSER_INVALID_STAGE_LOCATIONER_LOCAL_PATH_NOT_DIRECTORY)DatabaseErrorErrorInternalErrorOperationalErrorProgrammingError)CompressionTypeslookup_by_mime_sub_type)SnowflakeGCSRestClient)SnowflakeLocalStorageClient)SnowflakeS3RestClient)SnowflakeFileEncryptionMaterialSnowflakeStorageClient)SnowflakeConnection)SnowflakeCursor)CompressionTypenamestrreturndict[str, Any]c                 C  s   | ddd d ddS )Ntexti   Fr6   typelength	precisionscalenullable r6   rA   rA   `/var/www/Datamplify/venv/lib/python3.10/site-packages/snowflake/connector/file_transfer_agent.pyresult_text_column_descH      rD   c                 C  s   | ddddddS )Nfixed   r   Fr;   rA   rB   rA   rA   rC   result_fixed_column_descS   rE   rH   c                   @  s  e Zd ZU dZded< ded< ded< dZded< dZd	ed
< dZded< dZded< dZ	ded< dZ
ded< dZded< dZded< dZded< dZded< dZded< dZded< dZded< dZded< dZded< dZd ed!< dZd ed"< dZded#< dZded$< dZded%< dZd&ed'< dZd(ed)< d*Zded+< dZd,ed-< dZd.ed/< dZded0< dZded1< d2Z ded3< dZ!d4ed5< dZ"d4ed6< dZ#ded7< dZ$ded8< dS )9SnowflakeFileMetazAClass to keep track of information necessary for file operations.r7   r6   src_file_namestage_location_typeNzResultStatus | Noneresult_statusz!SnowflakeFileTransferAgent | Nonesfagent(type[SnowflakeProgressPercentage] | Noneput_callbackput_azure_callbackzIO[str] | Noneput_callback_output_streamget_callbackget_azure_callbackget_callback_output_streamFboolshow_progress_bar   intmultipart_thresholdz
str | Nonepresigned_url	overwritesha256_digest
int | Noneupload_sizereal_src_file_namezException | Noneerror_details
last_errorno_sleeping_timegcs_file_header_digestgcs_file_header_content_lengthzdict[str, Any] | None#gcs_file_header_encryption_metadataz&SnowflakeFileEncryptionMaterial | Noneencryption_materialr   src_file_sizezCompressionType | Nonesrc_compression_typer5   dst_compression_typerequire_compressdst_file_namedst_file_sizeIO[bytes] | Noneintermediate_stream
src_streamskip_upload_on_content_matchlocal_location)%__name__
__module____qualname____doc____annotations__rL   rM   rO   rP   rQ   rR   rS   rT   rV   rY   rZ   r[   r\   r^   r_   r`   ra   rb   rc   rd   re   rf   rg   rh   ri   rj   rk   rm   ro   rp   rq   rr   rA   rA   rA   rC   rI   e   sJ   
 rI   T	file_name
start_timefloat
total_sizeprogressfloat | intoutput_stream	IO | NonerV   bool | NonerU   c                 C  s2  d}|t  }d}t | }|dkr|| nd}	t|tr t|}t|ts)d}d}|dk r1d}d}|dkrBd}d|d	d
|	dd}|sQ|rQd|d	d
|	dd}|rtt|| }
d|  d|ddd|
 d||
    d|d dd| 
}|| |  t	d|  d| d| d| d| 
 |dkS )N
    g        r   z#error: progress var must be float
z	Halt...
r   zDone (z.3fzs, z.2fzMB/s).
(zMB/s)zMB): [#-z] g      Y@z% z
filename: z, start_time: z, total_size: z, progress: z, show_progress_bar:       ?)
r   r   
isinstancerX   rz   roundwriteflushloggerdebug)rx   ry   r{   r|   r~   rV   
bar_lengthstatuselapsed_time
throughputblockr:   rA   rA   rC   _update_progress   sL   



r   	file_sizerX   c                 C  sl   | t krtd|  dt  dt| t tkr"tt| t tnt}|tkr4t	d| dt d |S )Nz
File size z exceeds the maximum file size z allowed in S3.zSetting chunksize to z instead of the default .)
r   
ValueErrormathceilr   r   maxr   r   r   )r   calculated_chunk_sizerA   rA   rC   _chunk_size_calculator   s   r   seen_so_farsizec                 C  s    | |ks|dkr
dS t | | S )Nr   r   )rz   )r   r   rA   rA   rC   percent   s    r   c                   @  s,   e Zd ZdZejdfdddZdddZdS )SnowflakeProgressPercentagez'Built-in Progress bar for PUT commands.Tfilenamer7   filesizeint | floatr~   r   rV   r   r8   Nonec                 C  sh   | d}|dk rt|}tj|d| | _|| _|| _t|| _	d| _
d| _t | _t | _d S )Nr   r   F)rfindlenospathbasename	_filename_output_stream_show_progress_barrz   _size_seen_so_far_doner   _start_time	threadingLock_lock)selfr   r   r~   rV   last_pound_charrA   rA   rC   __init__   s   

z$SnowflakeProgressPercentage.__init__bytes_amountrX   c                 C  s   t N)NotImplementedError)r   r   rA   rA   rC   __call__   s   z$SnowflakeProgressPercentage.__call__N
r   r7   r   r   r~   r   rV   r   r8   r   )r   rX   )rs   rt   ru   rv   sysstdoutr   r   rA   rA   rA   rC   r      s    r   c                      0   e Zd Zejdfd fddZdddZ  ZS )SnowflakeS3ProgressPercentageTr   r7   r   r   r~   r   rV   r   r8   r   c                      t  j||||d d S N)r~   rV   superr   r   r   r   r~   rV   	__class__rA   rC   r         
z&SnowflakeS3ProgressPercentage.__init__r   rX   c              	   C  s   | j = | jr0|  j|7  _t| j| j}| js8t| j| j| j|| j| j	d| _W d    d S W d    d S W d    d S 1 sCw   Y  d S r   
r   r   r   r   r   r   r   r   r   r   )r   r   
percentagerA   rA   rC   r      s$   
"z&SnowflakeS3ProgressPercentage.__call__r   )r   rX   r8   r   rs   rt   ru   r   r   r   r   __classcell__rA   rA   r   rC   r      
    r   c                      r   ) SnowflakeAzureProgressPercentageTr   r7   r   r   r~   r   rV   r   r8   r   c                   r   r   r   r   r   rA   rC   r     r   z)SnowflakeAzureProgressPercentage.__init__currentrX   c              	   C  s   | j 9 | jr,|| _t| j| j}| js4t| j| j| j|| j| j	d| _W d    d S W d    d S W d    d S 1 s?w   Y  d S r   r   )r   r   r   rA   rA   rC   r     s$   
"z)SnowflakeAzureProgressPercentage.__call__r   )r   rX   r8   r   r   rA   rA   r   rC   r     r   r   c                   @  s    e Zd Zdd	d
ZdddZdS )StorageCredentialcredentialsr9   
connectionr3   commandr7   r8   r   c                 C  s(   || _ t | _t | _|| _|| _d S r   )credsr   	timestampr   r   lockr   _command)r   r   r   r   rA   rA   rC   r   0  s
   

zStorageCredential.__init__c                 C  s   | j 7 || jk rtd 	 W d    d S td | j | j}|d d d | _t	 | _W d    d S 1 s=w   Y  d S )Nz:Omitting renewal of storage token, as it already happened.zRenewing expired storage token.data	stageInfor   )
r   r   r   r   r   cursor_execute_helperr   r   r   )r   cur_timestampretrA   rA   rC   update<  s   


"zStorageCredential.updateN)r   r9   r   r3   r   r7   r8   r   r8   r   )rs   rt   ru   r   r   rA   rA   rA   rC   r   /  s    
r   c                   @  s2   e Zd ZU dZded< dZded< dZded< dS )TransferMetadatar   rX   num_files_startednum_files_completedchunks_in_queueN)rs   rt   ru   r   rw   r   r   rA   rA   rA   rC   r   I  s   
 r   c                   @  s   e Zd ZdZddejddejdddddddddedfdAd"d#ZdBd$d%ZdCd(d)Z	dDd-d.Z
dBd/d0ZdEd1d2ZdFd5d6ZdBd7d8ZdBd9d:ZdBd;d<ZdBd=d>Zd?d@ ZdS )GSnowflakeFileTransferAgentzkSnowflake File Transfer Agent provides cloud provider independent implementation for putting/getting files.NTFr   r4   r   r7   r   r9   rO   rN   rP   rQ   IO[str]rR   rS   rT   rV   rU   raise_put_get_errorforce_put_overwriterq   rY   r]   source_from_streamrn   use_s3_regional_urliobound_tpe_limitunsafe_file_write,reraise_error_in_file_transfer_work_functionr8   r   c                 C  s   || _ || _|| _|| _|r|n|| _|| _|| _|r|n|| _|	| _d | _	|| _
|
| _|| _|| _|| _g | _g | _|p>d| _|| _d | _|| _|| _|| _|| _d S )NrW   )_cursorr   _ret_put_callback_put_azure_callback_put_callback_output_stream_get_callback_get_azure_callback_get_callback_output_stream_use_accelerate_endpoint_raise_put_get_errorr   _force_put_overwrite_skip_upload_on_content_match_source_from_stream_file_metadata_results_multipart_threshold_use_s3_regional_url_credentials_iobound_tpe_limit_unsafe_file_write+_snowflake_server_dop_cap_for_file_transfer-_reraise_error_in_file_transfer_work_function)r   r   r   r   rO   rP   rQ   rR   rS   rT   rV   r   r   rq   rY   r   r   r   r   *snowflake_server_dop_cap_for_file_transferr   rA   rA   rC   r   S  s8   


z#SnowflakeFileTransferAgent.__init__c                 C  sZ  |    |   | jtkr|   | jD ]}| |_q|   | jtkr0t	j
| js0t	| j | jtkrFt	j
| jd sFt	| jd  | jD ]2}| j|_| j|_| |_| jtkr{| j|_| j|_| j|_| j|_| j|_| j|_ | j!|_"| j#|_$qIt%&d| j' d | j(r| jst)*| j+j,| j+t-dt.d | /| j | j0D ]}|j1j2|_1qd S )Nlocationz
parallel=[]zBWhile getting file(s) there was an error: the file does not exist.msgerrno)3_parse_command_init_file_metadata_command_typer   _process_file_compression_typer   rM   _transfer_accelerate_configr   r   r   isdir_local_locationmakedirs_stage_location_typer   _stage_info
_overwriter[   r   rq   r   rO   r   rP   r   rQ   r   rR   r   rS   r   rT   r   rV   r   rY   r   r   	_parallelr   r(   errorhandler_wrapperr   r   r*   r"   transferr   rL   value)r   mresultrA   rA   rC   execute  sR   







z"SnowflakeFileTransferAgent.executemetaslist[SnowflakeFileMeta]c           	   	     s  t t|t j}td| jd ur#tdj t |j}t jjt	t	|}t	|td  t
 t
 fdd|D }t|}t jtkd d tdrfdndt| d(fdd	d)	fdd}d*	
fddd)	fdd
td}d+ fd$d%|D ]&}t||d& r||j |j n
||j |j  jd'7  _q" j|k r  d urd urj|k sW d    n1 sw   Y  |_d S ),NzDecided IO-bound TPE size: %dz#IO-bound TPE size is limited to: %dzChunk ThreadPoolExecutor size: c                   s   g | ]}  |qS rA   )_create_file_transfer_client).0r  )r   rA   rC   
<listcomp>  s    z7SnowflakeFileTransferAgent.transfer.<locals>.<listcomp>zGoing to %sload %d filesupdownr8   r   c                     s@      j d7  _    W d    d S 1 sw   Y  d S )Nr   )r   notifyrA   )cv_main_threadtransfer_metadatarA   rC   notify_file_completed  s   
"zBSnowflakeFileTransferAgent.transfer.<locals>.notify_file_completedsuccessrU   r  r
   	file_metarI   done_clientr2   c              	     s(  | s!t d|jj d r|  |  n|    d S |jjtj	kr-  d S t d|jj   Q j
d krQt d    j
d ksAt|jD ]'}t||d}rl|j||| n
|j|||  j
d7  _
qV   W d    d S 1 sw   Y  d S )NzFailed to prepare r   zFinished preparing file    z2Chunk queue busy, waiting in file done callback...)r&  chunk_idr   )r   r   metar6   finish_uploaddelete_client_datafinish_downloadrL   r   SKIPPEDr   waitrangenum_of_chunksr   submitupload_chunkdownload_chunkr   )r$  r  r%  r&  	_chunk_id	_callback)cv_chunk_processfunction_and_callback_wrapper	is_uploadmax_concurrencynetwork_tper#  transfer_done_cbr"  rA   rC   preprocess_done_cb  sR   




"z?SnowflakeFileTransferAgent.transfer.<locals>.preprocess_done_cbr(  rX   c                   s  t d| d|d  d|j d|jj d	    jd8  _   W d    n1 s/w   Y  |j | sR| jd7  _t d| d|jj d|  n| j	d7  _	t d	|jj d
|j	 d|j d|j  |j	|j |jkrr|
  |    n(|jt|d t d|jj d W d    d S W d    d S W d    d S 1 sw   Y  d S )Nz
Chunk(id: z) r   /z	 of file z reached callbackzChunk - failed to transfer for unexpected exception zChunk progress: z: completed: z	 failed: z total: r&  zsubmitting z to done_postprocess)r   r   r0  r)  r6   r   r   r   failed_transferssuccessful_transfersr*  r+  r1  r,  r   )r$  r  r%  r&  r(  )r6  r7  r8  r#  postprocess_done_cbpostprocess_tper"  rA   rC   r;  !  sL   &
$

"z=SnowflakeFileTransferAgent.transfer.<locals>.transfer_done_cbc                   st   t d|jj d |j" | s%| jd7  _t d|jj d|     W d    d S 1 s3w   Y  d S )NzFile z reached postprocess callbackr   r>  )r   r   r)  r6   r   r@  )r$  r  r%  r&  )r#  rA   rC   rB  Q  s   "z@SnowflakeFileTransferAgent.transfer.<locals>.postprocess_done_cb_TworkCallable[..., _T]r5  9Callable[[bool, _T | Exception, SnowflakeFileMeta], None]argskwargsc                   s6  zd| |i |f}W n@ t yK } z4tjdt|  dd ||_d|f}jrA |  W d   n1 s<w   Y  W Y d}~nd}~ww z|g ||R   W dS  t y } z6 |  W d   n1 ssw   Y  |d stjdt  dd W Y d}~dS W Y d}~dS d}~ww )a'  This wrapper makes sure that callbacks are called from the TPEs.

            If the main thread adds a callback to a future that has already been
            fulfilled then the callback is executed by the main thread. This can
            lead to unexpected slowdowns and behavior.
            TzAn exception was raised in )exc_infoFNr   )	Exceptionr   errorreprr`   r   r   )rE  r5  r%  rH  rI  r  e)callbackr!  exception_caught_in_callbackexception_caught_in_workr   rA   rC   r7  d  s:   

zJSnowflakeFileTransferAgent.transfer.<locals>.function_and_callback_wrapperr?  r   r   )
r$  rU   r  r
   r%  rI   r&  r2   r8   r   )r$  rU   r  r
   r%  rI   r&  r2   r(  rX   r8   r   )rE  rF  r5  rG  r%  rI   rH  r
   rI  r
   r8   r   )minr   r   	cpu_countr   r   r   r   r  r   r   	Conditionr   r	  r   r   r   r1  prepare_uploadr)  prepare_downloadr   r   r.  r   )	r   r  r   preprocess_tpefilesnum_total_filesr<  rD  file_clientrA   )rO  r6  r!  rP  rQ  r7  r8  r9  r:  r#  rB  rC  r   r;  r"  rC   r    sv   

804	



z#SnowflakeFileTransferAgent.transferr)  rI   r2   c              	   C  s   | j tkrt|| jdt | jdS | j tkr"t|| jt	| j| jdS | j t
kr:t|| j| jt|j| j| j| jdS | j tkrOt|| j| j| jj| j| jdS t| j  d)N   )r   )use_accelerate_endpointr   r   z is an unknown stage type)r  r   r/   r  r   r   r   r   r   r   r   r0   r   rg   r   r   r   r.   r   _connectionr   rK  )r   r)  rA   rA   rC   r    sF   



	z7SnowflakeFileTransferAgent._create_file_transfer_clientc                 C  s6   | j tkr| jr| | jd }| | _d S d S d S )Nr   )r  r   r   r  transfer_accelerate_configr   )r   clientrA   rA   rC   r    s   z6SnowflakeFileTransferAgent._transfer_accelerate_configc           	      C  s  | j jj}g }| jtkrt| dr| jD ]}|jd ur |jj}nd}|j	d ur,|j	j}nd}|j
d ur8t|j
nd}|tkrA|jnt|j}|tkrM|jnt|j}td| j d|j dt|j dtj dttj 
 | jr|rt| j j| j td| dtd ||j|j|||||j|g qtd	td
tdtdtdtdtdtdgt |dS t| dr| jD ]<}|tkr|jnt|j}|j
d urt|j
nd}| jr|rt| j j| j td| dt!d ||j||j|g qtdtdtdtdgt |dS )Nr   NONEr   zraise_put_get_error: z, z+While putting file(s) there was an error: 'zU', this might be caused by your access to the blob storage provider, or by Snowflake.r  sourcetargetsource_sizetarget_sizesource_compressiontarget_compressionr   message)rowtyperowsetz+While getting file(s) there was an error: 'filer   )"r   r]  converter_classr	  r   hasattrr   rh   r6   ri   r`   rM  r   rg   r7   rm   r   r   r   rL   r<   r   ERRORr(   r  r   r*   r!   appendrk   rD   rH   sortedr    )	r   rk  ri  r)  rh   ri   r`   rg   rm   rA   rA   rC   r    s   


















z!SnowflakeFileTransferAgent.result	locations	list[str]c                 C  s   g }|D ]B}| j tkrAtj|}tr+t|dkr+|d dkr+|d dkr+|dd  }tj|s7tj|}t		|}||7 }q|
| q|S )Nr'  r   r=  :r   )r	  r   r   r   
expanduserr   r   isabsabspathglobrn  )r   rp  canonical_locationsrx   rX  rA   rA   rC   _expand_filenamesb  s    


z,SnowflakeFileTransferAgent._expand_filenamesc                 C  s   g | _ | jd dd urL| jd d }t| j | jtkr2| j t|d |d |d d d S |D ]}|d urK| j t|d |d |d d q4d S d S )Nr   encryptionMaterialqueryStageMasterKeyqueryIdsmkId)query_stage_master_keyquery_idsmk_id)	_encryption_materialr   getr   r   r	  r   rn  r1   )r   	root_nodeelemrA   rA   rC   _init_encryption_materialy  s0   
z4SnowflakeFileTransferAgent._init_encryption_materialc              
   C  sl  d| j vrt| jj| jtdtd | j d }|d | _|   t	|
dts7t| jj| jtdtd |d | _|
dd| _| jpK|
d	d
| _|d d  | _| jtvrmt| jj| jtd| j td |d d | _|d | _t| jd | jj| j| _| j d 
d| _| jtkr| jr| j| _n	t| | j| _d|vp|d | _d|v r|d   | _!d S d| _!d S t| j| _i | _"t#|d t#| j$krt%| jD ]\}}t&'| | j$| | j"|< qn#t#| j$dkrt| jj| jt(dt#|d  dt#| j$ t)d t*j+,|d | _-t*j+.| j-s4t| jj| jt/d| j- t0d d S d S )Nr   z!Failed to parse server's responser  r   src_locationszFailed to parse the locationparallelr   r[   Fr   locationTypez(Destination location type is not valid: r  r   presignedUrlsautoCompresssourceCompressionr   r   zNThe number of downloading files doesn't match the encryption materials: files=z	, encmat=localLocationz#The local path is not a directory: )1r   r(   r  r   r   r'   r%   r	  r  r   r  list_src_locationsr  r   r  upperr  VALID_STORAGEr*   r$   _stage_locationr  r   r   r   _presigned_urlsr   r   
_src_filesrx  _auto_compresslower_source_compression _src_file_to_encryption_materialr   r  	enumerater   r   r)   r#   r   r   rs  r  r  r+   r&   )r   responseidxsrc_filerA   rA   rC   r    s   
















z)SnowflakeFileTransferAgent._parse_commandc           
      C  s  t d| j  | jtkr| js7d| jv r$d| jd v r$| jd d nd}t| jj	| jt
d| td | jrn| jttj| jd | jd | j| jdtj| jt| jdkra| jd nd d | jd d S | jD ]U}tj|st| jj	| jt
d| td ntj|rt| jj	| jt
d	| td t|}| jttj|||j| jt| jdkr| jd nd d
 qqd S | jtkrLt }t| jD ]^\}}|sq| |}| d}|dkr||d d  n|}d }| j!r	|t| j!k r	| j!| }tj|}|"|dd ||< | jt|||| j| j#||| j$v r0| j$| nd d qdd |% D }	|	rNt &d|	  d S d S d S )Nzcommand type: r   r  r   zFile doesn't exist: r  r   )r6   rJ   ro   rg   rK   rf   zNot a file but a directory: )r6   rJ   rg   rK   rf   r=  r   )r6   rJ   rk   rK   rr   rZ   rf   c                 S  s   g | ]
\}}|d kr|qS )r   rA   )r  kvrA   rA   rC   r  g  s    zBSnowflakeFileTransferAgent._init_file_metadata.<locals>.<listcomp>zfDownloading multiple files with the same name could cause failures. File names with multiple entries: )'r   r   r	  r   r  r   r(   r  r   r   r+   r"   r   r   rn  rI   r   r   r   seekSEEK_ENDr  r   r  existsr  statst_sizer   dictr  3_strip_stage_prefix_from_dst_file_name_for_downloadfindr  r  r  r  itemswarning)
r   rx   statinfobasename_countsr  rk   first_path_sepurlr   duplicate_basenamesrA   rA   rC   r    s   
	

	
	

#


'z.SnowflakeFileTransferAgent._init_file_metadatac           
   	   C  sN  d }| j dkr
d}n'| j dkrd}nt| j }|d u s|js/t| jj| jtd| td d}| j	D ]}|j
}d }|rt  t|\}}|d u rd }| jslt|d}	|	d}W d    n1 sfw   Y  n| jd}| jd	 |d
rd}n$|r|d d dkrd}n|r|dkrd}n|rtt|ddkrd}|d urtd| d|  t|}ntd|  |d ur|jst| jj| jtd| td n|}|d ur||_|jr||_d|_|j|_q4t| jj| jtd| td q4| j|_d |_| jr|jtd j |_td |_q4|j|_d |_q4d S )Nauto_detectTnoneFzFeature is not supported: r  rbr[  r   z.brbr   s   ORCorcs   PAR1parquet   i/(zstdzdetected the encoding z: file=z$no file encoding was detected: file=GZIP) r  r-   is_supportedr(   r  r   r   r+   r   r   rJ   	mimetypesinit
guess_typer   openreadr  endswithrX   binasciihexlifyr   r   rh   ri   rj   r6   rk   r  r,   file_extension)
r   !user_specified_source_compressionr  r  rx   current_file_compression_type_encodingtestfrA   rA   rC   r
  n  s   








z9SnowflakeFileTransferAgent._process_file_compression_typec                 C  s   |S )zStrips the stage prefix from dst_file_name for download.

        Note that this is no-op in most cases, and therefore we return as is.
        But for some workloads they will monkeypatch this method to add their
        stripping logic.
        rA   )r   rk   rA   rA   rC   r    s   zNSnowflakeFileTransferAgent._strip_stage_prefix_from_dst_file_name_for_download)(r   r4   r   r7   r   r9   rO   rN   rP   rN   rQ   r   rR   rN   rS   rN   rT   r   rV   rU   r   rU   r   rU   rq   rU   rY   r]   r   rn   r   rU   r   r]   r   rU   r   rU   r8   r   r   )r  r  r8   r   )r)  rI   r8   r2   )r8   r9   )rp  rq  r8   rq  )rs   rt   ru   rv   r   r   r   r   r  r  r  r  r  rx  r  r  r  r
  r  rA   rA   rA   rC   r   P  s@    
=
6 
u
'

|


h
tgr   )r6   r7   r8   r9   )rx   r7   ry   rz   r{   rz   r|   r}   r~   r   rV   r   r8   rU   )r   rX   r8   rX   )r   rX   r   rz   r8   rz   )_
__future__r   r  rv  r   r  r   r   r   concurrent.futures.threadr   dataclassesr   	functoolsr   loggingr   r   typingr   r	   r
   r   r   _utilsr   azure_storage_clientr   compatr   	constantsr   r   r   r   r   r   r   r   r   r   r   r   r   converter_snowsqlr   	errorcoder   r    r!   r"   r#   r$   r%   r&   errorsr'   r(   r)   r*   r+   file_compression_typer,   r-   gcs_storage_clientr.   local_storage_clientr/   s3_storage_clientr0   storage_clientr1   r2   r   r3   r   r4   r5   r  INJECT_WAIT_IN_PUTrs   r   rD   rH   rI   r   r   r   r   r   r   r   r   r   r   rA   rA   rA   rC   <module>   sb   <(


1
+
