pybutt 2.0.0__py3-none-any.whl

pybutt/cli/import_command.py

@@ -0,0 +1,286 @@
+from pathlib import Path
+
+import typer
+
+from pybutt.cli.app import (
+    app,
+    build_sql_config,
+)
+from pybutt.core.config import (
+    BATCH_SIZE_DEFAULT,
+    CCI_DEFAULT,
+    DRIVER_DEFAULT,
+    ENCRYPT_DEFAULT,
+    IMPORT_ENGINE_DEFAULT,
+    MEM_COOLDOWN_DEFAULT,
+    MEM_HEARTBEAT_DEFAULT,
+    MEM_MAX_WAIT_DEFAULT,
+    MEM_SLEEP_DEFAULT,
+    MEM_THRESHOLD_DEFAULT,
+    PACKET_SIZE_DEFAULT,
+    RETRIES_DEFAULT,
+    SCHEMA_DEFAULT,
+    TRANSACTION_MODE_DEFAULT,
+    TRUST_CERT_DEFAULT,
+    TRUSTED_CONNECTION_DEFAULT,
+    TransactionMode,
+)
+from pybutt.core.logobs import configure_logging, get_logger
+from pybutt.exceptions import PyButtError
+from pybutt.io.importer import Importer
+
+logger = get_logger("cli.import")
+
+
+@app.command(
+    "import",
+    help=("Import Parquet files into a SQL Server table using a manifest file."),
+)
+def import_data(
+    manifest_path: Path = typer.Argument(  # noqa: B008
+        ..., help="Path to the input manifest file"
+    ),
+    verbose: bool = typer.Option(  # noqa: B008
+        False,
+        "--verbose",
+        "-V",
+        help="Show verbose logging output.",
+    ),
+    server: str = typer.Option(  # noqa: B008
+        ...,
+        "--server",
+        "-s",
+        help="SQL Server hostname or instance.",
+        rich_help_panel="Server Connection Options",
+    ),
+    database: str = typer.Option(  # noqa: B008
+        ...,
+        "--database",
+        "-d",
+        help="Target SQL Server database.",
+        rich_help_panel="Server Connection Options",
+    ),
+    driver: str = typer.Option(  # noqa: B008
+        DRIVER_DEFAULT,
+        "--driver",
+        "-D",
+        help="ODBC driver name.",
+        rich_help_panel="Server Connection Options",
+    ),
+    engine: str = typer.Option(  # noqa: B008
+        IMPORT_ENGINE_DEFAULT,
+        "--engine",
+        "-e",
+        help="Import engine to use: duckdb, pyodbc, or mssql-python.",
+        rich_help_panel="Server Connection Options",
+        case_sensitive=False,
+    ),
+    transaction_mode: TransactionMode = typer.Option(  # noqa: B008
+        TRANSACTION_MODE_DEFAULT,
+        "--transaction-mode",
+        "-M",
+        help=(
+            "Transaction scope: batch (per batch), rowgroup (per row group, "
+            "recommended), file (entire file)."
+        ),
+        rich_help_panel="Server Connection Options",
+    ),
+    schema: str = typer.Option(  # noqa: B008
+        SCHEMA_DEFAULT,
+        "--schema",
+        "-S",
+        help="Target table schema.",
+        rich_help_panel="SQL Data Object Options",
+    ),
+    table: str = typer.Option(  # noqa: B008
+        ...,
+        "--table",
+        "-t",
+        help="Target table name.",
+        rich_help_panel="SQL Data Object Options",
+    ),
+    cci: bool = typer.Option(  # noqa: B008
+        CCI_DEFAULT,
+        "--cci/--no-cci",
+        help=(
+            "Create a clustered columnstore index on the per-worker temp tables "
+            "used during multi-worker import. Use --no-cci to keep the previous "
+            "heap behaviour. Enabled by default."
+        ),
+        rich_help_panel="SQL Data Object Options",
+    ),
+    username: str | None = typer.Option(  # noqa: B008
+        None,
+        "--username",
+        "-u",
+        help="SQL Server username when not using trusted connection.",
+        rich_help_panel="Server Security Options",
+    ),
+    password: str | None = typer.Option(  # noqa: B008
+        None,
+        "--password",
+        "-p",
+        help="SQL Server password when not using trusted connection.",
+        rich_help_panel="Server Security Options",
+    ),
+    trusted_connection: bool = typer.Option(  # noqa: B008
+        TRUSTED_CONNECTION_DEFAULT,
+        "--trusted-connection",
+        "-T",
+        help="Use integrated Windows authentication instead of username/password.",
+        rich_help_panel="Server Security Options",
+    ),
+    trust_cert: bool = typer.Option(  # noqa: B008
+        TRUST_CERT_DEFAULT,
+        "--trust-cert",
+        "-c",
+        help="Trust the SQL Server TLS certificate.",
+        rich_help_panel="Server Security Options",
+    ),
+    encrypt: bool = typer.Option(  # noqa: B008
+        ENCRYPT_DEFAULT,
+        "--encrypt/--no-encrypt",
+        help="Enable or disable SQL Server encrypted transport.",
+        rich_help_panel="Server Security Options",
+    ),
+    temp_manifest_filename: str | None = typer.Option(  # noqa: B008
+        None,
+        "--imported-manifest-filename",
+        "-o",
+        help=(
+            "Override the import worker manifest filename written during "
+            "multi-worker import. Defaults to <schema>_<table>_import_manifest.json."
+        ),
+        rich_help_panel="File Options",
+    ),
+    batch_size: int | None = typer.Option(  # noqa: B008
+        BATCH_SIZE_DEFAULT,
+        "--batch-size",
+        "-b",
+        help="Rows per batch insert.",
+        rich_help_panel="Transport Tuning Options",
+        min=1,
+    ),
+    retries: int = typer.Option(  # noqa: B008
+        RETRIES_DEFAULT,
+        "--retries",
+        "-r",
+        help="Number of retry attempts for transient SQL errors.",
+        rich_help_panel="Transport Tuning Options",
+        min=1,
+    ),
+    packet_size: int = typer.Option(  # noqa: B008
+        PACKET_SIZE_DEFAULT,
+        "--packet-size",
+        help=(
+            "TDS packet size in bytes (512-32767). "
+            "Note: encrypted connections are capped at 16383."
+        ),
+        rich_help_panel="Transport Tuning Options",
+        min=512,
+        max=32767,
+    ),
+    worker_count: int = typer.Option(  # noqa: B008
+        1,
+        "--worker-count",
+        "-w",
+        help="Number of parallel import threads.",
+        rich_help_panel="Transport Tuning Options",
+        min=1,
+    ),
+    mem_heartbeat: float = typer.Option(  # noqa: B008
+        MEM_HEARTBEAT_DEFAULT,
+        "--mem-heartbeat",
+        help=("Log process memory (RSS + system %) every N seconds."),
+        rich_help_panel="Memory Tuning Options",
+        min=0,
+    ),
+    mem_threshold: float = typer.Option(  # noqa: B008
+        MEM_THRESHOLD_DEFAULT,
+        "--mem-threshold",
+        help=(
+            "System memory % at which workers are throttled. "
+            "Set to 0 to disable throttling."
+        ),
+        rich_help_panel="Memory Tuning Options",
+        min=0,
+        max=100,
+    ),
+    mem_sleep: float = typer.Option(  # noqa: B008
+        MEM_SLEEP_DEFAULT,
+        "--mem-sleep",
+        help=("Seconds to sleep per throttle check when memory is high. "),
+        rich_help_panel="Memory Tuning Options",
+        min=0.1,
+    ),
+    mem_max_wait: float = typer.Option(  # noqa: B008
+        MEM_MAX_WAIT_DEFAULT,
+        "--mem-max-wait",
+        help=("Max total seconds to wait during memory throttling before giving up."),
+        rich_help_panel="Memory Tuning Options",
+        min=0,
+    ),
+    mem_cooldown: float = typer.Option(  # noqa: B008
+        MEM_COOLDOWN_DEFAULT,
+        "--mem-cooldown",
+        help=(
+            "Seconds after a throttle event before re-checking. Prevents "
+            "the gate from serialising workers"
+        ),
+        rich_help_panel="Memory Tuning Options",
+        min=0,
+    ),
+) -> None:
+    """Import one or more Parquet files into SQL Server tables.
+
+    The command reads the manifest file and imports each Parquet file into the
+    target table. If the number of workers is greater than 1, the import will be
+    done using multiple tables created to the same data schema as the target table.
+    """
+
+    configure_logging(verbose)
+
+    if mem_threshold > 0:
+        logger.info(
+            "Memory throttling enabled: threads will sleep when system "
+            f"memory exceeds {mem_threshold:.0f}%% "
+            f"(--mem-threshold 0 to disable)"
+        )
+
+    config = build_sql_config(
+        server=server,
+        database=database,
+        username=username,
+        password=password,
+        driver=driver,
+        trusted_connection=trusted_connection,
+        trust_cert=trust_cert,
+        encrypt=encrypt,
+        retries=retries,
+        packet_size=packet_size,
+    )
+
+    try:
+        importer = Importer(
+            config=config,
+            schema=schema,
+            table=table,
+            input_path=manifest_path.parent,
+            manifest_filename=manifest_path.name,
+            worker_count=worker_count,
+            batch_size=batch_size,
+            transaction_mode=transaction_mode,
+            engine=engine.lower(),
+            temp_manifest_filename=temp_manifest_filename,
+            create_cci=cci,
+            mem_heartbeat=mem_heartbeat,
+            mem_threshold=mem_threshold,
+            mem_sleep=mem_sleep,
+            mem_max_wait=mem_max_wait,
+            mem_cooldown=mem_cooldown,
+        )
+        importer.perform_work()
+    except PyButtError as exc:
+        typer.secho(f"Import failed: {exc}", fg=typer.colors.RED, err=True)
+        raise SystemExit(1) from exc
+    typer.secho("Import completed successfully.", fg=typer.colors.GREEN)

pybutt/cli/inspect_command.py

@@ -0,0 +1,30 @@
+from pathlib import Path
+
+import typer
+
+from pybutt.cli.app import app
+from pybutt.core.logobs import get_logger
+from pybutt.files import inspect_manifest
+
+logger = get_logger("cli.inspect")
+
+
+@app.command(
+    "inspect",
+    help=(
+        "Inspect Parquet files listed in a manifest. "
+        "Shows file-level metadata and optionally column-level details."
+    ),
+)
+def inspect(
+    manifest: Path = typer.Argument(  # noqa: B008
+        ..., help="Path to the input manifest file"
+    ),
+    verbose: bool = typer.Option(  # noqa: B008
+        False, "--verbose", "-V", help="Show column details"
+    ),
+):
+    """
+    Inspect parquet files listed in a manifest.
+    """
+    inspect_manifest(manifest, verbose)

pybutt/cli/purge_command.py

@@ -0,0 +1,235 @@
+from pathlib import Path
+
+import typer
+
+from pybutt.cli.app import app, build_sql_config
+from pybutt.core.config import (
+    DRIVER_DEFAULT,
+    ENCRYPT_DEFAULT,
+    PACKET_SIZE_DEFAULT,
+    RETRIES_DEFAULT,
+    TRUST_CERT_DEFAULT,
+    TRUSTED_CONNECTION_DEFAULT,
+)
+from pybutt.core.logobs import configure_logging, get_logger
+from pybutt.exceptions import PyButtError
+from pybutt.files import load_manifest
+from pybutt.io.purger import TablePurger
+
+logger = get_logger("cli.purge")
+
+
+@app.command(
+    "purge",
+    help=(
+        "Purge objects listed in a manifest. "
+        "For file manifests, deletes each Parquet file then removes the manifest. "
+        "For table manifests, drops each SQL table then removes the manifest."
+    ),
+)
+def purge(
+    manifest_path: Path = typer.Argument(  # noqa: B008
+        ..., help="Path to the input manifest file."
+    ),
+    verbose: bool = typer.Option(  # noqa: B008
+        False,
+        "--verbose",
+        "-V",
+        help="Show verbose logging output.",
+    ),
+    server: str | None = typer.Option(  # noqa: B008
+        None,
+        "--server",
+        "-s",
+        help="SQL Server host (required for table manifests).",
+        rich_help_panel="Server Connection Options",
+    ),
+    database: str | None = typer.Option(  # noqa: B008
+        None,
+        "--database",
+        "-d",
+        help="Target database (required for table manifests).",
+        rich_help_panel="Server Connection Options",
+    ),
+    driver: str = typer.Option(  # noqa: B008
+        DRIVER_DEFAULT,
+        "--driver",
+        "-D",
+        help="ODBC driver name.",
+        rich_help_panel="Server Connection Options",
+    ),
+    username: str | None = typer.Option(  # noqa: B008
+        None,
+        "--username",
+        "-u",
+        help="SQL Server username when not using trusted connection.",
+        rich_help_panel="Server Security Options",
+    ),
+    password: str | None = typer.Option(  # noqa: B008
+        None,
+        "--password",
+        "-p",
+        help="SQL Server password when not using trusted connection.",
+        rich_help_panel="Server Security Options",
+    ),
+    trusted_connection: bool = typer.Option(  # noqa: B008
+        TRUSTED_CONNECTION_DEFAULT,
+        "--trusted-connection",
+        "-T",
+        help="Use integrated Windows authentication instead of username/password.",
+        rich_help_panel="Server Security Options",
+    ),
+    trust_cert: bool = typer.Option(  # noqa: B008
+        TRUST_CERT_DEFAULT,
+        "--trust-cert",
+        "-c",
+        help="Trust the SQL Server TLS certificate.",
+        rich_help_panel="Server Security Options",
+    ),
+    encrypt: bool = typer.Option(  # noqa: B008
+        ENCRYPT_DEFAULT,
+        "--encrypt/--no-encrypt",
+        help="Enable or disable SQL Server encrypted transport.",
+        rich_help_panel="Server Security Options",
+    ),
+    retries: int = typer.Option(  # noqa: B008
+        RETRIES_DEFAULT,
+        "--retries",
+        "-r",
+        help="Number of retry attempts for transient SQL errors.",
+        rich_help_panel="Transport Tuning Options",
+        min=1,
+    ),
+    packet_size: int = typer.Option(  # noqa: B008
+        PACKET_SIZE_DEFAULT,
+        "--packet-size",
+        help=(
+            "TDS packet size in bytes (512-32767). "
+            "Note: encrypted connections are capped at 16383."
+        ),
+        rich_help_panel="Transport Tuning Options",
+        min=512,
+        max=32767,
+    ),
+) -> None:
+    """Purge objects listed in a manifest and delete the manifest file."""
+
+    configure_logging(verbose)
+
+    try:
+        manifest = load_manifest(manifest_path)
+    except PyButtError as exc:
+        typer.secho(f"Purge failed: {exc}", fg=typer.colors.RED, err=True)
+        raise SystemExit(1) from exc
+
+    if manifest["type"] == "files":
+        _purge_files(manifest, manifest_path)
+        return
+
+    if manifest["type"] == "tables":
+        _purge_tables(
+            manifest,
+            manifest_path,
+            server=server,
+            database=database,
+            driver=driver,
+            username=username,
+            password=password,
+            trusted_connection=trusted_connection,
+            trust_cert=trust_cert,
+            encrypt=encrypt,
+            retries=retries,
+            packet_size=packet_size,
+        )
+        return
+
+    typer.secho(
+        f"Purge failed: unsupported manifest type '{manifest['type']}'",
+        fg=typer.colors.RED,
+        err=True,
+    )
+    raise SystemExit(1)
+
+
+def _purge_files(manifest: dict, manifest_path: Path) -> None:
+    """Delete Parquet files listed in the manifest, then delete the manifest."""
+    base_dir = manifest_path.parent
+    entries = manifest["entries"]
+    deleted = 0
+    missing = 0
+
+    for entry in entries:
+        filepath = base_dir / entry
+        if filepath.exists():
+            filepath.unlink()
+            logger.info(f"Deleted file: {filepath}")
+            deleted += 1
+        else:
+            logger.warning(f"File not found (skipping): {filepath}")
+            missing += 1
+
+    manifest_path.unlink()
+    logger.info(f"Deleted manifest: {manifest_path}")
+
+    summary = f"Purge complete: {deleted} file(s) deleted"
+    if missing:
+        summary += f", {missing} file(s) not found (skipped)"
+    summary += ", manifest removed."
+    typer.secho(summary, fg=typer.colors.GREEN)
+
+
+def _purge_tables(
+    manifest: dict,
+    manifest_path: Path,
+    *,
+    server: str | None,
+    database: str | None,
+    driver: str,
+    username: str | None,
+    password: str | None,
+    trusted_connection: bool,
+    trust_cert: bool,
+    encrypt: bool,
+    retries: int,
+    packet_size: int,
+) -> None:
+    """Drop SQL tables listed in the manifest, then delete the manifest."""
+    if not (server and database):
+        raise typer.BadParameter(
+            "--server and --database are required for table manifests"
+        )
+
+    entries = manifest["entries"]
+    if not entries:
+        typer.secho("No tables to purge.", fg=typer.colors.YELLOW)
+        manifest_path.unlink()
+        logger.info(f"Deleted manifest: {manifest_path}")
+        return
+
+    config = build_sql_config(
+        server=server,
+        database=database,
+        username=username,
+        password=password,
+        driver=driver,
+        trusted_connection=trusted_connection,
+        trust_cert=trust_cert,
+        encrypt=encrypt,
+        retries=retries,
+        packet_size=packet_size,
+    )
+
+    try:
+        purger = TablePurger(config=config, sources=entries)
+        dropped = purger.purge()
+    except PyButtError as exc:
+        typer.secho(f"Purge failed: {exc}", fg=typer.colors.RED, err=True)
+        raise SystemExit(1) from exc
+
+    manifest_path.unlink()
+    logger.info(f"Deleted manifest: {manifest_path}")
+
+    typer.secho(
+        f"Purge complete: {len(dropped)} table(s) dropped, manifest removed.",
+        fg=typer.colors.GREEN,
+    )

pybutt/core/__init__.py

@@ -0,0 +1,30 @@
+from .config import (
+    SqlConfig,
+    TransactionMode,
+    coerce_transaction_mode,
+    quote_identifier,
+    sanitise_dsn_value,
+    validate_engine,
+    validate_identifier,
+    validate_parameters,
+)
+from .logobs import (
+    configure_logging,
+    get_logger,
+)
+
+__all__ = [
+    # Config - types
+    "SqlConfig",
+    "TransactionMode",
+    # Config - validators
+    "coerce_transaction_mode",
+    "quote_identifier",
+    "sanitise_dsn_value",
+    "validate_engine",
+    "validate_parameters",
+    "validate_identifier",
+    # Logging
+    "configure_logging",
+    "get_logger",
+]

pybutt/core/base.py

@@ -0,0 +1,124 @@
+import re
+import time
+
+import duckdb as d
+import mssql_python
+import pyodbc
+
+from pybutt.exceptions import ConfigurationError, RetryExceededError
+
+from .config import (
+    SqlConfig,
+    quote_identifier,
+    sanitise_dsn_value,
+)
+from .logobs import get_logger
+
+logger = get_logger("base")
+
+
+def rows_from_arrow(arrow_obj) -> list[tuple]:
+    """Convert an Arrow Table or RecordBatch to a list of row-tuples.
+
+    Works with both ``pyarrow.Table`` and ``pyarrow.RecordBatch`` (anything
+    with a ``.columns`` attribute whose elements support ``.to_pylist()``).
+    """
+    return list(zip(*[col.to_pylist() for col in arrow_obj.columns], strict=True))
+
+
+class SqlServerIOBase:
+    def __init__(self, config: SqlConfig):
+        self.config = config
+        self.dsn = self.build_dsn()
+
+    def _connection_parts(self, *, include_driver: bool = True) -> list[str]:
+        """Build the common connection-string parts shared by all drivers."""
+        cfg = self.config
+        parts: list[str] = []
+        if include_driver:
+            parts.append(f"Driver={{{cfg.driver}}}")
+        parts.append(f"Server={sanitise_dsn_value(cfg.server)}")
+        parts.append(f"Database={sanitise_dsn_value(cfg.database)}")
+
+        if cfg.trusted_connection:
+            parts.append("Trusted_Connection=Yes")
+        else:
+            if include_driver:
+                if not cfg.username or not cfg.password:
+                    raise ConfigurationError(
+                        "Username/password required when not using trusted connection"
+                    )
+                parts.append(f"Uid={sanitise_dsn_value(cfg.username)}")
+                parts.append(f"Pwd={sanitise_dsn_value(cfg.password)}")
+            else:
+                if cfg.username:
+                    parts.append(f"UID={sanitise_dsn_value(cfg.username)}")
+                if cfg.password:
+                    parts.append(f"PWD={sanitise_dsn_value(cfg.password)}")
+
+        parts.append(f"TrustServerCertificate={'Yes' if cfg.trust_cert else 'No'}")
+
+        if cfg.encrypt:
+            parts.append("Encrypt=Yes")
+
+        parts.append(f"PacketSize={cfg.packet_size}")
+
+        return parts
+
+    def build_dsn(self):
+        return ";".join(self._connection_parts(include_driver=True)) + ";"
+
+    def connection_d(self):
+        conn = d.connect()
+        conn.execute("INSTALL odbc_scanner; LOAD odbc_scanner;")
+        return conn
+
+    def connection_p(self, autocommit=False):
+        conn = pyodbc.connect(self.dsn)
+        conn.autocommit = autocommit
+        return conn
+
+    def connection_m(self, autocommit=False):
+        conn_str = ";".join(self._connection_parts(include_driver=False)) + ";"
+        conn = mssql_python.connect(conn_str)
+        conn.setautocommit(autocommit)
+        return conn
+
+    def full_table_name(self):
+        return f"{quote_identifier(self.schema)}.{quote_identifier(self.table)}"
+
+    def safe_error_message(self, e: Exception) -> str:
+        msg = str(e)
+
+        # redact common sensitive tokens
+        msg = re.sub(r"(Pwd|Password)=[^;]+", r"\1=***", msg, flags=re.IGNORECASE)
+        msg = re.sub(r"(Uid|User ID)=[^;]+", r"\1=***", msg, flags=re.IGNORECASE)
+
+        return msg
+
+    def retry(self, fn, context="operation"):
+        last_error: Exception | None = None
+        for attempt in range(self.config.retries):
+            try:
+                return fn()
+            except MemoryError:
+                logger.error(f"{context} out of memory - not retrying (fatal)")
+                raise
+            except Exception as e:
+                last_error = e
+                safe_msg = self.safe_error_message(e)
+                logger.warning(
+                    f"{context} attempt {attempt + 1}/{self.config.retries} "
+                    f"failed: {safe_msg}"
+                )
+                time.sleep(2**attempt)
+        if last_error is not None:
+            raise RetryExceededError(
+                f"{context} failed after max retries: "
+                f"{self.safe_error_message(last_error)}"
+            ) from last_error
+        raise RetryExceededError(f"{context} failed after max retries")
+
+
+if __name__ == "__main__":
+    pass