pin saved selections

pin_saved_selections.py

"""CLI for managing Dagster+ saved selections (catalog views) and homepage pins.

A *saved selection* is a ``CatalogView`` — a named, filterable view of the
asset catalog (e.g. "all assets owned by alice", "assets tagged core_kpis").
A pinned saved selection appears on a user's Dagster+ homepage and in the
left-nav saved-selections list.

Besides asset selections, the homepage can also pin **jobs** and asset
groups. A pinned job is stored as a (public) ``CatalogView`` of type ``JOB``
and surfaces as a ``JobPinnableItem``.

Subcommands
-----------
``create``        Create a saved selection from a Dagster+ asset selection
                  query. With ``--pin``, also pin it to the (owner's)
                  homepage in one step.
``pin``           Pin an existing saved selection to a user's homepage.
``pin-jobs-by-tag``  Find jobs whose tags match a filter and pin each one to
                  one or more users' homepages.

Targeting another user (admin mode)
-----------------------------------
By default, actions are scoped to the user whose API token is in
``DAGSTER_CLOUD_API_TOKEN``. To create a view *for*, or pin into, another
user's list, pass ``--owner-email`` / ``--user-email``. That requires the
token to hold ``EDIT_OTHER_USERS_CATALOG_VIEWS`` (admin role). The
admin-on-behalf-of paths shipped in dagster-io/internal#24541.

Owners vs pin targets: a *public* view has no owner and can be pinned for
any user. A *private* view has an owner and (server-side ``_check_pinnable``)
can only be pinned for that owner. The ``create`` subcommand reflects this:

- ``--pin`` alone pins to ``--owner-email`` (or the caller) — always safe.
- ``--pin <email>`` pins to a specific user, distinct from the owner —
  only valid when the view is public (``--no-private``); the server will
  reject the pin otherwise.

The standalone ``pin`` subcommand follows the same rule via ``--user-email``.

Pinning jobs is simpler: ``addPinnedJob`` always creates a *public* JOB
view, so it is never subject to the private-owner restriction. Pinning a
job into another user's list still requires ``EDIT_OTHER_USERS_CATALOG_VIEWS``
(the same admin gate as everything else). Jobs have no server-side tag
filter, so ``pin-jobs-by-tag`` enumerates jobs and matches tags client-side.

Verification
------------
After each mutation the script issues the matching read query
(``catalogViews(ownerEmail=...)`` / ``pinnedItems(ownerEmail=...)``) and
confirms the item is visible from the target user's perspective. Exit
code is non-zero if verification fails.

Environment
-----------
    DAGSTER_CLOUD_ORGANIZATION   org slug (e.g. ``hooli``)
    DAGSTER_CLOUD_DEPLOYMENT     deployment name (e.g. ``data-eng-prod``)
    DAGSTER_CLOUD_API_TOKEN      ``user:...`` token; admin token required
                                 when targeting another user

Examples
--------
    # Create + pin in one step, for another user (admin token required)
    uv run --with requests pin_saved_selections.py create \\
        --name "Core KPIs" \\
        --description 'Assets tagged core_kpis' \\
        --query-selection 'tag:"core_kpis"' \\
        --owner-email christian@dagsterlabs.com \\
        --pin

    # Just create (no pin); print the id and stop
    uv run --with requests pin_saved_selections.py create \\
        --name "My KPIs" --query-selection 'tag:"core_kpis"'

    # Pin an existing view into another user's list
    uv run --with requests pin_saved_selections.py pin \\
        --catalog-view-id <id> \\
        --user-email christian@dagsterlabs.com

    # Create a *public* view owned by alice and pin it for bob
    uv run --with requests pin_saved_selections.py create \\
        --name "Shared KPIs" --query-selection 'tag:"core_kpis"' \\
        --no-private \\
        --owner-email alice@example.com \\
        --pin bob@example.com

    # Preview which jobs carry team:data — pin nothing
    uv run --with requests pin_saved_selections.py pin-jobs-by-tag \\
        --tag team=data --dry-run

    # Pin every job tagged team=data AND tier=critical to two users
    # (admin token required to target users other than the caller)
    uv run --with requests pin_saved_selections.py pin-jobs-by-tag \\
        --tag team=data --tag tier=critical --match all \\
        --user-email christian@dagsterlabs.com \\
        --user-email alice@example.com
"""

from __future__ import annotations

import argparse
import os
import sys
from typing import Any

import requests

CREATE_OR_UPDATE_CATALOG_VIEW = """
mutation CreateOrUpdateCatalogView(
    $id: String,
    $name: String!,
    $description: String!,
    $icon: String!,
    $isPrivate: Boolean!,
    $selection: CatalogViewSelectionInput!,
    $ownerEmail: String
) {
    createOrUpdateCatalogView(
        id: $id,
        name: $name,
        description: $description,
        icon: $icon,
        isPrivate: $isPrivate,
        selection: $selection,
        ownerEmail: $ownerEmail
    ) {
        __typename
        ... on CatalogView { id name }
        ... on PythonError { message stack }
        ... on UnauthorizedError { message }
        ... on SelectionNotResolvableError { message }
        ... on SelectionCantBeDeletedError { message }
    }
}
"""

ADD_PINNED_CATALOG_VIEW = """
mutation AddPinnedCatalogView($catalogViewId: String!, $userEmail: String) {
    addPinnedCatalogView(catalogViewId: $catalogViewId, userEmail: $userEmail) {
        __typename
        ... on AddPinnedItemSuccess {
            pinned
            item { ... on CatalogView { id name } }
        }
        ... on PythonError { message stack }
        ... on UnauthorizedError { message }
    }
}
"""

ADD_PINNED_JOB = """
mutation AddPinnedJob(
    $jobName: String!,
    $codeLocationName: String!,
    $repositoryName: String!,
    $userEmail: String
) {
    addPinnedJob(
        jobName: $jobName,
        codeLocationName: $codeLocationName,
        repositoryName: $repositoryName,
        userEmail: $userEmail
    ) {
        __typename
        ... on AddPinnedItemSuccess {
            pinned
            item {
                ... on JobPinnableItem {
                    id
                    name
                    selection { jobNames codeLocationName repositoryName }
                }
            }
        }
        ... on PythonError { message stack }
        ... on UnauthorizedError { message }
    }
}
"""

LIST_CATALOG_VIEWS = """
query CatalogViews($ownerEmail: String) {
    catalogViews(ownerEmail: $ownerEmail) {
        id
        name
        description
    }
}
"""

LIST_PINNED_ITEMS = """
query PinnedItems($ownerEmail: String) {
    pinnedItems(ownerEmail: $ownerEmail) {
        __typename
        ... on CatalogView { id name }
        ... on JobPinnableItem {
            id
            name
            selection { jobNames codeLocationName repositoryName }
        }
    }
}
"""

# Enumerate every job in the deployment along with its tags and the
# repository / code location it lives in. Jobs have no server-side tag
# filter, so ``pin-jobs-by-tag`` matches tags on the client.
LIST_JOBS_WITH_TAGS = """
query JobsWithTags {
    repositoriesOrError {
        __typename
        ... on RepositoryConnection {
            nodes {
                name
                location { name }
                jobs {
                    name
                    tags { key value }
                }
            }
        }
        ... on PythonError { message stack }
    }
}
"""


def _empty_selection(query_selection: str) -> dict[str, Any]:
    # ``CatalogViewSelectionInput`` requires every facet list, even when empty.
    return {
        "tags": [],
        "kinds": [],
        "owners": [],
        "groups": [],
        "codeLocations": [],
        "columns": [],
        "tableNames": [],
        "columnTags": [],
        "querySelection": query_selection,
    }


class DagsterPlusClient:
    def __init__(self, organization: str, deployment: str, token: str) -> None:
        self.url = f"https://{organization}.dagster.cloud/{deployment}/graphql"
        self.headers = {
            "Content-Type": "application/json",
            "Dagster-Cloud-Api-Token": token,
        }

    def execute(self, query: str, variables: dict[str, Any]) -> dict[str, Any]:
        resp = requests.post(
            self.url,
            json={"query": query, "variables": variables},
            headers=self.headers,
            timeout=30,
        )
        if resp.status_code >= 400:
            # Surface the server's response body — for a 500 it usually carries
            # the server-side error/traceback, which raise_for_status() hides.
            raise RuntimeError(
                f"HTTP {resp.status_code} from {self.url}\n{resp.text}"
            )
        payload = resp.json()
        if "errors" in payload:
            raise RuntimeError(f"GraphQL errors: {payload['errors']}")
        return payload["data"]

    def create_catalog_view(
        self,
        *,
        name: str,
        description: str,
        selection: dict[str, Any],
        icon: str = "bookmark",
        is_private: bool = True,
        owner_email: str | None = None,
    ) -> str:
        data = self.execute(
            CREATE_OR_UPDATE_CATALOG_VIEW,
            {
                "name": name,
                "description": description,
                "icon": icon,
                "isPrivate": is_private,
                "selection": selection,
                "ownerEmail": owner_email,
            },
        )
        result = data["createOrUpdateCatalogView"]
        if result["__typename"] != "CatalogView":
            raise RuntimeError(f"Failed to create catalog view {name!r}: {result}")
        return result["id"]

    def pin_catalog_view(self, catalog_view_id: str, user_email: str | None = None) -> str:
        data = self.execute(
            ADD_PINNED_CATALOG_VIEW,
            {"catalogViewId": catalog_view_id, "userEmail": user_email},
        )
        result = data["addPinnedCatalogView"]
        if result["__typename"] != "AddPinnedItemSuccess" or not result["pinned"]:
            raise RuntimeError(f"Failed to pin catalog view {catalog_view_id}: {result}")
        return result["item"]["name"]

    def pin_job(
        self,
        *,
        job_name: str,
        code_location_name: str,
        repository_name: str,
        user_email: str | None = None,
    ) -> str:
        """Pin a single job to a user's homepage; returns the pinned view id."""
        data = self.execute(
            ADD_PINNED_JOB,
            {
                "jobName": job_name,
                "codeLocationName": code_location_name,
                "repositoryName": repository_name,
                "userEmail": user_email,
            },
        )
        result = data["addPinnedJob"]
        if result["__typename"] != "AddPinnedItemSuccess" or not result["pinned"]:
            raise RuntimeError(f"Failed to pin job {job_name!r}: {result}")
        return result["item"]["id"]

    def list_catalog_views(self, owner_email: str | None = None) -> list[dict[str, Any]]:
        data = self.execute(LIST_CATALOG_VIEWS, {"ownerEmail": owner_email})
        return data["catalogViews"]

    def list_pinned_items(self, owner_email: str | None = None) -> list[dict[str, Any]]:
        data = self.execute(LIST_PINNED_ITEMS, {"ownerEmail": owner_email})
        return data["pinnedItems"]

    def list_jobs_with_tags(self) -> list[dict[str, Any]]:
        """Return every job in the deployment with its tags and location.

        Each item: ``{"name", "repository_name", "code_location_name",
        "tags": {key: value}}``.
        """
        data = self.execute(LIST_JOBS_WITH_TAGS, {})
        result = data["repositoriesOrError"]
        if result["__typename"] != "RepositoryConnection":
            raise RuntimeError(f"Failed to list jobs: {result}")
        jobs: list[dict[str, Any]] = []
        for repo in result["nodes"]:
            for job in repo["jobs"]:
                jobs.append(
                    {
                        "name": job["name"],
                        "repository_name": repo["name"],
                        "code_location_name": repo["location"]["name"],
                        "tags": {t["key"]: t["value"] for t in job["tags"]},
                    }
                )
        return jobs


def _pin_and_verify(
    client: DagsterPlusClient, catalog_view_id: str, user_email: str | None
) -> int:
    pinned_name = client.pin_catalog_view(catalog_view_id, user_email=user_email)
    target = user_email or "the authenticated user"
    print(f"Pinned saved selection {pinned_name!r} to {target}'s homepage")

    pinned = client.list_pinned_items(owner_email=user_email)
    matched = next(
        (
            p
            for p in pinned
            if p.get("__typename") == "CatalogView" and p.get("id") == catalog_view_id
        ),
        None,
    )
    if matched is None:
        print(
            f"FAIL: {catalog_view_id} not found in pinnedItems(ownerEmail={user_email!r}) "
            f"({len(pinned)} items returned)"
        )
        return 1
    print(
        f"Verified: {catalog_view_id} present in pinnedItems for {target} "
        f"({len(pinned)} items returned)"
    )
    return 0


def _job_is_pinned(
    pinned: list[dict[str, Any]],
    *,
    job_name: str,
    code_location_name: str,
    repository_name: str,
) -> bool:
    """True if ``pinnedItems`` contains a JobPinnableItem for this exact job."""
    for item in pinned:
        if item.get("__typename") != "JobPinnableItem":
            continue
        sel = item.get("selection") or {}
        if (
            job_name in (sel.get("jobNames") or [])
            and sel.get("codeLocationName") == code_location_name
            and sel.get("repositoryName") == repository_name
        ):
            return True
    return False


def _pin_job_and_verify(
    client: DagsterPlusClient, job: dict[str, Any], user_email: str | None
) -> int:
    target = user_email or "the authenticated user"
    label = f"{job['name']} ({job['code_location_name']}/{job['repository_name']})"
    client.pin_job(
        job_name=job["name"],
        code_location_name=job["code_location_name"],
        repository_name=job["repository_name"],
        user_email=user_email,
    )
    print(f"Pinned job {label} to {target}'s homepage")

    pinned = client.list_pinned_items(owner_email=user_email)
    if not _job_is_pinned(
        pinned,
        job_name=job["name"],
        code_location_name=job["code_location_name"],
        repository_name=job["repository_name"],
    ):
        print(
            f"FAIL: job {label} not found in pinnedItems(ownerEmail={user_email!r}) "
            f"({len(pinned)} items returned)"
        )
        return 1
    print(f"Verified: job {label} present in pinnedItems for {target}")
    return 0


def _parse_tag_filters(raw_tags: list[str]) -> list[tuple[str, str | None]]:
    """Parse ``--tag`` values into (key, value-or-None) pairs.

    ``KEY=VALUE`` matches that exact key/value. ``KEY`` (no ``=``) matches
    any job carrying that key, regardless of value.
    """
    filters: list[tuple[str, str | None]] = []
    for raw in raw_tags:
        key, sep, value = raw.partition("=")
        key = key.strip()
        if not key:
            raise ValueError(f"Invalid --tag {raw!r}: key must be non-empty.")
        filters.append((key, value if sep else None))
    return filters


def _job_matches(
    job_tags: dict[str, str],
    filters: list[tuple[str, str | None]],
    match_all: bool,
) -> bool:
    def matches_one(key: str, value: str | None) -> bool:
        if key not in job_tags:
            return False
        return value is None or job_tags[key] == value

    results = (matches_one(k, v) for k, v in filters)
    return all(results) if match_all else any(results)


def cmd_create(args: argparse.Namespace) -> int:
    client = _client_from_env()
    selection = _empty_selection(args.query_selection)
    view_id = client.create_catalog_view(
        name=args.name,
        description=args.description,
        selection=selection,
        icon=args.icon,
        is_private=args.private,
        owner_email=args.owner_email,
    )
    print(f"Created catalog view {args.name!r}: {view_id}")

    target = args.owner_email or "the authenticated user"
    views = client.list_catalog_views(owner_email=args.owner_email)
    matched = next((v for v in views if v["id"] == view_id), None)
    if matched is None:
        print(
            f"FAIL: {view_id} not found in catalogViews(ownerEmail={args.owner_email!r}) "
            f"({len(views)} views returned)"
        )
        return 1
    print(
        f"Verified: {view_id} visible in catalogViews for {target} "
        f"({len(views)} views returned)"
    )

    if args.pin is not None:
        # --pin alone (args.pin == "") → pin to the owner / caller.
        # --pin EMAIL → pin to EMAIL (only valid for public views).
        pin_target = args.pin or args.owner_email
        return _pin_and_verify(client, view_id, pin_target)
    return 0


def cmd_pin(args: argparse.Namespace) -> int:
    client = _client_from_env()
    return _pin_and_verify(client, args.catalog_view_id, args.user_email)


def cmd_pin_jobs_by_tag(args: argparse.Namespace) -> int:
    filters = _parse_tag_filters(args.tag)
    match_all = args.match == "all"
    client = _client_from_env()

    all_jobs = client.list_jobs_with_tags()
    matching = [j for j in all_jobs if _job_matches(j["tags"], filters, match_all)]

    pretty = " AND ".join(
        f"{k}={v}" if v is not None else f"{k}=*" for k, v in filters
    ) if match_all else " OR ".join(
        f"{k}={v}" if v is not None else f"{k}=*" for k, v in filters
    )
    print(
        f"Matched {len(matching)} of {len(all_jobs)} jobs against tag filter [{pretty}]:"
    )
    for job in matching:
        print(f"  - {job['name']} ({job['code_location_name']}/{job['repository_name']})")

    if not matching:
        print("No jobs matched; nothing to pin.")
        return 0

    if args.dry_run:
        print("\n--dry-run: no pins were created.")
        return 0

    # --user-email may be repeated; default ([]) means pin for the caller.
    targets: list[str | None] = list(args.user_email) if args.user_email else [None]

    exit_code = 0
    for user_email in targets:
        for job in matching:
            exit_code |= _pin_job_and_verify(client, job, user_email)
    return exit_code


def _client_from_env() -> DagsterPlusClient:
    return DagsterPlusClient(
        organization=os.environ["DAGSTER_CLOUD_ORGANIZATION"],
        deployment=os.environ["DAGSTER_CLOUD_DEPLOYMENT"],
        token=os.environ["DAGSTER_CLOUD_API_TOKEN"],
    )


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Manage Dagster+ saved selections (catalog views) and homepage pins.",
    )
    sub = parser.add_subparsers(dest="cmd", required=True)

    p_create = sub.add_parser("create", help="Create a saved selection (catalog view).")
    p_create.add_argument("--name", required=True)
    p_create.add_argument("--description", default="")
    p_create.add_argument(
        "--query-selection",
        required=True,
        help='Dagster+ asset selection query, e.g. \'owner:"a@b.com"\' or \'tag:"core_kpis"\'.',
    )
    p_create.add_argument(
        "--private",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Private view (default) or --no-private for a public view.",
    )
    p_create.add_argument("--icon", default="bookmark")
    p_create.add_argument(
        "--owner-email",
        default=None,
        help="Email of the user who should own the view. Targeting another user requires an admin token.",
    )
    p_create.add_argument(
        "--pin",
        nargs="?",
        const="",
        default=None,
        metavar="EMAIL",
        help=(
            "After creating, also pin the new view. With no value, pins to --owner-email "
            "(or the caller). With an email value, pins to that user — only valid for "
            "public views (private views can only be pinned for their owner)."
        ),
    )
    p_create.set_defaults(func=cmd_create)

    p_pin = sub.add_parser("pin", help="Pin a saved selection to a user's homepage.")
    p_pin.add_argument("--catalog-view-id", required=True, help="ID returned by `create`.")
    p_pin.add_argument(
        "--user-email",
        default=None,
        help="Email of the user to pin for. Targeting another user requires an admin token.",
    )
    p_pin.set_defaults(func=cmd_pin)

    p_jobs = sub.add_parser(
        "pin-jobs-by-tag",
        help="Find jobs whose tags match a filter and pin each to users' homepages.",
    )
    p_jobs.add_argument(
        "--tag",
        action="append",
        required=True,
        metavar="KEY[=VALUE]",
        help=(
            "Tag filter; repeatable. 'KEY=VALUE' matches that exact tag; bare 'KEY' "
            "matches any value for that key. Combine multiple with --match."
        ),
    )
    p_jobs.add_argument(
        "--match",
        choices=("any", "all"),
        default="any",
        help="With multiple --tag filters, require ANY (default) or ALL to match.",
    )
    p_jobs.add_argument(
        "--user-email",
        action="append",
        default=None,
        metavar="EMAIL",
        help=(
            "User to pin matching jobs for; repeatable. Omit to pin for the caller. "
            "Targeting other users requires an admin token (EDIT_OTHER_USERS_CATALOG_VIEWS)."
        ),
    )
    p_jobs.add_argument(
        "--dry-run",
        action="store_true",
        help="List matching jobs without pinning anything.",
    )
    p_jobs.set_defaults(func=cmd_pin_jobs_by_tag)

    return parser


def main(argv: list[str] | None = None) -> int:
    args = _build_parser().parse_args(argv)
    return args.func(args)


if __name__ == "__main__":
    sys.exit(main())