playground-ls-cli 4.14.1.dev8__py3-none-any.whl

localstack_cli/utils/archives.py

@@ -0,0 +1,271 @@
+import glob
+import io
+import logging
+import os
+import re
+import tarfile
+import tempfile
+import time
+import zipfile
+from subprocess import Popen
+from typing import IO, Literal
+
+from localstack_cli.constants import MAVEN_REPO_URL
+from localstack_cli.utils.files import load_file, mkdir, new_tmp_file, rm_rf, save_file
+from localstack_cli.utils.http import download
+from localstack_cli.utils.run import run
+
+from .checksum import verify_local_file_with_checksum_url
+from .run import is_command_available
+from .strings import truncate
+
+LOG = logging.getLogger(__name__)
+
+
+StrPath = str | os.PathLike
+
+
+def is_zip_file(content):
+    stream = io.BytesIO(content)
+    return zipfile.is_zipfile(stream)
+
+
+def get_unzipped_size(zip_file: str | IO[bytes]):
+    """Returns the size of the unzipped file."""
+    with zipfile.ZipFile(zip_file, "r") as zip_ref:
+        return sum(f.file_size for f in zip_ref.infolist())
+
+
+def unzip(path: str, target_dir: str, overwrite: bool = True) -> str | Popen | None:
+    from localstack_cli.utils.platform import is_debian
+
+    use_native_cmd = is_debian() or is_command_available("unzip")
+    if use_native_cmd:
+        # Running the native command can be an order of magnitude faster in the container. Also, `unzip`
+        #  is capable of extracting zip files with incorrect CRC codes (sometimes happens, e.g., with some
+        #  Node.js/Serverless versions), which can fail with Python's `zipfile` (extracting empty files).
+        flags = ["-o"] if overwrite else []
+        flags += ["-q"]
+        try:
+            cmd = ["unzip"] + flags + [path]
+            return run(cmd, cwd=target_dir, print_error=False)
+        except Exception as e:
+            error_str = truncate(str(e), max_length=200)
+            LOG.info(
+                'Unable to use native "unzip" command (using fallback mechanism): %s', error_str
+            )
+
+    try:
+        zip_ref = zipfile.ZipFile(path, "r")
+    except Exception as e:
+        LOG.warning("Unable to open zip file: %s: %s", path, e)
+        raise e
+
+    def _unzip_file_entry(zip_ref, file_entry, target_dir):
+        """Extracts a Zipfile entry and preserves permissions"""
+        out_path = os.path.join(target_dir, file_entry.filename)
+        if use_native_cmd and os.path.exists(out_path) and os.path.getsize(out_path) > 0:
+            # this can happen under certain circumstances if the native "unzip" command
+            # fails with a non-zero exit code, yet manages to extract parts of the zip file
+            return
+        zip_ref.extract(file_entry.filename, path=target_dir)
+        perm = file_entry.external_attr >> 16
+        # Make sure to preserve file permissions in the zip file
+        # https://www.burgundywall.com/post/preserving-file-perms-with-python-zipfile-module
+        os.chmod(out_path, perm or 0o777)
+
+    try:
+        for file_entry in zip_ref.infolist():
+            _unzip_file_entry(zip_ref, file_entry, target_dir)
+    finally:
+        zip_ref.close()
+
+
+def untar(path: str, target_dir: str):
+    mode = "r:gz" if path.endswith("gz") else "r"
+    with tarfile.open(path, mode) as tar:
+        tar.extractall(path=target_dir)
+
+
+def create_zip_file_cli(source_path: StrPath, base_dir: StrPath, zip_file: StrPath):
+    """
+    Creates a zip archive by using the native zip command. The native command can be an order of magnitude faster in CI
+    """
+    source = "." if source_path == base_dir else os.path.basename(source_path)
+    run(["zip", "-r", zip_file, source], cwd=base_dir)
+
+
+def create_zip_file_python(
+    base_dir: StrPath,
+    zip_file: StrPath,
+    mode: Literal["r", "w", "x", "a"] = "w",
+    content_root: str | None = None,
+):
+    with zipfile.ZipFile(zip_file, mode) as zip_file:
+        for root, dirs, files in os.walk(base_dir):
+            for name in files:
+                full_name = os.path.join(root, name)
+                relative = os.path.relpath(root, start=base_dir)
+                if content_root:
+                    dest = os.path.join(content_root, relative, name)
+                else:
+                    dest = os.path.join(relative, name)
+                zip_file.write(full_name, dest)
+
+
+def add_file_to_jar(class_file, class_url, target_jar, base_dir=None):
+    base_dir = base_dir or os.path.dirname(target_jar)
+    patch_class_file = os.path.join(base_dir, class_file)
+    if not os.path.exists(patch_class_file):
+        download(class_url, patch_class_file)
+        run(["zip", target_jar, class_file], cwd=base_dir)
+
+
+def update_jar_manifest(
+    jar_file_name: str, parent_dir: str, search: str | re.Pattern, replace: str
+):
+    manifest_file_path = "META-INF/MANIFEST.MF"
+    jar_path = os.path.join(parent_dir, jar_file_name)
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_manifest_file = os.path.join(tmp_dir, manifest_file_path)
+        run(["unzip", "-o", jar_path, manifest_file_path], cwd=tmp_dir)
+        manifest = load_file(tmp_manifest_file)
+
+    # return if the search pattern does not match (for idempotence, to avoid file permission issues further below)
+    if isinstance(search, re.Pattern):
+        if not search.search(manifest):
+            return
+        manifest = search.sub(replace, manifest, 1)
+    else:
+        if search not in manifest:
+            return
+        manifest = manifest.replace(search, replace, 1)
+
+    manifest_file = os.path.join(parent_dir, manifest_file_path)
+    save_file(manifest_file, manifest)
+    run(["zip", jar_file_name, manifest_file_path], cwd=parent_dir)
+
+
+def upgrade_jar_file(base_dir: str, file_glob: str, maven_asset: str):
+    """
+    Upgrade the matching Java JAR file in a local directory with the given Maven asset
+    :param base_dir: base directory to search the JAR file to replace in
+    :param file_glob: glob pattern for the JAR file to replace
+    :param maven_asset: name of Maven asset to download, in the form "<qualified_name>:<version>"
+    """
+
+    local_path = os.path.join(base_dir, file_glob)
+    parent_dir = os.path.dirname(local_path)
+    maven_asset = maven_asset.replace(":", "/")
+    parts = maven_asset.split("/")
+    maven_asset_url = f"{MAVEN_REPO_URL}/{maven_asset}/{parts[-2]}-{parts[-1]}.jar"
+    target_file = os.path.join(parent_dir, os.path.basename(maven_asset_url))
+    if os.path.exists(target_file):
+        # avoid re-downloading the newer JAR version if it already exists locally
+        return
+    matches = glob.glob(local_path)
+    if not matches:
+        return
+    for match in matches:
+        os.remove(match)
+    download(maven_asset_url, target_file)
+
+
+def download_and_extract(
+    archive_url: str,
+    target_dir: str,
+    retries: int | None = 0,
+    sleep: int | None = 3,
+    tmp_archive: str | None = None,
+    checksum_url: str | None = None,
+) -> None:
+    """
+    Download and extract an archive to a target directory with optional checksum verification.
+
+    Checksum verification is only performed if a `checksum_url` is provided.
+    Else, the archive is downloaded and extracted without verification.
+
+    :param archive_url: URL of the archive to download
+    :param target_dir: Directory to extract the archive contents to
+    :param retries: Number of download retries (default: 0)
+    :param sleep: Sleep time between retries in seconds (default: 3)
+    :param tmp_archive: Optional path for the temporary archive file
+    :param checksum_url: Optional URL of the checksum file for verification
+    :return: None
+    """
+    mkdir(target_dir)
+
+    _, ext = os.path.splitext(tmp_archive or archive_url)
+    tmp_archive = tmp_archive or new_tmp_file()
+    if not os.path.exists(tmp_archive) or os.path.getsize(tmp_archive) <= 0:
+        # create temporary placeholder file, to avoid duplicate parallel downloads
+        save_file(tmp_archive, "")
+
+        for i in range(retries + 1):
+            try:
+                download(archive_url, tmp_archive)
+                break
+            except Exception as e:
+                LOG.warning(
+                    "Attempt %d. Failed to download archive from %s: %s",
+                    i + 1,
+                    archive_url,
+                    e,
+                )
+                # only sleep between retries, not after the last one
+                if i < retries:
+                    time.sleep(sleep)
+
+    # if the temporary file we created above hasn't been replaced, we assume failure
+    if os.path.getsize(tmp_archive) <= 0:
+        raise Exception("Failed to download archive from %s: . Retries exhausted", archive_url)
+
+    # Verify checksum if provided
+    if checksum_url:
+        LOG.info("Verifying archive integrity...")
+        try:
+            verify_local_file_with_checksum_url(
+                file_path=tmp_archive,
+                checksum_url=checksum_url,
+            )
+        except Exception as e:
+            # clean up the corrupted download
+            rm_rf(tmp_archive)
+            raise e
+
+    if ext in (".zip", ".whl"):
+        unzip(tmp_archive, target_dir)
+    elif ext in (
+        ".bz2",
+        ".gz",
+        ".tgz",
+        ".xz",
+    ):
+        untar(tmp_archive, target_dir)
+    else:
+        raise Exception(f"Unsupported archive format: {ext}")
+
+
+def download_and_extract_with_retry(
+    archive_url,
+    tmp_archive,
+    target_dir,
+    checksum_url: str | None = None,
+):
+    try:
+        download_and_extract(
+            archive_url,
+            target_dir,
+            tmp_archive=tmp_archive,
+            checksum_url=checksum_url,
+        )
+    except Exception as e:
+        # try deleting and re-downloading the zip file
+        LOG.info("Unable to extract file, re-downloading ZIP archive %s: %s", tmp_archive, e)
+        rm_rf(tmp_archive)
+        download_and_extract(
+            archive_url,
+            target_dir,
+            tmp_archive=tmp_archive,
+            checksum_url=checksum_url,
+        )

localstack_cli/utils/batching.py

@@ -0,0 +1,258 @@
+import copy
+import logging
+import threading
+import time
+from typing import Generic, Protocol, TypeVar, overload
+
+LOG = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+# alias to signify whether a batch policy has been triggered
+BatchPolicyTriggered = bool
+
+
+# TODO: Add batching on bytes as well.
+class Batcher(Generic[T]):
+    """
+    A utility for collecting items into batches and flushing them when one or more batch policy conditions are met.
+
+    The batch policy can be created to trigger on:
+    - max_count: Maximum number of items added
+    - max_window: Maximum time window (in seconds)
+
+    If no limits are specified, the batcher is always in triggered state.
+
+    Example usage:
+
+        import time
+
+        # Triggers when 2 (or more) items are added
+        batcher = Batcher(max_count=2)
+        assert batcher.add(["item1", "item2", "item3"])
+        assert batcher.flush() == ["item1", "item2", "item3"]
+
+        # Triggers partially when 2 (or more) items are added
+        batcher = Batcher(max_count=2)
+        assert batcher.add(["item1", "item2", "item3"])
+        assert batcher.flush(partial=True) == ["item1", "item2"]
+        assert batcher.add("item4")
+        assert batcher.flush(partial=True) == ["item3", "item4"]
+
+        # Trigger 2 seconds after the first add
+        batcher = Batcher(max_window=2.0)
+        assert not batcher.add(["item1", "item2", "item3"])
+        time.sleep(2.1)
+        assert not batcher.add(["item4"])
+        assert batcher.flush() == ["item1", "item2", "item3", "item4"]
+    """
+
+    max_count: int | None
+    """
+    Maximum number of items, must be None or positive.
+    """
+
+    max_window: float | None
+    """
+    Maximum time window in seconds, must be None or positive.
+    """
+
+    _triggered: bool
+    _last_batch_time: float
+    _batch: list[T]
+
+    def __init__(self, max_count: int | None = None, max_window: float | None = None):
+        """
+        Initialize a new Batcher instance.
+
+        :param max_count: Maximum number of items that be None or positive.
+        :param max_window: Maximum time window in seconds that must be None or positive.
+        """
+        self.max_count = max_count
+        self.max_window = max_window
+
+        self._triggered = False
+        self._last_batch_time = time.monotonic()
+        self._batch = []
+
+    @property
+    def period(self) -> float:
+        return time.monotonic() - self._last_batch_time
+
+    def _check_batch_policy(self) -> bool:
+        """Check if any batch policy conditions are met"""
+        if self.max_count is not None and len(self._batch) >= self.max_count:
+            self._triggered = True
+        elif self.max_window is not None and self.period >= self.max_window:
+            self._triggered = True
+        elif not self.max_count and not self.max_window:
+            # always return true
+            self._triggered = True
+
+        return self._triggered
+
+    @overload
+    def add(self, item: T, *, deep_copy: bool = False) -> BatchPolicyTriggered: ...
+
+    @overload
+    def add(self, items: list[T], *, deep_copy: bool = False) -> BatchPolicyTriggered: ...
+
+    def add(self, item_or_items: T | list[T], *, deep_copy: bool = False) -> BatchPolicyTriggered:
+        """
+        Add an item or list of items to the collected batch.
+
+        Returns:
+            BatchPolicyTriggered: True if the batch policy was triggered during addition, False otherwise.
+        """
+        if deep_copy:
+            item_or_items = copy.deepcopy(item_or_items)
+
+        if isinstance(item_or_items, list):
+            self._batch.extend(item_or_items)
+        else:
+            self._batch.append(item_or_items)
+
+        # Check if the last addition triggered the batch policy
+        return self.is_triggered()
+
+    def flush(self, *, partial=False) -> list[T]:
+        result = []
+        if not partial or not self.max_count:
+            result = self._batch.copy()
+            self._batch.clear()
+        else:
+            batch_size = min(self.max_count, len(self._batch))
+            result = self._batch[:batch_size].copy()
+            self._batch = self._batch[batch_size:]
+
+        self._last_batch_time = time.monotonic()
+        self._triggered = False
+        self._check_batch_policy()
+
+        return result
+
+    def duration_until_next_batch(self) -> float:
+        if not self.max_window:
+            return -1
+        return max(self.max_window - self.period, -1)
+
+    def get_current_size(self) -> int:
+        return len(self._batch)
+
+    def is_triggered(self):
+        return self._triggered or self._check_batch_policy()
+
+
+class BatchHandler(Protocol[T]):
+    """
+    A BatchHandler is a callable that processes a list of items handed down by the AsyncBatcher.
+    """
+
+    def __call__(self, batch: list[T]) -> None: ...
+
+
+class AsyncBatcher(Generic[T]):
+    """
+    Class for managing asynchronous batching of items.
+
+    This class allows for efficient buffering and processing of items in batches by
+    periodically flushing the buffer to a given handler, or by automatically flushing
+    when the maximum batch size is reached. It is designed to be used in asynchronous
+    scenarios where the caller does not execute the flushing IO call itself, like with ``Batcher``.
+
+    :ivar max_flush_interval: Maximum time interval in seconds between
+        automatic flushes, regardless of the batch size.
+    :ivar max_batch_size: Maximum number of items in a batch. When reached,
+        the batch is flushed automatically.
+    :ivar handler: Callable handler that processes each flushed batch. The handler must
+        be provided during initialization and must accept a list of items as input.
+    """
+
+    max_flush_interval: float
+    max_batch_size: int
+    handler: BatchHandler[T]
+
+    _buffer: list[T]
+    _flush_lock: threading.Condition
+    _closed: bool
+
+    def __init__(
+        self,
+        handler: BatchHandler[T],
+        max_flush_interval: float = 10,
+        max_batch_size: int = 20,
+    ):
+        self.handler = handler
+        self.max_flush_interval = max_flush_interval
+        self.max_batch_size = max_batch_size
+
+        self._buffer = []
+        self._flush_lock = threading.Condition()
+        self._closed = False
+
+    def add(self, item: T):
+        """
+        Adds an item to the buffer.
+
+        :param item: the item to add
+        """
+        with self._flush_lock:
+            if self._closed:
+                raise ValueError("Batcher is stopped, can no longer add items")
+
+            self._buffer.append(item)
+
+            if len(self._buffer) >= self.max_batch_size:
+                self._flush_lock.notify_all()
+
+    @property
+    def current_batch_size(self) -> int:
+        """
+        Returns the current number of items in the buffer waiting to be flushed.
+        """
+        return len(self._buffer)
+
+    def run(self):
+        """
+        Runs the event loop that flushes the buffer to the handler based on the configured rules, and blocks until
+        ``close()`` is called. This method is meant to be run in a separate thread.
+        """
+        while not self._closed:
+            with self._flush_lock:
+                # wait returns once either the condition is notified (in which case wait returns True, indicating that
+                # something has triggered a flush manually), or the timeout expires (in which case wait returns False)
+                self._flush_lock.wait(self.max_flush_interval)
+
+                # if _flush_condition was notified because close() was called, we should still make sure we flush the
+                # last batch
+
+                # perform the flush, if there are any items in the buffer
+                if not self._buffer:
+                    continue
+
+                batch = self._buffer.copy()
+                self._buffer.clear()
+
+            # we can call the processor outside the lock so we can continue adding items into the next batch without
+            # waiting on the processor to return.
+            try:
+                self.handler(batch)
+            except Exception as e:
+                LOG.error(
+                    "Unhandled exception while processing a batch: %s",
+                    e,
+                    exc_info=LOG.isEnabledFor(logging.DEBUG),
+                )
+
+        # this marks that the main control loop is done
+        return
+
+    def close(self):
+        """
+        Triggers a close of the batcher, which will cause one last flush, and then end the main event loop.
+        """
+        with self._flush_lock:
+            if self._closed:
+                return
+            self._closed = True
+            self._flush_lock.notify_all()