oss2/select_response.py

import platform
import struct
import requests

from .compat import to_bytes
from .exceptions import RequestError
from .exceptions import SelectOperationFailed
from .exceptions import SelectOperationClientError
from .exceptions import InconsistentError
from . import utils
import logging

logger = logging.getLogger(__name__)
"""
The adapter class for Select object's response.
The response consists of frames. Each frame has the following format:

Type  |   Payload Length |  Header Checksum | Payload | Payload Checksum

|<4-->|  <--4 bytes------><---4 bytes-------><-n/a-----><--4 bytes--------->
And we have three kind of frames.
Data Frame:
Type:8388609
Payload:   Offset    |    Data
           <-8 bytes>

Continuous Frame
Type:8388612
Payload: Offset  (8-bytes)

End Frame
Type:8388613
Payload: Offset | total scanned bytes | http status code | error message
    <-- 8bytes--><-----8 bytes--------><---4 bytes-------><---variabe--->

"""
class SelectResponseAdapter(object):

    _CHUNK_SIZE = 8 * 1024
    _CONTINIOUS_FRAME_TYPE=8388612
    _DATA_FRAME_TYPE = 8388609
    _END_FRAME_TYPE = 8388613
    _META_END_FRAME_TYPE = 8388614
    _JSON_META_END_FRAME_TYPE = 8388615
    _FRAMES_FOR_PROGRESS_UPDATE = 10

    def __init__(self, response, progress_callback = None, content_length = None, enable_crc = False):
       self.response = response
       self.frame_off_set = 0
       self.frame_length = 0
       self.frame_data = b''
       self.check_sum_flag = 0
       self.file_offset = 0
       self.finished = 0
       self.raw_buffer = b''
       self.raw_buffer_offset = 0
       #self.resp_content_iter = response.__iter__()
       self.callback = progress_callback
       self.frames_since_last_progress_report = 0
       self.content_length = content_length
       self.resp_content_iter = response.__iter__()
       self.enable_crc = enable_crc
       self.payload = b''
       self.output_raw_data = response.headers.get("x-oss-select-output-raw", '') == "true"
       self.request_id = response.headers.get("x-oss-request-id",'')
       self.splits = 0
       self.rows = 0
       self.columns = 0

    def read(self):
        if self.finished:
            return b''
        
        content=b''
        for data in self:
            content += data
        
        return content
    
    def __iter__(self):
        return self

    def __next__(self):
        return self.next()
    
    def next(self):
        if self.output_raw_data == True:
             data = next(self.resp_content_iter) 
             if len(data) != 0:
                 return data
             else: raise StopIteration

        while self.finished == 0:
            if self.frame_off_set < self.frame_length:
                data = self.frame_data[self.frame_off_set : self.frame_length]
                self.frame_length = self.frame_off_set = 0
                return data
            else:
                self.read_next_frame()
                self.frames_since_last_progress_report += 1
                if (self.frames_since_last_progress_report >= SelectResponseAdapter._FRAMES_FOR_PROGRESS_UPDATE and self.callback is not None):
                    self.callback(self.file_offset, self.content_length)
                    self.frames_since_last_progress_report = 0
        
        raise StopIteration

    def read_raw(self, amt):
        ret = b''
        read_count = 0
        while amt > 0 and self.finished == 0:
            size = len(self.raw_buffer)
            if size == 0:
                self.raw_buffer = next(self.resp_content_iter)
                self.raw_buffer_offset = 0
                size = len(self.raw_buffer)
                if size == 0:
                    break

            if size - self.raw_buffer_offset >= amt:
                data = self.raw_buffer[self.raw_buffer_offset:self.raw_buffer_offset + amt]
                data_size = len(data)
                self.raw_buffer_offset += data_size
                ret += data
                read_count += data_size
                amt -= data_size
            else:
                data = self.raw_buffer[self.raw_buffer_offset:]
                data_len = len(data)
                ret += data
                read_count += data_len
                amt -= data_len 
                self.raw_buffer = b''
        
        return ret

    def read_next_frame(self):
        frame_type = bytearray(self.read_raw(4))
        payload_length = bytearray(self.read_raw(4))
        utils.change_endianness_if_needed(payload_length) # convert to little endian
        payload_length_val = struct.unpack("I", bytes(payload_length))[0]
        header_checksum = bytearray(self.read_raw(4))

        frame_type[0] = 0 #mask the version bit
        utils.change_endianness_if_needed(frame_type) # convert to little endian
        frame_type_val = struct.unpack("I", bytes(frame_type))[0]
        if (frame_type_val != SelectResponseAdapter._DATA_FRAME_TYPE and
            frame_type_val != SelectResponseAdapter._CONTINIOUS_FRAME_TYPE and
            frame_type_val != SelectResponseAdapter._END_FRAME_TYPE and
            frame_type_val != SelectResponseAdapter._META_END_FRAME_TYPE and
            frame_type_val != SelectResponseAdapter._JSON_META_END_FRAME_TYPE):
                logger.warning("Unexpected frame type: {0}. RequestId:{1}. This could be due to the old version of client.".format(frame_type_val, self.request_id))
                raise SelectOperationClientError(self.request_id, "Unexpected frame type:" + str(frame_type_val))

        self.payload = self.read_raw(payload_length_val)
        file_offset_bytes = bytearray(self.payload[0:8])
        utils.change_endianness_if_needed(file_offset_bytes)
        self.file_offset = struct.unpack("Q", bytes(file_offset_bytes))[0]
        if frame_type_val == SelectResponseAdapter._DATA_FRAME_TYPE:
            self.frame_length = payload_length_val - 8
            self.frame_off_set = 0
            self.check_sum_flag=1
            self.frame_data = self.payload[8:]
            checksum = bytearray(self.read_raw(4))  #read checksum crc32
            utils.change_endianness_if_needed(checksum)
            checksum_val = struct.unpack("I", bytes(checksum))[0]
            if self.enable_crc:
                crc32 = utils.Crc32()
                crc32.update(self.payload)
                checksum_calc = crc32.crc
                if checksum_val != checksum_calc:
                    logger.warning("Incorrect checksum: Actual {0} and calculated {1}. RequestId:{2}".format(checksum_val, checksum_calc, self.request_id))
                    raise InconsistentError("Incorrect checksum: Actual" + str(checksum_val) + ". Calculated:" + str(checksum_calc), self.request_id)
            
        elif frame_type_val == SelectResponseAdapter._CONTINIOUS_FRAME_TYPE:
            self.frame_length = self.frame_off_set = 0
            self.check_sum_flag=1
            self.read_raw(4)
        elif frame_type_val == SelectResponseAdapter._END_FRAME_TYPE:
            self.frame_off_set = 0
            scanned_size_bytes = bytearray(self.payload[8:16])
            status_bytes = bytearray(self.payload[16:20])
            utils.change_endianness_if_needed(status_bytes)
            status = struct.unpack("I", bytes(status_bytes))[0]
            error_msg_size = payload_length_val - 20
            error_msg=b''
            error_code = b''
            if error_msg_size > 0:
                error_msg = self.payload[20:error_msg_size + 20]
                error_code_index = error_msg.find(b'.')
                if error_code_index >= 0 and error_code_index < error_msg_size - 1:
                    error_code = error_msg[0:error_code_index]
                    error_msg = error_msg[error_code_index + 1:]

            if status // 100 != 2:
                raise SelectOperationFailed(status, error_code, error_msg)
            self.frame_length = 0
            if self.callback is not None:
                self.callback(self.file_offset, self.content_length)
            self.read_raw(4) # read the payload checksum
            self.frame_length = 0
            self.finished = 1
        elif frame_type_val == SelectResponseAdapter._META_END_FRAME_TYPE or frame_type_val == SelectResponseAdapter._JSON_META_END_FRAME_TYPE:
            self.frame_off_set = 0
            scanned_size_bytes = bytearray(self.payload[8:16])
            status_bytes = bytearray(self.payload[16:20])
            utils.change_endianness_if_needed(status_bytes)
            status = struct.unpack("I", bytes(status_bytes))[0]
            splits_bytes = bytearray(self.payload[20:24])
            utils.change_endianness_if_needed(splits_bytes)
            self.splits = struct.unpack("I", bytes(splits_bytes))[0]
            lines_bytes = bytearray(self.payload[24:32])
            utils.change_endianness_if_needed(lines_bytes)
            self.rows =  struct.unpack("Q", bytes(lines_bytes))[0]

            error_index = 36
            if frame_type_val == SelectResponseAdapter._META_END_FRAME_TYPE:
                column_bytes = bytearray(self.payload[32:36])
                utils.change_endianness_if_needed(column_bytes)
                self.columns = struct.unpack("I", bytes(column_bytes))[0]
            else:
                error_index = 32
            
            error_size = payload_length_val - error_index
            error_msg = b''
            error_code = b''
            if (error_size > 0):
                error_msg = self.payload[error_index:error_index + error_size]
                error_code_index = error_msg.find(b'.')
                if error_code_index >= 0 and error_code_index < error_size - 1:
                    error_code = error_msg[0:error_code_index]
                    error_msg = error_msg[error_code_index + 1:]

            self.read_raw(4) # read the payload checksum
            self.final_status = status
            self.frame_length = 0
            self.finished = 1
            if (status / 100 != 2):
                raise SelectOperationFailed(status, error_code, error_msg)