floe-python 0.8.0__py3-none-any.whl

floe/manifest_run.py

@@ -0,0 +1,299 @@
+"""Rebuild the `RootConfig` a run needs from a manifest (Rust: `manifest/reconstruct.rs`).
+
+The manifest was validated and resolved when it was generated: nothing is re-validated here. As in
+Rust, a malformed option block is dropped and an unknown enum value falls back to its default.
+"""
+
+import json
+import os
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any, TypeVar, cast, get_args
+
+from pydantic import BaseModel, ValidationError
+
+from floe.config.project import (
+    LineageConfig,
+    ProjectMetadata,
+    ReportConfig,
+    RootConfig,
+    config_version,
+)
+from floe.config.storage import CatalogsConfig, StorageDefinition, StoragesConfig
+from floe.config.types import (
+    EntityConfig,
+    IcebergPartitionField,
+    IcebergSinkOptions,
+    MergeOptions,
+    NormalizeColumnsConfig,
+    PiiConfig,
+    SchemaEvolutionConfig,
+    SchemaMismatchConfig,
+    SinkOptions,
+    SinkTarget,
+    SourceConfig,
+    SourceOptions,
+    SourcePartitioning,
+)
+from floe.config.validate import FLOE_VERSION
+from floe.errors import FloeConfigError, FloeIoError
+from floe.storage.paths import normalize_local_path, parse_s3_uri
+from floe.storage.s3 import S3Client
+
+Json = dict[str, Any]
+_M = TypeVar("_M", bound=BaseModel)
+_CLOUD = ("s3://", "gs://", "gcs://", "az://", "abfss://", "abfs://")
+_WRITE_MODES = ("overwrite", "append", "merge_scd1", "merge_scd2")
+
+
+def _omitted_as_none(model: type[BaseModel], value: object) -> object:
+    """Rust's `Option`: an omitted nullable field is `None`, in nested blocks too."""
+    if not isinstance(value, dict):
+        return value
+    filled = dict(cast(Json, value))
+    for name, field in model.model_fields.items():
+        args = get_args(field.annotation)
+        nested = next((a for a in args if isinstance(a, type) and issubclass(a, BaseModel)), None)  # fmt: skip
+        if name in filled and nested:
+            filled[name] = _omitted_as_none(nested, filled[name])
+        elif name not in filled and field.is_required() and type(None) in args:
+            filled[name] = None
+    return filled
+
+
+def _lenient(model: type[_M], value: object) -> _M | None:
+    """`serde_json::from_value(v).ok()`."""
+    try:
+        return (
+            None
+            if value is None
+            else model.model_validate(_omitted_as_none(model, value))
+        )
+    except ValidationError:
+        return None
+
+
+def _pick(value: str | None, allowed: tuple[str, ...], default: str) -> str:
+    return value if value in allowed else default
+
+
+def _storage(name: str, named_local: bool) -> str | None:
+    """`local` is the implicit storage, unless a definition is named so (Rust drops it either way)."""
+    return None if name == "local" and not named_local else name
+
+
+def _sink(t: Json, write_mode: str, named_local: bool) -> Json:
+    return {
+        "format": t["format"],
+        "path": t["path"],
+        "storage": _storage(t["storage"], named_local),
+        "options": _lenient(SinkOptions, t.get("options")),
+        "merge": _lenient(MergeOptions, t.get("merge")),
+        "iceberg": _lenient(IcebergSinkOptions, t.get("iceberg")),
+        "partition_by": t.get("partition_by"),
+        # #911: validated when generated, so not lenient; `num_buckets` may be omitted
+        "partition_spec": (spec := t.get("partition_spec"))
+        and [_omitted_as_none(IcebergPartitionField, f) for f in spec],
+        "write_mode": _pick(t.get("write_mode"), _WRITE_MODES, write_mode),
+    }
+
+
+def _schema(s: Json) -> Json:
+    keys = ("source", "nullable", "unique", "width", "trim")
+    return {
+        "columns": [
+            {"name": c["name"], "column_type": c["column_type"]}
+            | {k: c.get(k) for k in keys}
+            for c in s["columns"]
+        ],
+        "primary_key": s["primary_key"] or None,
+        "unique_keys": s["unique_keys"] or None,
+        "normalize_columns": _lenient(
+            NormalizeColumnsConfig, s.get("normalize_columns")
+        ),
+        "mismatch": _lenient(SchemaMismatchConfig, s.get("mismatch")),
+        "schema_evolution": _lenient(SchemaEvolutionConfig, s.get("schema_evolution")),
+    }
+
+
+def _entity(m: Json, named_local: bool) -> EntityConfig:
+    src, sinks = m["source"], m["sinks"]
+    write_mode = _pick(m.get("write_mode"), _WRITE_MODES, "overwrite")
+    rejected, archive = sinks.get("rejected"), sinks.get("archive")
+    source = {
+        "format": src["format"],
+        "path": src["path"],
+        "storage": _storage(src["storage"], named_local),
+        "options": _lenient(SourceOptions, src.get("options")) or SourceOptions(),
+        "cast_mode": src.get("cast_mode") or "strict",
+        "source_ref": src.get("ref"),
+        "resource": src.get("resource"),
+        "partitioning": _omitted_as_none(
+            SourcePartitioning, src.get("partitioning")
+        ),  # D28: validated when generated, so not lenient
+    }
+    sink = {
+        "write_mode": write_mode,
+        "accepted": _sink(sinks["accepted"], write_mode, named_local),
+        "rejected": rejected and _sink(rejected, write_mode, named_local),
+        "archive": archive
+        and {
+            "path": archive["path"],
+            "storage": _storage(archive["storage"], named_local),
+        },
+    }
+    incremental = _pick(m.get("incremental_mode"), ("archive", "file", "row"), "none")
+    return EntityConfig.model_validate(
+        {
+            "id": m["name"],
+            "name": (m.get("asset_key") or [m["name"]])[-1],
+            "metadata": None,
+            "domain": m.get("domain"),
+            "incremental_mode": incremental,
+            "state": {"path": m["state_path"]} if m.get("state_path") else None,
+            "source": source,
+            "sink": sink,
+            "policy": {
+                "severity": _pick(m.get("policy_severity"), ("reject", "abort"), "warn")
+            },
+            "schema": _schema(m["schema"]),
+            "pii": _lenient(PiiConfig, m.get("pii")),
+        }
+    )
+
+
+def _section(
+    model: type[_M], m: Json, name: str, shape: Callable[[Json], Json] | None = None
+) -> _M | None:
+    """An embedded block, back in config shape; a malformed one is an error, as in Rust."""
+    value = m.get(name)
+    try:
+        if value is not None and shape:
+            value = value | {"definitions": [shape(d) for d in value["definitions"]]}
+        return None if value is None else model.model_validate(value)
+    except (ValidationError, TypeError, KeyError, AttributeError) as e:
+        raise FloeConfigError(f"manifest {name} block is malformed: {e}") from e
+
+
+def _storage_def(d: Json) -> Json:
+    return {k: v for k, v in d.items() if k != "fs_type"} | {"type": d["fs_type"]}
+
+
+def _catalog_def(d: Json) -> Json:
+    rest = {k: v for k, v in d["type_config"].items() if k != "type"}  # D1: REST only
+    return {k: v for k, v in d.items() if k != "type_config"} | rest
+
+
+def _report(base: str, storages: StoragesConfig | None) -> tuple[ReportConfig, StoragesConfig | None]:  # fmt: skip
+    """Rust `RunContext::from_config`: an `s3://` base uses the storage on its bucket, else a synthetic one."""
+    if not base.startswith("s3://"):
+        local = storages and next((d.name for d in storages.definitions if d.type == "local"), None)  # fmt: skip
+        return ReportConfig(path=base.removeprefix("local://"), formatter=None, storage=local), storages  # fmt: skip
+    bucket = parse_s3_uri(base)[0]
+    defs = storages.definitions if storages else []
+    found = next((d.name for d in defs if d.type == "s3" and d.bucket == bucket), None)  # fmt: skip
+    # the report storage never becomes the default: unnamed paths stay local
+    if found is None:
+        found = "__report_s3__"
+        none = dict.fromkeys(("region", "account", "container", "prefix", "endpoint", "path_style_access"))  # fmt: skip
+        synthetic = StorageDefinition(name=found, type="s3", bucket=bucket, **none)
+        local = [] if storages else [StorageDefinition(name="local", type="local", bucket=None, **none)]  # fmt: skip
+        storages = StoragesConfig(default=storages and storages.default, definitions=[*defs, *local, synthetic])  # fmt: skip
+    return ReportConfig(path=base, formatter=None, storage=found), storages
+
+
+def _cloud(target: _M, uri: str | None) -> _M:
+    """A cloud `uri` replaces the path, so lineage names the dataset by its bucket."""
+    is_cloud = uri is not None and uri.startswith(_CLOUD)
+    return target.model_copy(update={"path": uri}) if is_cloud else target
+
+
+def _lineage_entity(entity: EntityConfig, m: Json) -> EntityConfig:
+    """Rust `lineage_inputs_from_manifest_json`."""
+    sink, sinks = entity.sink, m["sinks"]
+    accepted: SinkTarget = _cloud(sink.accepted, sinks["accepted"].get("uri"))
+    rejected = sink.rejected and _cloud(sink.rejected, sinks["rejected"].get("uri"))
+    source: SourceConfig = _cloud(entity.source, m["source"]["uri"])
+    update = {"accepted": accepted, "rejected": rejected}
+    return entity.model_copy(update={"source": source, "sink": sink.model_copy(update=update)})  # fmt: skip
+
+
+def _named_local(m: Json) -> bool:
+    try:
+        return any(d["name"] == "local" for d in m["storages"]["definitions"])
+    except (KeyError, TypeError):  # no storages; a malformed block fails in `_section`
+        return False
+
+
+def _read(location: str) -> bytes:
+    # S3 settings come from the environment (AWS_ENDPOINT_URL, ...), as for `-o s3://`
+    if location.startswith("s3://"):
+        found = S3Client(parse_s3_uri(location)[0]).read(location)
+        if found is None:
+            raise FloeIoError(f"manifest not found: {location}")
+        return found.body
+    if "://" in location:  # D1: gcs and adls are dropped
+        raise FloeConfigError(f"remote manifest is not supported: {location}")
+    try:
+        return Path(location).read_bytes()
+    except OSError as e:  # as the CLI's `_realpath`
+        raise FloeIoError(f"{e.strerror} (os error {e.errno})") from e
+
+
+def shown(location: str) -> str:
+    """The manifest as the run reports it: a remote one as typed."""
+    return location if "://" in location else str(normalize_local_path(Path(location)))
+
+
+def _replay_base(location: str, m: Json) -> Path:
+    """Rust `manifest_replay_config_base`: a remote manifest resolves local paths under its work root.
+
+    A local one resolves them beside `config_uri`, where `manifest generate` resolved them (#1008).
+    """
+    if not location.startswith("s3://"):
+        config = m.get("config_uri")
+        if isinstance(config, str) and config.startswith("local:///"):
+            return Path(config.removeprefix("local://"))
+        return Path(location)
+    env = os.environ.get("FLOE_WORK_ROOT", "")
+    image = str(m.get("runtime_env") or "").lower() == "image"
+    root = m.get("work_root") or (env if env.strip() else None) or ("/work" if image else os.getcwd())  # fmt: skip
+    return Path(root) / "manifest.json"  # the run resolves against its parent
+
+
+def load_manifest(location: str | Path) -> tuple[RootConfig, list[EntityConfig], Path]:
+    """The run's config, the entities its lineage observes, and the path its relative paths resolve against."""
+    location = str(location)
+    data = _read(location)
+    try:
+        m = json.loads(data.decode())
+        spec = m["spec_version"]
+        entities = [_entity(e, _named_local(m)) for e in m["entities"]]
+        report_base = m["report_base_uri"]
+        meta = m.get("metadata")  # D20: dumped without its None fields
+        metadata = None if meta is None else ProjectMetadata.model_validate(_omitted_as_none(ProjectMetadata, meta))  # fmt: skip
+    except (ValueError, KeyError, TypeError, AttributeError) as e:
+        why = f"missing field `{e.args[0]}`" if isinstance(e, KeyError) else e
+        raise FloeConfigError(f"manifest parse error: {why}") from e
+    try:
+        supported = config_version(spec) <= FLOE_VERSION
+    except FloeConfigError:
+        supported = False
+    if not supported:
+        floe = f"{FLOE_VERSION[0]}.{FLOE_VERSION[1]}"
+        raise FloeConfigError(f"manifest spec_version={spec} is unsupported (newer than Floe {floe})")  # fmt: skip
+    report, storages = _report(report_base, _section(StoragesConfig, m, "storages", _storage_def))  # fmt: skip
+    root = RootConfig(
+        version=spec,
+        metadata=metadata,
+        storages=storages,
+        catalogs=_section(CatalogsConfig, m, "catalogs", _catalog_def),
+        env=None,
+        domains=[],
+        report=report,
+        lineage=_section(LineageConfig, m, "lineage"),
+        entities=entities,
+        config_files=[],
+    )
+    observed = [_lineage_entity(e, me) for e, me in zip(entities, m["entities"])]
+    return root, observed, _replay_base(location, m)

floe/pii.py

@@ -0,0 +1,134 @@
+import hashlib
+import hmac
+import os
+import re
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Literal
+
+import polars as pl
+
+from floe.errors import FloeConfigError, FloeRunError
+
+Strategy = Literal["hash", "drop", "nullify", "redact", "mask", "tokenize"]
+
+
+@dataclass(frozen=True)
+class PiiColumn:
+    name: str
+    strategy: Strategy
+    mask_pattern: str | None = None
+    redact_value: str | None = None
+    key: str | None = None
+
+
+def apply_pii_masking(
+    df: pl.DataFrame,
+    columns: Sequence[PiiColumn],
+    schema_to_runtime: Mapping[str, str],
+) -> pl.DataFrame:
+    for col in columns:
+        name = schema_to_runtime.get(col.name, col.name)
+        if name not in df.columns:
+            continue
+        match col.strategy:
+            case "hash":
+                df = df.with_columns(_hash(_as_string(df[name]), name, col.key))
+            case "drop":
+                df = df.drop(name)
+            case "nullify":
+                df = df.with_columns(pl.lit(None, dtype=df.schema[name]).alias(name))
+            case "redact":
+                text = _as_string(df[name])
+                df = df.with_columns(
+                    text.to_frame().select(
+                        _when_set(name, pl.lit(col.redact_value or "[REDACTED]"))
+                    )
+                )
+            case "mask":
+                text = _as_string(df[name])
+                df = df.with_columns(
+                    text.to_frame().select(_mask(name, col.mask_pattern or ""))
+                )
+            case "tokenize":
+                pass  # rejected at config validation; never reached at runtime
+    return df
+
+
+def _as_string(series: pl.Series) -> pl.Series:
+    try:
+        return series.cast(pl.String)
+    except pl.exceptions.PolarsError as e:
+        raise FloeRunError(
+            f"PII: failed to cast column {series.name} to string: {e}"
+        ) from e
+
+
+def _when_set(name: str, value: pl.Expr) -> pl.Expr:
+    return pl.when(pl.col(name).is_not_null()).then(value).alias(name)
+
+
+def token_n(pattern: str, token: str) -> int:
+    start = pattern.find(token)
+    if start < 0:
+        return 0
+    rest = pattern[start + len(token) :]
+    end = rest.find("}")
+    digits = rest[:end]
+    if end < 0 or not re.fullmatch(r"\+?[0-9]+", digits) or int(digits) >= 2**64:
+        return 0
+    return int(digits)
+
+
+def _mask(name: str, pattern: str) -> pl.Expr:
+    first_n = token_n(pattern, "{first")
+    last_n = token_n(pattern, "{last")
+    text = pl.col(name)
+    chars = text.str.len_chars()
+    last = pl.min_horizontal(pl.lit(last_n, pl.UInt32), chars)
+    first = pl.min_horizontal(pl.lit(first_n, pl.UInt32), chars - last)
+    out = _when_set(name, pl.lit(pattern))  # broadcast the pattern to one row per value
+    if first_n:
+        out = out.str.replace_all(
+            f"{{first{first_n}}}", text.str.slice(0, first), literal=True
+        )
+    if last_n:
+        out = out.str.replace_all(
+            f"{{last{last_n}}}", text.str.slice(chars - last, last), literal=True
+        )
+    return out.alias(name)
+
+
+def _resolve_key(raw: str, name: str) -> str:
+    if "${" not in raw:
+        return raw
+    if not (raw.startswith("${") and raw.endswith("}")):
+        raise FloeConfigError(
+            f"pii column {name}: key must be a plain value or a single ${{VAR_NAME}} reference; "
+            "mixing literal text with ${...} is not supported"
+        )
+    var = raw[2:-1]
+    if not var or "{" in var or "}" in var:
+        raise FloeConfigError(f"pii column {name}: key has invalid placeholder syntax")
+    value = os.environ.get(var)
+    if value is None:
+        raise FloeRunError(
+            f"pii column {name}: key references env var {var} which is not set"
+        )
+    return value
+
+
+def _hash(values: pl.Series, name: str, key: str | None) -> pl.Series:
+    # ponytail: ~0.8M rows/s single-threaded; upgrade to vectorised HMAC if it becomes a bottleneck
+    if key is None:
+        out = [
+            None if v is None else hashlib.sha256(v.encode()).hexdigest()
+            for v in values
+        ]
+    else:
+        raw = _resolve_key(key, name).encode()
+        out = [
+            None if v is None else hmac.digest(raw, v.encode(), "sha256").hex()
+            for v in values
+        ]
+    return pl.Series(name, out, dtype=pl.String)

floe/read/avro.py

@@ -0,0 +1,40 @@
+"""Avro reader: native column types, in chunks (Rust: `io/read/avro.rs`)."""
+
+from collections.abc import Collection, Iterator
+from dataclasses import dataclass
+from pathlib import Path
+
+import polars as pl
+
+from floe.errors import FloeIoError
+
+
+@dataclass(frozen=True)
+class AvroChunk:
+    """`offset` is the whole-file index of the chunk's first row."""
+
+    offset: int
+    df: pl.DataFrame
+
+
+def read_avro(
+    path: Path, columns: Collection[str] = (), chunk_rows: int = 1_000_000
+) -> Iterator[AvroChunk]:
+    """Yield the file in chunks of at most `chunk_rows` rows; always at least one (maybe empty).
+
+    Columns keep their native types (casting them is #758, D15). `columns` projects to the
+    declared columns that exist in the file; none of them present keeps everything.
+    """
+    try:
+        # ponytail: Polars reads the whole file (no streaming Avro scan); chunks bound only what
+        # the checks hold at once. Switch to a streaming reader if Avro files outgrow memory.
+        df = pl.read_avro(path)
+    except (pl.exceptions.PolarsError, OSError) as err:
+        raise FloeIoError(f"avro read failed: {err}") from err
+    picked = [name for name in df.columns if name in columns]
+    if picked:
+        df = df.select(picked)
+    for offset in range(0, df.height, chunk_rows):
+        yield AvroChunk(offset, df.slice(offset, chunk_rows))
+    if df.height == 0:
+        yield AvroChunk(0, df)

floe/read/csv.py

@@ -0,0 +1,138 @@
+"""CSV reader: every column as a raw string, in bounded chunks (Rust: `io/read/csv.rs`)."""
+
+import csv
+import json
+from collections.abc import Iterator, Sequence
+from dataclasses import dataclass
+from itertools import islice
+from pathlib import Path
+from typing import Literal
+
+import polars as pl
+
+from floe.config.types import SourceOptions
+from floe.errors import FloeConfigError, FloeIoError
+
+RAGGED_MESSAGE = "found more fields than defined in 'Schema'"
+_ENCODINGS: dict[str, Literal["utf8", "utf8-lossy"]] = {
+    "utf8": "utf8",
+    "lossyutf8": "utf8-lossy",
+}
+
+
+@dataclass(frozen=True)
+class CsvChunk:
+    """`offset` is the whole-file index of the chunk's first data row (add it to the row index).
+
+    `ragged` marks rows with more fields than the header (decision D3, #603): their values are
+    truncated and the caller rejects just those rows with `RAGGED_MESSAGE`.
+    """
+
+    offset: int
+    df: pl.DataFrame
+    ragged: pl.Series
+
+
+@dataclass(frozen=True)
+class _Args:
+    separator: str
+    encoding: Literal["utf8", "utf8-lossy"]
+    null_values: list[str] | None
+    has_header: bool
+
+    def scan(
+        self,
+        path: Path,
+        schema: dict[str, type[pl.String]] | None = None,
+        n_rows: int | None = None,
+        ragged_from: int | None = None,
+    ) -> pl.LazyFrame:
+        """`ragged_from`: truncate over-long rows and skip that many data rows."""
+        return pl.scan_csv(
+            path,
+            has_header=self.has_header,
+            separator=self.separator,
+            encoding=self.encoding,
+            null_values=self.null_values,
+            infer_schema=False,
+            schema=schema,
+            n_rows=n_rows,
+            truncate_ragged_lines=ragged_from is not None,
+            skip_rows_after_header=ragged_from or 0,
+        )
+
+
+def _scan_args(options: SourceOptions) -> _Args:
+    separator = ";" if options.separator is None else options.separator
+    if len(separator.encode()) != 1:
+        raise FloeConfigError(
+            f"separator must be a single byte, got {json.dumps(separator, ensure_ascii=False)}"
+        )
+    raw = "utf8" if options.encoding is None else options.encoding
+    encoding = _ENCODINGS.get(raw.lower().replace("-", "").replace("_", ""))
+    if encoding is None:
+        raise FloeConfigError(f"unsupported encoding: {raw}")
+    return _Args(
+        separator, encoding, options.null_values or None, options.header is not False
+    )
+
+
+def _columns(path: Path, args: _Args, declared: Sequence[str]) -> list[str]:
+    try:
+        probe = args.scan(path, n_rows=0 if args.has_header else 1).collect()
+    except (pl.exceptions.PolarsError, OSError) as err:
+        # Rust's polars prints the error kind: `no data: empty CSV`
+        kind = "no data: " if isinstance(err, pl.exceptions.NoDataError) else ""
+        raise FloeIoError(f"csv header read failed: {kind}{err}") from err
+    if args.has_header:
+        return probe.columns
+    names = list(declared[: probe.width])
+    return names + [f"extra_column_{i + 1}" for i in range(len(declared), probe.width)]
+
+
+def _field_counts(path: Path, args: _Args, skip: int) -> Iterator[int]:
+    errors = "replace" if args.encoding == "utf8-lossy" else "strict"
+    with path.open(newline="", encoding="utf-8-sig", errors=errors) as handle:
+        rows = csv.reader(handle, delimiter=args.separator)
+        for _ in islice(rows, skip):
+            pass
+        for row in rows:
+            yield len(row)
+
+
+def read_csv(
+    path: Path,
+    options: SourceOptions,
+    declared: Sequence[str] = (),
+    chunk_rows: int = 1_000_000,
+) -> Iterator[CsvChunk]:
+    """Yield the file in chunks of at most `chunk_rows` rows; always at least one (maybe empty).
+
+    `declared` names the headless columns, as Rust. Nulls follow `null_values` and Polars: an
+    unquoted empty field is always null.
+    """
+    args = _scan_args(options)
+    names = _columns(path, args, declared)
+    schema = dict.fromkeys(names, pl.String)
+    done = 0
+    try:
+        for df in args.scan(path, schema).collect_batches(chunk_size=chunk_rows):
+            yield CsvChunk(done, df, pl.Series([False] * df.height, dtype=pl.Boolean))
+            done += df.height
+    except pl.exceptions.ComputeError as err:
+        if RAGGED_MESSAGE not in str(err):
+            raise FloeIoError(f"csv read failed: {err}") from err
+        # ponytail: the ragged-row pass reads the file again in Python's csv module, only for files
+        # that have such rows; a Polars option to flag them would replace it.
+        counts = _field_counts(path, args, done + int(args.has_header))
+        for df in args.scan(path, schema, ragged_from=done).collect_batches(
+            chunk_size=chunk_rows
+        ):
+            mask = [n > len(names) for n in islice(counts, df.height)]
+            yield CsvChunk(done, df, pl.Series(mask, dtype=pl.Boolean))
+            done += df.height
+    except pl.exceptions.PolarsError as err:
+        raise FloeIoError(f"csv read failed: {err}") from err
+    if done == 0:
+        empty = pl.DataFrame(schema=dict.fromkeys(names, pl.String))
+        yield CsvChunk(0, empty, pl.Series([], dtype=pl.Boolean))