Implement decompression lz4

Author: danielebriggi

src/sqlitecloud/driver.py

@@ -279,12 +279,16 @@ def _internal_parse_buffer(
 
         # check for compressed result
         if cmd == SQCLOUD_CMD.COMPRESSED.value:
-            buffer = self._internal_uncompress_data(buffer, blen)
+            buffer = self._internal_uncompress_data(buffer)
             if buffer is None:
                 raise SQCloudException(
                     f"An error occurred while decompressing the input buffer of len {blen}."
                 )
 
+            # buffer after decompression
+            blen = len(buffer)
+            cmd = chr(buffer[0])
+
         # first character contains command type
         if cmd in [
             SQCLOUD_CMD.ZEROSTRING.value,
@@ -338,7 +342,10 @@ def _internal_parse_buffer(
 
         elif cmd in [SQCLOUD_CMD.ROWSET.value, SQCLOUD_CMD.ROWSET_CHUNK.value]:
             # CMD_ROWSET:          *LEN 0:VERSION ROWS COLS DATA
+            # - When decompressed, LEN for ROWSET is *0
+            #
             # CMD_ROWSET_CHUNK:    /LEN IDX:VERSION ROWS COLS DATA
+            #
             rowset_signature = self._internal_parse_rowset_signature(buffer)
             if rowset_signature.start < 0:
                 raise SQCloudException("Cannot parse rowset signature")
@@ -361,7 +368,7 @@ def _internal_parse_buffer(
             # continue parsing next chunk in the buffer
             sign_len = rowset_signature.len
             buffer = buffer[sign_len + len(f"/{sign_len} ") :]
-            if buffer:
+            if cmd == SQCLOUD_CMD.ROWSET_CHUNK.value and buffer:
                 return self._internal_parse_buffer(connection, buffer, len(buffer))
 
             return rowset
@@ -387,7 +394,7 @@ def _internal_parse_buffer(
         # TODO: exception here?
         return SQCloudResult(None)
 
-    def _internal_uncompress_data(self, buffer: bytes, blen: int) -> Optional[bytes]:
+    def _internal_uncompress_data(self, buffer: bytes) -> Optional[bytes]:
         """
         %LEN COMPRESSED UNCOMPRESSED BUFFER
 
@@ -398,51 +405,34 @@ def _internal_uncompress_data(self, buffer: bytes, blen: int) -> Optional[bytes]
         Returns:
             str: The uncompressed data.
         """
-        tlen = 0  # total length
-        clen = 0  # compressed length
-        ulen = 0  # uncompressed length
-        hlen = 0  # raw header length
-        seek1 = 0
+        space_index = buffer.index(b" ")
+        buffer = buffer[space_index + 1 :]
 
-        start = 1
-        counter = 0
-        for i in range(blen):
-            if chr(buffer[i]) != " ":
-                continue
-            counter += 1
+        # extract compressed size
+        space_index = buffer.index(b" ")
+        compressed_size = int(buffer[:space_index].decode("utf-8"))
+        buffer = buffer[space_index + 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
+        # extract decompressed size
+        space_index = buffer.index(b" ")
+        uncompressed_size = int(buffer[:space_index].decode("utf-8"))
+        buffer = buffer[space_index + 1 :]
 
-        # copy raw header
-        hlen = start - seek1
-        header = buffer[start : start + hlen]
+        # extract data header
+        header = buffer[:-compressed_size]
 
-        # compute index of the first compressed byte
-        start += hlen
+        # extract compressed data
+        compressed_buffer = buffer[-compressed_size:]
 
-        # perform real decompression
-        # clone = header + lz4.block.decompress(buffer[start:])
-        clone = lz4decode(buffer, start, header)
+        decompressed_buffer = header + lz4.block.decompress(
+            compressed_buffer, uncompressed_size
+        )
 
         # sanity check result
-        if len(clone) != ulen + hlen:
+        if len(decompressed_buffer) != uncompressed_size + len(header):
             return None
 
-        return clone
+        return decompressed_buffer
 
     def _internal_parse_array(self, buffer: bytes) -> list:
         start = 0

src/sqlitecloud/resultset.py

@@ -75,4 +75,4 @@ def get_name(self, col: int) -> Optional[str]:
         return self._result.colname[col]
 
     def get_result(self) -> Optional[any]:
-        return self.get_value(0, 0)
+        return self.get_value(0, 0)

src/tests/integration/test_client.py

@@ -584,8 +584,7 @@ def test_download_database(self, sqlitecloud_connection):
         assert cursor.description[0][0] == "AlbumId"
         assert cursor.description[1][0] == "Title"
 
-    # TODO
-    def test_compression(self):
+    def test_compression_single_column(self):
         account = SqliteCloudAccount()
         account.hostname = os.getenv("SQLITE_HOST")
         account.apikey = os.getenv("SQLITE_API_KEY")
@@ -595,12 +594,34 @@ def test_compression(self):
         client.config.compression = True
 
         # min compression size for rowset set by default to 20400 bytes
-        rowset = client.exec_query("SELECT '" + "a" * (1024 * 20) + "' AS DDD")
+        blob_size = 20 * 1024
+        # rowset = client.exec_query("SELECT * from albums inner join albums a2 on albums.AlbumId = a2.AlbumId")
+        rowset = client.exec_query(
+            f"SELECT hex(randomblob({blob_size})) AS 'someColumnName'"
+        )
 
         assert rowset.nrows == 1
         assert rowset.ncols == 1
-        assert rowset.get_value(0, 0).startswith("aaaaa")
-        assert len(rowset.get_value(0, 0)) == 100
+        assert rowset.get_name(0) == "someColumnName"
+        assert len(rowset.get_value(0, 0)) == blob_size * 2
+
+    def test_compression_multiple_columns(self):
+        account = SqliteCloudAccount()
+        account.hostname = os.getenv("SQLITE_HOST")
+        account.apikey = os.getenv("SQLITE_API_KEY")
+        account.database = os.getenv("SQLITE_DB")
+
+        client = SqliteCloudClient(cloud_account=account)
+        client.config.compression = True
+
+        # min compression size for rowset set by default to 20400 bytes
+        rowset = client.exec_query(
+            "SELECT * from albums inner join albums a2 on albums.AlbumId = a2.AlbumId"
+        )
+
+        assert rowset.nrows > 0
+        assert rowset.ncols > 0
+        assert rowset.get_name(0) == "AlbumId"
 
     # def test_send_blob(self, sqlitecloud_connection):
     #     connection, client = sqlitecloud_connection
@@ -615,19 +636,18 @@ def test_compression(self):
     #     blob = b""
     #     result = client.sendblob(blob, connection)
     #     assert result is not None
-    #     # Add additional assertions as needed
+    #
 
     # def test_send_large_blob(self, sqlitecloud_connection):
     #     connection, client = sqlitecloud_connection
     #     blob = b"A" * 1024 * 1024  # 1MB blob
     #     result = client.sendblob(blob, connection)
     #     assert result is not None
-    #     # Add additional assertions as needed
+    #
 
     # def test_send_blob_with_connection_closed(self, sqlitecloud_connection):
     #     connection, client = sqlitecloud_connection
     #     client.disconnect(connection)
     #     blob = b"Hello, this is a test blob"
     #     with pytest.raises(Exception):
     #         client.sendblob(blob, connection)
-    # Add additional assertions as needed

src/tests/unit/test_resultset.py

@@ -78,7 +78,7 @@ def test_get_value_array(self):
         result = SQCloudResult(result=[1, 2, 3])
         result_set = SqliteCloudResultSet(result)
 
-        assert [1,2,3] == result_set.get_value(0, 0)
+        assert [1, 2, 3] == result_set.get_value(0, 0)
 
     def test_get_colname(self):
         result = SQCloudResult()