pybutt 2.0.0__py3-none-any.whl

pybutt/core/config.py

@@ -0,0 +1,144 @@
+import re
+from dataclasses import dataclass
+from enum import StrEnum
+
+from pybutt.exceptions import (
+    EngineSelectionError,
+    InvalidIdentifierError,
+    InvalidParameterError,
+)
+
+ENGINE_CHOICES = frozenset({"duckdb", "pyodbc", "mssql-python"})
+
+
+class TransactionMode(StrEnum):
+    """Control how transactions are handled during import."""
+
+    BATCH = "batch"  # Each batch of batch_size rows in its own transaction
+    ROWGROUP = "rowgroup"  # Each row group in the parquet file in its own transaction
+    FILE = "file"  # Entire file in one transaction
+
+
+# Global defaults
+DRIVER_DEFAULT = "ODBC Driver 18 for SQL Server"
+SCHEMA_DEFAULT = "dbo"
+TRUSTED_CONNECTION_DEFAULT = False
+TRUST_CERT_DEFAULT = False
+ENCRYPT_DEFAULT = True
+RETRIES_DEFAULT = 3
+
+# Default memory heartbeat interval in seconds. Set to 30 so operators always
+# have a recent RSS breadcrumb trail when a worker is OOM-killed.
+MEM_HEARTBEAT_DEFAULT: float = 30.0
+
+# Default memory-pressure throttle threshold (% system memory used). When system
+# memory exceeds this %, workers sleep until pressure drops. Set to 85% so OOM
+# kill is avoided without throttling during normal operation.
+MEM_THRESHOLD_DEFAULT: float = 85.0
+
+# Seconds to sleep per throttle cycle and max total wait before giving up.
+MEM_SLEEP_DEFAULT: float = 5.0
+MEM_MAX_WAIT_DEFAULT: float = 300.0
+
+# Cooldown seconds after a throttle event before the gate re-checks. Prevents
+# the gate from firing on every loop iteration and serialising workers.
+MEM_COOLDOWN_DEFAULT: float = 30.0
+
+# Default TDS packet size in bytes. 16383 is the maximum for encrypted
+# connections (SQL Server caps encrypted packets at this size). Valid range
+# for all drivers is 512–32767.
+PACKET_SIZE_DEFAULT: int = 4_096
+
+# Import specific defaults
+IMPORT_ENGINE_DEFAULT = "mssql-python"
+BATCH_SIZE_DEFAULT = 1_000
+TRANSACTION_MODE_DEFAULT = TransactionMode.ROWGROUP
+CCI_DEFAULT = True
+
+# Export specific defaults
+EXPORT_ENGINE_DEFAULT = "pyodbc"
+FETCH_SIZE_DEFAULT = 1_000
+ROWGROUP_SIZE_DEFAULT = 1_048_576
+
+
+def validate_engine(engine: str, allowed: frozenset[str] | None = None) -> str:
+    """Raise :class:`EngineSelectionError` if *engine* is not in *allowed*."""
+    choices = allowed if allowed is not None else ENGINE_CHOICES
+    if engine not in choices:
+        raise EngineSelectionError(f"engine must be one of {sorted(choices)}")
+    return engine
+
+
+def coerce_transaction_mode(mode: TransactionMode | str) -> TransactionMode:
+    """Accept a :class:`TransactionMode` or its string value and return the enum."""
+    if isinstance(mode, str):
+        return TransactionMode(mode)
+    return mode
+
+
+IDENTIFIER_REGEX = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+
+# TVF parameters must be a comma-separated list of literals:
+# integers/decimals, single-quoted strings (no nested quotes), NULLs.
+_PARAM_TOKEN_RE = re.compile(
+    r"\s*(?:" r"NULL" r"|[+-]?\d+(?:\.\d+)?" r"|'[^']*'" r")\s*",
+    re.IGNORECASE,
+)
+
+
+def validate_parameters(params: str) -> str:
+    """Reject parameter strings that could contain SQL injection payloads.
+
+    Accepts only comma-separated SQL literals: numbers, single-quoted
+    strings (no embedded quotes), and NULL.
+    """
+    tokens = params.split(",")
+    for token in tokens:
+        if not _PARAM_TOKEN_RE.fullmatch(token):
+            raise InvalidParameterError(
+                f"Unsafe TVF parameter token: {token.strip()!r}. "
+                "Only numeric literals, single-quoted strings, and NULL are allowed."
+            )
+    return params
+
+
+def validate_identifier(name: str) -> str:
+    if not IDENTIFIER_REGEX.match(name):
+        raise InvalidIdentifierError(f"Invalid identifier: {name}")
+    return name
+
+
+def quote_identifier(name: str) -> str:
+    return f"[{name.replace(']', ']]')}]"
+
+
+def sanitise_dsn_value(value: str) -> str:
+    """Escape ODBC connection-string metacharacters in a value.
+
+    Braces and semicolons are special in ODBC DSN strings. If the value
+    contains any of them, wrap it in ``{…}`` (doubling any literal
+    ``}`` inside) so the driver interprets the whole token as one value.
+    """
+    if not value:
+        return value
+    if any(ch in value for ch in (";", "{", "}", "=")):
+        return "{" + value.replace("}", "}}") + "}"
+    return value
+
+
+@dataclass
+class SqlConfig:
+    server: str
+    database: str
+    username: str | None = None
+    password: str | None = None
+    driver: str = DRIVER_DEFAULT
+    trusted_connection: bool = TRUSTED_CONNECTION_DEFAULT
+    trust_cert: bool = TRUST_CERT_DEFAULT
+    encrypt: bool = ENCRYPT_DEFAULT
+    retries: int = RETRIES_DEFAULT
+    packet_size: int = PACKET_SIZE_DEFAULT
+
+
+if __name__ == "__main__":
+    pass

pybutt/core/logobs.py

@@ -0,0 +1,445 @@
+"""Centralised logging/observability helpers for PyButt.
+
+All PyButt modules log through the ``pybutt`` logger (via :func:`get_logger`)
+rather than the root logger. The CLI calls :func:`configure_logging` once at
+startup; spawned export worker processes call it again through the pool
+initialiser (see ``Exporter.perform_work``) so their output is formatted
+identically on every platform (``spawn`` is the default on Windows/macOS and is
+forced here on all OSes).
+
+Library/API users who want PyButt's formatted output should call
+:func:`configure_logging` themselves; otherwise standard ``logging`` rules apply.
+"""
+
+import logging
+import threading
+import time as _time
+
+import psutil
+
+LOGGER_NAME = "pybutt"
+
+# Timestamp + level + process/thread identity so concurrent workers' lines can be
+# told apart and ordered. Identity matters because a single import run fans out
+# across threads and an export run across (spawned) processes.
+LOG_FORMAT = (
+    "%(asctime)s %(levelname)s [%(processName)s/%(threadName)s] %(name)s: %(message)s"
+)
+DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
+
+
+def get_logger(name: str | None = None) -> logging.Logger:
+    """Return a child of the ``pybutt`` logger (or the root pybutt logger)."""
+    if name is None:
+        return logging.getLogger(LOGGER_NAME)
+    return logging.getLogger(f"{LOGGER_NAME}.{name}")
+
+
+def configure_logging(verbose: bool = False) -> logging.Logger:
+    """Configure the ``pybutt`` logger. Idempotent and safe to call repeatedly.
+
+    Adds a single stderr handler with :data:`LOG_FORMAT`, sets the level
+    (``DEBUG`` when ``verbose`` else ``INFO``), and disables propagation so we
+    don't double-emit through the root logger or fight library handlers.
+    """
+    logger = logging.getLogger(LOGGER_NAME)
+    logger.setLevel(logging.DEBUG if verbose else logging.INFO)
+
+    if not any(getattr(h, "_pybutt_handler", False) for h in logger.handlers):
+        handler = logging.StreamHandler()
+        handler.setFormatter(logging.Formatter(LOG_FORMAT, DATE_FORMAT))
+        handler._pybutt_handler = True  # marker so we never add a duplicate
+        logger.addHandler(handler)
+
+    logger.propagate = False
+    return logger
+
+
+def init_worker_logging(level: int) -> None:
+    """Pool initialiser: configure logging inside a spawned worker process."""
+    configure_logging(verbose=level <= logging.DEBUG)
+    logging.getLogger(LOGGER_NAME).setLevel(level)
+
+
+def context(**fields: object) -> str:
+    """Render structured ``key=value`` context, skipping ``None`` values.
+
+    Example: ``context(file="a.parquet", rg="3/40", batch=12)`` ->
+    ``"file=a.parquet rg=3/40 batch=12"``.
+    """
+    return " ".join(
+        f"{key}={value}" for key, value in fields.items() if value is not None
+    )
+
+
+# --- memory observability --------------------------------------------------
+#
+# psutil gives a uniform *current* RSS on Windows/Linux/BSD/macOS (stdlib
+# ``resource`` is Unix-only and its units differ by OS). There is no portable
+# "peak RSS", so we track a running peak ourselves, per process. Export workers
+# are separate processes, so each tracks (and reports) its own peak.
+
+_process = psutil.Process()
+_peak_rss = 0
+
+
+def _human_bytes(num: float) -> str:
+    """Render a byte count compactly, e.g. ``1.8GB`` / ``512.0MB`` / ``900B``."""
+    value = float(num)
+    for unit in ("B", "KB", "MB", "GB"):
+        if value < 1024 or unit == "GB":
+            return f"{int(value)}{unit}" if unit == "B" else f"{value:.1f}{unit}"
+        value /= 1024
+    return f"{value:.1f}GB"
+
+
+def rss_bytes() -> int:
+    """Return current process RSS in bytes, updating the per-process peak.
+
+    Returns 0 if the platform/process info is unavailable, so logging never
+    fails because of a memory probe.
+    """
+    global _peak_rss
+    try:
+        rss = _process.memory_info().rss
+    except Exception:
+        return 0
+    if rss > _peak_rss:
+        _peak_rss = rss
+    return rss
+
+
+def peak_rss_bytes() -> int:
+    """Return the highest RSS observed in this process (refreshes first)."""
+    rss_bytes()
+    return _peak_rss
+
+
+def sys_mem_fields() -> dict[str, str]:
+    """System-wide memory fields for :func:`context`.
+
+    Returns ``{"sys_pct": "78%", "sys_avail": "4.2GB"}`` so log lines show
+    how close the *machine* is to the OOM-kill threshold — not just this
+    process's own RSS.
+    """
+    try:
+        vm = psutil.virtual_memory()
+        return {
+            "sys_pct": f"{vm.percent:.0f}%",
+            "sys_avail": _human_bytes(vm.available),
+        }
+    except Exception:
+        return {}
+
+
+def mem_fields() -> dict[str, str]:
+    """RSS + system-wide memory fields for :func:`context`.
+
+    Splat into ``context`` at boundary log points so the last line before an
+    OOM-kill shows the memory trend and exactly where it died, e.g.::
+
+        context(file=fn, rows=n, **mem_fields())
+    """
+    rss = rss_bytes()
+    return {
+        "rss": _human_bytes(rss),
+        "peak": _human_bytes(_peak_rss),
+        **sys_mem_fields(),
+    }
+
+
+def log_memory_budget(
+    *,
+    operation: str,
+    workers: int,
+    total_rows: int | None = None,
+    threshold_pct: float = 0,
+) -> None:
+    """Log a pre-flight memory budget so operators can gauge headroom.
+
+    Called once before ``perform_work()`` begins real processing. The
+    estimate is deliberately rough (and labelled as such) — it exists to
+    surface an immediate "this probably won't fit" signal, not to be
+    precise.
+    """
+    log = get_logger("budget")
+    try:
+        vm = psutil.virtual_memory()
+    except Exception:
+        return
+
+    avail = vm.available
+    total = vm.total
+    pct = vm.percent
+
+    parts = [
+        f"operation={operation}",
+        f"workers={workers}",
+        f"sys_total={_human_bytes(total)}",
+        f"sys_avail={_human_bytes(avail)}",
+        f"sys_pct={pct:.0f}%",
+    ]
+    if total_rows is not None:
+        parts.append(f"total_rows={total_rows}")
+    if threshold_pct > 0:
+        headroom_bytes = int(total * (1 - threshold_pct / 100)) - (total - avail)
+        parts.append(f"threshold={threshold_pct:.0f}%")
+        parts.append(f"headroom={_human_bytes(max(headroom_bytes, 0))}")
+
+    log.info("Memory budget " + " ".join(parts))
+
+
+def log_failure_summary(
+    *,
+    operation: str,
+    workers: int,
+    completed: list[str] | None = None,
+    failed_error: str = "",
+) -> None:
+    """Log a structured post-mortem when a pool/executor fails.
+
+    Gives the operator a concise picture of what finished before the
+    failure so they know how much progress was lost.
+    """
+    log = get_logger("postmortem")
+    try:
+        vm = psutil.virtual_memory()
+        sys_info = f"sys_pct={vm.percent:.0f}% sys_avail={_human_bytes(vm.available)}"
+    except Exception:
+        sys_info = ""
+
+    completed = completed or []
+    parts = [
+        f"operation={operation}",
+        f"workers={workers}",
+        f"completed={len(completed)}/{workers}",
+    ]
+    if sys_info:
+        parts.append(sys_info)
+    if failed_error:
+        parts.append(f"error={failed_error}")
+
+    log.error("FAILURE SUMMARY " + " ".join(parts))
+    if completed:
+        log.error(f"  Completed units: {', '.join(completed)}")
+
+
+class MemoryHeartbeat:
+    """Periodically log process RSS while a long operation runs.
+
+    Use as a context manager. A no-op when ``interval <= 0`` so callers can pass
+    a user-configured value unconditionally. The thread is a daemon and is
+    stopped/joined on exit. Runs in whichever process enters it, so for export
+    it must be entered inside the worker (where the memory actually lives).
+    """
+
+    def __init__(
+        self,
+        interval: float,
+        unit: str | None = None,
+        progress: dict[str, object] | None = None,
+    ):
+        self.interval = interval or 0
+        self.progress = progress
+        self.unit = unit
+        self._stop = threading.Event()
+        self._thread: threading.Thread | None = None
+
+    def __enter__(self) -> "MemoryHeartbeat":
+        if self.interval > 0:
+            self._thread = threading.Thread(
+                target=self._run, name="mem-heartbeat", daemon=True
+            )
+            self._thread.start()
+        return self
+
+    def __exit__(self, *exc: object) -> bool:
+        self._stop.set()
+        if self._thread is not None:
+            self._thread.join(timeout=self.interval + 1)
+        return False
+
+    def _run(self) -> None:
+        log = get_logger("mem")
+        while not self._stop.wait(self.interval):
+            extra = dict(self.progress) if self.progress else {}
+            log.info(
+                "Memory heartbeat " + context(unit=self.unit, **extra, **mem_fields())
+            )
+
+
+class WorkerMonitor:
+    """Monitor child worker processes from the parent and log their RSS.
+
+    Runs a daemon thread that polls each worker PID via ``psutil``. When a
+    worker disappears (e.g. OOM-killed by SIGKILL), the monitor logs the last
+    known RSS and system memory state so the operator has a breadcrumb trail
+    even though the child had no chance to log anything itself.
+
+    Use as a context manager. A no-op when ``interval <= 0``.
+    """
+
+    def __init__(self, pids: list[int], interval: float):
+        self.interval = interval or 0
+        self._pids = list(pids)
+        self._stop = threading.Event()
+        self._thread: threading.Thread | None = None
+        self._last_rss: dict[int, int] = {}
+
+    def __enter__(self) -> "WorkerMonitor":
+        if self.interval > 0 and self._pids:
+            self._thread = threading.Thread(
+                target=self._run, name="worker-monitor", daemon=True
+            )
+            self._thread.start()
+        return self
+
+    def __exit__(self, *exc: object) -> bool:
+        self._stop.set()
+        if self._thread is not None:
+            self._thread.join(timeout=self.interval + 1)
+        return False
+
+    def _run(self) -> None:
+        log = get_logger("monitor")
+        procs: dict[int, psutil.Process] = {}
+        for pid in self._pids:
+            try:
+                procs[pid] = psutil.Process(pid)
+            except (psutil.NoSuchProcess, psutil.AccessDenied):
+                pass
+
+        gone: set[int] = set()
+
+        while not self._stop.wait(self.interval):
+            sys_fields = sys_mem_fields()
+            for pid in self._pids:
+                if pid in gone:
+                    continue
+                proc = procs.get(pid)
+                if proc is None:
+                    gone.add(pid)
+                    log.warning(
+                        "Worker vanished "
+                        + context(
+                            pid=pid,
+                            last_rss=_human_bytes(self._last_rss.get(pid, 0)),
+                            status="GONE",
+                            **sys_fields,
+                        )
+                        + " — likely OOM-killed"
+                    )
+                    continue
+                try:
+                    rss = proc.memory_info().rss
+                    self._last_rss[pid] = rss
+                    log.debug(
+                        "Worker health "
+                        + context(
+                            pid=pid,
+                            rss=_human_bytes(rss),
+                            status="alive",
+                            **sys_fields,
+                        )
+                    )
+                except psutil.NoSuchProcess:
+                    gone.add(pid)
+                    log.warning(
+                        "Worker vanished "
+                        + context(
+                            pid=pid,
+                            last_rss=_human_bytes(self._last_rss.get(pid, 0)),
+                            status="GONE",
+                            **sys_fields,
+                        )
+                        + " — likely OOM-killed"
+                    )
+                except (psutil.AccessDenied, Exception):
+                    pass
+
+
+class MemoryGate:
+    """Cooperative throttle: sleep the caller when system memory is high.
+
+    Call :meth:`check` at natural pause points in hot loops (before a
+    ``fetchmany``, ``read_row_group``, etc.). When system memory exceeds
+    *threshold_pct*, the caller sleeps in increments of *sleep_seconds*
+    until memory drops or *max_wait* is exhausted.
+
+    After a throttle event (or max_wait timeout), a cooldown of
+    *cooldown_seconds* prevents the gate from re-triggering on every
+    subsequent loop iteration — allowing workers to make real progress
+    between checks.
+
+    A no-op when ``threshold_pct <= 0``.
+    """
+
+    def __init__(
+        self,
+        threshold_pct: float = 0.0,
+        sleep_seconds: float = 5.0,
+        max_wait: float = 300.0,
+        cooldown_seconds: float = 30.0,
+    ):
+        self.threshold_pct = threshold_pct
+        self.sleep_seconds = sleep_seconds
+        self.max_wait = max_wait
+        self.cooldown_seconds = cooldown_seconds
+        self._log = get_logger("gate")
+        self._enabled = threshold_pct > 0
+        self._last_release: float = 0.0
+
+    def check(self, context_msg: str = "") -> float:
+        """Block while system memory exceeds the threshold.
+
+        Returns the total seconds waited (0.0 if no throttling occurred).
+        Skips the check entirely if still within the cooldown window from
+        the last throttle event.
+        """
+        if not self._enabled:
+            return 0.0
+
+        now = _time.monotonic()
+        if now - self._last_release < self.cooldown_seconds:
+            return 0.0
+
+        waited = 0.0
+        vm = psutil.virtual_memory()
+        if vm.percent <= self.threshold_pct:
+            return 0.0
+
+        self._log.warning(
+            "Memory pressure — throttling "
+            + context(
+                reason=context_msg or "gate",
+                threshold=f"{self.threshold_pct:.0f}%",
+                **mem_fields(),
+            )
+        )
+
+        while vm.percent > self.threshold_pct and waited < self.max_wait:
+            _time.sleep(self.sleep_seconds)
+            waited += self.sleep_seconds
+            vm = psutil.virtual_memory()
+            self._log.info(
+                "Throttle wait "
+                + context(
+                    waited=f"{waited:.0f}s",
+                    sys_pct=f"{vm.percent:.0f}%",
+                    threshold=f"{self.threshold_pct:.0f}%",
+                    sys_avail=_human_bytes(vm.available),
+                )
+            )
+
+        if waited > 0:
+            self._log.info(
+                "Throttle released "
+                + context(
+                    total_waited=f"{waited:.0f}s",
+                    sys_pct=f"{vm.percent:.0f}%",
+                    sys_avail=_human_bytes(vm.available),
+                )
+            )
+
+        self._last_release = _time.monotonic()
+        return waited

pybutt/exceptions.py

@@ -0,0 +1,82 @@
+class PyButtError(Exception):
+    """Base class for all PyButt-specific errors."""
+
+
+class ConfigurationError(PyButtError, ValueError):
+    """Raised for invalid application configuration."""
+
+
+class EngineSelectionError(ConfigurationError):
+    """Raised when an unsupported engine is selected."""
+
+
+class InvalidIdentifierError(ConfigurationError):
+    """Raised when a SQL identifier is invalid."""
+
+
+class InvalidParameterError(ConfigurationError):
+    """Raised when a TVF parameter string contains unsafe content."""
+
+
+class ManifestError(PyButtError, ValueError):
+    """Base class for manifest validation errors."""
+
+
+class ManifestNotFoundError(FileNotFoundError, ManifestError):
+    """Raised when a manifest file cannot be found."""
+
+
+class InvalidManifestError(ManifestError):
+    """Raised when a manifest file contains invalid data."""
+
+
+class InvalidManifestEntryError(InvalidManifestError):
+    """Raised when a manifest entry is malformed."""
+
+
+class DuplicateManifestEntryError(InvalidManifestError):
+    """Raised when a manifest contains duplicate file entries."""
+
+
+class UnsupportedManifestVersionError(InvalidManifestError):
+    """Raised when a manifest has an unsupported version."""
+
+
+class UnsupportedManifestTypeError(InvalidManifestError):
+    """Raised when a manifest type is not supported."""
+
+
+class MissingManifestEntryError(FileNotFoundError, InvalidManifestError):
+    """Raised when a manifest references a missing Parquet file."""
+
+
+class PathTraversalError(InvalidManifestError):
+    """Raised when a manifest entry resolves outside its base directory."""
+
+
+class SchemaMismatchError(PyButtError, ValueError):
+    """Raised when Parquet schema does not match the destination table schema."""
+
+
+class DataExportError(PyButtError, RuntimeError):
+    """Raised when exporting data fails."""
+
+
+class DataImportError(PyButtError, RuntimeError):
+    """Raised when importing data fails."""
+
+
+class BatchImportError(DataImportError):
+    """Raised when a batch import fails after retries."""
+
+
+class RowGroupImportError(DataImportError):
+    """Raised when a row group import fails after retries."""
+
+
+class RetryExceededError(PyButtError, RuntimeError):
+    """Raised when retry logic exhausts all attempts."""
+
+
+class TableEmptyError(DataExportError):
+    """Raised when the source table is empty or missing."""

pybutt/files/__init__.py

@@ -0,0 +1,28 @@
+from .combine import combine_parquet_files
+from .inspect import inspect_manifest, inspect_parquet_file
+from .manifest import (
+    MANIFEST_VERSION_1,
+    MANIFEST_VERSION_2,
+    SUPPORTED_MANIFEST_TYPES,
+    default_import_manifest_filename,
+    default_manifest_filename,
+    load_file_manifest,
+    load_manifest,
+    validate_manifest_entries,
+    write_manifest,
+)
+
+__all__ = [
+    "MANIFEST_VERSION_1",
+    "MANIFEST_VERSION_2",
+    "SUPPORTED_MANIFEST_TYPES",
+    "default_manifest_filename",
+    "default_import_manifest_filename",
+    "load_file_manifest",
+    "load_manifest",
+    "validate_manifest_entries",
+    "write_manifest",
+    "inspect_manifest",
+    "inspect_parquet_file",
+    "combine_parquet_files",
+]