flyteplugins-union 0.4.0__py3-none-win_amd64.whl

flyteplugins/union/__init__.py

@@ -0,0 +1,4 @@
+"""Union SDK - Proprietary extensions for Flyte.
+
+This package provides Union-specific functionality on top of the open-source Flyte SDK.
+"""

flyteplugins/union/cli/__init__.py

@@ -0,0 +1,48 @@
+import tempfile
+
+import rich_click as click
+
+
+def edit_with_retry(yaml_text: str, apply_fn, *, console, noun: str = "resource"):
+    """Open an editor and retry or save to file on failure.
+
+    Args:
+        yaml_text: Initial YAML content to edit.
+        apply_fn: Callable that takes the edited YAML string and applies it.
+                  Should raise on failure.
+        console: Rich console for output.
+        noun: Name of the resource for messages (e.g. "role", "policy").
+
+    Returns:
+        The result of apply_fn on success, or None if cancelled.
+    """
+    edited = click.edit(yaml_text, extension=".yaml")
+    if edited is None:
+        console.print(f"[yellow]Edit cancelled, {noun} not modified.[/yellow]")
+        return None
+
+    while True:
+        try:
+            return apply_fn(edited)
+        except Exception as e:
+            console.print(f"[red]Error:[/red] {e}\n")
+            action = click.prompt(
+                "Would you like to (r)etry editing, (s)ave to file, or (a)bort?",
+                type=click.Choice(["r", "s", "a"], case_sensitive=False),
+                default="r",
+            )
+            if action == "r":
+                edited = click.edit(edited, extension=".yaml")
+                if edited is None:
+                    console.print(f"[yellow]Edit cancelled, {noun} not modified.[/yellow]")
+                    return None
+            elif action == "s":
+                tmp = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", prefix=f"{noun}-", delete=False)
+                tmp.write(edited)
+                tmp.close()
+                console.print(f"[cyan]Saved to {tmp.name}[/cyan]")
+                console.print(f"You can retry with: [cyan]--file {tmp.name}[/cyan]")
+                return None
+            else:
+                console.print(f"[yellow]Aborted, {noun} not modified.[/yellow]")
+                return None

flyteplugins/union/cli/_tui/__init__.py

@@ -0,0 +1,67 @@
+"""Lazy entrypoint for the volume explore TUI.
+
+Importing ``textual`` is gated behind the ``[tui]`` extra so headless
+installs don't pay the cost. The launcher below catches the
+:class:`ImportError` and re-raises with a clear install hint, matching the
+ergonomic of :mod:`flyte.cli._tui` upstream.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from flyteplugins.union.cli._volume_index import IndexReader
+
+
+async def launch_volume_explore(
+    *,
+    reader: "IndexReader",
+    title: str,
+    subtitle: str = "",
+    lineage: "list[dict] | None" = None,
+    provenance: "dict | None" = None,
+    can_navigate_parent: bool = False,
+) -> "str | None":
+    """Run the volume explore TUI against an already-open index.
+
+    ``lineage`` is the version chain rendered in the right-pane Lineage
+    box (PRD: walk ``parent`` pointers). Each entry is a dict with at
+    least ``label`` (e.g. ``current`` / ``parent``); optional keys
+    ``name``, ``index_path``, ``total_inodes``, ``used_bytes``, ``error``.
+
+    ``provenance`` carries ``produced_by`` / ``parent_produced_by`` action
+    refs for the Provenance box; ``can_navigate_parent`` enables the "open
+    parent" keybinding.
+
+    Returns the navigation request the user made before exiting — ``"parent"``
+    to re-open explore against the parent version, or ``None`` to stop.
+
+    Async: the :class:`IndexReader` is aiosqlite-backed, so we ``await`` the
+    volume summary up front (``compose`` can't await) and drive the Textual app
+    via ``run_async`` inside the caller's event loop.
+
+    Raises a helpful :class:`ImportError` when ``textual`` is not installed.
+    """
+    try:
+        from flyteplugins.union.cli._tui._volume_explore import VolumeExploreApp, summary_header_text
+    except ImportError as exc:
+        raise ImportError(
+            "The volume explore TUI requires the 'textual' package. "
+            "Install it with:  pip install 'flyteplugins-union[tui]'"
+        ) from exc
+
+    # Precompute the header summary now (async read) so compose() stays sync.
+    summary = await reader.summary()
+    app = VolumeExploreApp(
+        reader=reader,
+        title=title,
+        subtitle=subtitle,
+        root_inode=summary.root_inode,
+        header_text=summary_header_text(summary, subtitle),
+        lineage=lineage or [],
+        provenance=provenance,
+        can_navigate_parent=can_navigate_parent,
+    )
+    await app.run_async()
+    return app.navigate

flyteplugins/union/cli/_tui/_volume_explore.py

@@ -0,0 +1,484 @@
+"""Textual TUI for browsing a Volume metadata index.
+
+Mirrors the shape of :mod:`flyte.cli._tui._explore` in flyte-sdk:
+
+* Same color palette (Flyte purple).
+* Header + ``Horizontal`` body + Footer skeleton.
+* Left pane: a ``Tree`` widget (vim ``j``/``k`` bindings + arrow keys).
+* Right pane: ``TabbedContent`` with a ``DetailPanel`` of bordered
+  ``_DetailBox`` cards (Volume Header / Selected Node / Lineage).
+* Top-level key bindings: ``q`` quit, ``r`` refresh.
+
+The Tree is lazy: directory nodes are added with a placeholder child so
+they render as expandable; the children are fetched from the
+:class:`IndexReader` on first expansion. This keeps the load proportional
+to what the user actually looks at rather than walking a million-inode
+volume on launch.
+"""
+
+from __future__ import annotations
+
+import datetime
+import stat as _stat
+from typing import Any, ClassVar
+
+from textual.app import App, ComposeResult
+from textual.binding import Binding, BindingType
+from textual.containers import Horizontal, VerticalScroll
+from textual.widgets import Footer, Header, Static, TabbedContent, TabPane, Tree
+from textual.widgets.tree import TreeNode
+
+from flyteplugins.union.cli._volume_index import IndexEntry, IndexReader, VolumeSummary
+
+# Match flyte-sdk's _explore palette exactly so the experience feels
+# native to anyone who's run `flyte start tui`.
+_FLYTE_PURPLE = "#7652a2"
+_FLYTE_PURPLE_LIGHT = "#f7f5fd"
+_FLYTE_PURPLE_DARK = "#171020"
+
+_KIND_ICON = {
+    "dir": "📁",
+    "file": "📄",
+    "symlink": "🔗",
+}
+
+
+def _fmt_size(n: int) -> str:
+    units = ["B", "KB", "MB", "GB", "TB", "PB"]
+    f = float(n)
+    for u in units:
+        if f < 1024 or u == units[-1]:
+            return f"{f:,.1f} {u}" if u != "B" else f"{int(f):,} {u}"
+        f /= 1024
+    return f"{n} B"
+
+
+def _fmt_mtime(mtime_ns: int) -> str:
+    if not mtime_ns:
+        return ""
+    try:
+        dt = datetime.datetime.fromtimestamp(mtime_ns / 1e9)
+    except (OSError, OverflowError, ValueError):
+        return ""
+    return dt.strftime("%Y-%m-%d %H:%M:%S")
+
+
+def _fmt_mode(mode: int, kind: str) -> str:
+    if not mode:
+        return f"<{kind}>"
+    try:
+        return _stat.filemode(mode)
+    except (TypeError, ValueError):
+        return oct(mode)
+
+
+def _label(entry: IndexEntry) -> str:
+    icon = _KIND_ICON.get(entry.kind, "•")
+    if entry.kind == "dir":
+        return f"{icon} {entry.name}/"
+    if entry.kind == "symlink" and entry.symlink_target:
+        return f"{icon} {entry.name} → {entry.symlink_target}"
+    return f"{icon} {entry.name}  ({_fmt_size(entry.size)})"
+
+
+class _DetailBox(Static):
+    """Bordered card; markup off so raw paths/JSON render literally."""
+
+    DEFAULT_CSS = """
+    _DetailBox {
+        border: solid $accent;
+        padding: 0 1;
+        margin-bottom: 1;
+        height: auto;
+    }
+    """
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        kwargs.setdefault("markup", False)
+        super().__init__(*args, **kwargs)
+
+
+class VolumeFileTree(Tree[int]):
+    """Lazy directory tree backed by the :class:`IndexReader`.
+
+    Each node's ``data`` carries the inode number. Directories are added
+    with a stub child so the disclosure arrow shows up; the stub is
+    replaced with real entries on first expand.
+    """
+
+    BINDINGS: ClassVar[list[BindingType]] = [
+        Binding("down,j", "cursor_down", "Cursor Down", show=False),
+        Binding("up,k", "cursor_up", "Cursor Up", show=False),
+    ]
+
+    def __init__(self, reader: IndexReader, root_inode: int, **kwargs: Any) -> None:
+        super().__init__("/", data=root_inode, **kwargs)
+        self._reader = reader
+        self._populated: set[int] = set()
+
+    def on_mount(self) -> None:
+        self.root.allow_expand = True
+        # Reader I/O is async (aiosqlite); populate + expand the root on a
+        # worker so the event loop stays responsive while it loads.
+        self.run_worker(self._populate_and_expand(self.root), exclusive=False)
+
+    def reload(self) -> None:
+        """Reset the tree to its root and re-populate from scratch (the ``r``
+        refresh). Must clear ``_populated`` — otherwise the post-reset root is
+        still flagged as populated and :meth:`_populate` would no-op, leaving an
+        empty tree.
+        """
+        self.reset(self.root.label, data=self.root.data)
+        self._populated.clear()
+        self.root.allow_expand = True
+        self.run_worker(self._populate_and_expand(self.root), exclusive=False)
+
+    async def _populate_and_expand(self, node: TreeNode[int]) -> None:
+        await self._populate(node)
+        node.expand()
+
+    async def _populate(self, node: TreeNode[int]) -> None:
+        inode = node.data
+        if inode is None or inode in self._populated:
+            return
+        self._populated.add(inode)
+        node.remove_children()
+        for entry in await self._reader.children(inode):
+            child = node.add(_label(entry), data=entry.inode, allow_expand=(entry.kind == "dir"))
+            if entry.kind == "dir":
+                # Placeholder so the disclosure arrow shows up; replaced
+                # on first expand.
+                child.add_leaf("…")
+                self._populated.discard(entry.inode)
+
+    def on_tree_node_expanded(self, event: Tree.NodeExpanded) -> None:
+        node = event.node
+        if node.data is not None and node.data not in self._populated:
+            self.run_worker(self._populate(node), exclusive=False)
+
+
+class VolumeDetailPanel(VerticalScroll):
+    """Right pane — volume summary + selected node details.
+
+    Volume Header is populated once on mount; Selected Node is rewritten
+    on each tree-cursor change; Provenance and Lineage are static text the
+    launcher fills in.
+    """
+
+    def __init__(
+        self,
+        reader: IndexReader,
+        header_text: str,
+        provenance_text: str = "",
+        lineage_text: str = "",
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self._reader = reader
+        self._header_text = header_text
+        self._provenance_text = provenance_text
+        self._lineage_text = lineage_text
+
+    def compose(self) -> ComposeResult:
+        yield _DetailBox(id="box-volume")
+        yield _DetailBox(id="box-selected")
+        yield _DetailBox(id="box-provenance")
+        yield _DetailBox(id="box-lineage")
+
+    def on_mount(self) -> None:
+        vbox = self.query_one("#box-volume", _DetailBox)
+        vbox.border_title = "Volume"
+        vbox.update(self._header_text)
+
+        sbox = self.query_one("#box-selected", _DetailBox)
+        sbox.border_title = "Selected"
+        sbox.update("Pick a file or directory in the tree.")
+
+        pbox = self.query_one("#box-provenance", _DetailBox)
+        pbox.border_title = "Provenance"
+        pbox.update(self._provenance_text or "(no provenance recorded)")
+
+        lbox = self.query_one("#box-lineage", _DetailBox)
+        lbox.border_title = "Lineage"
+        lbox.update(self._lineage_text or "(no lineage — load a Volume value to walk parent indexes)")
+
+    def show_selected(self, entry: IndexEntry, chunk_count: int) -> None:
+        sbox = self.query_one("#box-selected", _DetailBox)
+        sbox.border_title = "Selected"
+        lines: list[str] = []
+        lines.append(f"name:    {entry.name or '/'}")
+        lines.append(f"inode:   {entry.inode}")
+        lines.append(f"kind:    {entry.kind}")
+        if entry.kind == "file":
+            lines.append(f"size:    {_fmt_size(entry.size)} ({entry.size:,} bytes)")
+            if chunk_count >= 0:
+                lines.append(f"chunks:  {chunk_count}")
+            else:
+                lines.append("chunks:  — (not available for this backend)")
+        if entry.kind == "symlink" and entry.symlink_target:
+            lines.append(f"target:  {entry.symlink_target}")
+        lines.append(f"mode:    {_fmt_mode(entry.mode, entry.kind)}  ({oct(entry.mode)})")
+        lines.append(f"owner:   {entry.uid}:{entry.gid}")
+        lines.append(f"nlink:   {entry.nlink}")
+        mtime = _fmt_mtime(entry.mtime_ns)
+        if mtime:
+            lines.append(f"mtime:   {mtime}")
+        sbox.update("\n".join(lines))
+
+    def set_lineage(self, text: str) -> None:
+        lbox = self.query_one("#box-lineage", _DetailBox)
+        lbox.update(text)
+
+
+def _lineage_text(lineage: "list[dict]") -> str:
+    """Render the lineage chain, newest first.
+
+    The chain is walked to its root via each version's metadata object, so
+    every entry is a real published version — including intra-action
+    intermediates. Each shows its producing action, message, and stats.
+    """
+    if not lineage:
+        return ""
+    lines: list[str] = []
+    for entry in lineage:
+        label = entry.get("label", "?")
+        name = entry.get("name", "")
+        line = f"[{label}]"
+        if name:
+            line += f" {name}"
+        if entry.get("message"):
+            # Commit message inline with the label (PRD UC2: 'epoch N' /
+            # 'post-augmentation' / etc.) — git-log shape.
+            line += f"  — {entry['message']}"
+        lines.append(line)
+        if "error" in entry:
+            lines.append(f"  error: {entry['error']}")
+            continue
+        produced_by = entry.get("produced_by")
+        if produced_by:
+            who = produced_by.get("action") or "?"
+            if produced_by.get("run"):
+                who = f"{produced_by['run']} / {who}"
+            lines.append(f"  by:      {who}")
+        if entry.get("total_inodes") is not None:
+            lines.append(f"  inodes:  {entry['total_inodes']:,}")
+        if entry.get("used_bytes") is not None:
+            lines.append(f"  size:    {entry['used_bytes']:,} bytes")
+        path = entry.get("index_path")
+        if path:
+            lines.append(f"  index:   {path}")
+    return "\n".join(lines)
+
+
+def _fmt_action_ref(ref: "dict | None") -> list[str]:
+    """Render an action-ref dict (run / action / output / scope) as lines."""
+    if not ref:
+        return ["  (unknown)"]
+    out: list[str] = []
+    if ref.get("run"):
+        out.append(f"  run:     {ref['run']}")
+    if ref.get("action"):
+        out.append(f"  action:  {ref['action']}")
+    if ref.get("op"):
+        out.append(f"  output:  {ref['op']}")
+    if ref.get("project") or ref.get("domain"):
+        out.append(f"  scope:   {ref.get('project') or '?'}/{ref.get('domain') or '?'}")
+    return out or ["  (unknown)"]
+
+
+def _provenance_text(provenance: "dict | None", can_navigate_parent: bool) -> str:
+    """Render the producing action for this version and its parent.
+
+    ``provenance`` carries ``produced_by`` / ``parent_produced_by`` action-ref
+    dicts (or None). When the parent can be opened, append the keybinding hint.
+    """
+    if not provenance:
+        return ""
+    lines: list[str] = ["produced by:"]
+    lines += _fmt_action_ref(provenance.get("produced_by"))
+    parent = provenance.get("parent_produced_by")
+    if parent or can_navigate_parent:
+        lines.append("")
+        hint = "  ·  press 'p' to open" if can_navigate_parent else ""
+        lines.append(f"parent produced by:{hint}")
+        lines += _fmt_action_ref(parent)
+    return "\n".join(lines)
+
+
+def summary_header_text(s: VolumeSummary, subtitle: str) -> str:
+    """Render the Volume Header box text from a precomputed summary.
+
+    Takes a :class:`VolumeSummary` (read by the launcher before the app starts)
+    rather than the reader, so the app's ``compose`` stays synchronous while the
+    underlying reads are async (aiosqlite).
+    """
+    lines: list[str] = []
+    if subtitle:
+        lines.append(subtitle)
+    lines.append(f"inodes:  {s.total_inodes:,}")
+    lines.append(f"size:    {_fmt_size(s.used_bytes)} ({s.used_bytes:,} bytes)")
+    interesting = ("nextInode", "nextChunk", "nextSession", "nextinode", "nextchunk", "nextsession")
+    for key in interesting:
+        if key in s.counters:
+            lines.append(f"{key}: {s.counters[key]:,}")
+    return "\n".join(lines)
+
+
+class VolumeExploreApp(App[None]):
+    """Standalone Textual app for browsing a Volume metadata index."""
+
+    BINDINGS: ClassVar[list[BindingType]] = [
+        Binding("q", "quit_app", "Quit"),
+        Binding("r", "refresh_tree", "Refresh"),
+        Binding("p", "open_parent", "Open parent"),
+    ]
+
+    CSS = f"""
+    Screen {{
+        background: {_FLYTE_PURPLE_DARK};
+    }}
+    Header {{
+        background: {_FLYTE_PURPLE};
+        color: {_FLYTE_PURPLE_LIGHT};
+    }}
+    Footer {{
+        background: {_FLYTE_PURPLE};
+        color: {_FLYTE_PURPLE_LIGHT};
+    }}
+    Horizontal {{
+        height: 1fr;
+    }}
+    VolumeFileTree {{
+        width: 1fr;
+        min-width: 32;
+        border: solid {_FLYTE_PURPLE};
+        border-title-color: {_FLYTE_PURPLE_LIGHT};
+        background: {_FLYTE_PURPLE_DARK};
+        color: {_FLYTE_PURPLE_LIGHT};
+    }}
+    #right-tabs {{
+        width: 2fr;
+    }}
+    VolumeDetailPanel {{
+        background: {_FLYTE_PURPLE_DARK};
+    }}
+    TabPane {{
+        padding: 0;
+    }}
+    Tabs {{
+        background: {_FLYTE_PURPLE_DARK};
+        color: {_FLYTE_PURPLE_LIGHT};
+    }}
+    Tab {{
+        background: {_FLYTE_PURPLE_DARK};
+        color: {_FLYTE_PURPLE_LIGHT};
+    }}
+    Tab.-active {{
+        background: {_FLYTE_PURPLE};
+        color: {_FLYTE_PURPLE_LIGHT};
+    }}
+    _DetailBox {{
+        border: solid {_FLYTE_PURPLE};
+        border-title-color: {_FLYTE_PURPLE_LIGHT};
+        padding: 0 1;
+        margin-bottom: 1;
+        height: auto;
+        color: {_FLYTE_PURPLE_LIGHT};
+    }}
+    """
+
+    def __init__(
+        self,
+        *,
+        reader: IndexReader,
+        title: str,
+        subtitle: str = "",
+        root_inode: int,
+        header_text: str,
+        lineage: "list[dict] | None" = None,
+        provenance: "dict | None" = None,
+        can_navigate_parent: bool = False,
+    ) -> None:
+        super().__init__()
+        self._reader = reader
+        self._title = title
+        self._subtitle = subtitle
+        # Precomputed by the launcher (async reader.summary()) so compose()
+        # — which can't await — stays synchronous.
+        self._root_inode = root_inode
+        self._header_text = header_text
+        self._lineage = lineage or []
+        self._provenance = provenance
+        self._can_navigate_parent = can_navigate_parent
+        # Set to "parent" by ``action_open_parent`` before exit so the
+        # launcher knows to re-open explore against this volume's parent.
+        self.navigate: "str | None" = None
+
+    def compose(self) -> ComposeResult:
+        yield Header()
+        with Horizontal():
+            tree = VolumeFileTree(self._reader, root_inode=self._root_inode, id="file-tree")
+            tree.border_title = self._title
+            tree.show_root = True
+            yield tree
+            with TabbedContent(initial="tab-details", id="right-tabs"):
+                with TabPane("Details", id="tab-details"):
+                    yield VolumeDetailPanel(
+                        self._reader,
+                        header_text=self._header_text,
+                        provenance_text=_provenance_text(self._provenance, self._can_navigate_parent),
+                        lineage_text=_lineage_text(self._lineage),
+                        id="detail-panel",
+                    )
+        yield Footer()
+
+    def on_mount(self) -> None:
+        self.title = self._title
+        self.query_one("#file-tree", VolumeFileTree).focus()
+
+    def on_tree_node_highlighted(self, event: Tree.NodeHighlighted) -> None:
+        self.run_worker(self._render_selected(event.node), exclusive=True)
+
+    def on_tree_node_selected(self, event: Tree.NodeSelected) -> None:
+        self.run_worker(self._render_selected(event.node), exclusive=True)
+
+    async def _render_selected(self, node: TreeNode[int]) -> None:
+        inode = node.data
+        if inode is None:
+            return
+        entry = await self._reader.get(inode)
+        if entry is None:
+            return
+        chunks = await self._reader.chunk_count(inode) if entry.kind == "file" else 0
+        panel = self.query_one("#detail-panel", VolumeDetailPanel)
+        panel.show_selected(entry, chunks)
+
+    def check_action(self, action: str, parameters: "tuple[object, ...]") -> "bool | None":
+        """Disable the ``p`` (open-parent) binding when there's no navigable
+        parent, so it shows greyed in the footer and the keypress is a true
+        no-op — instead of exiting, re-downloading, and landing back on the
+        same volume.
+        """
+        if action == "open_parent" and not self._can_navigate_parent:
+            return False
+        return True
+
+    def action_quit_app(self) -> None:
+        self.exit()
+
+    def action_open_parent(self) -> None:
+        """Exit and signal the launcher to re-open explore on the parent
+        version. Guarded by :meth:`check_action`, so this only fires when a
+        navigable parent exists.
+        """
+        if not self._can_navigate_parent:
+            self.notify("No parent version to open.", severity="warning")
+            return
+        self.navigate = "parent"
+        self.exit()
+
+    def action_refresh_tree(self) -> None:
+        # Re-populate the tree against the same reader. The IndexReader is
+        # snapshot-based (already-downloaded index), so this is purely
+        # a state reset — useful after an accidental collapse-all.
+        self.query_one("#file-tree", VolumeFileTree).reload()