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

dagflow_cli/pex.py

@@ -0,0 +1,523 @@
+from __future__ import annotations
+
+import importlib.metadata
+import json
+import os
+import re
+import shutil
+import subprocess
+import sys
+import tempfile
+import zipfile
+from contextlib import contextmanager
+from dataclasses import dataclass
+from importlib import resources
+from pathlib import Path
+from typing import ContextManager, Iterator, Sequence
+
+from packaging.requirements import InvalidRequirement, Requirement
+from packaging.utils import InvalidWheelFilename, canonicalize_name, parse_wheel_filename
+
+try:
+    import tomllib
+except ModuleNotFoundError:  # pragma: no cover - Python < 3.11
+    import tomli as tomllib  # type: ignore[no-redef]
+
+
+@dataclass(frozen=True)
+class PexArtifacts:
+    source_pex: Path
+    deps_pex: Path
+
+
+DEFAULT_PYTHON_VERSION = "3.11"
+SUPPORTED_PYTHON_VERSIONS = {"3.11", "3.12"}
+
+
+@dataclass(frozen=True)
+class PexBuildConfig:
+    project_path: Path
+    output_dir: Path
+    python_version: str = DEFAULT_PYTHON_VERSION
+    python: str = sys.executable
+
+
+@dataclass(frozen=True)
+class DependencyBuildInputs:
+    requirements: tuple[str, ...]
+    direct_requirements: tuple[str, ...]
+    local_packages: tuple[Path, ...]
+
+
+@dataclass(frozen=True)
+class RequirementEntry:
+    value: str
+    relative_to: Path
+
+
+SOURCE_IGNORED_PATTERNS = (
+    ".git",
+    ".venv",
+    ".pytest_cache",
+    ".tox",
+    ".nox",
+    ".mypy_cache",
+    ".cache",
+    "__pycache__",
+    "*.pyc",
+    "*.egg-info",
+    ".dagflow-*",
+    "build",
+    "dist",
+)
+
+
+def build_pex_artifacts(config: PexBuildConfig) -> PexArtifacts:
+    project_path = config.project_path.expanduser().resolve()
+    if not project_path.exists():
+        raise ValueError(f"project path does not exist: {project_path}")
+    target_python = _target_python(config.python_version)
+
+    output_dir = config.output_dir.expanduser().resolve()
+    output_dir.mkdir(parents=True, exist_ok=True)
+    source_pex = output_dir / "source.pex"
+    deps_pex = output_dir / "deps.pex"
+
+    inputs = _collect_dependency_inputs(project_path, target_python)
+    _validate_runtime_dependencies(inputs.direct_requirements)
+    with _requirements_file(inputs.requirements, project_path) as requirements:
+        with _staged_source_directories(
+            project_path,
+            inputs.local_packages,
+            target_python,
+            output_dir,
+        ) as source_directories:
+            _run_pex(
+                _source_pex_command(
+                    config.python,
+                    target_python,
+                    source_directories,
+                    source_pex,
+                ),
+                cwd=project_path,
+            )
+        with _complete_platform_file(config.python_version) as complete_platform:
+            _run_pex(
+                _deps_pex_command(config.python, complete_platform, requirements, deps_pex),
+                cwd=project_path,
+            )
+        _validate_deps_pex(deps_pex)
+
+    return PexArtifacts(source_pex=source_pex, deps_pex=deps_pex)
+
+
+def default_output_dir(location_name: str, version: str) -> Path:
+    return Path("dist") / "dagflow" / location_name / version
+
+
+def default_version() -> str:
+    try:
+        result = subprocess.run(
+            ["git", "rev-parse", "--short=12", "HEAD"],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+    except (OSError, subprocess.CalledProcessError):
+        from datetime import datetime, timezone
+
+        return datetime.now(timezone.utc).strftime("%Y%m%d%H%M%S")
+    return result.stdout.strip()
+
+
+def _target_python(python_version: str) -> str:
+    if python_version not in SUPPORTED_PYTHON_VERSIONS:
+        raise ValueError(
+            f"unsupported python version {python_version}; supported versions are "
+            + ", ".join(sorted(SUPPORTED_PYTHON_VERSIONS))
+        )
+    current_version = f"{sys.version_info.major}.{sys.version_info.minor}"
+    if python_version == current_version:
+        return sys.executable
+    executable = f"python{python_version}"
+    resolved = shutil.which(executable)
+    if resolved is None:
+        raise ValueError(
+            f"{executable} was not found on PATH; install Python {python_version} "
+            f"or choose a runtime version available in this environment"
+        )
+    return resolved
+
+
+@contextmanager
+def _requirements_file(requirements: Sequence[str], project_path: Path) -> Iterator[Path]:
+    generated = tempfile.NamedTemporaryFile(
+        "w",
+        encoding="utf-8",
+        suffix="-requirements.txt",
+        prefix=".dagflow-",
+        dir=project_path,
+        delete=False,
+    )
+    path = Path(generated.name)
+    try:
+        generated.write("\n".join(requirements))
+        generated.write("\n")
+        generated.close()
+        yield path
+    finally:
+        generated.close()
+        path.unlink(missing_ok=True)
+
+
+@contextmanager
+def _temporary_json_file(value: object) -> Iterator[Path]:
+    generated = tempfile.NamedTemporaryFile(
+        "w", encoding="utf-8", suffix="-complete-platform.json", delete=False
+    )
+    path = Path(generated.name)
+    try:
+        json.dump(value, generated)
+        generated.close()
+        yield path
+    finally:
+        generated.close()
+        path.unlink(missing_ok=True)
+
+
+def _runtime_complete_platform(python_version: str) -> dict[str, object]:
+    manifest = json.loads(
+        resources.files("dagflow_cli")
+        .joinpath("runtime-platforms.json")
+        .read_text(encoding="utf-8")
+    )[python_version]
+    return manifest["complete_platform"]
+
+
+def _complete_platform_file(python_version: str) -> ContextManager[Path]:
+    return _temporary_json_file(_runtime_complete_platform(python_version))
+
+
+def _deps_pex_command(python: str, complete_platform: Path, requirements: Path, output: Path) -> list[str]:
+    return [
+        python,
+        "-m",
+        "pex",
+        "--complete-platform",
+        str(complete_platform),
+        "-r",
+        str(requirements),
+        "--pip-version",
+        "latest-compatible",
+        "--resolver-version",
+        "pip-2020-resolver",
+        "--inherit-path=false",
+        "--no-strip-pex-env",
+        "-o",
+        str(output),
+    ]
+
+
+def _source_pex_command(
+    python: str,
+    target_python: str,
+    source_directories: Sequence[Path],
+    output: Path,
+) -> list[str]:
+    source_args = [arg for path in source_directories for arg in ("-D", str(path))]
+    return [
+        python,
+        "-m",
+        "pex",
+        "--python",
+        target_python,
+        *source_args,
+        "--inherit-path=false",
+        "--no-strip-pex-env",
+        "-o",
+        str(output),
+    ]
+
+
+def _run_pex(command: Sequence[str], cwd: Path) -> None:
+    env = {**os.environ, "PEX_ROOT": os.environ.get("PEX_ROOT", str(Path.home() / ".pex"))}
+    subprocess.run(command, cwd=cwd, env=env, check=True)
+
+
+@contextmanager
+def _staged_source_directories(
+    project_path: Path,
+    local_packages: Sequence[Path],
+    python: str,
+    output_dir: Path | None = None,
+) -> Iterator[tuple[Path, ...]]:
+    with tempfile.TemporaryDirectory(prefix="dagflow-source-") as temp_dir:
+        staging_root = Path(temp_dir)
+        build_dir = staging_root / "installed"
+        source_dir = staging_root / "source"
+        build_dir.mkdir()
+
+        for local_package in (*local_packages, project_path):
+            if _is_buildable_project(local_package):
+                _install_project_source(local_package, build_dir, python)
+
+        base_ignore = shutil.ignore_patterns(*SOURCE_IGNORED_PATTERNS)
+
+        def ignore(directory: str, names: list[str]) -> set[str]:
+            ignored = set(base_ignore(directory, names))
+            if output_dir is not None:
+                for name in names:
+                    if (Path(directory) / name).resolve() == output_dir.resolve():
+                        ignored.add(name)
+            return ignored
+
+        shutil.copytree(
+            project_path,
+            source_dir,
+            ignore=ignore,
+            symlinks=True,
+        )
+        _remove_installed_source_collisions(source_dir, build_dir)
+        directories = [source_dir]
+        if any(build_dir.iterdir()):
+            directories.insert(0, build_dir)
+        yield tuple(directories)
+
+
+def _is_buildable_project(project_path: Path) -> bool:
+    return (project_path / "setup.py").exists() or (project_path / "pyproject.toml").exists()
+
+
+def _remove_installed_source_collisions(source_dir: Path, build_dir: Path) -> None:
+    for installed_path in build_dir.rglob("*"):
+        if not installed_path.is_file() and not installed_path.is_symlink():
+            continue
+        source_path = source_dir / installed_path.relative_to(build_dir)
+        if source_path.is_file() or source_path.is_symlink():
+            source_path.unlink()
+
+
+def _install_project_source(project_path: Path, build_dir: Path, python: str) -> None:
+    try:
+        subprocess.run(
+            [
+                python,
+                "-m",
+                "pip",
+                "install",
+                "--target",
+                str(build_dir),
+                "--no-deps",
+                ".",
+            ],
+            cwd=project_path,
+            check=True,
+        )
+    except subprocess.CalledProcessError as exc:
+        raise ValueError(f"could not build local project source at {project_path}") from exc
+
+
+def _collect_dependency_inputs(project_path: Path, python: str) -> DependencyBuildInputs:
+    direct_entries = _project_dependency_entries(project_path, python)
+    pending = [(project_path, direct_entries)]
+    seen_projects: set[Path] = set()
+    local_packages: list[Path] = []
+    requirements: list[str] = []
+
+    while pending:
+        current_project, entries = pending.pop()
+        current_project = current_project.resolve()
+        if current_project in seen_projects:
+            continue
+        seen_projects.add(current_project)
+
+        for entry in entries:
+            local_path = _local_package_path(entry)
+            if local_path is None:
+                if entry.value not in requirements:
+                    requirements.append(entry.value)
+                continue
+            if local_path == current_project:
+                continue
+            if not local_path.is_dir():
+                raise ValueError(
+                    f"local package {local_path} referenced by {entry.value!r} is not a directory"
+                )
+            if local_path not in local_packages:
+                local_packages.append(local_path)
+            pending.append((local_path, _project_dependency_entries(local_path, python)))
+
+    return DependencyBuildInputs(
+        requirements=tuple(requirements),
+        direct_requirements=tuple(entry.value for entry in direct_entries),
+        local_packages=tuple(local_packages),
+    )
+
+
+def _project_dependency_entries(project_path: Path, python: str) -> list[RequirementEntry]:
+    entries: list[RequirementEntry] = []
+    requirements = project_path / "requirements.txt"
+    if requirements.exists():
+        entries.extend(_requirements_file_entries(requirements))
+    entries.extend(
+        RequirementEntry(value=value, relative_to=project_path)
+        for value in _setup_py_dependencies(project_path, python)
+    )
+    entries.extend(
+        RequirementEntry(value=value, relative_to=project_path)
+        for value in _pyproject_dependencies(project_path)
+    )
+    return entries
+
+
+def _validate_runtime_dependencies(dependencies: Sequence[str]) -> None:
+    names = {_requirement_name(dependency) for dependency in dependencies}
+    names.discard(None)
+    missing = [name for name in ("dagster", "dagster-postgres") if name not in names]
+    if missing:
+        required = ", ".join(missing)
+        raise ValueError(
+            f"deps.pex must include direct project requirement(s): {required}. "
+            "The code-location image does not provide Dagster or PostgreSQL storage classes."
+        )
+
+
+def _requirements_file_entries(
+    path: Path,
+    seen: set[Path] | None = None,
+) -> list[RequirementEntry]:
+    resolved = path.resolve()
+    visited = seen or set()
+    if resolved in visited:
+        return []
+    visited.add(resolved)
+
+    entries: list[RequirementEntry] = []
+    for raw_line in _logical_requirement_lines(resolved):
+        line = re.sub(r"(^#|\s+#).*", "", raw_line).strip()
+        if not line:
+            continue
+        parts = line.split(maxsplit=1)
+        if parts[0] in {"-r", "--requirement"} and len(parts) == 2:
+            entries.extend(_requirements_file_entries(resolved.parent / parts[1], visited))
+        elif parts[0].startswith("-r") and len(parts[0]) > 2:
+            entries.extend(_requirements_file_entries(resolved.parent / parts[0][2:], visited))
+        elif line != ".":
+            entries.append(RequirementEntry(value=line, relative_to=resolved.parent))
+    return entries
+
+
+def _logical_requirement_lines(path: Path) -> list[str]:
+    lines: list[str] = []
+    current = ""
+    for raw_line in path.read_text(encoding="utf-8").splitlines():
+        stripped = raw_line.rstrip()
+        if stripped.endswith("\\"):
+            current += stripped[:-1]
+        else:
+            lines.append(current + stripped)
+            current = ""
+    if current:
+        lines.append(current)
+    return lines
+
+
+def _local_package_path(entry: RequirementEntry) -> Path | None:
+    value = entry.value.strip()
+    if value.startswith("-e ") or value.startswith("--editable "):
+        value = value.split(maxsplit=1)[1]
+    if not value.startswith(("./", "../", "/")):
+        return None
+    return (entry.relative_to / value).resolve()
+
+
+def _requirement_name(value: str) -> str | None:
+    try:
+        return canonicalize_name(Requirement(value).name)
+    except InvalidRequirement:
+        return None
+
+
+def _validate_deps_pex(path: Path) -> None:
+    try:
+        with zipfile.ZipFile(path) as archive:
+            pex_info = json.loads(archive.read("PEX-INFO"))
+    except (OSError, KeyError, json.JSONDecodeError, zipfile.BadZipFile) as exc:
+        raise ValueError(f"could not inspect built deps.pex: {exc}") from exc
+
+    distribution_names: set[str] = set()
+    for filename in pex_info.get("distributions", {}):
+        try:
+            name, _version, _build, _tags = parse_wheel_filename(filename)
+        except InvalidWheelFilename:
+            continue
+        distribution_names.add(canonicalize_name(name))
+
+    missing = [name for name in ("dagster", "dagster-postgres") if name not in distribution_names]
+    if missing:
+        raise ValueError(
+            "built deps.pex does not contain required distribution(s): " + ", ".join(missing)
+        )
+
+
+def _setup_py_dependencies(project_path: Path, python: str) -> list[str]:
+    setup_py = project_path / "setup.py"
+    if not setup_py.exists():
+        return []
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        result = subprocess.run(
+            [python, str(setup_py), "egg_info", f"--egg-base={temp_dir}"],
+            cwd=project_path,
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+        if result.returncode:
+            raise ValueError(
+                "could not inspect setup.py dependencies: "
+                + result.stdout
+                + result.stderr
+            )
+
+        distributions = list(importlib.metadata.distributions(path=[temp_dir]))
+        if len(distributions) != 1:
+            raise ValueError(f"could not find generated distribution metadata for {setup_py}")
+        dependencies: list[str] = []
+        for value in distributions[0].requires or []:
+            requirement = Requirement(value)
+            if requirement.marker and re.search(r"\bextra\b", str(requirement.marker)):
+                continue
+            dependencies.append(value)
+        return dependencies
+
+
+def _pyproject_dependencies(project_path: Path) -> list[str]:
+    pyproject = project_path / "pyproject.toml"
+    if not pyproject.exists():
+        return []
+
+    try:
+        data = tomllib.loads(pyproject.read_text(encoding="utf-8"))
+    except Exception as exc:
+        raise ValueError(f"could not parse pyproject.toml: {exc}") from exc
+
+    deps: list[str] = []
+    project = data.get("project", {})
+    deps.extend(str(dep) for dep in project.get("dependencies", []))
+
+    poetry = data.get("tool", {}).get("poetry", {})
+    for dep_name, dep_spec in poetry.get("dependencies", {}).items():
+        if dep_name == "python":
+            continue
+        deps.append(_poetry_dependency(dep_name, dep_spec))
+    return [dep for dep in deps if dep]
+
+
+def _poetry_dependency(dep_name: str, dep_spec: object) -> str:
+    if isinstance(dep_spec, str):
+        return f"{dep_name}{dep_spec}"
+    if isinstance(dep_spec, dict):
+        version = dep_spec.get("version", "")
+        return f"{dep_name}{version}" if version else dep_name
+    return dep_name

dagflow_cli/plugin_spec.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from snowflake.cli.api.plugins.command import (
+    SNOWCLI_ROOT_COMMAND_PATH,
+    CommandSpec,
+    CommandType,
+    plugin_hook_impl,
+)
+
+from dagflow_cli import commands
+
+
+@plugin_hook_impl
+def command_spec() -> CommandSpec:
+    return CommandSpec(
+        parent_command_path=SNOWCLI_ROOT_COMMAND_PATH,
+        command_type=CommandType.COMMAND_GROUP,
+        typer_instance=commands.app.create_instance(),
+    )