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

dagflow_cli/manager.py

@@ -0,0 +1,283 @@
+from __future__ import annotations
+
+import json
+import re
+import urllib.error
+import urllib.request
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+try:
+    from snowflake.cli.api.sql_execution import SqlExecutionMixin
+except ModuleNotFoundError:  # pragma: no cover - exercised only without snowflake-cli installed
+    class SqlExecutionMixin:  # type: ignore[no-redef]
+        def execute_query(self, query: str) -> Any:
+            raise RuntimeError("snowflake-cli is required to execute Dagflow CLI commands")
+
+
+CODE_LOCATION_RE = re.compile(r"^[A-Za-z][A-Za-z0-9_]*$")
+VERSION_RE = re.compile(r"^[A-Za-z0-9_.-]+$")
+MODULE_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_.]*$")
+PYTHON_FILE_RE = re.compile(r"^[A-Za-z0-9_./-]+\.py$")
+IDENTIFIER_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_$]*$")
+
+
+@dataclass(frozen=True)
+class Entrypoint:
+    type: str
+    value: str
+
+    def as_payload(self) -> dict[str, str]:
+        if self.type == "python_file":
+            return {"type": self.type, "path": self.value}
+        return {"type": self.type, "name": self.value}
+
+
+class DagflowCliManager(SqlExecutionMixin):
+    def dagster_endpoint_sql(self, app_name: str | None) -> str:
+        return f"SHOW ENDPOINTS IN SERVICE {self._dagster_service_name(app_name)}"
+
+    def dagster_endpoint_url(self, app_name: str | None) -> str:
+        rows = self._rows_from_cursor(self.execute_query(self.dagster_endpoint_sql(app_name)))
+        ingress_url = next(
+            (
+                str(row.get("ingress_url") or row.get("INGRESS_URL"))
+                for row in rows
+                if str(row.get("name") or row.get("NAME")).lower() == "dagster-ui"
+            ),
+            "",
+        )
+        if not ingress_url:
+            raise RuntimeError("Dagster endpoint dagster-ui was not found")
+        if ingress_url.startswith("http://") or ingress_url.startswith("https://"):
+            return ingress_url.rstrip("/")
+        return f"https://{ingress_url.rstrip('/')}"
+
+    def session_token(self) -> str:
+        try:
+            self.execute_query("ALTER SESSION SET PYTHON_CONNECTOR_QUERY_RESULT_FORMAT = 'json'")
+            return str(self._conn._rest._token_request("ISSUE")["data"]["sessionToken"])
+        except Exception as exc:
+            raise RuntimeError("Could not issue a Snowflake session token for the Dagster endpoint") from exc
+
+    def put_source_pex_sql(self, app_name: str | None, location_name: str, version: str, path: Path) -> str:
+        return self._put_sql(app_name, location_name, version, path, "source.pex")
+
+    def put_deps_pex_sql(self, app_name: str | None, location_name: str, version: str, path: Path) -> str:
+        return self._put_sql(app_name, location_name, version, path, "deps.pex")
+
+    def list_code_locations_sql(self, app_name: str | None) -> str:
+        return f"SHOW SERVICES IN SCHEMA {self._code_location_schema(app_name)}"
+
+    def describe_code_location_sql(self, app_name: str | None, location_name: str) -> str:
+        self._validate_code_location(location_name)
+        return f"DESCRIBE SERVICE {self._service_name(app_name, location_name)}"
+
+    def show_code_location_containers_sql(self, app_name: str | None, location_name: str) -> str:
+        self._validate_code_location(location_name)
+        return f"SHOW SERVICE CONTAINERS IN SERVICE {self._service_name(app_name, location_name)}"
+
+    def code_location_status(self, app_name: str | None, location_name: str) -> dict[str, Any]:
+        self._validate_code_location(location_name)
+        rows = self._rows_from_cursor(self.execute_query(self.list_code_locations_sql(app_name)))
+        expected = location_name.upper()
+        row = next(
+            (row for row in rows if str(row.get("name") or row.get("NAME")).upper() == expected),
+            None,
+        )
+        if row is None:
+            raise RuntimeError(f"Code location {location_name} was not found")
+        comment = row.get("comment") or row.get("COMMENT")
+        desired: Any = None
+        if comment:
+            try:
+                desired = json.loads(str(comment))
+            except json.JSONDecodeError:
+                desired = str(comment)
+        return {
+            "name": row.get("name") or row.get("NAME"),
+            "status": row.get("status") or row.get("STATUS"),
+            "database_name": row.get("database_name") or row.get("DATABASE_NAME"),
+            "schema_name": row.get("schema_name") or row.get("SCHEMA_NAME"),
+            "dns_name": row.get("dns_name") or row.get("DNS_NAME"),
+            "desired": desired,
+        }
+
+    def upload_artifacts(
+        self,
+        app_name: str | None,
+        location_name: str,
+        version: str,
+        source_pex: Path,
+        deps_pex: Path,
+    ) -> None:
+        self.execute_query(self.put_source_pex_sql(app_name, location_name, version, source_pex))
+        self.execute_query(self.put_deps_pex_sql(app_name, location_name, version, deps_pex))
+
+    def deploy_code_location_api(
+        self,
+        app_name: str | None,
+        location_name: str,
+        version: str,
+        entrypoint: Entrypoint,
+        python_version: str,
+        blocking: bool,
+    ) -> dict[str, Any]:
+        self._validate_code_location(location_name)
+        self._validate_version(version)
+        self._validate_entrypoint(entrypoint)
+        return self._request_code_location_api(
+            app_name,
+            "PUT",
+            f"/code-locations/{location_name}",
+            {
+                "version": version,
+                "entrypoint": entrypoint.as_payload(),
+                "python_version": python_version,
+                "blocking": blocking,
+            },
+        )
+
+    def remove_code_location_api(self, app_name: str | None, location_name: str) -> dict[str, Any]:
+        self._validate_code_location(location_name)
+        return self._request_code_location_api(
+            app_name,
+            "DELETE",
+            f"/code-locations/{location_name}",
+            None,
+        )
+
+    def list_code_locations(self, app_name: str | None) -> Any:
+        return self.execute_query(self.list_code_locations_sql(app_name))
+
+    def describe_code_location(self, app_name: str | None, location_name: str) -> Any:
+        return self.execute_query(self.describe_code_location_sql(app_name, location_name))
+
+    def code_location_service_exists(self, app_name: str | None, location_name: str) -> bool:
+        try:
+            self.execute_query(self.describe_code_location_sql(app_name, location_name))
+        except Exception as exc:
+            if self._is_missing_service_error(exc):
+                return False
+            raise
+        return True
+
+    def show_code_location_containers(self, app_name: str | None, location_name: str) -> Any:
+        return self.execute_query(self.show_code_location_containers_sql(app_name, location_name))
+
+    def _put_sql(
+        self,
+        app_name: str | None,
+        location_name: str,
+        version: str,
+        path: Path,
+        artifact_name: str,
+    ) -> str:
+        self._validate_code_location(location_name)
+        self._validate_version(version)
+        if path.name != artifact_name:
+            raise ValueError(f"{artifact_name} path must end with {artifact_name}")
+        stage_path = f"{self._pex_stage(app_name)}/{location_name.lower()}/{version}/"
+        return f"PUT file://{path.expanduser().resolve().as_posix()} {stage_path} AUTO_COMPRESS=FALSE OVERWRITE=TRUE"
+
+    def _request_code_location_api(
+        self,
+        app_name: str | None,
+        method: str,
+        path: str,
+        payload: dict[str, Any] | None,
+    ) -> dict[str, Any]:
+        body = json.dumps(payload).encode("utf-8") if payload is not None else None
+        request = urllib.request.Request(
+            f"{self.dagster_endpoint_url(app_name).rstrip('/')}{path}",
+            data=body,
+            headers={
+                "Authorization": f'Snowflake Token="{self.session_token()}"',
+                "Content-Type": "application/json",
+            },
+            method=method,
+        )
+        try:
+            with urllib.request.urlopen(request, timeout=360) 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"Dagflow code-location request failed: HTTP {exc.code}: {response_body}"
+            ) from exc
+
+    @classmethod
+    def _validate_code_location(cls, value: str) -> None:
+        if not CODE_LOCATION_RE.fullmatch(value):
+            raise ValueError("code location names must start with a letter and contain only letters, numbers, and underscores")
+
+    @classmethod
+    def _validate_version(cls, value: str) -> None:
+        if not VERSION_RE.fullmatch(value):
+            raise ValueError("versions must contain only letters, numbers, underscores, dots, and dashes")
+
+    @classmethod
+    def _validate_entrypoint(cls, entrypoint: Entrypoint) -> None:
+        if entrypoint.type not in {"module", "package", "python_file"}:
+            raise ValueError(f"unsupported entrypoint type: {entrypoint.type}")
+        if entrypoint.type == "python_file":
+            if not PYTHON_FILE_RE.fullmatch(entrypoint.value) or ".." in entrypoint.value:
+                raise ValueError("python_file entrypoints must be relative .py files without '..'")
+            return
+        if not MODULE_RE.fullmatch(entrypoint.value):
+            raise ValueError("module and package entrypoints must be valid dotted Python names")
+
+    @classmethod
+    def _identifier(cls, value: str) -> str:
+        if not IDENTIFIER_RE.fullmatch(value):
+            raise ValueError(f"invalid Snowflake identifier: {value}")
+        return value.upper()
+
+    @classmethod
+    def _code_location_schema(cls, app_name: str | None) -> str:
+        if app_name:
+            return f"{cls._identifier(app_name)}.code_locations"
+        return "code_locations"
+
+    @classmethod
+    def _service_name(cls, app_name: str | None, location_name: str) -> str:
+        return f"{cls._code_location_schema(app_name)}.{cls._identifier(location_name)}"
+
+    @classmethod
+    def _dagster_service_name(cls, app_name: str | None) -> str:
+        if app_name:
+            return f"{cls._identifier(app_name)}.core.dagster"
+        return "core.dagster"
+
+    @classmethod
+    def _pex_stage(cls, app_name: str | None) -> str:
+        if app_name:
+            return f"@{cls._identifier(app_name)}.core.pex"
+        return "@core.pex"
+
+    @staticmethod
+    def _rows_from_cursor(cursor: Any) -> list[dict[str, Any]]:
+        rows = cursor.fetchall()
+        if not rows:
+            return []
+        first = rows[0]
+        if isinstance(first, dict):
+            return list(rows)
+        description = getattr(cursor, "description", None)
+        if not description:
+            return []
+        columns = [str(column[0]).lower() for column in description]
+        return [dict(zip(columns, row)) for row in rows]
+
+    @staticmethod
+    def _is_missing_service_error(exc: Exception) -> bool:
+        message = str(exc).lower()
+        if "not authorized" in message and "not exist" not in message and "not found" not in message:
+            return False
+        return (
+            "does not exist" in message
+            or "not exist" in message
+            or "not found" in message
+            or "doesn't exist" in message
+        )