pybutt 2.0.0__py3-none-any.whl

pybutt/files/combine.py

@@ -0,0 +1,93 @@
+from pathlib import Path
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from pybutt.exceptions import InvalidManifestError
+from pybutt.files.manifest import (
+    load_file_manifest,
+    validate_manifest_entries,
+    write_manifest,
+)
+
+
+def _write_table_chunks(writer, table, rowgroup_size: int):
+    if table.num_rows < rowgroup_size:
+        return table
+
+    offset = 0
+    while offset + rowgroup_size <= table.num_rows:
+        chunk = table.slice(offset, rowgroup_size)
+        writer.write_table(chunk, row_group_size=rowgroup_size)
+        offset += rowgroup_size
+
+    if offset < table.num_rows:
+        return table.slice(offset)
+
+    return None
+
+
+def combine_parquet_files(
+    manifest_path: Path,
+    output_file: Path,
+    rowgroup_size: int = 1_048_576,
+    delete_originals: bool = False,
+    new_manifest_name: str | None = None,
+):
+    """combine all parquet files listed in a manifest into a single Parquet file.
+
+    The resulting file will use the schema of the first file. All subsequent
+    files must be schema-compatible (column names/types) or behavior is undefined.
+    """
+    manifest_path = Path(manifest_path)
+    base_dir = manifest_path.parent
+
+    manifest = load_file_manifest(manifest_path, operation="combine")
+
+    entries = validate_manifest_entries(manifest, base_dir)
+    if not entries:
+        raise InvalidManifestError("Manifest contains no entries to combine")
+
+    output_file = Path(output_file)
+    output_file.parent.mkdir(parents=True, exist_ok=True)
+
+    # Use schema from first file
+    first_path = base_dir / entries[0]
+    first_pf = pq.ParquetFile(first_path)
+    schema = first_pf.schema_arrow
+
+    with pq.ParquetWriter(output_file, schema, compression="snappy") as writer:
+        buffered_table = None
+
+        for entry in entries:
+            src = base_dir / entry
+            pf = pq.ParquetFile(src)
+            # If schema differs, let pyarrow handle or raise downstream
+            for batch in pf.iter_batches():
+                table = pa.Table.from_batches([batch])
+                if buffered_table is None:
+                    buffered_table = table
+                else:
+                    buffered_table = pa.concat_tables([buffered_table, table])
+
+                buffered_table = _write_table_chunks(
+                    writer, buffered_table, rowgroup_size
+                )
+
+        if buffered_table is not None and buffered_table.num_rows > 0:
+            writer.write_table(buffered_table, row_group_size=rowgroup_size)
+
+    if delete_originals:
+        for entry in entries:
+            src = base_dir / entry
+            if src.exists() and src != output_file:
+                src.unlink()
+        if manifest_path.exists():
+            manifest_path.unlink()
+
+    new_manifest_name = (
+        new_manifest_name or f"{manifest_path.stem}_combined{manifest_path.suffix}"
+    )
+    write_manifest(base_dir / new_manifest_name, [output_file.name])
+
+    return output_file

pybutt/files/inspect.py

@@ -0,0 +1,51 @@
+from pathlib import Path
+
+import pyarrow.parquet as pq
+
+from pybutt.files.manifest import load_file_manifest
+
+
+def inspect_parquet_file(filepath: Path, verbose: bool = False) -> dict:
+    pf = pq.ParquetFile(filepath)
+
+    info = {
+        "file": filepath.name,
+        "rows": pf.metadata.num_rows,
+        "row_groups": pf.metadata.num_row_groups,
+        "row_group_sizes": {
+            pf.metadata.row_group(i).num_rows for i in range(pf.metadata.num_row_groups)
+        },
+    }
+
+    if verbose:
+        schema = pf.schema_arrow
+        info["columns"] = {field.name: str(field.type) for field in schema}
+
+    return info
+
+
+def inspect_manifest(manifest_path: str | Path, verbose: bool = False):
+    manifest_path = Path(manifest_path)
+    base_dir = manifest_path.parent
+
+    manifest = load_file_manifest(manifest_path, operation="Inspect")
+
+    for filename in manifest["entries"]:
+        filepath = base_dir / filename
+        if not filepath.exists():
+            print(f"Missing file: {filepath}")
+            continue
+
+        info = inspect_parquet_file(filepath, verbose=verbose)
+
+        print(info["file"])
+        print(f"  rows: {info['rows']}")
+        print(f"  row groups: {info['row_groups']}")
+        print(f"  group sizes: {info['row_group_sizes']}")
+
+        if verbose:
+            print("  columns:")
+            for col, typ in info["columns"].items():
+                print(f"    {col}: {typ}")
+
+        print()

pybutt/files/manifest.py

@@ -0,0 +1,160 @@
+import json
+from pathlib import Path
+
+from pybutt.core.config import validate_identifier
+from pybutt.exceptions import (
+    DuplicateManifestEntryError,
+    InvalidManifestEntryError,
+    InvalidManifestError,
+    ManifestNotFoundError,
+    MissingManifestEntryError,
+    PathTraversalError,
+    UnsupportedManifestTypeError,
+    UnsupportedManifestVersionError,
+)
+
+MANIFEST_VERSION_1 = 1
+MANIFEST_VERSION_2 = 2
+SUPPORTED_MANIFEST_TYPES = frozenset({"files", "tables"})
+
+
+def _parse_manifest_dict(data):
+    if not isinstance(data, dict):
+        raise InvalidManifestError(
+            "Manifest must be a list or an object with version, type, and entries"
+        )
+
+    version = data.get("version")
+    if version not in {MANIFEST_VERSION_1, MANIFEST_VERSION_2}:
+        raise UnsupportedManifestVersionError(
+            f"Unsupported manifest version: {version}"
+        )
+
+    manifest_type = data.get("type")
+    if version == MANIFEST_VERSION_1:
+        manifest_type = manifest_type or "files"
+        if manifest_type != "files":
+            raise UnsupportedManifestTypeError(
+                "Version 1 manifests support only type 'files'"
+            )
+    else:
+        if manifest_type not in SUPPORTED_MANIFEST_TYPES:
+            raise InvalidManifestError(
+                "Manifest type must be 'files' or 'tables' for version 2"
+            )
+
+    entries = data.get("entries")
+    if not isinstance(entries, list):
+        raise InvalidManifestError("Manifest entries must be a list")
+
+    if manifest_type == "tables":
+        return {
+            "version": version,
+            "type": manifest_type,
+            "entries": [_validate_table_name(e) for e in entries],
+        }
+
+    return {"version": version, "type": manifest_type, "entries": entries}
+
+
+def _validate_table_name(value):
+    if not isinstance(value, str):
+        raise InvalidManifestEntryError(
+            f"Invalid manifest table entry (not string): {value}"
+        )
+
+    parts = value.split(".")
+    if len(parts) != 2:
+        raise InvalidManifestEntryError(
+            f"Invalid table name format, expected schema.table: {value}"
+        )
+
+    schema, table = parts
+    validate_identifier(schema)
+    validate_identifier(table)
+    return value
+
+
+def default_manifest_filename(
+    schema: str, table: str, op_type: str = "", suffix: str = "manifest"
+) -> str:
+    return f"{schema}_{table}_{op_type}{suffix}.json"
+
+
+def default_import_manifest_filename(
+    schema: str,
+    table: str,
+) -> str:
+    return default_manifest_filename(schema=schema, table=table, op_type="import_")
+
+
+def write_manifest(
+    path: str | Path,
+    entries: list[str],
+    manifest_type: str = "files",
+    version: int = MANIFEST_VERSION_2,
+) -> Path:
+    """Write a versioned manifest JSON file and return its :class:`Path`."""
+    path = Path(path)
+    with open(path, "w") as f:
+        json.dump(
+            {"version": version, "type": manifest_type, "entries": entries},
+            f,
+            indent=4,
+        )
+    return path
+
+
+def load_manifest(manifest_path: str | Path) -> dict:
+    manifest_path = Path(manifest_path)
+
+    if not manifest_path.exists():
+        raise ManifestNotFoundError(f"Manifest not found: {manifest_path}")
+
+    with open(manifest_path) as f:
+        data = json.load(f)
+
+    if isinstance(data, list):
+        return {"version": MANIFEST_VERSION_1, "type": "files", "entries": data}
+
+    return _parse_manifest_dict(data)
+
+
+def load_file_manifest(
+    manifest_path: str | Path, *, operation: str = "Operation"
+) -> dict:
+    """Load a manifest and raise if it is not a file manifest."""
+    manifest = load_manifest(manifest_path)
+    if manifest["type"] != "files":
+        raise UnsupportedManifestTypeError(
+            f"{operation} only supports file manifests, got: {manifest['type']}"
+        )
+    return manifest
+
+
+def validate_manifest_entries(manifest: dict, base_dir: Path) -> list[str]:
+    seen = set()
+    validated = []
+
+    for item in manifest["entries"]:
+        if not isinstance(item, str):
+            raise InvalidManifestEntryError(
+                f"Invalid manifest entry (not string): {item}"
+            )
+
+        if item in seen:
+            raise DuplicateManifestEntryError(f"Duplicate file in manifest: {item}")
+
+        if manifest["type"] == "files":
+            filepath = (base_dir / item).resolve()
+            if not filepath.is_relative_to(base_dir.resolve()):
+                raise PathTraversalError(
+                    f"Manifest entry escapes base directory: {item}"
+                )
+            if not filepath.exists():
+                raise MissingManifestEntryError(f"Missing file: {filepath}")
+
+        seen.add(item)
+        validated.append(item)
+
+    return validated

pybutt/io/__init__.py

@@ -0,0 +1,6 @@
+from .combiner import TableCombine
+from .exporter import Exporter
+from .importer import Importer
+from .purger import TablePurger
+
+__all__ = ["Exporter", "Importer", "TableCombine", "TablePurger"]

pybutt/io/combiner.py

@@ -0,0 +1,119 @@
+from collections.abc import Iterable
+
+from pybutt.core.base import SqlServerIOBase
+from pybutt.core.config import (
+    SCHEMA_DEFAULT,
+    TransactionMode,
+    coerce_transaction_mode,
+    quote_identifier,
+    validate_engine,
+    validate_identifier,
+)
+from pybutt.core.logobs import context, get_logger
+from pybutt.exceptions import SchemaMismatchError
+
+logger = get_logger("combiner")
+
+
+class TableCombine(SqlServerIOBase):
+    """Combine multiple SQL tables into a single target table.
+
+    Sources should be provided as fully-qualified schema.table strings.
+    """
+
+    def __init__(
+        self,
+        config,
+        table: str,
+        sources: Iterable[str],
+        schema: str = SCHEMA_DEFAULT,
+        transaction_mode: TransactionMode = TransactionMode.BATCH,
+        engine: str = "pyodbc",
+    ):
+        super().__init__(config)
+        self.schema = validate_identifier(schema)
+        self.table = validate_identifier(table)
+        self.sources: list[str] = list(sources)
+        self.transaction_mode = coerce_transaction_mode(transaction_mode)
+        validate_engine(engine, allowed=frozenset({"pyodbc", "duckdb"}))
+        self.engine = engine
+
+    def _parse_schema_table(self, fq: str) -> tuple[str, str]:
+        parts = fq.split(".")
+        if len(parts) != 2:
+            raise ValueError(f"Invalid source table name: {fq}")
+        schema, table = parts
+        validate_identifier(schema)
+        validate_identifier(table)
+        return schema, table
+
+    def _ensure_target_exists_and_schema(self, cur, first_source: str):
+        # If target exists, validate schema equality;
+        # otherwise create from first source (no rows)
+        cur.execute("SELECT OBJECT_ID(?)", (f"{self.schema}.{self.table}",))
+        exists = cur.fetchone()[0] is not None
+
+        # Get column list for source
+        src_schema, src_table = self._parse_schema_table(first_source)
+        q_src = f"{quote_identifier(src_schema)}.{quote_identifier(src_table)}"
+        q_tgt = f"{quote_identifier(self.schema)}.{quote_identifier(self.table)}"
+        cur.execute(f"SELECT TOP 0 * FROM {q_src}")
+        src_cols = [c[0] for c in cur.description]
+
+        if not exists:
+            # Create target table with same schema as source
+            cur.execute(f"SELECT TOP 0 * INTO {q_tgt} FROM {q_src}")
+            cur.connection.commit()
+            return src_cols
+
+        # Target exists: get target columns
+        cur.execute(f"SELECT TOP 0 * FROM {q_tgt}")
+        tgt_cols = [c[0] for c in cur.description]
+
+        if set(src_cols) != set(tgt_cols):
+            raise SchemaMismatchError(
+                "Source and target schemas differ for "
+                f"{first_source} vs {self.schema}.{self.table}"
+            )
+
+        return src_cols
+
+    def combine(self):
+        """Combine all source tables into the target table.
+
+        Implementation: create target if missing using first source schema, then
+        run `INSERT INTO target SELECT * FROM source` for each source.
+        """
+        with self.connection_p(autocommit=False) as conn:
+            with conn.cursor() as cur:
+                # Ensure first source schema compatible / create target
+                first_source = self.sources[0]
+                self._ensure_target_exists_and_schema(
+                    cur, first_source, self.schema, self.table
+                )
+
+                # Insert from each source
+                q_tgt = (
+                    f"{quote_identifier(self.schema)}"
+                    f".{quote_identifier(self.table)}"
+                )
+                for src in self.sources:
+                    src_schema, src_table = self._parse_schema_table(src)
+                    q_src = (
+                        f"{quote_identifier(src_schema)}.{quote_identifier(src_table)}"
+                    )
+                    logger.info(
+                        "Combining "
+                        + context(
+                            source=f"{src_schema}.{src_table}",
+                            target=f"{self.schema}.{self.table}",
+                        )
+                    )
+                    try:
+                        cur.execute(f"INSERT INTO {q_tgt} SELECT * FROM {q_src}")
+                        conn.commit()
+                    except Exception:
+                        conn.rollback()
+                        raise
+
+        logger.info("Table combine completed")