mercuto-client 0.1.0__py3-none-any.whl

mercuto_client/ingester/processor.py

@@ -0,0 +1,291 @@
+import logging
+import os
+import shutil
+import sqlite3
+from datetime import datetime, timezone
+from typing import Callable, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class FileProcessor:
+    """
+    System for processing files in a strict order with retry logic.
+
+    Keeps an SQLite database to track files and their processing status.
+    Keeps old files in the buffer and only deletes them once max_files is reached (whether processed or not).
+
+    :param buffer_dir: Directory where files are stored
+    :param db_path: Path to the SQLite database file for tracking file processing status.
+    :param process_callback: Callable that processes a file. Should return True if processing is successful
+    :param max_attempts: Maximum number of attempts to process a file before marking it as failed.
+    :param max_files: Maximum number of files to keep in the buffer directory. If None, no limit is enforced.
+    :param free_space_mb: Optional minimum free space in MB to keep on the partition where the buffer directory is located.
+        This is combined with max_files to determine when to delete old files and takes precedence over max_files.
+    :param clock: Optional callable that returns the timestamp for the file based on the filename. Takes in the file name as an argument
+        and should return a float representing the timestamp in seconds since the epoch. Defaults to datetime.now().timestamp().
+        This clock is NOT used when scanning existing files, it takes its own clock function.
+
+
+    Provides a callback for processing files, which should return True if successful.
+    If processing fails, it retries up to max_attempts times before marking the file as failed.
+
+    Periodically call `cleanup_old_files()` to remove old files from the buffer directory.
+    Use `scan_existing_files()` to register files that were added while the system was offline.
+    Use `process_next_file()` to process the next file in the buffer in strict order.
+    Add files to the buffer using `add_file_to_db()`, or use `start_watching()` to automatically watch a directory for new files.
+
+    Example usage:
+    ```python
+    processor = FileProcessor(...)
+    processor.scan_existing_files()
+
+    while True:
+        processor.add_file_to_db("path/to/file.txt")
+        processor.process_next_file()
+    ```
+    """
+
+    def __init__(self, buffer_dir: str, db_path: str,
+                 process_callback: Callable[[str], bool],
+                 max_attempts: int,
+                 max_files: Optional[int] = None,
+                 free_space_mb: Optional[float] = None,
+                 clock: Optional[Callable[[str], float]] = None
+                 ) -> None:
+        self._buffer_dir = buffer_dir
+        self._db_path = db_path
+        self._max_files = max_files
+        self._max_attempts = max_attempts
+        self._process_callback = process_callback
+        self._free_space_mb = free_space_mb
+        if clock is None:
+            # Default clock function to get current time in seconds since epoch
+            def clock(_): return datetime.now(timezone.utc).timestamp()
+        self._clock = clock
+
+        os.makedirs(self._buffer_dir, exist_ok=True)
+        self._init_db()
+
+    def get_db_path(self) -> str:
+        """Returns the path to the SQLite database."""
+        return self._db_path
+
+    def get_buffer_dir(self) -> str:
+        """Returns the path to the buffer directory."""
+        return self._buffer_dir
+
+    def _init_db(self) -> None:
+        """Initialize SQLite database with attempt tracking."""
+        conn = sqlite3.connect(self._db_path)
+        cursor = conn.cursor()
+        cursor.execute("""
+        CREATE TABLE IF NOT EXISTS file_buffer (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            filename TEXT UNIQUE,
+            filepath TEXT UNIQUE,
+            status TEXT DEFAULT 'pending',
+            attempts INTEGER DEFAULT 0,
+            timestamp REAL
+        )
+        """)
+        conn.commit()
+        conn.close()
+        logger.info("Database initialized.")
+
+    def scan_existing_files(self, clock: Optional[Callable[[str], float]] = None) -> None:
+        """
+        Detect files added while offline and process them.
+        :param clock: Optional callable that returns the timestamp for the file based on the filename.
+            If not provided, uses the file's creation time.
+            This may be innaccurate if multiple files are added at once, or if the file system does not support accurate timestamps.
+
+        """
+        if clock is None:
+            # Default clock function to get current time in seconds since epoch
+            def clock(filepath: str) -> float:
+                return os.path.getctime(filepath)
+
+        conn = sqlite3.connect(self._db_path)
+        cursor = conn.cursor()
+
+        files_with_timestamps = [
+            (filename, clock(os.path.join(self._buffer_dir, filename)))
+            for filename in os.listdir(self._buffer_dir)
+            if os.path.isfile(os.path.join(self._buffer_dir, filename))
+        ]
+
+        for filename, timestamp in sorted(files_with_timestamps, key=lambda x: x[1]):
+            filepath = os.path.join(self._buffer_dir, filename)
+            timestamp = clock(filepath)
+
+            cursor.execute(
+                "SELECT COUNT(*) FROM file_buffer WHERE filename = ?", (filename,))
+            exists: int = cursor.fetchone()[0]
+
+            if not exists:
+                logger.info(f"Registering existing {filename} for processing...")
+                cursor.execute("INSERT INTO file_buffer (filename, filepath, status, attempts, timestamp) VALUES (?, ?, 'pending', 0, ?)",
+                               (filename, filepath, timestamp))
+
+        conn.commit()
+        conn.close()
+
+    def process_next_file(self) -> Optional[str]:
+        """
+        Attempt to process the next file in the sequence (if exists), ensuring strict order.
+        Returns the filepath of the processed file if successful or None if no pending files are found or failed.
+        """
+        conn = sqlite3.connect(self._db_path)
+        cursor = conn.cursor()
+        cursor.execute(
+            "SELECT filepath, attempts FROM file_buffer WHERE status = 'pending' ORDER BY timestamp ASC LIMIT 1")
+
+        pending_files: list[tuple[str, int]] = cursor.fetchall()
+        conn.close()
+
+        assert len(pending_files) <= 1, "More than one pending file found, which violates strict order."
+
+        for (filepath, attempts) in pending_files:
+            if self._process_file(filepath, attempts):
+                return filepath
+        return None
+
+    def _process_file(self, filepath: str, attempts: int) -> bool:
+        if not os.path.exists(filepath):
+            logger.warning(f"File {filepath} does not exist. Skipping.")
+            self._mark_as_failed(filepath)
+        try:
+            success: bool = self._process_callback(filepath)
+        except Exception as e:
+            logger.error(f"Processing error for {filepath}: {e}")
+            success = False
+
+        if success:
+            self._mark_as_processed(filepath)
+            return True
+
+        if attempts >= self._max_attempts:
+            logger.warning(
+                f"Max retries reached for {filepath}. Moving to next file.")
+
+            self._mark_as_failed(filepath)
+            return True  # Give up and move to next file
+        else:
+            self._increment_attempt(filepath)
+            return False
+
+    def _increment_attempt(self, filepath: str) -> None:
+        # Reopen DB to update attempt count
+        conn = sqlite3.connect(self._db_path)
+        cursor = conn.cursor()
+        cursor.execute(
+            "UPDATE file_buffer SET attempts = attempts + 1 WHERE filepath = ?", (filepath,))
+        conn.commit()
+        conn.close()
+
+    def _mark_as_failed(self, filepath: str) -> None:
+        """Marks a file as failed in the database."""
+        conn = sqlite3.connect(self._db_path)
+        cursor = conn.cursor()
+        cursor.execute(
+            "UPDATE file_buffer SET status = 'failed' WHERE filepath = ?", (filepath,))
+        conn.commit()
+        conn.close()
+        logger.info(f"File {filepath} marked as failed.")
+
+    def _mark_as_processed(self, filepath: str) -> None:
+        """Marks a file as processed in the database."""
+        conn = sqlite3.connect(self._db_path)
+        cursor = conn.cursor()
+        cursor.execute(
+            "UPDATE file_buffer SET status = 'processed' WHERE filepath = ?", (filepath,))
+        conn.commit()
+        conn.close()
+        logger.info(f"File {filepath} marked as processed.")
+
+    def cleanup_old_files_with_max_files(self) -> None:
+        """Remove old files beyond the max file count."""
+        if self._max_files is None:
+            return
+        conn = sqlite3.connect(self._db_path)
+        cursor = conn.cursor()
+        cursor.execute(
+            "SELECT filepath FROM file_buffer ORDER BY timestamp DESC LIMIT -1 OFFSET ?", (self._max_files,))
+        files_to_delete: list[tuple[str]] = cursor.fetchall()
+        conn.close()
+
+        for (filepath,) in reversed(files_to_delete):
+            if os.path.exists(filepath):
+                os.remove(filepath)
+                conn = sqlite3.connect(self._db_path)
+                cursor = conn.cursor()
+                cursor.execute(
+                    "DELETE FROM file_buffer WHERE filepath = ?", (filepath,))
+                conn.commit()
+                conn.close()
+                logger.info(f"Deleted old file {filepath}")
+
+    def cleanup_old_files(self) -> None:
+        """Remove old files based on max_files and free space."""
+        if self._max_files is not None:
+            self.cleanup_old_files_with_max_files()
+
+        if self._free_space_mb is not None:
+            self.cleanup_old_files_with_free_space()
+
+    def cleanup_old_files_with_free_space(self) -> None:
+        """Remove old files ensuring free space is maintained."""
+        if self._free_space_mb is None:
+            return
+
+        def free_space_mb() -> float:
+            """Returns the free space in MB on the partition where the buffer directory is located."""
+            total, used, free = shutil.disk_usage(self._buffer_dir)
+            return free / (1024 * 1024)
+
+        while free_space_mb() < self._free_space_mb:
+            if not self._delete_oldest_file():
+                logger.warning("No more files to delete to free up space.")
+                break
+
+    def _delete_oldest_file(self) -> bool:
+        """Deletes the oldest file in the buffer directory."""
+        conn = sqlite3.connect(self._db_path)
+        cursor = conn.cursor()
+        cursor.execute(
+            "SELECT filepath FROM file_buffer ORDER BY timestamp ASC LIMIT 1")
+        oldest_file: Optional[tuple[str]] = cursor.fetchone()
+        if not oldest_file:
+            logger.info("No files to delete.")
+            return False
+
+        filepath = oldest_file[0]
+        removed = False
+        if os.path.exists(filepath):
+            os.remove(filepath)
+            cursor.execute(
+                "DELETE FROM file_buffer WHERE filepath = ?", (filepath,))
+            conn.commit()
+            removed = True
+            logger.info(f"Deleted oldest file {filepath}")
+        else:
+            logger.warning(f"Oldest file {filepath} does not exist.")
+
+        conn.close()
+        return removed
+
+    def add_file_to_db(self, filepath: str) -> None:
+        """Adds a new file to database and triggers processing."""
+        timestamp = self._clock(filepath)
+        filename: str = os.path.basename(filepath)
+        conn = sqlite3.connect(self._db_path)
+        cursor = conn.cursor()
+
+        cursor.execute("""
+        INSERT OR IGNORE INTO file_buffer (filename, filepath, status, attempts, timestamp)
+        VALUES (?, ?, 'pending', 0, ?)
+        """, (filename, filepath, timestamp))
+
+        conn.commit()
+        conn.close()

mercuto_client/ingester/util.py

@@ -0,0 +1,64 @@
+import itertools
+import shutil
+from pathlib import Path
+from typing import Iterable, Iterator, TypeVar
+
+import requests
+
+
+def get_my_public_ip() -> str:
+    """
+    Fetches the public IP address of the machine making the request.
+    Uses the 'checkip.amazonaws.com' service to retrieve the IP address.
+    :return The public IP address as a string in the form 'x.x.x.x'.
+    :raises
+        requests.RequestException: If the request to the IP service fails.
+        requests.Timeout: If the request times out.
+    """
+    r = requests.get('https://checkip.amazonaws.com', timeout=30)
+    r.raise_for_status()
+    return r.content.decode().strip()
+
+
+def get_directory_size(directory: str) -> int:
+    """
+    Returns the total size (in bytes) of the target directory, including all subdirectories.
+
+    :param directory: Path to the target directory.
+    :return: Total size in bytes.
+    """
+    dir_path = Path(directory)
+    return sum(f.stat().st_size for f in dir_path.rglob("*") if f.is_file())
+
+
+def get_free_space_excluding_files(directory: str) -> int:
+    """
+    Returns the number of free bytes on the partition of the target directory,
+    excluding the total size of files in that directory.
+
+    :param directory: Path to the target directory.
+    :return: Free bytes available in the partition after subtracting file sizes.
+    """
+    # Get partition's free space
+    total, used, free = shutil.disk_usage(directory)
+
+    # Calculate the total size of files in the directory
+    files_size = get_directory_size(directory)
+
+    # Exclude file sizes from free space
+    return max(0, free - files_size)
+
+
+T = TypeVar('T')
+
+
+def batched(iterable: Iterable[T], n: int) -> Iterator[tuple[T, ...]]:
+    """
+    Implementation of itertools.batched for < Python 3.12
+    """
+    it = iter(iterable)
+    while True:
+        chunk = tuple(itertools.islice(it, n))
+        if not chunk:
+            break
+        yield chunk

mercuto_client/mocks.py

@@ -0,0 +1,203 @@
+import contextlib
+import dataclasses
+import re
+import uuid
+from typing import Any, Optional, Protocol
+
+from .exceptions import MercutoClientException
+from .types import Tenant, User, UserDetails, VerifyMeResult
+
+"""
+This module provides a context manager for mocking the Mercuto Client.
+
+Within the mock_client() context, any calls to the Mercuto Client will be intercepted and handled by the MercutoMocker.
+
+
+Usage:
+
+client = MercutoClient()
+
+with mock_client() as mock:
+    # Create a user in the system
+    api_key = mock.add_user()
+
+    # Login using the generated user key
+    client.connect(api_key = api_key)
+
+    # Call Mercuto Client endpoints as normal
+    client.identity().verify_me()
+
+"""
+
+
+class MakeRequestHookType(Protocol):
+    def __call__(self, method: str, url: str, *args, **kwargs) -> Any:
+        """
+        *args and **kwargs are arguments passed to requests.request(method, url, *args, **kwargs)
+        """
+        pass
+
+
+@contextlib.contextmanager
+def mock_client():
+    from .client import MercutoClient
+    original = MercutoClient._make_request
+    mocker = MercutoMocker()
+    try:
+        setattr(MercutoClient, '_make_request', mocker._on_make_request)
+        yield mocker
+    finally:
+        setattr(MercutoClient, '_make_request', original)
+
+
+class MercutoMocker:
+    @dataclasses.dataclass
+    class MockedUser:
+        code: str
+        username: str
+        description: str
+        tenant: str
+        permission_group: str
+        mobile_number: Optional[str] = None
+
+    def __init__(self) -> None:
+        self._hooks: dict[tuple[str, str], MakeRequestHookType] = {}
+
+        # Users based on Api-Key
+        self._known_users: dict[str, MercutoMocker.MockedUser] = {}
+        self._known_tenants: dict[str, Tenant] = {}
+
+        self._setup_default_hooks()
+
+    def on(self, method: str, path: str, callback: MakeRequestHookType):
+        self._hooks[(method, path)] = callback
+
+    def add_tenant(self, code: Optional[str] = None) -> None:
+        if code is None:
+            code = str(uuid.uuid4())
+        self._known_tenants[code] = Tenant(
+            code=code,
+            name=f"Tenant {code}",
+            description=f"Tenant {code}",
+            logo_url=None,
+        )
+
+    def add_user(self,
+                 user: Optional[str] = None,
+                 tenant: Optional[str] = None,
+                 permission_group: Optional[str] = None,
+                 key: Optional[str] = None,
+                 username: Optional[str] = None,
+                 description: Optional[str] = None,
+                 mobile_number: Optional[str] = None,
+                 ):
+        if user is None:
+            # User code
+            user = str(uuid.uuid4())
+        if tenant is None:
+            tenant = str(uuid.uuid4())
+        if permission_group is None:
+            permission_group = str(uuid.uuid4())
+        if key is None:
+            key = str(uuid.uuid4())
+        if username is None:
+            username = f"{user.replace(' ', '')}@example.com"
+        if description is None:
+            description = f"User {user}"
+
+        mocked = self.MockedUser(user, username, description, tenant, permission_group,
+                                 mobile_number=mobile_number)
+        self._known_users[key] = mocked
+        return key
+
+    def delete_user(self, key: Optional[str] = None, code: Optional[str] = None):
+        if key is not None and key in self._known_users:
+            del self._known_users[key]
+
+        if code is not None:
+            key = next((k for k, u in self._known_users.items()
+                       if u.code == code), None)
+            if key is not None:
+                del self._known_users[key]
+
+    def _mocked_get_tenant(self, method: str, url: str, *args, **kwargs) -> Tenant:
+        key = kwargs.get('headers', {}).get('X-Api-Key', None)
+        if key is None:
+            raise MercutoClientException("No X-Api-Key header provided")
+        tenant = self._known_tenants.get(key, None)
+        if tenant is None:
+            raise MercutoClientException("Tenant not found")
+        return tenant
+
+    def _mocked_verify_me(self, method: str, url: str, *args, **kwargs) -> VerifyMeResult:
+        key = kwargs.get('headers', {}).get('X-Api-Key', None)
+        if key is None:
+            raise MercutoClientException("No X-Api-Key header provided")
+        user = self._known_users.get(key, None)
+        if user is None:
+            raise MercutoClientException(f"User {user} not found")
+        return VerifyMeResult(
+            user=user.code,
+            tenant=user.tenant,
+            permission_group=user.permission_group,
+            acl_policy='{"version": 1, "permissions": []}'
+        )
+
+    def _mocked_get_user(self, method: str, url: str, *args, **kwargs) -> User:
+        apikey = kwargs.get('headers', {}).get('X-Api-Key', None)
+        servicekey = kwargs.get('headers', {}).get('X-Service-Token', None)
+        if apikey is None and servicekey is None:
+            raise MercutoClientException(
+                "No X-Api-Key or X-Service-Token header provided")
+        user_code = url.split('/')[-1]
+        user = next((u for u in self._known_users.values()
+                    if u.code == user_code), None)
+        if user is None:
+            raise MercutoClientException("User not found")
+        return User(
+            code=user.code,
+            username=user.username,
+            description=user.description,
+            tenant=user.tenant,
+            permission_group=user.permission_group,
+        )
+
+    def _mocked_get_user_details(self, method: str, url: str, *args, **kwargs) -> UserDetails:
+        apikey = kwargs.get('headers', {}).get('X-Api-Key', None)
+        servicekey = kwargs.get('headers', {}).get('X-Service-Token', None)
+        if apikey is None and servicekey is None:
+            raise MercutoClientException(
+                "No X-Api-Key or X-Service-Token header provided")
+        url_parts = url.split('/')
+        assert url_parts[-1] == 'details'
+        user_code = url_parts[-2]
+        user = next((u for u in self._known_users.values()
+                    if u.code == user_code), None)
+        if user is None:
+            raise MercutoClientException("User not found")
+        return UserDetails(
+            code=user.code,
+            username=user.username,
+            mobile_number=user.mobile_number,
+            email_address=None,
+            first_name=None,
+            last_name=None,
+            api_keys=[]
+        )
+
+    def _setup_default_hooks(self) -> None:
+        self.on('GET', r'\/identity\/verify\/me', self._mocked_verify_me)
+        self.on('GET', r'\/identity\/users\/[^\/]+', self._mocked_get_user)
+        self.on('GET', r'\/identity\/tenant\/[^\/]+', self._mocked_get_tenant)
+        self.on(
+            'GET', r'\/identity\/users\/[^\/]+\/details', self._mocked_get_user_details)
+
+    def _on_make_request(self, method: str, url: str, *args, **kwargs) -> Any:
+
+        # First check any custom hooks
+        for (hook_method, pattern), callback in self._hooks.items():
+            if method == hook_method and re.fullmatch(pattern, url) is not None:
+                return callback(method, url, *args, **kwargs)
+
+        raise NotImplementedError(
+            "Mocking is not supported for this endpoint: %s %s" % (method, url))