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

dagflow_cli/workspace.py

@@ -0,0 +1,215 @@
+from __future__ import annotations
+
+import time
+from typing import Any, Callable
+
+from dagflow_cli.launch import _post_json
+
+
+RELOAD_REPOSITORY_LOCATION_MUTATION = """
+mutation ReloadRepositoryLocation($repositoryLocationName: String!) {
+  reloadRepositoryLocation(repositoryLocationName: $repositoryLocationName) {
+    __typename
+    ... on WorkspaceLocationEntry { name loadStatus }
+    ... on RepositoryLocationNotFound { message }
+    ... on ReloadNotSupported { message }
+    ... on PythonError { message stack }
+    ... on UnauthorizedError { message }
+  }
+}
+"""
+
+RELOAD_WORKSPACE_MUTATION = """
+mutation ReloadWorkspace {
+  reloadWorkspace {
+    __typename
+    ... on Workspace {
+      locationEntries {
+        name
+        loadStatus
+      }
+    }
+    ... on PythonError { message stack }
+    ... on UnauthorizedError { message }
+  }
+}
+"""
+
+WORKSPACE_LOCATION_STATUS_QUERY = """
+query WorkspaceLocationStatus {
+  workspaceOrError {
+    __typename
+    ... on Workspace {
+      locationEntries {
+        name
+        loadStatus
+        locationOrLoadError {
+          __typename
+          ... on PythonError { message stack }
+        }
+      }
+    }
+    ... on PythonError { message stack }
+  }
+}
+"""
+
+
+class ReloadRepositoryLocationNotFound(RuntimeError):
+    pass
+
+
+def reload_repository_location(
+    endpoint_url: str,
+    token: str,
+    location_name: str,
+    post_json: Callable[[str, dict[str, Any], str], dict[str, Any]] | None = None,
+) -> dict[str, Any]:
+    dagster_location_name = _dagster_location_name(location_name)
+    response = (post_json or _post_json)(
+        f"{endpoint_url.rstrip('/')}/graphql",
+        {
+            "query": RELOAD_REPOSITORY_LOCATION_MUTATION,
+            "variables": {"repositoryLocationName": dagster_location_name},
+        },
+        token,
+    )
+    return _result_or_raise(response, "reloadRepositoryLocation")
+
+
+def reload_workspace(
+    endpoint_url: str,
+    token: str,
+    post_json: Callable[[str, dict[str, Any], str], dict[str, Any]] | None = None,
+) -> dict[str, Any]:
+    response = (post_json or _post_json)(
+        f"{endpoint_url.rstrip('/')}/graphql",
+        {"query": RELOAD_WORKSPACE_MUTATION, "variables": {}},
+        token,
+    )
+    return _result_or_raise(response, "reloadWorkspace")
+
+
+def reload_after_code_location_deploy(
+    endpoint_url: str,
+    token: str,
+    location_name: str,
+    service_existed: bool,
+    post_json: Callable[[str, dict[str, Any], str], dict[str, Any]] | None = None,
+) -> dict[str, Any]:
+    dagster_location_name = _dagster_location_name(location_name)
+    if service_existed:
+        try:
+            reload_result = reload_repository_location(
+                endpoint_url, token, dagster_location_name, post_json=post_json
+            )
+        except ReloadRepositoryLocationNotFound:
+            reload_result = reload_workspace(endpoint_url, token, post_json=post_json)
+    else:
+        reload_result = reload_workspace(endpoint_url, token, post_json=post_json)
+    wait_for_code_location_loaded(endpoint_url, token, dagster_location_name, post_json=post_json)
+    return reload_result
+
+
+def reload_after_code_location_delete(
+    endpoint_url: str,
+    token: str,
+    location_name: str,
+    post_json: Callable[[str, dict[str, Any], str], dict[str, Any]] | None = None,
+) -> dict[str, Any]:
+    dagster_location_name = _dagster_location_name(location_name)
+    reload_result = reload_workspace(endpoint_url, token, post_json=post_json)
+    wait_for_code_location_removed(endpoint_url, token, dagster_location_name, post_json=post_json)
+    return reload_result
+
+
+def wait_for_code_location_loaded(
+    endpoint_url: str,
+    token: str,
+    location_name: str,
+    *,
+    timeout_s: float = 300,
+    interval_s: float = 5,
+    post_json: Callable[[str, dict[str, Any], str], dict[str, Any]] | None = None,
+) -> dict[str, Any]:
+    deadline = time.monotonic() + timeout_s
+    last_entry: dict[str, Any] | None = None
+    while time.monotonic() < deadline:
+        entry = code_location_entry(endpoint_url, token, location_name, post_json=post_json)
+        if entry is not None:
+            last_entry = entry
+            error = entry.get("locationOrLoadError")
+            if isinstance(error, dict) and error.get("__typename") == "PythonError":
+                raise RuntimeError(f"Dagster code location {location_name} failed to load: {error}")
+            if entry.get("loadStatus") == "LOADED":
+                return entry
+        time.sleep(interval_s)
+    raise TimeoutError(
+        f"Timed out waiting for Dagster code location {location_name} to load; "
+        f"last entry was {last_entry}"
+    )
+
+
+def wait_for_code_location_removed(
+    endpoint_url: str,
+    token: str,
+    location_name: str,
+    *,
+    timeout_s: float = 300,
+    interval_s: float = 5,
+    post_json: Callable[[str, dict[str, Any], str], dict[str, Any]] | None = None,
+) -> None:
+    deadline = time.monotonic() + timeout_s
+    last_entry: dict[str, Any] | None = None
+    while time.monotonic() < deadline:
+        entry = code_location_entry(endpoint_url, token, location_name, post_json=post_json)
+        if entry is None:
+            return
+        last_entry = entry
+        time.sleep(interval_s)
+    raise TimeoutError(
+        f"Timed out waiting for Dagster code location {location_name} to be removed; "
+        f"last entry was {last_entry}"
+    )
+
+
+def code_location_entry(
+    endpoint_url: str,
+    token: str,
+    location_name: str,
+    post_json: Callable[[str, dict[str, Any], str], dict[str, Any]] | None = None,
+) -> dict[str, Any] | None:
+    dagster_location_name = _dagster_location_name(location_name)
+    response = (post_json or _post_json)(
+        f"{endpoint_url.rstrip('/')}/graphql",
+        {"query": WORKSPACE_LOCATION_STATUS_QUERY, "variables": {}},
+        token,
+    )
+    if response.get("errors"):
+        raise RuntimeError(f"Dagster GraphQL workspace query failed: {response['errors']}")
+    workspace = response.get("data", {}).get("workspaceOrError")
+    if not isinstance(workspace, dict):
+        raise RuntimeError(f"Dagster GraphQL workspace query returned an unexpected response: {response}")
+    if workspace.get("__typename") != "Workspace":
+        raise RuntimeError(f"Dagster GraphQL workspace query failed: {workspace}")
+    for entry in workspace.get("locationEntries", []):
+        if entry.get("name") == dagster_location_name:
+            return entry
+    return None
+
+
+def _result_or_raise(response: dict[str, Any], field_name: str) -> dict[str, Any]:
+    if response.get("errors"):
+        raise RuntimeError(f"Dagster GraphQL reload failed: {response['errors']}")
+    result = response.get("data", {}).get(field_name)
+    if not isinstance(result, dict):
+        raise RuntimeError(f"Dagster GraphQL reload returned an unexpected response: {response}")
+    if result.get("__typename") in {"Workspace", "WorkspaceLocationEntry"}:
+        return result
+    if result.get("__typename") == "RepositoryLocationNotFound":
+        raise ReloadRepositoryLocationNotFound(f"Dagster GraphQL reload failed: {result}")
+    raise RuntimeError(f"Dagster GraphQL reload failed: {result}")
+
+
+def _dagster_location_name(location_name: str) -> str:
+    return location_name.lower()

dagflow_cli_plugin-0.1.0.dist-info/METADATA

@@ -0,0 +1,30 @@
+Metadata-Version: 2.4
+Name: dagflow-cli-plugin
+Version: 0.1.0
+Summary: Snowflake CLI plugin for managing Dagflow deployments.
+Requires-Python: <3.13,>=3.11
+Description-Content-Type: text/markdown
+Requires-Dist: dagster-dg-core>=1.12.1
+Requires-Dist: packaging<27,>=24
+Requires-Dist: pex<3,>=2.33
+Requires-Dist: pyyaml<7,>=6
+Requires-Dist: snowflake-cli<4,>=3
+
+# Dagflow CLI
+
+Snowflake CLI plugin for building, deploying, and managing Dagflow code locations.
+
+Install it in the same Python environment as Snowflake CLI:
+
+```bash
+pip install dagflow-cli-plugin
+```
+
+Enable the external plugin in the Snowflake CLI configuration:
+
+```toml
+[cli.plugins.dg]
+enabled = true
+```
+
+Dagflow commands are then available under `snow dg`.

dagflow_cli_plugin-0.1.0.dist-info/RECORD

@@ -0,0 +1,13 @@
+dagflow_cli/__init__.py,sha256=wsTNGr7bQbNQDQeLV7mDb2lCtW550s52LSR7RkZojr4,48
+dagflow_cli/commands.py,sha256=19XXfmw8CNaq91A-BI_ypzDB5emRk4VGjloDTRtlmLA,16526
+dagflow_cli/launch.py,sha256=5aNcMwi0jBA8Encz47jxhbI8xV7onXyYjtsOvRwmBWc,3998
+dagflow_cli/manager.py,sha256=N1SyuF7-gjApjeMADA84f3iU5ccnEXEdWtnn79XSz2w,11499
+dagflow_cli/pex.py,sha256=_3Q8I3p_dFTiq0sglJKX30kxfZxPwySMoZqGJ090gPM,16867
+dagflow_cli/plugin_spec.py,sha256=WcNplrJZ5fqMfL_drydVw3r_hjjkr0jUMGKI46HKf7g,450
+dagflow_cli/runtime-platforms.json,sha256=fCQznKQsGwcQndqgdplky-qmla3Xy_cl3aGDE3D2IxI,52331
+dagflow_cli/workspace.py,sha256=Hh29w-g9EV5omdKVVdU1Qy7kUz6DI5Hg9lO5sFvREPk,7238
+dagflow_cli_plugin-0.1.0.dist-info/METADATA,sha256=KO0W3idaxOpTiyvHdG0IT0yfJlOF8cwusf_4AtXZ7Ss,726
+dagflow_cli_plugin-0.1.0.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
+dagflow_cli_plugin-0.1.0.dist-info/entry_points.txt,sha256=oDpBkyH-Bhp8JuQv74ZrTWMm8WO1jymR8Xem-0aw8KE,60
+dagflow_cli_plugin-0.1.0.dist-info/top_level.txt,sha256=iCuyZvhHP_1btyzkW7izMxij84empD5tyk2LlkAfxuU,12
+dagflow_cli_plugin-0.1.0.dist-info/RECORD,,

dagflow_cli_plugin-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (83.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

dagflow_cli_plugin-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[snowflake.cli.plugin.command]
+dg = dagflow_cli.plugin_spec

dagflow_cli_plugin-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+dagflow_cli