o
    0j                     @   s   d Z ddlmZ ddlZddlmZ ddlmZ ddlm	Z	m
Z
 ddlmZ ddlmZ eejd	d
ZG dd de	ZdefddZddededee fddZdedefddZdedefddZdS )uo   
本文件实现了对bos的封装, 首先安装 bce-python-sdk

Authors: suoyi@baidu.com
Date:    2024/01/03

    )logN)BceClientConfiguration)BceCredentials)	BosClientBceClientError)utils)ListZAISTUDIO_BOS_RETRY_TIMES
   c                       sH   e Zd ZdZ	d
ddZ									d fdd	Zdd	 Z  ZS )MyBosClientu>   
    重写BosClient的_upload方法，增加重试功能
    Nc                 C   s  |	  rtd| d d S d}ttD ]S}z&| j||||||||
|d	}|||jjd td| d d}W  n+ t	yh } zt
d| d t| t
d	|d
  d W Y d }~qd }~ww |s|	  t
d| d td| dd S )Nz%upload task canceled with partNumber=!F)progress_callbacktraffic_limit)Z
partNumberZeTagz$upload task success with partNumber=Tz#upload task failed with partNumber=u	   重试第   u   次u   !已取消上传)Z	is_cancelr   debugrangeRETRY_TIMESZupload_part_from_fileappendmetadataetag	Exceptionerrorcancelr   )selfbucket_nameZ
object_keyZ	upload_idZpart_numberZ	part_size	file_nameoffsetZ	part_listuploadTaskHandler   r   successiresponsee r!   [/var/www/html/Deteccion_Ine/venv/lib/python3.10/site-packages/aistudio_sdk/utils/bos_sdk.py_upload_task   s:   

"zMyBosClient._upload_task   c                    s"   t  |||||||||	|
||S )u[   调用原始的 put_super_obejct_from_file，但这里会使用上面定义的 _upload_task)superput_super_obejct_from_file)r   r   keyr   
chunk_size
thread_numr   content_typeZstorage_classZuser_headersr   r   config	__class__r!   r"   r&   <   s   
z&MyBosClient.put_super_obejct_from_filec                 C   s   dS )u-   需要覆盖父类的方法，否则会报错Zbosr!   )r   r!   r!   r"   _compute_service_idL   s   zMyBosClient._compute_service_id)NN)	r$   NNNNNNNN)__name__
__module____qualname____doc__r#   r&   r.   __classcell__r!   r!   r,   r"   r
      s    
r
   returnc                 C   s   t tt||| |d}|S )u   
    获取sts client
    )credentialsZendpointZsecurity_token)r
   r   r   )Zbos_hostZsts_akZsts_skZsession_token
bos_clientr!   r!   r"   
sts_clientP   s   r7    r6   bucketfilesc                 C   s&   |D ]}| j ||| |ddd qdS )u<   
    上传文件
    key_prefix: 上传文件的前缀
    r$   N)r(   r)   )r&   )r6   r9   r:   Z
key_prefixfiler!   r!   r"   upload_files\   s   r<   c                 C   s   |  ||t|S ),   
    上传文件
    key: 存储路径
    )Zput_object_from_filestr)r6   r9   r;   r'   r!   r!   r"   upload_filed   s   r?   c                 C   sX   t tjdd}tjdd}|rt |}| j||t|||dd}|s*td |S )r=   ZAISTUDIO_UPLOAD_CHUNK_SIZE_MBr$   ZAISTUDIO_UPLOAD_THREAD_NUMN)r(   r)   r   u  upload file failed: 已经取消或者上传失败，如果上传失败，请配置环境变量 AISTUDIO_UPLOAD_CHUNK_SIZE_MB (int类型，默认为5，单位MB)，减小分块大小后重试，例如：export AISTUDIO_UPLOAD_CHUNK_SIZE_MB=3 后重新执行如果带宽过小，需要配置环境变量 AISTUDIO_UPLOAD_THREAD_NUM 减少线程数，防止部分分块上传超时，例如：export AISTUDIO_UPLOAD_THREAD_NUM=1 后重新执行)intosenvirongetr&   r>   r   r   )r6   r9   r;   r'   r(   r)   resr!   r!   r"   upload_super_filem   s   
rE   )r8   )r2   Zaistudio_sdkr   rA   Z!baidubce.bce_client_configurationr   Zbaidubce.auth.bce_credentialsr   Z baidubce.services.bos.bos_clientr   r   Zbaidubcer   typingr   r@   rB   rC   r   r
   r7   r>   r<   r?   rE   r!   r!   r!   r"   <module>   s   8	