prefect-cloud 0.1.0__py3-none-any.whl

prefect_cloud/auth.py

@@ -0,0 +1,419 @@
+import json
+import os
+import socket
+import threading
+import webbrowser
+from contextlib import asynccontextmanager, contextmanager
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from pathlib import Path
+from queue import Queue
+from typing import Any, AsyncGenerator, Generator, Sequence
+from urllib.parse import parse_qs, quote, urlparse
+from uuid import UUID
+
+import toml
+from pydantic import BaseModel, TypeAdapter
+
+from prefect_cloud.client import PrefectCloudClient, SyncPrefectCloudClient
+from prefect_cloud.utilities.tui import prompt_select_from_list
+
+
+def _get_cloud_urls() -> tuple[str, str]:
+    """Get the appropriate Cloud UI and API URLs based on environment"""
+
+    env = os.environ.get("CLOUD_ENV")
+    from_env = os.environ.get("PREFECT_CLOUD_API_URL")
+
+    if from_env:
+        cloud_api_url = from_env
+    elif env in ("prd", "prod", None):
+        cloud_api_url = "https://api.prefect.cloud/api"
+    elif env == "stg":
+        cloud_api_url = "https://api.stg.prefect.dev/api"
+    elif env == "dev":
+        cloud_api_url = "https://api.prefect.dev/api"
+    elif env == "lcl":
+        cloud_api_url = "http://localhost:8000/api"
+    else:
+        raise ValueError(f"Invalid CLOUD_ENV: {env}")
+
+    if os.environ.get("PREFECT_CLOUD_UI_URL"):
+        cloud_ui_url = os.environ["PREFECT_CLOUD_UI_URL"]
+    else:
+        cloud_ui_url = cloud_api_url.replace("https://api.", "https://app.").replace(
+            "/api", ""
+        )
+
+    return cloud_ui_url, cloud_api_url
+
+
+CLOUD_UI_URL, CLOUD_API_URL = _get_cloud_urls()
+
+PREFECT_HOME = Path.home() / ".prefect"
+
+
+class Account(BaseModel):
+    account_id: UUID
+    account_name: str
+    account_handle: str
+
+
+class Workspace(BaseModel):
+    account_id: UUID
+    account_name: str
+    account_handle: str
+    workspace_id: UUID
+    workspace_name: str
+    workspace_handle: str
+
+    @property
+    def full_handle(self) -> str:
+        return f"{self.account_handle}/{self.workspace_handle}"
+
+    @property
+    def api_url(self) -> str:
+        return (
+            f"{CLOUD_API_URL}/accounts/{self.account_id}/workspaces/{self.workspace_id}"
+        )
+
+
+class Me(BaseModel):
+    id: UUID
+    email: str
+    first_name: str
+    last_name: str
+    handle: str
+
+
+async def login(api_key: str | None = None, workspace_id_or_slug: str | None = None):
+    """Logs the user into Prefect Cloud interactively, setting their active profile"""
+    if not api_key:
+        api_key = get_api_key()
+
+    if not api_key or not await key_is_valid(api_key):
+        api_key = login_interactively()
+        if not api_key:
+            return
+
+    workspace: Workspace | None = None
+    if workspace_id_or_slug:
+        workspace = await lookup_workspace(api_key, workspace_id_or_slug)
+
+    if not workspace:
+        workspace = await prompt_for_workspace(api_key)
+        if not workspace:
+            return
+
+    set_cloud_profile(api_key, workspace)
+
+
+def logout():
+    """Logs the user out of Prefect Cloud"""
+    remove_cloud_profile()
+
+
+@asynccontextmanager
+async def cloud_client(api_key: str) -> AsyncGenerator[PrefectCloudClient, None]:
+    """Creates a client for the Prefect Cloud API"""
+
+    async with PrefectCloudClient(api_url=CLOUD_API_URL, api_key=api_key) as client:
+        yield client
+
+
+@contextmanager
+def sync_cloud_client(api_key: str) -> Generator[SyncPrefectCloudClient, None, None]:
+    """Creates a client for the Prefect Cloud API"""
+    with SyncPrefectCloudClient(api_url=CLOUD_API_URL, api_key=api_key) as client:
+        yield client
+
+
+async def get_prefect_cloud_client() -> PrefectCloudClient:
+    _, api_url, api_key = await get_cloud_urls_or_login()
+    return PrefectCloudClient(
+        api_url=api_url,
+        api_key=api_key,
+    )
+
+
+async def get_cloud_urls_or_login() -> tuple[str, str, str]:
+    """Gets the cloud UI URL, API URL, and API key"""
+    ui_url, api_url, api_key = get_cloud_urls_without_login()
+    if not ui_url or not api_url or not api_key:
+        await login()
+
+    ui_url, api_url, api_key = get_cloud_urls_without_login()
+    if not ui_url or not api_url or not api_key:
+        raise ValueError("No cloud profile found")
+
+    return ui_url, api_url, api_key
+
+
+def get_cloud_urls_without_login() -> tuple[str | None, str | None, str | None]:
+    """Gets the cloud UI URL, API URL, and API key"""
+    api_url: str | None = get_api_url()
+    if not api_url:
+        return None, None, None
+
+    ui_url = (
+        api_url.replace("https://api.", "https://app.")
+        .replace("/api", "")
+        .replace("/accounts/", "/account/")
+        .replace("/workspaces/", "/workspace/")
+    )
+
+    api_key = get_api_key()
+    if not api_key:
+        return None, None, None
+
+    return ui_url, api_url, api_key
+
+
+def get_api_key_or_login() -> str:
+    """Gets a validated API key or logs the user in if no API key is available"""
+    api_key = get_api_key()
+    if not api_key:
+        api_key = login_interactively()
+        if not api_key:
+            raise ValueError("No API key found")
+    return api_key
+
+
+async def key_is_valid(api_key: str) -> bool:
+    """Checks if the given API key is valid"""
+    async with cloud_client(api_key) as client:
+        response = await client.request("GET", "/me/")
+        return response.status_code == 200
+
+
+def login_interactively() -> str | None:
+    """Logs the user into Prefect Cloud interactively"""
+    result_queue: Queue[str | None] = Queue()
+    with login_server(result_queue) as uri:
+        login_url = f"{CLOUD_UI_URL}/auth/client?callback={quote(uri)}&g=true"
+
+        threading.Thread(
+            target=webbrowser.open_new_tab, args=(login_url,), daemon=True
+        ).start()
+
+        return result_queue.get()
+
+
+@contextmanager
+def login_server(result_queue: Queue[str | None]) -> Generator[str, None, None]:
+    """Runs a local server to handle the callback from Prefect Cloud"""
+
+    class LoginHandler(BaseHTTPRequestHandler):
+        def log_message(self, format: str, *args: Any) -> None:
+            pass
+
+        def add_cors_headers(self) -> None:
+            self.send_header("Access-Control-Allow-Origin", "*")
+            self.send_header("Access-Control-Allow-Headers", "*")
+            self.send_header("Access-Control-Allow-Methods", "POST, OPTIONS")
+
+        def do_OPTIONS(self) -> None:
+            self.send_response(200)
+            self.add_cors_headers()
+            self.end_headers()
+
+        def do_GET(self) -> None:
+            query_params = parse_qs(urlparse(self.path).query)
+            api_key = query_params.get("key", [""])[0] or None
+
+            self.send_response(200)
+            self.send_header("Content-type", "text/html")
+            self.end_headers()
+            self.wfile.write(b"""
+                <script>
+                if (window.opener) {
+                    window.opener.postMessage('auth-complete', '*')
+                }
+                </script>
+                <p>You can close this window.</p>
+            """)
+            result_queue.put(api_key or None)
+
+        def do_POST(self):
+            content_length = int(self.headers.get("Content-Length", 0))
+            body = self.rfile.read(content_length)
+            data = json.loads(body)
+            api_key = data.get("api_key")
+
+            self.send_response(200)
+            self.add_cors_headers()
+            self.end_headers()
+            self.wfile.write(b"{}")
+            result_queue.put(api_key or None)
+
+    with socket.socket() as sock:
+        sock.bind(("", 0))
+        port = sock.getsockname()[1]
+
+    server = HTTPServer(("localhost", port), LoginHandler)
+    threading.Thread(target=server.serve_forever, daemon=True).start()
+
+    try:
+        yield f"http://localhost:{port}"
+    finally:
+        server.server_close()
+
+
+async def me(api_key: str) -> Me:
+    """Gets the current user's information"""
+    async with cloud_client(api_key) as client:
+        response = await client.request("GET", "/me/")
+        response.raise_for_status()
+
+    return Me.model_validate_json(response.text)
+
+
+async def get_accounts(api_key: str) -> Sequence[Account]:
+    """Gets the list of accounts for the current user"""
+    async with cloud_client(api_key) as client:
+        response = await client.request("GET", "/me/accounts")
+        response.raise_for_status()
+
+        return TypeAdapter(list[Account]).validate_json(response.text)
+
+
+async def get_workspaces(api_key: str) -> Sequence[Workspace]:
+    """Gets the list of workspaces for the current user"""
+    async with cloud_client(api_key) as client:
+        response = await client.request("GET", "/me/workspaces")
+        response.raise_for_status()
+
+        workspaces = TypeAdapter(list[Workspace]).validate_json(response.text)
+        workspaces.sort(key=lambda w: w.full_handle)
+
+        return workspaces
+
+
+async def lookup_workspace(api_key: str, workspace_id_or_slug: str) -> Workspace | None:
+    """Looks up a workspace by ID or slug"""
+    workspace_id: UUID | None = None
+    try:
+        workspace_id = UUID(workspace_id_or_slug)
+    except ValueError:
+        pass
+
+    for workspace in await get_workspaces(api_key):
+        if workspace_id and workspace.workspace_id == workspace_id:
+            return workspace
+        if workspace.full_handle == workspace_id_or_slug:
+            return workspace
+        if workspace.workspace_handle == workspace_id_or_slug:
+            return workspace
+
+    return None
+
+
+async def prompt_for_workspace(api_key: str) -> Workspace | None:
+    """Prompts the user to select a workspace from the list of available workspaces"""
+    workspaces = await get_workspaces(api_key)
+
+    if not workspaces:
+        return None
+
+    if len(workspaces) == 1:
+        return workspaces[0]
+
+    selected = prompt_select_from_list(
+        "Select a workspace",
+        sorted([workspace.full_handle for workspace in workspaces]),
+    )
+    return next(
+        workspace for workspace in workspaces if workspace.full_handle == selected
+    )
+
+
+def get_from_env_or_profile(key: str) -> str | None:
+    """Gets value from the environment or the current cloud profile"""
+    profile = get_cloud_profile() or {}
+
+    if from_env := os.environ.get(key):
+        return from_env
+    elif from_profile := profile.get(key):
+        return from_profile
+    else:
+        return None
+
+
+def get_api_key() -> str | None:
+    return get_from_env_or_profile("PREFECT_API_KEY")
+
+
+def get_api_url() -> str | None:
+    return get_from_env_or_profile("PREFECT_API_URL")
+
+
+def load_profiles() -> dict[str, Any]:
+    """Loads the profiles from the profiles file"""
+    profile_path = PREFECT_HOME / "profiles.toml"
+    if profile_path.exists():
+        profiles = toml.load(profile_path)
+    else:
+        profiles = {}
+
+    if "profiles" not in profiles:
+        profiles["profiles"] = {}
+
+    return profiles
+
+
+def cloud_profile_name() -> str:
+    """Returns the name of the current cloud profile"""
+    if CLOUD_API_URL == "https://api.stg.prefect.dev/api":
+        return "prefect-cloud-stg"
+    elif CLOUD_API_URL == "https://api.prefect.dev/api":
+        return "prefect-cloud-dev"
+    elif CLOUD_API_URL.startswith("http://localhost"):
+        return "prefect-cloud-lcl"
+    else:
+        return "prefect-cloud"
+
+
+def get_cloud_profile() -> dict[str, str] | None:
+    """Returns the current cloud profile"""
+    profile_path = PREFECT_HOME / "profiles.toml"
+    if not profile_path.exists():
+        return None
+
+    profiles = toml.load(profile_path)
+    return profiles.get("profiles", {}).get(cloud_profile_name())
+
+
+def set_cloud_profile(api_key: str, workspace: Workspace) -> None:
+    """Writes the current cloud profile"""
+    profile_name = cloud_profile_name()
+    profile_path = PREFECT_HOME / "profiles.toml"
+
+    profiles = load_profiles()
+
+    profiles["active"] = profile_name
+    if profile_name not in profiles["profiles"]:
+        profiles["profiles"][profile_name] = {}
+
+    profiles["profiles"][profile_name].update(
+        {
+            "PREFECT_API_KEY": api_key,
+            "PREFECT_API_URL": workspace.api_url,
+        }
+    )
+
+    profile_path.parent.mkdir(parents=True, exist_ok=True)
+    profile_path.write_text(toml.dumps(profiles))
+
+
+def remove_cloud_profile() -> None:
+    """Removes the current cloud profile"""
+    profile_path = PREFECT_HOME / "profiles.toml"
+    if not profile_path.exists():
+        return
+
+    profiles = load_profiles()
+    profiles["profiles"].pop(cloud_profile_name(), None)
+    try:
+        profiles["active"] = list(profiles["profiles"].keys())[0]
+        profile_path.write_text(toml.dumps(profiles))
+    except IndexError:
+        profile_path.unlink()

prefect_cloud/cli/completions.py

@@ -0,0 +1,54 @@
+import json
+import time
+from pathlib import Path
+
+from prefect_cloud.auth import get_cloud_urls_without_login, sync_cloud_client
+
+COMPLETION_CACHE = Path.home() / ".prefect" / "prefect-cloud-completions.json"
+CACHE_TTL = 86400
+
+
+def clear_cache():
+    if COMPLETION_CACHE.exists():
+        COMPLETION_CACHE.unlink()
+
+
+def complete_deployment(incomplete: str) -> list[str]:
+    _, api_url, api_key = get_cloud_urls_without_login()
+    if not api_url or not api_key:
+        return []
+
+    deployment_names = None
+    if (
+        COMPLETION_CACHE.exists()
+        and time.time() - COMPLETION_CACHE.stat().st_mtime < CACHE_TTL
+    ):
+        try:
+            with open(COMPLETION_CACHE) as f:
+                cache = json.load(f)
+                deployment_names = cache["deployment_names"]
+        except (IOError, json.JSONDecodeError, KeyError):
+            deployment_names = []
+
+    if deployment_names is None:
+        with sync_cloud_client(api_key) as client:
+            response = client.post(f"{api_url}/deployments/filter")
+            response.raise_for_status()
+            deployments = response.json()
+
+            response = client.post(f"{api_url}/flows/filter")
+            response.raise_for_status()
+            flows = response.json()
+
+        flow_names = {flow["id"]: flow["name"] for flow in flows}
+
+        deployment_names = [
+            f"{flow_names[deployment['flow_id']]}/{deployment['name']}"
+            for deployment in deployments
+        ]
+
+        COMPLETION_CACHE.parent.mkdir(parents=True, exist_ok=True)
+        with open(COMPLETION_CACHE, "w") as f:
+            json.dump({"deployment_names": deployment_names}, f)
+
+    return [name for name in deployment_names if name.startswith(incomplete)]