dagflow-cli-plugin 0.1.0__py3-none-any.whl

dagflow_cli/__init__.py

@@ -0,0 +1 @@
+"""Snowflake CLI plugin support for Dagflow."""

dagflow_cli/commands.py

@@ -0,0 +1,423 @@
+from __future__ import annotations
+
+import tempfile
+import tomllib
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Iterator
+
+import click
+import typer
+from snowflake.cli.api.commands.decorators import with_output
+from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
+from snowflake.cli.api.output.types import CommandResult, MultipleResults, ObjectResult, SingleQueryResult
+
+from dagflow_cli.launch import build_launch_execution_params, launch_job, load_run_config, parse_tags
+from dagflow_cli.manager import DagflowCliManager, Entrypoint
+from dagflow_cli.pex import (
+    DEFAULT_PYTHON_VERSION,
+    PexBuildConfig,
+    build_pex_artifacts,
+    default_output_dir,
+    default_version,
+)
+from dagflow_cli.workspace import (
+    reload_repository_location,
+    reload_workspace,
+    wait_for_code_location_loaded,
+)
+
+
+app = SnowTyperFactory(name="dg", help="Manage Dagflow deployments")
+code_location_app = SnowTyperFactory(name="code-location", help="Manage Dagflow code locations")
+job_app = SnowTyperFactory(name="job", help="Manage Dagflow jobs")
+
+
+@dataclass(frozen=True)
+class DgProjectDefaults:
+    path: Path
+    location_name: str
+    module_name: str
+
+
+def _entrypoint(module_name: str | None, package_name: str | None, python_file: Path | None) -> Entrypoint:
+    selected = [value is not None for value in (module_name, package_name, python_file)]
+    if sum(selected) != 1:
+        raise typer.BadParameter("pass exactly one of --module-name, --package-name, or --python-file")
+    if module_name is not None:
+        return Entrypoint("module", module_name)
+    if package_name is not None:
+        return Entrypoint("package", package_name)
+    assert python_file is not None
+    return Entrypoint("python_file", python_file.as_posix())
+
+
+def _looks_like_dg_project(path: Path) -> bool:
+    project_path = path.expanduser().resolve()
+    if (project_path / "dg.toml").exists():
+        return True
+    pyproject = project_path / "pyproject.toml"
+    if not pyproject.exists():
+        return False
+    try:
+        config = tomllib.loads(pyproject.read_text())
+    except tomllib.TOMLDecodeError:
+        return False
+    return "dg" in config.get("tool", {})
+
+
+def _dg_project_defaults(path: Path) -> DgProjectDefaults | None:
+    project_path = path.expanduser().resolve()
+    if not _looks_like_dg_project(project_path):
+        return None
+
+    try:
+        from dagster_dg_core.context import DgContext
+    except ImportError as error:
+        raise click.ClickException(
+            "This looks like a Dagster dg project, but dagster-dg-core is not installed. "
+            'Install the Dagflow CLI plugin: pip install "dagflow-cli-plugin".'
+        ) from error
+
+    try:
+        context = DgContext.for_project_environment(
+            project_path,
+            {"suppress_warnings": ["project_and_activated_venv_mismatch"]},
+        )
+    except SystemExit as error:
+        raise click.ClickException(f"Could not load Dagster dg project at {project_path}") from error
+    except Exception as error:
+        raise click.ClickException(f"Could not load Dagster dg project at {project_path}: {error}") from error
+
+    build_path = context.root_path
+    build_config = context.build_config
+    if build_config and build_config.get("directory"):
+        build_path = Path(build_config["directory"])
+
+    return DgProjectDefaults(
+        path=build_path,
+        location_name=context.code_location_name,
+        module_name=context.code_location_target_module_name,
+    )
+
+
+def _resolve_code_location_inputs(
+    path: Path,
+    location_name: str | None,
+    module_name: str | None,
+    package_name: str | None,
+    python_file: Path | None,
+) -> tuple[Path, str, Entrypoint]:
+    defaults = _dg_project_defaults(path)
+    resolved_path = defaults.path if defaults else path
+    resolved_location_name = location_name or (defaults.location_name if defaults else None)
+    resolved_module_name = module_name or (
+        defaults.module_name if defaults and package_name is None and python_file is None else None
+    )
+
+    if resolved_location_name is None:
+        raise typer.BadParameter(
+            "missing code location name; pass NAME or run from a Dagster dg project with code_location_name"
+        )
+
+    return (
+        resolved_path,
+        resolved_location_name,
+        _entrypoint(resolved_module_name, package_name, python_file),
+    )
+
+
+def _build_pex_artifacts_for_cli(config: PexBuildConfig):
+    try:
+        return build_pex_artifacts(config)
+    except ValueError as error:
+        raise click.ClickException(str(error)) from None
+
+
+@contextmanager
+def _build_artifacts(
+    path: Path,
+    python_version: str,
+) -> Iterator[tuple[Path, Path]]:
+    with tempfile.TemporaryDirectory(prefix="dagflow-pex-") as temp_dir:
+        artifacts = _build_pex_artifacts_for_cli(
+            PexBuildConfig(
+                project_path=path,
+                output_dir=Path(temp_dir),
+                python_version=python_version,
+            )
+        )
+        yield artifacts.source_pex, artifacts.deps_pex
+
+
+def _require_existing_pex(path: Path, option_name: str) -> Path:
+    resolved = path.expanduser().resolve()
+    if not resolved.exists():
+        raise click.ClickException(f"{option_name} does not exist: {resolved}")
+    if not resolved.is_file():
+        raise click.ClickException(f"{option_name} is not a file: {resolved}")
+    return resolved
+
+
+@code_location_app.command("package")
+def package(
+    path: Path = typer.Argument(..., help="Path to the Dagster code location project"),
+    location_name: str = typer.Option(..., "--location-name", help="Dagflow code location name"),
+    version: str | None = typer.Option(None, "--version", help="PEX artifact version"),
+    module_name: str | None = typer.Option(None, "--module-name", help="Dagster definitions module"),
+    package_name: str | None = typer.Option(None, "--package-name", help="Dagster definitions package"),
+    python_file: Path | None = typer.Option(None, "--python-file", help="Dagster definitions Python file"),
+    python_version: str = typer.Option(DEFAULT_PYTHON_VERSION, "--python-version", help="Target runtime Python version"),
+    output_dir: Path | None = typer.Option(None, "--output-dir", help="Directory for source.pex and deps.pex"),
+) -> None:
+    """Build source.pex and deps.pex for a Dagflow code location."""
+    _entrypoint(module_name, package_name, python_file)
+    resolved_version = version or default_version()
+    artifacts = _build_pex_artifacts_for_cli(
+        PexBuildConfig(
+            project_path=path,
+            output_dir=output_dir or default_output_dir(location_name, resolved_version),
+            python_version=python_version,
+        )
+    )
+    typer.echo(f"source_pex={artifacts.source_pex}")
+    typer.echo(f"deps_pex={artifacts.deps_pex}")
+
+
+def _upload_and_deploy_code_location(
+    app_name: str | None,
+    location_name: str,
+    version: str,
+    entrypoint: Entrypoint,
+    python_version: str,
+    source_pex: Path,
+    deps_pex: Path,
+    blocking: bool,
+) -> CommandResult:
+    manager = DagflowCliManager()
+    manager.upload_artifacts(app_name, location_name, version, source_pex, deps_pex)
+    result = manager.deploy_code_location_api(
+        app_name,
+        location_name,
+        version,
+        entrypoint,
+        python_version,
+        blocking,
+    )
+    return ObjectResult(result)
+
+
+def _deploy_code_location(
+    path: Path,
+    app_name: str | None,
+    location_name: str | None,
+    version: str | None,
+    module_name: str | None,
+    package_name: str | None,
+    python_file: Path | None,
+    python_version: str,
+    blocking: bool,
+) -> CommandResult:
+    resolved_path, resolved_location_name, entrypoint = _resolve_code_location_inputs(
+        path, location_name, module_name, package_name, python_file
+    )
+    resolved_version = version or default_version()
+    with _build_artifacts(
+        path=resolved_path,
+        python_version=python_version,
+    ) as (built_source_pex, built_deps_pex):
+        return _upload_and_deploy_code_location(
+            app_name=app_name,
+            location_name=resolved_location_name,
+            version=resolved_version,
+            entrypoint=entrypoint,
+            python_version=python_version,
+            source_pex=built_source_pex,
+            deps_pex=built_deps_pex,
+            blocking=blocking,
+        )
+
+
+@app.command("deploy", requires_connection=True, requires_global_options=True)
+@with_output
+def deploy(
+    name: str | None = typer.Argument(None, help="Dagflow code location name"),
+    path: Path = typer.Option(Path("."), "--path", help="Path to the Dagster code location project"),
+    app_name: str | None = typer.Option(None, "--app-name", help="Installed Dagflow application name"),
+    version: str | None = typer.Option(None, "--version", help="PEX artifact version"),
+    module_name: str | None = typer.Option(None, "--module-name", help="Dagster definitions module"),
+    package_name: str | None = typer.Option(None, "--package-name", help="Dagster definitions package"),
+    python_file: Path | None = typer.Option(None, "--python-file", help="Dagster definitions Python file"),
+    python_version: str = typer.Option(DEFAULT_PYTHON_VERSION, "--python-version", help="Target runtime Python version"),
+    blocking: bool = typer.Option(True, "--blocking/--no-blocking", help="Wait for code location readiness"),
+    **options: Any,
+) -> CommandResult:
+    """Package and deploy a Dagflow code location."""
+    return _deploy_code_location(
+        path=path,
+        app_name=app_name,
+        location_name=name,
+        version=version,
+        module_name=module_name,
+        package_name=package_name,
+        python_file=python_file,
+        python_version=python_version,
+        blocking=blocking,
+    )
+
+
+@code_location_app.command("add", requires_connection=True, requires_global_options=True)
+@with_output
+def add(
+    name: str = typer.Argument(..., help="Dagflow code location name"),
+    app_name: str | None = typer.Option(None, "--app-name", help="Installed Dagflow application name"),
+    version: str | None = typer.Option(None, "--version", help="PEX artifact version"),
+    module_name: str | None = typer.Option(None, "--module-name", help="Dagster definitions module"),
+    package_name: str | None = typer.Option(None, "--package-name", help="Dagster definitions package"),
+    python_file: Path | None = typer.Option(None, "--python-file", help="Dagster definitions Python file"),
+    python_version: str = typer.Option(DEFAULT_PYTHON_VERSION, "--python-version", help="Target runtime Python version"),
+    source_pex: Path = typer.Option(..., "--source-pex", help="Existing source.pex to upload"),
+    deps_pex: Path = typer.Option(..., "--deps-pex", help="Existing deps.pex to upload"),
+    blocking: bool = typer.Option(True, "--blocking/--no-blocking", help="Wait for code location readiness"),
+    **options: Any,
+) -> CommandResult:
+    """Add or update a Dagflow code location from existing PEX artifacts."""
+    resolved_version = version or default_version()
+    entrypoint = _entrypoint(module_name, package_name, python_file)
+    return _upload_and_deploy_code_location(
+        app_name=app_name,
+        location_name=name,
+        version=resolved_version,
+        entrypoint=entrypoint,
+        python_version=python_version,
+        source_pex=_require_existing_pex(source_pex, "--source-pex"),
+        deps_pex=_require_existing_pex(deps_pex, "--deps-pex"),
+        blocking=blocking,
+    )
+
+
+def _delete_code_location(name: str, app_name: str | None, yes: bool) -> CommandResult:
+    if not yes:
+        typer.confirm(f"Delete code location {name}?", abort=True)
+    manager = DagflowCliManager()
+    return ObjectResult(manager.remove_code_location_api(app_name, name))
+
+
+def _reload_code_location(name: str, app_name: str | None, blocking: bool) -> CommandResult:
+    manager = DagflowCliManager()
+    endpoint_url = manager.dagster_endpoint_url(app_name)
+    token = manager.session_token()
+    result = reload_repository_location(
+        endpoint_url=endpoint_url,
+        token=token,
+        location_name=name,
+    )
+    if blocking:
+        wait_for_code_location_loaded(
+            endpoint_url=endpoint_url,
+            token=token,
+            location_name=name,
+        )
+    return ObjectResult(result)
+
+
+def _refresh_code_locations(app_name: str | None) -> CommandResult:
+    manager = DagflowCliManager()
+    result = reload_workspace(
+        endpoint_url=manager.dagster_endpoint_url(app_name),
+        token=manager.session_token(),
+    )
+    return ObjectResult(result)
+
+
+@code_location_app.command("delete", requires_connection=True, requires_global_options=True)
+@with_output
+def delete(
+    name: str = typer.Argument(..., help="Dagflow code location name"),
+    app_name: str | None = typer.Option(None, "--app-name", help="Installed Dagflow application name"),
+    yes: bool = typer.Option(False, "--yes", help="Skip confirmation"),
+    **options: Any,
+) -> CommandResult:
+    """Delete a Dagflow code location."""
+    return _delete_code_location(name, app_name, yes)
+
+
+@code_location_app.command("reload", requires_connection=True, requires_global_options=True)
+@with_output
+def reload(
+    name: str = typer.Argument(..., help="Dagflow code location name"),
+    app_name: str | None = typer.Option(None, "--app-name", help="Installed Dagflow application name"),
+    blocking: bool = typer.Option(True, "--blocking/--no-blocking", help="Wait for code location readiness"),
+    **options: Any,
+) -> CommandResult:
+    """Reload a Dagflow code location in Dagster."""
+    return _reload_code_location(name, app_name, blocking)
+
+
+@code_location_app.command("reload-all", requires_connection=True, requires_global_options=True)
+@with_output
+def reload_all(
+    app_name: str | None = typer.Option(None, "--app-name", help="Installed Dagflow application name"),
+    **options: Any,
+) -> CommandResult:
+    """Reload all Dagflow code locations in Dagster."""
+    return _refresh_code_locations(app_name)
+
+
+@code_location_app.command("list", requires_connection=True, requires_global_options=True)
+@with_output
+def list_locations(
+    app_name: str | None = typer.Option(None, "--app-name", help="Installed Dagflow application name"),
+    **options: Any,
+) -> CommandResult:
+    """List Dagflow code locations."""
+    cursor = DagflowCliManager().list_code_locations(app_name)
+    return SingleQueryResult(cursor)
+
+
+def _get_code_location(name: str, app_name: str | None) -> CommandResult:
+    return ObjectResult(DagflowCliManager().code_location_status(app_name, name))
+
+
+@code_location_app.command("get", requires_connection=True, requires_global_options=True)
+@with_output
+def get(
+    name: str = typer.Argument(..., help="Dagflow code location name"),
+    app_name: str | None = typer.Option(None, "--app-name", help="Installed Dagflow application name"),
+    **options: Any,
+) -> CommandResult:
+    """Get a Dagflow code location."""
+    return _get_code_location(name, app_name)
+
+
+@job_app.command("launch", requires_connection=True, requires_global_options=True)
+@with_output
+def launch(
+    job_name: str = typer.Option(..., "--job", "-j", help="Dagster job name"),
+    location_name: str = typer.Option(..., "--location", "-l", help="Dagflow code location name"),
+    app_name: str | None = typer.Option(None, "--app-name", help="Installed Dagflow application name"),
+    repository_name: str = typer.Option("__repository__", "--repository", help="Dagster repository name"),
+    config: Path | None = typer.Option(None, "--config", help="Run config YAML file"),
+    tag: list[str] | None = typer.Option(None, "--tag", help="Run tag as KEY=VALUE; may be repeated"),
+    **options: Any,
+) -> CommandResult:
+    """Launch a Dagster job using Dagflow's configured run launcher."""
+    manager = DagflowCliManager()
+    execution_params = build_launch_execution_params(
+        job_name=job_name,
+        location_name=location_name,
+        repository_name=repository_name,
+        run_config=load_run_config(config),
+        tags=parse_tags(tuple(tag or ())),
+    )
+    result = launch_job(
+        endpoint_url=manager.dagster_endpoint_url(app_name),
+        token=manager.session_token(),
+        execution_params=execution_params,
+    )
+    return ObjectResult({"run_id": result.run_id, "status": "LAUNCHED"})
+
+
+app.add_typer(code_location_app)
+app.add_typer(job_app)

dagflow_cli/launch.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+import json
+import urllib.error
+import urllib.request
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable
+
+import yaml
+
+
+DEFAULT_REPOSITORY_NAME = "__repository__"
+LAUNCH_MUTATION = """
+mutation Launch($executionParams: ExecutionParams!) {
+  launchPipelineExecution(executionParams: $executionParams) {
+    __typename
+    ... on LaunchRunSuccess { run { runId } }
+    ... on PythonError { message stack }
+    ... on RunConfigValidationInvalid { errors { message } }
+    ... on PipelineNotFoundError { message }
+  }
+}
+"""
+
+
+@dataclass(frozen=True)
+class LaunchResult:
+    run_id: str
+    response_type: str
+
+
+def build_launch_execution_params(
+    job_name: str,
+    location_name: str,
+    repository_name: str = DEFAULT_REPOSITORY_NAME,
+    run_config: dict[str, Any] | None = None,
+    tags: dict[str, str] | None = None,
+) -> dict[str, Any]:
+    if not job_name:
+        raise ValueError("job name is required")
+    if not location_name:
+        raise ValueError("code location name is required")
+    return {
+        "selector": {
+            "repositoryLocationName": location_name,
+            "repositoryName": repository_name,
+            "pipelineName": job_name,
+        },
+        "runConfigData": run_config or {},
+        "mode": "default",
+        "executionMetadata": {
+            "tags": [{"key": key, "value": value} for key, value in (tags or {}).items()]
+        },
+    }
+
+
+def load_run_config(path: Path | None) -> dict[str, Any]:
+    if path is None:
+        return {}
+    loaded = yaml.safe_load(path.expanduser().read_text())
+    if loaded is None:
+        return {}
+    if not isinstance(loaded, dict):
+        raise ValueError("run config YAML must contain a mapping at the document root")
+    return loaded
+
+
+def parse_tags(values: tuple[str, ...]) -> dict[str, str]:
+    tags: dict[str, str] = {}
+    for value in values:
+        if "=" not in value:
+            raise ValueError("tags must use KEY=VALUE syntax")
+        key, tag_value = value.split("=", 1)
+        if not key:
+            raise ValueError("tag keys cannot be empty")
+        tags[key] = tag_value
+    return tags
+
+
+def launch_job(
+    endpoint_url: str,
+    token: str,
+    execution_params: dict[str, Any],
+    post_json: Callable[[str, dict[str, Any], str], dict[str, Any]] | None = None,
+) -> LaunchResult:
+    payload = {"query": LAUNCH_MUTATION, "variables": {"executionParams": execution_params}}
+    response = (post_json or _post_json)(f"{endpoint_url.rstrip('/')}/graphql", payload, token)
+    if response.get("errors"):
+        raise RuntimeError(f"Dagster GraphQL launch failed: {response['errors']}")
+    result = response.get("data", {}).get("launchPipelineExecution")
+    if not isinstance(result, dict):
+        raise RuntimeError(f"Dagster GraphQL launch returned an unexpected response: {response}")
+    if result.get("__typename") == "LaunchRunSuccess":
+        run_id = result.get("run", {}).get("runId")
+        if not run_id:
+            raise RuntimeError(f"Dagster GraphQL launch response did not include a run id: {result}")
+        return LaunchResult(run_id=str(run_id), response_type="LaunchRunSuccess")
+    raise RuntimeError(f"Dagster run launch failed: {result}")
+
+
+def _post_json(url: str, payload: dict[str, Any], token: str) -> dict[str, Any]:
+    body = json.dumps(payload).encode("utf-8")
+    request = urllib.request.Request(
+        url,
+        data=body,
+        headers={
+            "Authorization": f'Snowflake Token="{token}"',
+            "Content-Type": "application/json",
+        },
+        method="POST",
+    )
+    try:
+        with urllib.request.urlopen(request, timeout=120) as response:
+            return json.loads(response.read().decode("utf-8"))
+    except urllib.error.HTTPError as exc:
+        response_body = exc.read().decode("utf-8", errors="replace")
+        raise RuntimeError(f"Dagster GraphQL launch request failed: HTTP {exc.code}: {response_body}") from exc