sunstone-data-platform 1.0.0__py3-none-any.whl

sunstone_data/import_/parser.py

@@ -0,0 +1,492 @@
+"""CSV and Excel parsing to PyArrow, with Iceberg table creation."""
+
+import json
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import fastexcel
+import pyarrow as pa
+import pyarrow.csv as pa_csv
+from pyiceberg.exceptions import BadRequestError
+from pyiceberg.io.fsspec import FsspecFileIO
+from pyiceberg.table import Table
+
+from sunstone_data.catalog import create_namespace_if_not_exists, get_catalog
+from sunstone_data.config import Settings, get_settings
+from sunstone_data.import_.fetcher import FileType, SourceInfo
+from sunstone_data.import_.pipeline import ImportResult
+from sunstone_data.metadata.provenance import (
+    make_activity_uri,
+    make_file_uri,
+    make_table_uri,
+    make_url_uri,
+)
+from sunstone_data.utils import sanitize_table_name
+
+if TYPE_CHECKING:
+    from sunstone_data.metadata import MetadataBuilder
+
+
+def _fix_table_io_endpoint(table: Table, settings: Settings) -> Table:
+    """Fix table IO endpoint for local development.
+
+    When running PyIceberg on the host against Nessie in Docker, Nessie vends
+    the internal Docker endpoint. This function replaces the table's IO with
+    one using the host-accessible endpoint.
+
+    Args:
+        table: PyIceberg Table object loaded from REST catalog.
+        settings: Platform settings with correct endpoint.
+
+    Returns:
+        The same table with IO endpoint fixed.
+    """
+    props = dict(table.io.properties)
+    changed = False
+
+    if settings.storage_backend == "gcs":
+        # Fix fake-gcs-server internal endpoint
+        gcs_host = props.get("gcs.service.host", "")
+        if "fake-gcs-server" in gcs_host and settings.gcs_endpoint:
+            props["gcs.service.host"] = settings.gcs_endpoint
+            changed = True
+    else:
+        # Fix the object store's internal endpoint
+        s3_endpoint = props.get("s3.endpoint", "")
+        if "object-store:9000" in s3_endpoint:
+            props["s3.endpoint"] = settings.s3_endpoint
+            changed = True
+
+    if changed:
+        new_io = FsspecFileIO(properties=props)
+        object.__setattr__(table, "io", new_io)
+
+    return table
+
+
+def parse_csv(file_path: Path, encoding: str | None = None) -> pa.Table:
+    """Parse CSV file to PyArrow table with streaming.
+
+    Args:
+        file_path: Path to CSV file.
+        encoding: File encoding (auto-detect if None).
+
+    Returns:
+        PyArrow Table with parsed data.
+    """
+    read_options = pa_csv.ReadOptions(
+        encoding=encoding or "utf-8",
+        block_size=64 * 1024 * 1024,  # 64MB chunks for type inference
+    )
+    parse_options = pa_csv.ParseOptions()
+    convert_options = pa_csv.ConvertOptions()
+
+    # Use streaming reader
+    reader = pa_csv.open_csv(
+        file_path,
+        read_options=read_options,
+        parse_options=parse_options,
+        convert_options=convert_options,
+    )
+
+    return reader.read_all()
+
+
+def parse_excel(file_path: Path) -> dict[str, pa.Table]:
+    """Parse Excel file, returning dict of sheet_name -> PyArrow table.
+
+    Args:
+        file_path: Path to Excel file.
+
+    Returns:
+        Dict mapping sheet names to PyArrow tables.
+    """
+    excel_file = fastexcel.read_excel(file_path)
+    sheet_names = excel_file.sheet_names
+
+    tables = {}
+    for sheet_name in sheet_names:
+        sheet = excel_file.load_sheet_by_name(sheet_name)
+        arrow_data = sheet.to_arrow()
+
+        # fastexcel may return RecordBatch - convert to Table
+        if isinstance(arrow_data, pa.RecordBatch):
+            arrow_table = pa.Table.from_batches([arrow_data])
+        else:
+            arrow_table = arrow_data
+
+        # Skip empty sheets
+        if arrow_table.num_rows > 0:
+            tables[sheet_name] = arrow_table
+
+    return tables
+
+
+def create_iceberg_table(
+    namespace: str,
+    table_name: str,
+    arrow_table: pa.Table,
+    properties: dict[str, str],
+    branch: str = "main",
+    force_schema: bool = False,
+    settings: Settings | None = None,
+    builder: "MetadataBuilder | None" = None,
+) -> tuple[int | None, bool]:
+    """Create Iceberg table from Arrow table with metadata.
+
+    If table exists with same checksum, returns without modification.
+    If table exists with different checksum, appends as new snapshot.
+
+    Args:
+        namespace: Target namespace.
+        table_name: Target table name.
+        arrow_table: Data to write.
+        properties: Metadata properties to set.
+        branch: Branch to import to (default: "main").
+        force_schema: Allow unsafe schema changes (default: False).
+        settings: Optional settings override.
+        builder: Optional MetadataBuilder. On a same-checksum no-op its
+            ``activity_id`` is reset to the one the original committing import
+            recorded on the table, so a retry does not mint a second activity
+            for a snapshot that never changed (DP-27).
+
+    Returns:
+        ``(snapshot_id, already_existed)``. The snapshot id is the table's
+        current one after the operation — for an unchanged (same-checksum)
+        table, the existing one; None only if the table somehow carries no
+        snapshot. ``already_existed`` is True only when the same-checksum
+        short-circuit fired and nothing was written.
+
+    Raises:
+        ValueError: The schema is not flat, or an unsafe schema change was
+            detected without ``force_schema``.
+    """
+    settings = settings or get_settings()
+    catalog = get_catalog(branch=branch)
+
+    # Ensure namespace exists
+    create_namespace_if_not_exists(catalog, namespace)
+
+    identifier = f"{namespace}.{table_name}"
+
+    # Check if table exists
+    existing_tables = [t for _, t in catalog.list_tables(namespace)]
+
+    if table_name in existing_tables:
+        # Table exists - check checksum
+        table = catalog.load_table(identifier)
+        existing_checksum = table.metadata.properties.get("sunstone.import.checksum")
+
+        if existing_checksum == properties.get("sunstone.import.checksum"):
+            # Same checksum - already imported, skip. Adopt the committing
+            # import's activity id so provenance for this unchanged snapshot
+            # keeps pointing at the activity that actually produced it.
+            if builder is not None:
+                committed_activity_id = table.metadata.properties.get("sunstone.import.activity_id")
+                if committed_activity_id:
+                    builder.activity_id = committed_activity_id
+            return table.metadata.current_snapshot_id, True
+
+        # Fix IO endpoint for local development (the catalog vends object-store:9000)
+        table = _fix_table_io_endpoint(table, settings)
+
+        # Check for schema changes
+        existing_schema = table.schema().as_arrow()
+        new_schema = arrow_table.schema
+
+        if not existing_schema.equals(new_schema):
+            # Schema differs - validate evolution
+            from sunstone_data.versioning.schema_evolution import (
+                detect_schema_changes,
+                validate_schema_evolution,
+            )
+
+            changes = detect_schema_changes(existing_schema, new_schema)
+
+            try:
+                validate_schema_evolution(changes, force=force_schema)
+            except ValueError as e:
+                # Unsafe schema changes detected
+                raise ValueError(
+                    f"Schema change detected that could cause data loss. {e}\nUse --force-schema to override."
+                ) from e
+
+        # Different checksum - append as new snapshot (version)
+        table.append(arrow_table)
+
+        # Update properties (a property-only commit adds no snapshot)
+        with table.transaction() as txn:
+            for key, value in properties.items():
+                txn.set_properties({key: value})
+        return table.metadata.current_snapshot_id, False
+
+    # Create new table. The catalog enforces the flat-schema rule server-side.
+    try:
+        table = catalog.create_table(
+            identifier=identifier,
+            schema=arrow_table.schema,
+            properties=properties,
+        )
+    except BadRequestError as e:
+        raise ValueError(
+            f"Cannot create table {identifier!r}: Iceberg tables must have a flat schema "
+            f"(no struct, list, or map columns). {e}"
+        ) from e
+
+    # Fix IO endpoint for local development (the catalog vends object-store:9000)
+    table = _fix_table_io_endpoint(table, settings)
+
+    table.append(arrow_table)
+    return table.metadata.current_snapshot_id, False
+
+
+def schema_summary(arrow_table: pa.Table) -> str:
+    """Create human-readable schema summary.
+
+    Args:
+        arrow_table: PyArrow table.
+
+    Returns:
+        String like "col1 (int64), col2 (string), col3 (float64)"
+    """
+    parts = []
+    for field in arrow_table.schema:
+        parts.append(f"{field.name} ({field.type})")
+    return ", ".join(parts[:5]) + ("..." if len(parts) > 5 else "")
+
+
+def _extract_schema_columns(arrow_table: pa.Table) -> list[dict]:
+    """Extract column metadata from PyArrow schema.
+
+    Args:
+        arrow_table: PyArrow table.
+
+    Returns:
+        List of {name: str, type: str} for each column.
+    """
+    return [{"name": field.name, "type": str(field.type)} for field in arrow_table.schema]
+
+
+def _build_extended_properties(
+    base_properties: dict[str, str],
+    arrow_table: pa.Table,
+    namespace: str,
+    table_name: str,
+    builder: "MetadataBuilder | None",
+) -> dict[str, str]:
+    """Build extended Iceberg properties including PROV-DM URIs.
+
+    Args:
+        base_properties: Base properties dict to extend.
+        arrow_table: Parsed PyArrow table.
+        namespace: Target namespace.
+        table_name: Target table name.
+        builder: Optional MetadataBuilder with source info.
+
+    Returns:
+        Extended properties dict including sunstone.prov.* fields.
+    """
+    properties = dict(base_properties)
+
+    if builder is None:
+        return properties
+
+    # Schema columns (same for original and iceberg in direct import)
+    columns = _extract_schema_columns(arrow_table)
+
+    # Extended sunstone.import.* properties
+    properties.update(
+        {
+            "sunstone.import.activity_id": builder.activity_id,
+            "sunstone.import.user": builder.user,
+            "sunstone.import.system_version": builder.system_version,
+            "sunstone.import.row_count": str(arrow_table.num_rows),
+            "sunstone.import.column_count": str(arrow_table.num_columns),
+            "sunstone.import.schema_original": json.dumps(columns),
+            "sunstone.import.schema_iceberg": json.dumps(columns),
+        }
+    )
+
+    # Add source_normalized and file_size if source info is available
+    if builder.source_info is not None:
+        properties["sunstone.import.source_normalized"] = builder.source_info.normalized_path
+        properties["sunstone.import.file_size"] = str(builder.source_info.file_size)
+
+        # Determine source entity URI
+        original_path = builder.source_info.original_path
+        if original_path.startswith(("http://", "https://")):
+            source_uri = make_url_uri(original_path)
+        else:
+            source_uri = make_file_uri(builder.source_info.normalized_path)
+
+        # PROV-DM properties
+        properties.update(
+            {
+                "sunstone.prov.activity_uri": make_activity_uri(builder.activity_id),
+                "sunstone.prov.source_entity_uri": source_uri,
+                "sunstone.prov.table_entity_uri": make_table_uri(namespace, table_name),
+                "sunstone.prov.user_agent_uri": f"urn:user:{builder.user}",
+                "sunstone.prov.system_agent_uri": f"urn:system:sunstone-data:{builder.system_version}",
+            }
+        )
+
+    return properties
+
+
+def parse_and_create_tables(
+    source_info: SourceInfo,
+    base_name: str,
+    namespace: str,
+    checksum: str,
+    source: str,
+    tags: list[str] | None = None,
+    branch: str = "main",
+    force_schema: bool = False,
+    settings: Settings | None = None,
+    builder: "MetadataBuilder | None" = None,
+    owner_org: str | None = None,
+) -> list[ImportResult]:
+    """Parse file and create Iceberg tables.
+
+    Args:
+        source_info: Source information from fetcher.
+        base_name: Base table name.
+        namespace: Target namespace.
+        checksum: File checksum.
+        source: Original source path/URL.
+        tags: Optional tags.
+        branch: Branch to import to (default: "main").
+        force_schema: Allow unsafe schema changes (default: False).
+        settings: Optional settings override.
+        builder: Optional MetadataBuilder for extended metadata capture.
+
+    Returns:
+        List of ImportResult (one for CSV, multiple for Excel sheets).
+    """
+    settings = settings or get_settings()
+    now = datetime.now(timezone.utc).isoformat()
+
+    # Build base properties
+    base_properties = {
+        "sunstone.import.checksum": checksum,
+        "sunstone.import.source": source,
+        "sunstone.import.timestamp": now,
+        "sunstone.import.format": source_info.file_type.value,
+    }
+    if source_info.encoding:
+        base_properties["sunstone.import.encoding"] = source_info.encoding
+    if tags:
+        base_properties["sunstone.import.tags"] = ",".join(tags)
+    if owner_org:
+        base_properties["sunstone.owner.org"] = owner_org
+
+    results = []
+
+    if source_info.file_type == FileType.CSV:
+        # Parse CSV
+        arrow_table = parse_csv(source_info.local_path, source_info.encoding)
+
+        # Extract schema columns for metadata
+        columns = _extract_schema_columns(arrow_table)
+
+        # Capture schema info in builder
+        if builder is not None:
+            builder.set_schema_info(
+                original_columns=columns,
+                iceberg_columns=columns,
+                row_count=arrow_table.num_rows,
+                column_count=arrow_table.num_columns,
+            )
+
+        # Build properties with extended metadata
+        properties = {
+            **base_properties,
+            "sunstone.import.original_filename": source_info.filename,
+        }
+        properties = _build_extended_properties(properties, arrow_table, namespace, base_name, builder)
+
+        snapshot_id, already_existed = create_iceberg_table(
+            namespace=namespace,
+            table_name=base_name,
+            arrow_table=arrow_table,
+            properties=properties,
+            branch=branch,
+            force_schema=force_schema,
+            settings=settings,
+            builder=builder,
+        )
+
+        results.append(
+            ImportResult(
+                table_name=base_name,
+                namespace=namespace,
+                row_count=arrow_table.num_rows,
+                column_count=arrow_table.num_columns,
+                checksum=checksum,
+                original_path="",  # Will be filled by pipeline
+                schema_summary=schema_summary(arrow_table),
+                already_existed=already_existed,
+                snapshot_id=snapshot_id,
+            )
+        )
+
+    elif source_info.file_type == FileType.EXCEL:
+        # Parse Excel - each sheet becomes a table
+        sheets = parse_excel(source_info.local_path)
+
+        # Track if this is first sheet (for schema metadata)
+        first_sheet = True
+
+        for sheet_name, arrow_table in sheets.items():
+            # Sanitize sheet name for table suffix
+            safe_sheet = sanitize_table_name(sheet_name)
+            table_name = f"{base_name}_{safe_sheet}" if len(sheets) > 1 else base_name
+
+            # Extract schema columns
+            columns = _extract_schema_columns(arrow_table)
+
+            # Capture schema info from first sheet (representative)
+            if builder is not None and first_sheet:
+                builder.set_schema_info(
+                    original_columns=columns,
+                    iceberg_columns=columns,
+                    row_count=arrow_table.num_rows,
+                    column_count=arrow_table.num_columns,
+                )
+                first_sheet = False
+
+            # Build properties with extended metadata
+            properties = {
+                **base_properties,
+                "sunstone.import.original_filename": source_info.filename,
+                "sunstone.import.sheet_name": sheet_name,
+            }
+            properties = _build_extended_properties(properties, arrow_table, namespace, table_name, builder)
+
+            snapshot_id, already_existed = create_iceberg_table(
+                namespace=namespace,
+                table_name=table_name,
+                arrow_table=arrow_table,
+                properties=properties,
+                branch=branch,
+                force_schema=force_schema,
+                settings=settings,
+                builder=builder,
+            )
+
+            results.append(
+                ImportResult(
+                    table_name=table_name,
+                    namespace=namespace,
+                    row_count=arrow_table.num_rows,
+                    column_count=arrow_table.num_columns,
+                    checksum=checksum,
+                    original_path="",  # Will be filled by pipeline
+                    schema_summary=schema_summary(arrow_table),
+                    already_existed=already_existed,
+                    snapshot_id=snapshot_id,
+                )
+            )
+
+    return results