U
    d/eM8                     @   s0  d 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 dd	lmZmZmZ d
ZdZdZdZdZd
ZdZdZdZdZG dd dZG dd deZG dd deZG dd deZG dd dZ G dd dZ!G dd  d Z"d!d" Z#d#d$ Z$d%d& Z%d'd( Z&G d)d* d*Z'G d+d, d,Z(d-S ).z,Request/response of SelectObjectContent API.    )absolute_import)ABCMetacrc32)BytesIO)ElementTree   )MinioException)Element
SubElementfindtextNONEZGZIPZBZIP2ZUSEZIGNOREZDOCUMENTLINESZALWAYSZASNEEDEDc                   @   s$   e Zd ZdZeZdd Zdd ZdS )InputSerializationzInput serialization.c                 C   s:   |d k	r0|t ttfkr0tdt  dt dt || _d S )Nzcompression type must be ,  or )COMPRESSION_TYPE_NONECOMPRESSION_TYPE_GZIPCOMPRESSION_TYPE_BZIP2
ValueError_compression_type)selfcompression_type r   0/tmp/pip-unpacked-wheel-xery97c7/minio/select.py__init__1   s    zInputSerialization.__init__c                 C   s   | j dk	rt|d| j  |S )Convert to XML.NZCompressionType)r   r   r   elementr   r   r   toxml@   s    
zInputSerialization.toxmlN)__name__
__module____qualname____doc__r   __metaclass__r   r   r   r   r   r   r   ,   s   r   c                       s.   e Zd ZdZd fdd	Z fddZ  ZS )CSVInputSerializationzCSV input serialization.Nc	           	         sj   t  | || _|| _|| _|d k	rN|tttfkrNtdt dt dt || _	|| _
|| _|| _d S )Nzfile header info must be r   r   )superr   _allow_quoted_record_delimiter	_comments_field_delimiterFILE_HEADER_INFO_USEFILE_HEADER_INFO_IGNOREFILE_HEADER_INFO_NONEr   _file_header_info_quote_character_quote_escape_character_record_delimiter)	r   r   Zallow_quoted_record_delimitercommentsfield_delimiterZfile_header_infoquote_characterquote_escape_characterrecord_delimiter	__class__r   r   r   J   s&    zCSVInputSerialization.__init__c                    s   t  | t|d}| jdk	r.t|d| j | jdk	rFt|d| j | jdk	r^t|d| j | jdk	rvt|d| j | jdk	rt|d| j | jdk	rt|d| j | j	dk	rt|d	| j	 dS )
r   CSVNZAllowQuotedRecordDelimiterZCommentsFieldDelimiterZFileHeaderInfoQuoteCharacterQuoteEscapeCharacterRecordDelimiter)
r&   r   r   r'   r(   r)   r-   r.   r/   r0   r   r6   r   r   r   d   s0    







zCSVInputSerialization.toxml)NNNNNNNNr    r!   r"   r#   r   r   __classcell__r   r   r6   r   r%   G   s              r%   c                       s.   e Zd ZdZd fdd	Z fddZ  ZS )JSONInputSerializationzJSON input serialization.Nc                    s>   t  | |d k	r4|ttfkr4tdt dt || _d S )Nzjson type must be r   )r&   r   JSON_TYPE_DOCUMENTJSON_TYPE_LINESr   
_json_type)r   r   Z	json_typer6   r   r   r      s    
zJSONInputSerialization.__init__c                    s2   t  | t|d}| jdk	r.t|d| j dS )r   JSONNType)r&   r   r   rB   r   r6   r   r   r      s    

zJSONInputSerialization.toxml)NNr=   r   r   r6   r   r?      s   r?   c                       s,   e Zd ZdZ fddZ fddZ  ZS )ParquetInputSerializationzParquet input serialization.c                    s   t  d  d S N)r&   r   r   r6   r   r   r      s    z"ParquetInputSerialization.__init__c                    s   t  | t|dS )r   ZParquet)r&   r   r   r   r6   r   r   r      s    zParquetInputSerialization.toxmlr=   r   r   r6   r   rE      s   rE   c                   @   s"   e Zd ZdZdddZdd ZdS )CSVOutputSerializationzCSV output serialization.Nc                 C   sJ   || _ || _|| _|d k	r:|ttfkr:tdt dt || _|| _d S )Nzquote fields must be r   )r)   r.   r/   QUOTE_FIELDS_ALWAYSQUOTE_FIELDS_ASNEEDEDr   _quote_fieldsr0   )r   r2   r3   r4   Zquote_fieldsr5   r   r   r   r      s     zCSVOutputSerialization.__init__c                 C   s   t |d}| jdk	r"t |d| j | jdk	r:t |d| j | jdk	rRt |d| j | jdk	rjt |d| j | jdk	rt |d| j dS )r   r8   Nr9   r:   r;   ZQuoteFieldsr<   )r   r)   r.   r/   rK   r0   r   r   r   r   r      s    





zCSVOutputSerialization.toxml)NNNNNr    r!   r"   r#   r   r   r   r   r   r   rH      s        
rH   c                   @   s"   e Zd ZdZdddZdd ZdS )JSONOutputSerializationzJSON output serialization.Nc                 C   s
   || _ d S rF   )r0   )r   r5   r   r   r   r      s    z JSONOutputSerialization.__init__c                 C   s&   t |d}| jdk	r"t |d| j dS )r   rC   Nr<   )r   r0   r   r   r   r   r      s    

zJSONOutputSerialization.toxml)NrL   r   r   r   r   rM      s   
rM   c                   @   s"   e Zd ZdZdddZdd ZdS )	SelectRequestzSelect object content request.FNc                 C   sV   || _ t|tttfstd|| _t|ttfs:td|| _	|| _
|| _|| _d S )Nzkinput serialization must be CSVInputSerialization, JSONInputSerialization or ParquetInputSerialization typezSoutput serialization must be CSVOutputSerialization or JSONOutputSerialization type)_expression
isinstancer%   r?   rE   r   _input_serializationrH   rM   _output_serialization_request_progress_scan_start_range_scan_end_range)r   Z
expressionZinput_serializationZoutput_serializationZrequest_progressZscan_start_rangeZscan_end_ranger   r   r   r      s.    zSelectRequest.__init__c                 C   s   t d}t|d| j t|dd | jt|d | jt|d | jr^tt|ddd	 | jsj| jrt|d
}| jrt|d| j | jrt|d| j |S )r   ZSelectObjectContentRequestZ
ExpressionZExpressionTypeZSQLr   ZOutputSerializationZRequestProgressZEnabledtrueZ	ScanRangeZStartEnd)	r
   r   rO   rQ   r   rR   rS   rT   rU   )r   r   tagr   r   r   r      s,      
zSelectRequest.toxml)FNNrL   r   r   r   r   rN      s       
rN   c                 C   s"   |  |}t||krtd|S )z8Wrapper to RawIOBase.read() to error out on short reads.zinsufficient data)readlenIOError)readersizedatar   r   r   _read  s    
r_   c                 C   s   t j| ddS )z$Convert byte data to big-endian int.big)	byteorder)int
from_bytesr^   r   r   r   _int  s    re   c                 C   s   t | d@ S )zWrapper to binascii.crc32().l    r   rd   r   r   r   _crc32  s    rf   c                 C   sn   t | }i }|d}|sqjt|t|}tt|ddkrDtdt|tt|d}| || < q|S )zDecode header data.r      zheader value type is not 7   )r   rY   r_   re   r[   decode)r^   r\   headerslengthnamevaluer   r   r   _decode_header#  s    
rn   c                   @   s<   e Zd ZdZdd Zedd Zedd Zedd	 Zd
S )StatszProgress/Stats information.c                 C   s6   t | }t|d| _t|d| _t|d| _d S )NZBytesScannedZBytesProcessedZBytesReturned)ET
fromstringri   r   _bytes_scanned_bytes_processed_bytes_returned)r   r^   r   r   r   r   r   6  s    zStats.__init__c                 C   s   | j S )zGet bytes scanned.)rr   rG   r   r   r   bytes_scanned<  s    zStats.bytes_scannedc                 C   s   | j S )zGet bytes processed.)rs   rG   r   r   r   bytes_processedA  s    zStats.bytes_processedc                 C   s   | j S )zGet bytes returned.)rt   rG   r   r   r   bytes_returnedF  s    zStats.bytes_returnedN)	r    r!   r"   r#   r   propertyru   rv   rw   r   r   r   r   ro   3  s   

ro   c                   @   sZ   e Zd ZdZdd Zdd Zdd Zdd	 Zd
d Zdd Z	dd Z
dd ZdddZdS )SelectObjectReaderzm
    BufferedIOBase compatible reader represents response data of
    Minio.select_object_content() API.
    c                 C   s   || _ d | _d | _d S rF   )	_response_stats_payload)r   responser   r   r   r   R  s    zSelectObjectReader.__init__c                 C   s   | S rF   r   rG   r   r   r   	__enter__W  s    zSelectObjectReader.__enter__c                 C   s   |   S rF   )close)r   exc_type	exc_valueexc_tracebackr   r   r   __exit__Z  s    zSelectObjectReader.__exit__c                 C   s   dS )zReturn this is readable.Tr   rG   r   r   r   readable]  s    zSelectObjectReader.readablec                 C   s   dS )zReturn this is not writeable.Fr   rG   r   r   r   	writeablea  s    zSelectObjectReader.writeablec                 C   s   | j   | j   dS )z-Close response and release network resources.N)rz   r   release_connrG   r   r   r   r   e  s    
zSelectObjectReader.closec                 C   s   | j S )zGet stats information.)r{   rG   r   r   r   statsj  s    zSelectObjectReader.statsc           
      C   s  | j  rdS t| j d}t| j d}t|t|krRtdt| dt| t|dd }t| j |d d d }tt| j d}t|| | |krtdt|| |  d| t|dd }t|d| }|dd	krt|d
 d|d |ddkr dS || d }|ddksF|dk rN|  S ||||  }	|ddkrt	|	| _
|  S |ddkr|	| _t|	S td|d dS )zRead and decode response.r         z prelude CRC mismatch; expected: z, got: Nz message CRC mismatch; expected: z:message-typeerrorz:error-codez: z:error-messagez:event-typerW      ZContr   )Progressro   ZRecordszunknown event-type )rz   isclosedr_   rf   re   r[   rn   getr	   ro   r{   r|   rZ   )
r   ZpreludeZprelude_crctotal_lengthr^   Zmessage_crcheader_lengthrj   Zpayload_lengthpayloadr   r   r   r_   n  sH    

zSelectObjectReader._read   c                 c   sV   |   dkrR| jr | j}|t| jk r4| jd| }| jt|d | _|V  qq dS )z
        Stream extracted payload from response data. Upon completion, caller
        should call self.close() to release network resources.
        r   N)r_   r|   rZ   )r   	num_bytesresultr   r   r   stream  s    zSelectObjectReader.streamN)r   )r    r!   r"   r#   r   r~   r   r   r   r   r   r_   r   r   r   r   r   ry   L  s   4ry   N))r#   
__future__r   abcr   binasciir   ior   Z	xml.etreer   rp   r   r	   xmlr
   r   r   r   r   r   r*   r+   r,   r@   rA   rI   rJ   r   r%   r?   rE   rH   rM   rN   r_   re   rf   rn   ro   ry   r   r   r   r   <module>   s<   9)9