Add Upload /downoload database

Author: danielebriggi

README.md

@@ -13,15 +13,12 @@ You can install SqliteCloud Package using Python Package Index (PYPI):
 $ pip install SqliteCloud
 ```
 
-- Follow the instructions reported here https://github.com/sqlitecloud/sdk/tree/master/C to build the driver.
-
-- Set SQLITECLOUD_DRIVER_PATH environment variable to the path of the driver file build.
-
 ## Usage
 <hr>
 
 ```python
-from sqlitecloud.client import SqliteCloudClient, SqliteCloudAccount
+from sqlitecloud.client import SqliteCloudClient
+from sqlitecloud.types import SqliteCloudAccount
 ```
 
 ### _Init a connection_
@@ -45,9 +42,8 @@ conn = client.open_connection()
 ### _Execute a query_
 You can bind values to parametric queries: you can pass parameters as positional values in an array
 ```python
-result = client.exec_statement(
-    "SELECT * FROM table_name WHERE id = ?",
-    [1],
+result = client.exec_query(
+    "SELECT * FROM table_name WHERE id = 1"
     conn=conn
 )
 ```

src/sqlitecloud/client.py

@@ -92,7 +92,7 @@ def sendblob(self, blob: bytes, conn: SQCloudConnect) -> SqliteCloudResultSet:
             blob (bytes): The blob to be sent to the database.
             conn (SQCloudConnect): The connection to the database.
         """
-        return self.driver.sendblob(blob, conn)
+        return self.driver.send_blob(blob, conn)
 
     def _parse_connection_string(self, connection_string) -> SQCloudConfig:
         # URL STRING FORMAT

src/sqlitecloud/download.py

@@ -0,0 +1,23 @@
+from io import BufferedWriter
+import logging
+
+from sqlitecloud.driver import Driver
+from sqlitecloud.types import SQCloudConnect
+
+
+def xCallback(
+    fd: BufferedWriter, data: bytes, blen: int, ntot: int, nprogress: int
+) -> None:
+    fd.write(data)
+
+    if blen == 0:
+        logging.log(logging.DEBUG, "DOWNLOAD COMPLETE")
+    else:
+        logging.log(logging.DEBUG, f"{(nprogress + blen) / ntot * 100:.2f}%")
+
+
+def download_db(connection: SQCloudConnect, dbname: str, filename: str) -> None:
+    driver = Driver()
+
+    with open(filename, "wb") as fd:
+        driver.download_database(connection, dbname, fd, xCallback, False)

src/sqlitecloud/driver.py

@@ -1,9 +1,11 @@
+from io import BufferedReader, BufferedWriter
 import ssl
-from typing import Optional, Union
+from typing import Callable, Optional, Union
 import lz4.block
 from sqlitecloud.resultset import SQCloudResult
 from sqlitecloud.types import (
     SQCLOUD_CMD,
+    SQCLOUD_DEFAULT,
     SQCLOUD_INTERNAL_ERRCODE,
     SQCLOUD_ROWSET,
     SQCloudConfig,
@@ -17,6 +19,8 @@
 
 
 class Driver:
+    SQCLOUD_DEFAULT_UPLOAD_SIZE = 512 * 1024
+
     def __init__(self) -> None:
         # Used while parsing chunked rowset
         self._rowset: SQCloudResult = None
@@ -77,7 +81,7 @@ def disconnect(self, conn: SQCloudConnect):
     def execute(self, command: str, connection: SQCloudConnect) -> SQCloudResult:
         return self._internal_run_command(connection, command)
 
-    def sendblob(self, blob: bytes, conn: SQCloudConnect) -> SQCloudResult:
+    def send_blob(self, blob: bytes, conn: SQCloudConnect) -> SQCloudResult:
         try:
             conn.isblob = True
             return self._internal_run_command(conn, blob)
@@ -90,6 +94,113 @@ def _internal_reconnect(self, buffer: bytes) -> bool:
     def _internal_setup_pubsub(self, buffer: bytes) -> bool:
         return True
 
+    def upload_database(
+        self,
+        connection: SQCloudConnect,
+        dbname: str,
+        key: Optional[str],
+        is_file_transfer: bool,
+        snapshot_id: int,
+        is_internal_db: bool,
+        fd: BufferedReader,
+        dbsize: int,
+        xCallback: Callable[[BufferedReader, int, int, int], bytes],
+    ) -> None:
+        keyarg = "KEY " if key else ""
+        keyvalue = key if key else ""
+
+        # prepare command to execute
+        command = ""
+        if is_file_transfer:
+            internalarg = "INTERNAL" if is_internal_db else ""
+            command = f"TRANSFER DATABASE '{dbname}' {keyarg}{keyvalue} SNAPSHOT {snapshot_id} {internalarg}"
+        else:
+            command = f"UPLOAD DATABASE '{dbname}' {keyarg}{keyvalue}"
+
+        # execute command on server side
+        result = self._internal_run_command(connection, command)
+        if not result.data[0]:
+            raise SQCloudException(
+                "An error occurred while initializing the upload of the database."
+            )
+
+        buffer: bytes = b""
+        blen = 0
+        nprogress = 0
+        try:
+            while True:
+                # execute callback to read buffer
+                blen = SQCLOUD_DEFAULT.UPLOAD_SIZE.value
+                try:
+                    buffer = xCallback(fd, blen, dbsize, nprogress)
+                    blen = len(buffer)
+                except Exception as e:
+                    raise SQCloudException(
+                        "An error occurred while reading the file."
+                    ) from e
+
+                try:
+                    # send also the final confirmation blob of zero bytes
+                    self.send_blob(buffer, connection)
+                except Exception as e:
+                    raise SQCloudException(
+                        "An error occurred while uploading the file."
+                    ) from e
+
+                # update progress
+                nprogress += blen
+
+                if blen == 0:
+                    # Upload completed
+                    break
+        except Exception as e:
+            self._internal_run_command(connection, "UPLOAD ABORT")
+            raise e
+
+    def download_database(
+        self,
+        connection: SQCloudConnect,
+        dbname: str,
+        fd: BufferedWriter,
+        xCallback: Callable[[BufferedWriter, int, int, int], bytes],
+        if_exists: bool,
+    ) -> None:
+        exists_cmd = " IF EXISTS" if if_exists else ""
+        result = self._internal_run_command(
+            connection, f"DOWNLOAD DATABASE {dbname}{exists_cmd};"
+        )
+
+        if result.nrows == 0:
+            raise SQCloudException(
+                "An error occurred while initializing the download of the database."
+            )
+
+        # result is an ARRAY (database size, number of pages, raft_index)
+        download_info = result.data[0]
+        db_size = int(download_info[0])
+
+        # loop to download
+        progress_size = 0
+
+        try:
+            while progress_size < db_size:
+                result = self._internal_run_command(connection, "DOWNLOAD STEP")
+
+                # res is BLOB, decode it
+                data = result.data[0]
+                data_len = len(data)
+
+                # execute callback (with progress_size updated)
+                progress_size += data_len
+                xCallback(fd, data, data_len, db_size, progress_size)
+
+                # check exit condition
+                if data_len == 0:
+                    break
+        except Exception as e:
+            self._internal_run_command(connection, "DOWNLOAD ABORT")
+            raise e
+
     def _internal_config_apply(
         self, connection: SQCloudConnect, config: SQCloudConfig
     ) -> None:
@@ -136,7 +247,7 @@ def _internal_config_apply(
 
     def _internal_run_command(
         self, connection: SQCloudConnect, command: Union[str, bytes]
-    ) -> None:
+    ) -> SQCloudResult:
         self._internal_socket_write(connection, command)
         return self._internal_socket_read(connection)
 

src/sqlitecloud/types.py

@@ -2,6 +2,10 @@
 from typing import Optional
 from enum import Enum
 
+class SQCLOUD_DEFAULT(Enum):
+    PORT = 8860
+    TIMEOUT = 12
+    UPLOAD_SIZE = 512*1024
 
 class SQCLOUD_CMD(Enum):
     STRING = "+"
@@ -62,7 +66,7 @@ def __init__(
         password: Optional[str] = "",
         hostname: Optional[str] = "",
         dbname: Optional[str] = "",
-        port: Optional[int] = 8860,
+        port: Optional[int] = SQCLOUD_DEFAULT.PORT.value,
         apikey: Optional[str] = "",
     ) -> None:
         # User name is required unless connectionstring is provided
@@ -98,7 +102,7 @@ def __init__(self) -> None:
         # Optional query timeout passed directly to TLS socket
         self.timeout = 0
         # Socket connection timeout
-        self.connect_timeout = 20
+        self.connect_timeout = SQCLOUD_DEFAULT.TIMEOUT.value
 
         # Enable compression
         self.compression = False
@@ -132,7 +136,7 @@ def __init__(self) -> None:
 
 class SQCloudException(Exception):
     def __init__(
-        self, message: str, code: Optional[int] = -1, xerrcode: Optional[int] = 0
+        self, message: str, code: int = -1, xerrcode: int = 0
     ) -> None:
         self.errmsg = str(message)
         self.errcode = code

src/sqlitecloud/upload.py

@@ -0,0 +1,54 @@
+from io import BufferedReader
+import os
+from typing import Optional
+from sqlitecloud.driver import Driver
+from sqlitecloud.types import SQCloudConnect
+import logging
+
+def xCallback(fd: BufferedReader, blen: int, ntot: int, nprogress: int) -> bytes:
+    buffer = fd.read(blen)
+    nread = len(buffer)
+
+    if nread == 0:
+        logging.log(logging.DEBUG, "UPLOAD COMPLETE\n\n")
+    else:
+        logging.log(logging.DEBUG, f"{(nprogress + nread) / ntot * 100:.2f}%")
+
+    return buffer
+
+
+def upload_db(
+    connection: SQCloudConnect, dbname: str, key: Optional[str], filename: str
+) -> bool:
+    """
+    Uploads a SQLite database to the SQLite Cloud node using the provided connection.
+
+    Args:
+        connection (SQCloudConnect): The connection object used to connect to the node.
+        dbname (str): The name of the database in SQLite Cloud.
+        key (Optional[str]): The encryption key for the database. If None, no encryption is used.
+        filename (str): The path to the SQLite database file to be uploaded.
+
+    Returns:
+        bool: True if the upload is successful, SQCloudException in case of errors.
+    """
+    
+    # Create a driver object
+    driver = Driver()
+
+    with open(filename, 'rb') as fd:
+        dbsize = os.path.getsize(filename)
+
+        driver.upload_database(
+            connection,
+            dbname,
+            key,
+            False,
+            0,
+            False,
+            fd,
+            dbsize,
+            xCallback,
+        )
+
+    return True

src/tests/assets/test.db

@@ -554,36 +554,6 @@ def test_stress_test_20x_batched_selects(self, sqlitecloud_connection):
                         query_ms < self.EXPECT_SPEED_MS
                     ), f"{num_queries}x batched selects, {query_ms}ms per query"
 
-    def test_download_database(self, sqlitecloud_connection):
-        connection, client = sqlitecloud_connection
-
-        rowset = client.exec_query(
-            "DOWNLOAD DATABASE " + os.getenv("SQLITE_DB"), connection
-        )
-
-        result_array = rowset.get_result()
-
-        db_size = int(result_array[0])
-
-        tot_read = 0
-        data: bytes = b""
-        while tot_read < db_size:
-            result = client.exec_query("DOWNLOAD STEP;", connection)
-
-            data += result.get_result()
-            tot_read += len(data)
-
-        temp_file = tempfile.mkstemp(prefix="chinook")[1]
-        with open(temp_file, "wb") as f:
-            f.write(data)
-
-        db = sqlite3.connect(temp_file)
-        cursor = db.execute("SELECT * FROM albums")
-        rowset = cursor.fetchall()
-
-        assert cursor.description[0][0] == "AlbumId"
-        assert cursor.description[1][0] == "Title"
-
     def test_compression_single_column(self):
         account = SqliteCloudAccount()
         account.hostname = os.getenv("SQLITE_HOST")

src/tests/integration/test_client.py

@@ -0,0 +1,41 @@
+import os
+import sqlite3
+import tempfile
+
+import pytest
+
+from sqlitecloud import download
+from sqlitecloud.client import SqliteCloudClient
+from sqlitecloud.types import SQCloudConnect, SqliteCloudAccount
+
+
+class TestDownload:
+    @pytest.fixture()
+    def sqlitecloud_connection(self):
+        account = SqliteCloudAccount()
+        account.username = os.getenv("SQLITE_USER")
+        account.password = os.getenv("SQLITE_PASSWORD")
+        account.dbname = os.getenv("SQLITE_DB")
+        account.hostname = os.getenv("SQLITE_HOST")
+        account.port = 8860
+
+        client = SqliteCloudClient(cloud_account=account)
+
+        connection = client.open_connection()
+        assert isinstance(connection, SQCloudConnect)
+
+        yield (connection, client)
+
+        client.disconnect(connection)
+
+    def test_download_database(self, sqlitecloud_connection):
+        connection, _ = sqlitecloud_connection
+
+        temp_file = tempfile.mkstemp(prefix="chinook")[1]
+        download.download_db(connection, "chinook.sqlite", temp_file)
+
+        db = sqlite3.connect(temp_file)
+        cursor = db.execute("SELECT * FROM albums")
+
+        assert cursor.description[0][0] == "AlbumId"
+        assert cursor.description[1][0] == "Title"