parse value, number, array and lz4 compression

Author: danielebriggi

requirements-dev.txt

@@ -6,3 +6,4 @@ pytest==7.0.1
 pytest-mock==3.6.1
 black==22.8.0
 python-dotenv==0.20.0
+lz4==3.1.10

requirements.txt

@@ -0,0 +1 @@
+lz4==3.1.10

src/sqlitecloud/driver.py

@@ -1,18 +1,23 @@
 import ssl
+from termios import CSTART
 from typing import Optional
+import lz4.block
 from sqlitecloud.types import (
     SQCLOUD_CMD,
     SQCLOUD_INTERNAL_ERRCODE,
     SQCloudConfig,
     SQCloudConnect,
     SQCloudException,
     SQCloudNumber,
+    SQCloudValue,
 )
 import socket
 
 
 class Driver:
-    def connect(self, hostname: str, port: int, config: SQCloudConfig) -> SQCloudConnect:
+    def connect(
+        self, hostname: str, port: int, config: SQCloudConfig
+    ) -> SQCloudConnect:
         """
         Connects to the SQLite Cloud server.
 
@@ -150,24 +155,28 @@ def _internal_socket_read(self, connection: SQCloudConnect) -> any:
                 data = connection.socket.recv(buffer_size)
                 if not data:
                     break
-                
+
                 # update buffers
                 data = data.decode()
                 buffer += data
                 nread += len(data)
 
                 c = buffer[0]
 
-                if c == SQCLOUD_CMD.INT or c == SQCLOUD_CMD.FLOAT or c == SQCLOUD_CMD.NULL:
-                    if buffer[nread-1] != ' ':
+                if (
+                    c == SQCLOUD_CMD.INT.value
+                    or c == SQCLOUD_CMD.FLOAT.value
+                    or c == SQCLOUD_CMD.NULL.value
+                ):
+                    if buffer[nread - 1] != " ":
                         continue
-                elif c == SQCLOUD_CMD.ROWSET_CHUNK:
-                    isEndOfChunk = buffer.endswith(SQCLOUD_CMD.CHUNKS_END)
+                elif c == SQCLOUD_CMD.ROWSET_CHUNK.value:
+                    isEndOfChunk = buffer.endswith(SQCLOUD_CMD.CHUNKS_END.value)
                     if not isEndOfChunk:
                         continue
                 else:
                     n: SQCloudNumber = self._internal_parse_number(buffer)
-                    can_be_zerolength = c == SQCLOUD_CMD.BLOB or c == SQCLOUD_CMD.STRING
+                    can_be_zerolength = c == SQCLOUD_CMD.BLOB.value or c == SQCLOUD_CMD.STRING.value
                     if n.value == 0 and not can_be_zerolength:
                         continue
                     if n.value + n.cstart != nread:
@@ -183,7 +192,7 @@ def _internal_socket_read(self, connection: SQCloudConnect) -> any:
             )
 
     def _internal_parse_number(self, buffer: str, index: int = 1) -> SQCloudNumber:
-        sqlite_number = SQCloudNumber()
+        sqcloud_number = SQCloudNumber()
         extvalue = 0
         isext = False
         blen = len(buffer)
@@ -193,24 +202,261 @@ def _internal_parse_number(self, buffer: str, index: int = 1) -> SQCloudNumber:
             c = buffer[i]
 
             # check for optional extended error code (ERRCODE:EXTERRCODE)
-            if c == ':':
+            if c == ":":
                 isext = True
                 continue
 
             # check for end of value
-            if c == ' ':
-                sqlite_number.cstart = i + 1
-                sqlite_number.extcode = extvalue
-                return sqlite_number
+            if c == " ":
+                sqcloud_number.cstart = i + 1
+                sqcloud_number.extcode = extvalue
+                return sqcloud_number
+
+            v = int(c) if c.isdigit() else 0
 
             # compute numeric value
             if isext:
-                extvalue = (extvalue * 10) + int(buffer[i])
+                extvalue = (extvalue * 10) + v
             else:
-                sqlite_number.value = (sqlite_number.value * 10) + int(buffer[i])
+                sqcloud_number.value = (sqcloud_number.value * 10) + v
+
+        return sqcloud_number
 
-        return 0
-    
     def _internal_parse_buffer(self, buffer: str, blen: int) -> any:
-        # TODO
-        return
+        # possible return values:
+        # True 	=> OK
+        # False 	=> error
+        # integer
+        # double
+        # string
+        # list
+        # object
+        # None
+
+        # check OK value
+        if buffer == "+2 OK":
+            return True
+
+        cmd = buffer[0]
+
+        # check for compressed result
+        if cmd == SQCLOUD_CMD.COMPRESSED.value:
+            # TODO: use exception
+            buffer = self._internal_uncompress_data(buffer, blen)
+            if buffer is None:
+                raise SQCloudException(
+                    f"An error occurred while decompressing the input buffer of len {len}.",
+                    -1,
+                )
+
+        # first character contains command type
+        if cmd in [
+            SQCLOUD_CMD.ZEROSTRING.value,
+            SQCLOUD_CMD.RECONNECT.value,
+            SQCLOUD_CMD.PUBSUB.value,
+            SQCLOUD_CMD.COMMAND.value,
+            SQCLOUD_CMD.STRING.value,
+            SQCLOUD_CMD.ARRAY.value,
+            SQCLOUD_CMD.BLOB.value,
+            SQCLOUD_CMD.JSON.value,
+        ]:
+            cstart = 0
+            sqlite_number = self._internal_parse_number(buffer, cstart)
+            len = sqlite_number.value
+            if len == 0:
+                return ""
+
+            if cmd == SQCLOUD_CMD.ZEROSTRING.value:
+                len -= 1
+            clone = buffer[cstart : cstart + len]
+
+            if cmd == SQCLOUD_CMD.COMMAND.value:
+                return self._internal_run_command(clone)
+            elif cmd == SQCLOUD_CMD.PUBSUB.value:
+                # TODO
+                return self._internal_setup_pubsub(clone)
+            elif cmd == SQCLOUD_CMD.RECONNECT.value:
+                return self._internal_reconnect(clone)
+            elif cmd == SQCLOUD_CMD.ARRAY.value:
+                return self._internal_parse_array(clone)
+
+            return clone
+
+        elif cmd == SQCLOUD_CMD.ERROR.value:
+            # -LEN ERRCODE:EXTCODE ERRMSG
+            sqlite_number = self._internal_parse_number(buffer)
+            len = sqlite_number.value
+            cstart = sqlite_number.cstart
+            clone = buffer[cstart:]
+
+            sqlite_number = self._internal_parse_number(clone, 0)
+            cstart2 = sqlite_number.cstart
+
+            errcode = sqlite_number.value
+            xerrcode = sqlite_number.extcode
+
+            len -= cstart2
+            errmsg = clone[cstart2:]
+
+            raise SQCloudException(errmsg, errcode, xerrcode)
+
+        elif cmd in [SQCLOUD_CMD.ROWSET.value, SQCLOUD_CMD.ROWSET_CHUNK.value]:
+            # CMD_ROWSET:          *LEN 0:VERSION ROWS COLS DATA
+            # CMD_ROWSET_CHUNK:    /LEN IDX:VERSION ROWS COLS DATA
+            # TODO: return a custom object
+            start = self._internal_parse_rowset_signature(
+                buffer, len, idx, version, nrows, ncols
+            )
+            if start < 0:
+                return False
+
+            # check for end-of-chunk condition
+            if start == 0 and version == 0:
+                rowset = self.rowset
+                self.rowset = None
+                return rowset
+
+            rowset = self._internal_parse_rowset(
+                buffer, start, idx, version, nrows, ncols
+            )
+
+            # continue parsing next chunk in the buffer
+            buffer = buffer[len + len("/{} ".format(len)) :]
+            if buffer:
+                return self.internal_parse_buffer(buffer, len(buffer))
+
+            return rowset
+
+        elif cmd == SQCLOUD_CMD.NULL.value:
+            return None
+
+        elif cmd in [SQCLOUD_CMD.INT.value, SQCLOUD_CMD.FLOAT.value]:
+            # TODO
+            clone = self._internal_parse_value(buffer, blen)
+            if clone is None:
+                return 0
+            if cmd == SQCLOUD_CMD.INT.value:
+                return int(clone)
+            return float(clone)
+
+        elif cmd == SQCLOUD_CMD.RAWJSON.value:
+            return None
+
+        return None
+
+    def _internal_uncompress_data(self, buffer: str, blen: int) -> Optional[str]:
+        """
+        %LEN COMPRESSED UNCOMPRESSED BUFFER
+
+        Args:
+            buffer (str): The compressed data buffer.
+            blen (int): The length of the buffer.
+
+        Returns:
+            str: The uncompressed data.
+
+        Raises:
+            None
+        """
+        tlen = 0  # total length
+        clen = 0  # compressed length
+        ulen = 0  # uncompressed length
+        hlen = 0  # raw header length
+        seek1 = 0
+
+        start = 1
+        counter = 0
+        for i in range(blen):
+            if buffer[i] != " ":
+                continue
+            counter += 1
+
+            data = buffer[start:i]
+            start = i + 1
+
+            if counter == 1:
+                tlen = int(data)
+                seek1 = start
+            elif counter == 2:
+                clen = int(data)
+            elif counter == 3:
+                ulen = int(data)
+                break
+
+        # sanity check header values
+        if tlen == 0 or clen == 0 or ulen == 0 or start == 1 or seek1 == 0:
+            return None
+
+        # copy raw header
+        hlen = start - seek1
+        header = buffer[start : start + hlen]
+
+        # compute index of the first compressed byte
+        start += hlen
+
+        # perform real decompression
+        clone = header + str(lz4.block.decompress(buffer[start:]))
+
+        # sanity check result
+        if len(clone) != ulen + hlen:
+            return None
+
+        return clone
+
+    def _internal_reconnect(self, buffer: str) -> bool:
+        return True
+
+    def _internal_parse_array(self, buffer: str) -> list:
+        start = 0
+        sqlite_number = self._internal_parse_number(buffer, start)
+        n = sqlite_number.value
+        start = sqlite_number.cstart
+
+        r = []
+        for i in range(n):
+            sqcloud_value = self._internal_parse_value(buffer, start)
+            start += sqcloud_value.cellsize
+            r.append(sqcloud_value.value)
+
+        return r
+
+    def _internal_parse_value(self, buffer: str, index: int = 0) -> SQCloudValue:
+        sqcloud_value = SQCloudValue()
+        len = 0
+        cellsize = 0
+
+        # handle special NULL value case
+        if buffer is None or buffer[index] == SQCLOUD_CMD.NULL.value:
+            len = 0
+            if cellsize is not None:
+                cellsize = 2
+
+            sqcloud_value.len = len
+            sqcloud_value.cellsize = cellsize
+
+            return sqcloud_value
+
+        sqcloud_number = self._internal_parse_number(buffer, index + 1)
+        blen = sqcloud_number.value
+        cstart = sqcloud_number.cstart
+
+        # handle decimal/float cases
+        if buffer[index] == SQCLOUD_CMD.INT.value or buffer[index] == SQCLOUD_CMD.FLOAT.value:
+            nlen = cstart - index
+            len = nlen - 2
+            cellsize = nlen
+
+            sqcloud_value.value = buffer[index + 1 : index + 1 + len]
+            sqcloud_value.len
+            sqcloud_value.cellsize = cellsize
+
+            return sqcloud_value
+
+        len = blen - 1 if buffer[index] == SQCLOUD_CMD.ZEROSTRING.value else blen
+        cellsize = blen + cstart - index
+
+        sqcloud_value.value = buffer[cstart : cstart + len]
+        sqcloud_value.len = len
+        sqcloud_value.cellsize = cellsize
+
+        return sqcloud_value

src/sqlitecloud/types.py

@@ -2,6 +2,8 @@
 from typing import List, Optional
 from enum import Enum
 
+from click import Option
+
 
 class SQCLOUD_VALUE_TYPE(Enum):
     VALUE_INTEGER = 1
@@ -174,3 +176,14 @@ def __init__(self) -> None:
         self.value: int = 0
         self.cstart: int = 0
         self.extcode: int = None
+
+
+class SQCloudValue:
+    """
+    Represents the parse value.
+    """
+
+    def __init__(self) -> None:
+        self.value: Optional[str] = None
+        self.len: int = 0
+        self.cellsize: int = 0