lfss 0.7.14__py3-none-any.whl → 0.8.0__py3-none-any.whl

lfss/{client → api}/__init__.py

@@ -1,6 +1,8 @@
 import os, time, pathlib
 from threading import Lock
-from .api import Connector
+from .connector import Connector
+from ..src.datatype import FileRecord
+from ..src.utils import decode_uri_compnents
 from ..src.bounded_pool import BoundedThreadPoolExecutor
 
 def upload_file(
@@ -11,8 +13,9 @@ def upload_file(
     interval: float = 0, 
     verbose: bool = False, 
     **put_kwargs
-    ):
+    ) -> tuple[bool, str]:
     this_try = 0
+    error_msg = ""
     while this_try <= n_retries:
         try:
             with open(file_path, 'rb') as f:
@@ -24,6 +27,7 @@ def upload_file(
                 raise e
             if verbose:
                 print(f"Error uploading {file_path}: {e}, retrying...")
+                error_msg = str(e)
             this_try += 1
         finally:
             time.sleep(interval)
@@ -31,8 +35,8 @@ def upload_file(
     if this_try > n_retries:
         if verbose:
             print(f"Failed to upload {file_path} after {n_retries} retries.")
-        return False
-    return True
+        return False, error_msg
+    return True, error_msg
 
 def upload_directory(
     connector: Connector, 
@@ -43,7 +47,7 @@ def upload_directory(
     interval: float = 0,
     verbose: bool = False,
     **put_kwargs
-    ) -> list[str]:
+    ) -> list[tuple[str, str]]:
     assert path.endswith('/'), "Path must end with a slash."
     if path.startswith('/'):
         path = path[1:]
@@ -52,8 +56,8 @@ def upload_directory(
     _counter = 0
     _counter_lock = Lock()
 
-    faild_files = []
-    def put_file(file_path):
+    faild_items = []
+    def put_file(c: Connector, file_path):
         with _counter_lock:
             nonlocal _counter
             _counter += 1
@@ -62,18 +66,19 @@ def upload_directory(
             if verbose:
                 print(f"[{this_count}] Uploading {file_path} to {dst_path}")
 
-        if not upload_file(
-            connector, file_path, dst_path, 
+        if not (res:=upload_file(
+            c, file_path, dst_path, 
             n_retries=n_retries, interval=interval, verbose=verbose, **put_kwargs
-            ):
-            faild_files.append(file_path)
+            ))[0]:
+            faild_items.append((file_path, res[1]))
 
-    with BoundedThreadPoolExecutor(n_concurrent) as executor:
-        for root, dirs, files in os.walk(directory):
-            for file in files:
-                executor.submit(put_file, os.path.join(root, file))
+    with connector.session(n_concurrent) as c:
+        with BoundedThreadPoolExecutor(n_concurrent) as executor:
+            for root, dirs, files in os.walk(directory):
+                for file in files:
+                    executor.submit(put_file, c, os.path.join(root, file))
 
-    return faild_files
+    return faild_items
 
 def download_file(
     connector: Connector, 
@@ -83,17 +88,19 @@ def download_file(
     interval: float = 0, 
     verbose: bool = False, 
     overwrite: bool = False
-    ):
+    ) -> tuple[bool, str]:
     this_try = 0
+    error_msg = ""
     while this_try <= n_retries:
         if not overwrite and os.path.exists(file_path):
             if verbose:
                 print(f"File {file_path} already exists, skipping download.")
-            return True
+            return True, error_msg
         try:
             blob = connector.get(src_url)
-            if not blob:
-                return False
+            if blob is None:
+                error_msg = "File not found."
+                return False, error_msg
             pathlib.Path(file_path).parent.mkdir(parents=True, exist_ok=True)
             with open(file_path, 'wb') as f:
                 f.write(blob)
@@ -103,6 +110,7 @@ def download_file(
                 raise e
             if verbose:
                 print(f"Error downloading {src_url}: {e}, retrying...")
+                error_msg = str(e)
             this_try += 1
         finally:
             time.sleep(interval)
@@ -110,8 +118,8 @@ def download_file(
     if this_try > n_retries:
         if verbose:
             print(f"Failed to download {src_url} after {n_retries} retries.")
-        return False
-    return True
+        return False, error_msg
+    return True, error_msg
 
 def download_directory(
     connector: Connector,
@@ -122,7 +130,7 @@ def download_directory(
     interval: float = 0,
     verbose: bool = False,
     overwrite: bool = False
-    ) -> list[str]:
+    ) -> list[tuple[str, str]]:
 
     directory = str(directory)
 
@@ -133,23 +141,32 @@ def download_directory(
     
     _counter = 0
     _counter_lock = Lock()
-    failed_files = []
-    def get_file(src_url):
-        nonlocal _counter, failed_files
+    failed_items: list[tuple[str, str]] = []
+    def get_file(c, src_url):
+        nonlocal _counter, failed_items
         with _counter_lock:
             _counter += 1
             this_count = _counter
-            dst_path = f"{directory}{os.path.relpath(src_url, src_path)}"
+            dst_path = f"{directory}{os.path.relpath(decode_uri_compnents(src_url), decode_uri_compnents(src_path))}"
             if verbose:
                 print(f"[{this_count}] Downloading {src_url} to {dst_path}")
 
-        if not download_file(
-            connector, src_url, dst_path, 
+        if not (res:=download_file(
+            c, src_url, dst_path, 
             n_retries=n_retries, interval=interval, verbose=verbose, overwrite=overwrite
-            ):
-            failed_files.append(src_url)
+            ))[0]:
+            failed_items.append((src_url, res[1]))
         
-    with BoundedThreadPoolExecutor(n_concurrent) as executor:
-        for file in connector.list_path(src_path, flat=True).files:
-            executor.submit(get_file, file.url)
-    return failed_files
+    batch_size = 10000
+    file_list: list[FileRecord] = []
+    with connector.session(n_concurrent) as c:
+        file_count = c.count_files(src_path, flat=True)
+        for offset in range(0, file_count, batch_size):
+            file_list.extend(c.list_files(
+                src_path, offset=offset, limit=batch_size, flat=True
+            ))
+
+        with BoundedThreadPoolExecutor(n_concurrent) as executor:
+            for file in file_list:
+                executor.submit(get_file, c, file.url)
+    return failed_items

lfss/{client/api.py → api/connector.py}

@@ -1,9 +1,13 @@
-from typing import Optional, Literal
+from __future__ import annotations
+from typing import Optional, Literal, Iterator
 import os
 import requests
+import requests.adapters
 import urllib.parse
+from lfss.src.error import PathNotFoundError
 from lfss.src.datatype import (
-    FileReadPermission, FileRecord, DirectoryRecord, UserRecord, PathContents
+    FileReadPermission, FileRecord, DirectoryRecord, UserRecord, PathContents, 
+    FileSortKey, DirSortKey
     )
 from lfss.src.utils import ensure_uri_compnents
 
@@ -11,12 +15,41 @@ _default_endpoint = os.environ.get('LFSS_ENDPOINT', 'http://localhost:8000')
 _default_token = os.environ.get('LFSS_TOKEN', '')
 
 class Connector:
+    class Session:
+        def __init__(self, connector: Connector, pool_size: int = 10):
+            self.connector = connector
+            self.pool_size = pool_size
+        def open(self):
+            self.close()
+            if self.connector._session is None:
+                s = requests.Session()
+                adapter = requests.adapters.HTTPAdapter(pool_connections=self.pool_size, pool_maxsize=self.pool_size)
+                s.mount('http://', adapter)
+                s.mount('https://', adapter)
+                self.connector._session = s
+        def close(self):
+            if self.connector._session is not None:
+                self.connector._session.close()
+            self.connector._session = None
+        def __call__(self):
+            return self.connector
+        def __enter__(self):
+            self.open()
+            return self.connector
+        def __exit__(self, exc_type, exc_value, traceback):
+            self.close()
+
     def __init__(self, endpoint=_default_endpoint, token=_default_token):
         assert token, "No token provided. Please set LFSS_TOKEN environment variable."
         self.config = {
             "endpoint": endpoint,
             "token": token
         }
+        self._session: Optional[requests.Session] = None
+    
+    def session(self, pool_size: int = 10):
+        """ avoid creating a new session for each request.  """
+        return self.Session(self, pool_size)
     
     def _fetch_factory(
         self, method: Literal['GET', 'POST', 'PUT', 'DELETE'], 
@@ -31,9 +64,13 @@ class Connector:
             headers.update({
                 'Authorization': f"Bearer {self.config['token']}",
             })
-            with requests.Session() as s:
-                response = s.request(method, url, headers=headers, **kwargs)
+            if self._session is not None:
+                response = self._session.request(method, url, headers=headers, **kwargs)
                 response.raise_for_status()
+            else:
+                with requests.Session() as s:
+                    response = s.request(method, url, headers=headers, **kwargs)
+                    response.raise_for_status()
             return response
         return f
 
@@ -80,9 +117,9 @@ class Connector:
         )
         return response.json()
     
-    def _get(self, path: str) -> Optional[requests.Response]:
+    def _get(self, path: str, stream: bool = False) -> Optional[requests.Response]:
         try:
-            response = self._fetch_factory('GET', path)()
+            response = self._fetch_factory('GET', path)(stream=stream)
         except requests.exceptions.HTTPError as e:
             if e.response.status_code == 404:
                 return None
@@ -94,6 +131,12 @@ class Connector:
         response = self._get(path)
         if response is None: return None
         return response.content
+    
+    def get_stream(self, path: str) -> Iterator[bytes]:
+        """Downloads a file from the specified path, will raise PathNotFoundError if path not found."""
+        response = self._get(path, stream=True)
+        if response is None: raise PathNotFoundError("Path not found: " + path)
+        return response.iter_content(chunk_size=1024)
 
     def get_json(self, path: str) -> Optional[dict]:
         response = self._get(path)
@@ -118,12 +161,50 @@ class Connector:
                 return None
             raise e
     
-    def list_path(self, path: str, flat: bool = False) -> PathContents:
+    def list_path(self, path: str) -> PathContents:
+        """ 
+        shorthand list with limited options, 
+        for large directories / more options, use list_files and list_dirs instead.
+        """
         assert path.endswith('/')
-        response = self._fetch_factory('GET', path, {'flat': flat})()
+        response = self._fetch_factory('GET', path)()
         dirs = [DirectoryRecord(**d) for d in response.json()['dirs']]
         files = [FileRecord(**f) for f in response.json()['files']]
         return PathContents(dirs=dirs, files=files)
+    
+    def count_files(self, path: str, flat: bool = False) -> int:
+        assert path.endswith('/')
+        response = self._fetch_factory('GET', '_api/count-files', {'path': path, 'flat': flat})()
+        return response.json()['count']
+
+    def list_files(
+        self, path: str, offset: int = 0, limit: int = 1000,
+        order_by: FileSortKey = '', order_desc: bool = False, 
+        flat: bool = False
+    ) -> list[FileRecord]:
+        assert path.endswith('/')
+        response = self._fetch_factory('GET', "_api/list-files", {
+            'path': path,
+            'offset': offset, 'limit': limit, 'order_by': order_by, 'order_desc': order_desc, 'flat': flat
+        })()
+        return [FileRecord(**f) for f in response.json()]
+    
+    def count_dirs(self, path: str) -> int:
+        assert path.endswith('/')
+        response = self._fetch_factory('GET', '_api/count-dirs', {'path': path})()
+        return response.json()['count']
+        
+    def list_dirs(
+        self, path: str, offset: int = 0, limit: int = 1000,
+        order_by: DirSortKey = '', order_desc: bool = False, 
+        skim: bool = True
+    ) -> list[DirectoryRecord]:
+        assert path.endswith('/')
+        response = self._fetch_factory('GET', "_api/list-dirs", {
+            'path': path,
+            'offset': offset, 'limit': limit, 'order_by': order_by, 'order_desc': order_desc, 'skim': skim
+        })()
+        return [DirectoryRecord(**d) for d in response.json()]
 
     def set_file_permission(self, path: str, permission: int | FileReadPermission):
         """Sets the file permission for the specified path."""

lfss/cli/cli.py

@@ -1,4 +1,4 @@
-from lfss.client import Connector, upload_directory, upload_file, download_file, download_directory
+from lfss.api import Connector, upload_directory, upload_file, download_file, download_directory
 from pathlib import Path
 import argparse
 from lfss.src.datatype import FileReadPermission

lfss/cli/user.py

@@ -29,7 +29,7 @@ async def _main():
     sp_set.add_argument('username', type=str)
     sp_set.add_argument('-p', '--password', type=str, default=None)
     sp_set.add_argument('-a', '--admin', type=parse_bool, default=None)
-    sp_set.add_argument('--permission', type=int, default=None)
+    sp_set.add_argument('--permission', type=parse_permission, default=None)
     sp_set.add_argument('--max-storage', type=parse_storage_size, default=None)
     
     sp_list = sp.add_parser('list')

lfss/src/connection_pool.py

@@ -46,7 +46,7 @@ class SqlConnection:
 
 class SqlConnectionPool:
     _r_sem: Semaphore
-    _w_sem: Semaphore
+    _w_sem: Lock | Semaphore
     def __init__(self):
         self._readers: list[SqlConnection] = []
         self._writer: None | SqlConnection = None
@@ -57,7 +57,8 @@ class SqlConnectionPool:
         self._readers = []
 
         self._writer = SqlConnection(await get_connection(read_only=False))
-        self._w_sem = Semaphore(1)
+        self._w_sem = Lock()
+        # self._w_sem = Semaphore(1)
 
         for _ in range(n_read):
             conn = await get_connection(read_only=True)