dagster-cloud-cli 1.12.16__py2.py3-none-any.whl

dagster_cloud_cli/core/artifacts.py

@@ -0,0 +1,192 @@
+import os
+from pathlib import Path
+
+import requests
+
+from dagster_cloud_cli.config_utils import (
+    URL_ENV_VAR_NAME,
+    get_deployment,
+    get_organization,
+    get_user_token,
+)
+from dagster_cloud_cli.core.headers.auth import DagsterCloudInstanceScope
+from dagster_cloud_cli.core.headers.impl import get_dagster_cloud_api_headers
+from dagster_cloud_cli.gql import url_from_config
+
+
+def _dagster_cloud_http_client():
+    # indirection for mocking, could share a session
+    return requests
+
+
+def download_artifact(
+    *,
+    url: str,
+    api_token: str,
+    scope: DagsterCloudInstanceScope,
+    key: str,
+    path: Path | str,
+    deployment: str | None = None,
+):
+    response = _dagster_cloud_http_client().post(
+        url=f"{url}/gen_artifact_get",
+        headers=get_dagster_cloud_api_headers(
+            api_token,
+            scope,
+            deployment_name=deployment,
+        ),
+        json={"key": key},
+    )
+    response.raise_for_status()
+
+    payload = response.json()
+
+    response = requests.get(
+        payload["url"],
+    )
+    response.raise_for_status()
+    Path(path).write_bytes(response.content)
+
+
+def upload_artifact(
+    *,
+    url: str,
+    api_token: str,
+    scope: DagsterCloudInstanceScope,
+    key: str,
+    path: Path | str,
+    deployment: str | None = None,
+):
+    upload_file = Path(path).resolve(strict=True)
+
+    response = _dagster_cloud_http_client().post(
+        url=f"{url}/gen_artifact_post",
+        headers=get_dagster_cloud_api_headers(
+            api_token,
+            scope,
+            deployment_name=deployment,
+        ),
+        json={"key": key},
+    )
+    response.raise_for_status()
+    payload = response.json()
+
+    response = requests.post(
+        payload["url"],
+        data=payload["fields"],
+        files={
+            "file": upload_file.open(),
+        },
+    )
+    response.raise_for_status()
+
+
+def _resolve_org(passed_org: str | None) -> str:
+    org = passed_org or get_organization()
+    if org is None:
+        raise Exception(
+            "Unable to resolve organization, pass organization or set the "
+            "DAGSTER_CLOUD_ORGANIZATION environment variable."
+        )
+    return org
+
+
+def _resolve_token(passed_token: str | None) -> str:
+    api_token = get_user_token()
+    if api_token is None:
+        raise Exception(
+            "Unable to resolve api_token, pass api_token or set the "
+            "DAGSTER_CLOUD_API_TOKEN environment variable"
+        )
+    return api_token
+
+
+def _resolve_deploy(passed_deploy: str | None) -> str:
+    deploy = passed_deploy or get_deployment()
+    if deploy is None:
+        raise Exception(
+            "Unable to resolve deployment, pass deployment or set the "
+            "DAGSTER_CLOUD_DEPLOYMENT environment variable."
+        )
+    return deploy
+
+
+def _resolve_url(organization: str, deployment: str | None = None) -> str:
+    env_url = os.getenv(URL_ENV_VAR_NAME)
+    if env_url:
+        return env_url
+
+    return url_from_config(organization, deployment)
+
+
+def upload_organization_artifact(
+    key: str,
+    path: str | Path,
+    organization: str | None = None,
+    api_token: str | None = None,
+):
+    upload_artifact(
+        url=_resolve_url(
+            organization=_resolve_org(organization),
+        ),
+        api_token=_resolve_token(api_token),
+        scope=DagsterCloudInstanceScope.ORGANIZATION,
+        key=key,
+        path=path,
+    )
+
+
+def upload_deployment_artifact(
+    key: str,
+    path: str | Path,
+    organization: str | None = None,
+    deployment: str | None = None,
+    api_token: str | None = None,
+):
+    upload_artifact(
+        url=_resolve_url(
+            organization=_resolve_org(organization),
+            deployment=_resolve_deploy(deployment),
+        ),
+        api_token=_resolve_token(api_token),
+        scope=DagsterCloudInstanceScope.DEPLOYMENT,
+        key=key,
+        path=path,
+    )
+
+
+def download_organization_artifact(
+    key: str,
+    path: str | Path,
+    organization: str | None = None,
+    api_token: str | None = None,
+):
+    download_artifact(
+        url=_resolve_url(
+            organization=_resolve_org(organization),
+        ),
+        api_token=_resolve_token(api_token),
+        scope=DagsterCloudInstanceScope.ORGANIZATION,
+        key=key,
+        path=path,
+    )
+
+
+def download_deployment_artifact(
+    key: str,
+    path: str | Path,
+    organization: str | None = None,
+    deployment: str | None = None,
+    api_token: str | None = None,
+):
+    deployment = _resolve_deploy(deployment)
+    download_artifact(
+        url=_resolve_url(
+            organization=_resolve_org(organization),
+            deployment=deployment,
+        ),
+        api_token=_resolve_token(api_token),
+        scope=DagsterCloudInstanceScope.DEPLOYMENT,
+        key=key,
+        path=path,
+    )

dagster_cloud_cli/core/docker_runner.py

@@ -0,0 +1,57 @@
+import os
+import subprocess
+
+from dagster_cloud_cli import ui
+
+DAGSTER_CLOUD_BUILDER_IMAGE = "ghcr.io/dagster-io/dagster-manylinux-builder:manylinux_2_28_x86_64"
+
+
+def run_dagster_cloud(
+    map_folders: dict[str, str], run_args: list[str], env: dict[str, str | None]
+) -> subprocess.CompletedProcess:
+    proc = run_docker_image(
+        image_name=os.getenv("DAGSTER_CLOUD_BUILDER_IMAGE", DAGSTER_CLOUD_BUILDER_IMAGE),
+        map_folders=map_folders,
+        env=env,
+        entrypoint="dagster-cloud",
+        run_args=run_args,
+    )
+    return proc
+
+
+def run_docker_image(
+    image_name,
+    map_folders: dict[str, str],
+    env: dict[str, str | None],
+    entrypoint: str,
+    run_args: list[str],
+) -> subprocess.CompletedProcess:
+    mount_args = []
+    for target_folder, source_folder in map_folders.items():
+        mount_args.extend(["--mount", f"type=bind,source={source_folder},target={target_folder}"])
+
+    env_args = []
+    for env_name, env_value in env.items():
+        if env_value is None:
+            # specifying just the name copies the env from the calling environment
+            env_args.extend(["--env", env_name])
+        else:
+            env_args.extend(["--env", f"{env_name}={env_value}"])
+
+    ui.print(f"Running docker {image_name}")
+    ui.print("Mapped folders:")
+    for target_dir, source_dir in map_folders.items():
+        ui.print(f" - {source_dir} -> {target_dir}")
+
+    cmd = [
+        "docker",
+        "run",
+        "--platform=linux/amd64",
+        *env_args,
+        *mount_args,
+        "--entrypoint",
+        entrypoint,
+        image_name,
+        *run_args,
+    ]
+    return subprocess.run(cmd, capture_output=True, check=False)

dagster_cloud_cli/core/errors.py

@@ -0,0 +1,32 @@
+from requests import HTTPError
+
+
+class DagsterCloudAgentServerError(Exception):
+    """Raise this when there's an error in the GraphQL layer."""
+
+
+DEFAULT_MAINTENANCE_TIMEOUT = 3600 * 4 * 4
+DEFAULT_MAINTENANCE_RETRY_INTERVAL = 30
+
+
+class DagsterCloudMaintenanceException(Exception):
+    def __init__(self, message, timeout, retry_interval):
+        self.timeout = timeout or DEFAULT_MAINTENANCE_TIMEOUT
+        self.retry_interval = retry_interval or DEFAULT_MAINTENANCE_RETRY_INTERVAL
+        super().__init__(message)
+
+
+class DagsterCloudHTTPError(Exception):
+    """Clearer error message for exceptions hitting Dagster Cloud servers."""
+
+    def __init__(self, http_error: HTTPError):
+        self.response = http_error.response
+        error_content = http_error.response.content.decode("utf-8", errors="ignore")
+        super().__init__(http_error.__str__() + ": " + str(error_content))
+
+
+def raise_http_error(response):
+    try:
+        response.raise_for_status()
+    except HTTPError as e:
+        raise DagsterCloudHTTPError(e) from e

dagster_cloud_cli/core/graphql_client.py

@@ -0,0 +1,420 @@
+import logging
+import re
+import time
+from collections.abc import Callable, Mapping
+from contextlib import contextmanager
+from email.utils import mktime_tz, parsedate_tz
+from typing import Any
+
+import requests
+from dagster_shared import check
+from requests.adapters import HTTPAdapter
+from requests.exceptions import (
+    ChunkedEncodingError,
+    ConnectionError as RequestsConnectionError,
+    HTTPError,
+    ReadTimeout as RequestsReadTimeout,
+)
+
+from dagster_cloud_cli.core.errors import (
+    DagsterCloudAgentServerError,
+    DagsterCloudHTTPError,
+    DagsterCloudMaintenanceException,
+)
+from dagster_cloud_cli.core.headers.auth import DagsterCloudInstanceScope
+from dagster_cloud_cli.core.headers.impl import get_dagster_cloud_api_headers
+
+DEFAULT_RETRIES = 6
+DEFAULT_BACKOFF_FACTOR = 0.5
+DEFAULT_TIMEOUT = 60
+
+logger = logging.getLogger("dagster_cloud")
+
+
+RETRY_STATUS_CODES = [
+    502,
+    503,
+    504,
+    429,
+    409,
+]
+
+
+class DagsterCloudAgentHttpClient:
+    def __init__(
+        self,
+        session: requests.Session,
+        headers: dict[str, Any] | None = None,
+        verify: bool = True,
+        timeout: int = DEFAULT_TIMEOUT,
+        cookies: dict[str, Any] | None = None,
+        proxies: dict[str, Any] | None = None,
+        max_retries: int = 0,
+        backoff_factor: float = DEFAULT_BACKOFF_FACTOR,
+    ):
+        self.headers = headers or {}
+        self.verify = verify
+        self.timeout = timeout
+        self.cookies = cookies
+        self._session = session
+        self._proxies = proxies
+        self._max_retries = max_retries
+        self._backoff_factor = backoff_factor
+
+    @property
+    def session(self) -> requests.Session:
+        return self._session
+
+    def post(self, *args, **kwargs):
+        return self.execute("POST", *args, **kwargs)
+
+    def get(self, *args, **kwargs):
+        return self.execute("GET", *args, **kwargs)
+
+    def put(self, *args, **kwargs):
+        return self.execute("PUT", *args, **kwargs)
+
+    def execute(
+        self,
+        method: str,
+        url: str,
+        headers: Mapping[str, str] | None = None,
+        idempotent: bool = False,
+        **kwargs,
+    ):
+        retry_on_read_timeout = idempotent or bool(
+            headers.get("Idempotency-Key") if headers else False
+        )
+
+        return _retry_loop(
+            lambda: self._execute_retry(method, url, headers, **kwargs),
+            max_retries=self._max_retries,
+            backoff_factor=self._backoff_factor,
+            retry_on_read_timeout=retry_on_read_timeout,
+        )
+
+    def _execute_retry(
+        self,
+        method: str,
+        url: str,
+        headers: Mapping[str, Any] | None,
+        **kwargs,
+    ):
+        response = self._session.request(
+            method,
+            url,
+            headers={
+                **(self.headers if self.headers is not None else {}),
+                **(headers if headers is not None else {}),
+            },
+            cookies=self.cookies,
+            timeout=self.timeout,
+            verify=self.verify,
+            proxies=self._proxies,
+            **kwargs,
+        )
+        try:
+            result = response.json()
+            if not isinstance(result, dict):
+                result = {}
+        except ValueError:
+            result = {}
+
+        if "maintenance" in result:
+            maintenance_info = result["maintenance"]
+            raise DagsterCloudMaintenanceException(
+                message=maintenance_info.get("message"),
+                timeout=maintenance_info.get("timeout"),
+                retry_interval=maintenance_info.get("retry_interval"),
+            )
+
+        if "errors" in result:
+            raise DagsterCloudAgentServerError(f"Error in GraphQL response: {result['errors']}")
+
+        response.raise_for_status()
+
+        return result
+
+
+def _retry_loop(
+    execute_retry: Callable,
+    max_retries: int,
+    backoff_factor: float,
+    retry_on_read_timeout: bool,
+):
+    start_time = time.time()
+    retry_number = 0
+    error_msg_set = set()
+    requested_sleep_time = None
+    while True:
+        try:
+            return execute_retry()
+        except (HTTPError, RequestsConnectionError, RequestsReadTimeout, ChunkedEncodingError) as e:
+            retryable_error = False
+            if isinstance(e, HTTPError):
+                retryable_error = e.response.status_code in RETRY_STATUS_CODES
+                error_msg = e.response.status_code
+                requested_sleep_time = _get_retry_after_sleep_time(e.response.headers)
+            elif isinstance(e, RequestsConnectionError):
+                retryable_error = True
+                error_msg = str(e)
+            else:
+                retryable_error = retry_on_read_timeout
+                error_msg = str(e)
+
+            error_msg_set.add(error_msg)
+            if retryable_error and retry_number < max_retries:
+                retry_number += 1
+                sleep_time = 0
+                if requested_sleep_time:
+                    sleep_time = requested_sleep_time
+                else:
+                    sleep_time = backoff_factor * (2**retry_number)
+
+                logger.warning(
+                    f"Error in Dagster Cloud request ({error_msg}). Retrying in"
+                    f" {sleep_time} seconds..."
+                )
+                time.sleep(sleep_time)
+            else:
+                # Throw the error straight if no retries were involved
+                if max_retries == 0 or not retryable_error:
+                    if isinstance(e, HTTPError):
+                        raise DagsterCloudHTTPError(e) from e
+                    else:
+                        raise
+                else:
+                    if len(error_msg_set) == 1:
+                        status_code_msg = str(next(iter(error_msg_set)))
+                    else:
+                        status_code_msg = str(error_msg_set)
+                    raise DagsterCloudAgentServerError(
+                        f"Max retries ({max_retries}) exceeded, too many"
+                        f" {status_code_msg} error responses."
+                    ) from e
+        except DagsterCloudMaintenanceException as e:
+            if time.time() - start_time > e.timeout:
+                raise
+
+            logger.warning(
+                "Dagster Cloud is currently unavailable due to scheduled maintenance. Retrying"
+                f" in {e.retry_interval} seconds..."
+            )
+            time.sleep(e.retry_interval)
+        except DagsterCloudAgentServerError:
+            raise
+        except Exception as e:
+            raise DagsterCloudAgentServerError(str(e)) from e
+
+
+class DagsterCloudGraphQLClient:
+    def __init__(
+        self,
+        url: str,
+        session: requests.Session,
+        headers: dict[str, Any] | None = None,
+        verify: bool = True,
+        timeout: int = DEFAULT_TIMEOUT,
+        cookies: dict[str, Any] | None = None,
+        proxies: dict[str, Any] | None = None,
+        max_retries: int = 0,
+        backoff_factor: float = DEFAULT_BACKOFF_FACTOR,
+    ):
+        self.url = url
+        self.headers = headers
+        self.verify = verify
+        self.timeout = timeout
+        self.cookies = cookies
+        self._session = session
+        self._proxies = proxies
+        self._max_retries = max_retries
+        self._backoff_factor = backoff_factor
+
+    @property
+    def session(self) -> requests.Session:
+        return self._session
+
+    def execute(
+        self,
+        query: str,
+        variable_values: Mapping[str, Any] | None = None,
+        headers: Mapping[str, str] | None = None,
+        idempotent_mutation: bool = False,
+    ):
+        if "mutation " in query and not idempotent_mutation:
+            all_headers = {
+                **(self.headers if self.headers is not None else {}),
+                **(headers if headers is not None else {}),
+            }
+
+            # mutations can be made idempotent if they use Idempotency-Key header
+            retry_on_read_timeout = bool(all_headers.get("Idempotency-Key"))
+        else:
+            retry_on_read_timeout = True
+
+        return _retry_loop(
+            lambda: self._execute_retry(query, variable_values, headers),
+            max_retries=self._max_retries,
+            backoff_factor=self._backoff_factor,
+            retry_on_read_timeout=retry_on_read_timeout,
+        )
+
+    def _execute_retry(
+        self,
+        query: str,
+        variable_values: Mapping[str, Any] | None,
+        headers: Mapping[str, Any] | None,
+    ):
+        response = self._session.post(
+            self.url,
+            headers={
+                **(self.headers if self.headers is not None else {}),
+                **(headers if headers is not None else {}),
+                "Content-type": "application/json",
+            },
+            cookies=self.cookies,
+            timeout=self.timeout,
+            verify=self.verify,
+            json={
+                "query": query,
+                "variables": variable_values if variable_values else {},
+            },
+            proxies=self._proxies,
+        )
+        try:
+            result = response.json()
+            if not isinstance(result, dict):
+                result = {}
+        except ValueError:
+            result = {}
+
+        if "errors" not in result and "data" not in result and "maintenance" not in result:
+            response.raise_for_status()
+            raise requests.HTTPError("Unexpected GraphQL response", response=response)
+
+        if "maintenance" in result:
+            maintenance_info = result["maintenance"]
+            raise DagsterCloudMaintenanceException(
+                message=maintenance_info.get("message"),
+                timeout=maintenance_info.get("timeout"),
+                retry_interval=maintenance_info.get("retry_interval"),
+            )
+
+        if "errors" in result:
+            raise DagsterCloudAgentServerError(f"Error in GraphQL response: {result['errors']}")
+
+        return result
+
+
+def get_agent_headers(config_value: dict[str, Any], scope: DagsterCloudInstanceScope):
+    return get_dagster_cloud_api_headers(
+        config_value["agent_token"],
+        scope=scope,
+        deployment_name=config_value.get("deployment"),
+        additional_headers=config_value.get("headers"),
+    )
+
+
+class HTTPAdapterWithSocketOptions(HTTPAdapter):
+    def __init__(self, *args, **kwargs):
+        self.socket_options = kwargs.pop("socket_options", None)
+        super().__init__(*args, **kwargs)
+
+    def init_poolmanager(self, *args, **kwargs):
+        if self.socket_options is not None:
+            kwargs["socket_options"] = self.socket_options
+        super().init_poolmanager(*args, **kwargs)
+
+
+@contextmanager
+def create_graphql_requests_session(adapter_kwargs: Mapping[str, Any] | None = None):
+    with requests.Session() as session:
+        adapter = HTTPAdapterWithSocketOptions(**(adapter_kwargs or {}))
+        session.mount("https://", adapter)
+        session.mount("http://", adapter)
+        yield session
+
+
+def create_agent_http_client(
+    session: requests.Session,
+    config_value: dict[str, Any],
+    scope: DagsterCloudInstanceScope = DagsterCloudInstanceScope.DEPLOYMENT,
+):
+    return DagsterCloudAgentHttpClient(
+        headers=get_agent_headers(config_value, scope=scope),
+        verify=config_value.get("verify", True),
+        timeout=config_value.get("timeout", DEFAULT_TIMEOUT),
+        cookies=config_value.get("cookies", {}),
+        # Requests library modifies proxies dictionary so create a copy
+        proxies=(
+            check.is_dict(config_value.get("proxies")).copy() if config_value.get("proxies") else {}
+        ),
+        session=session,
+        max_retries=config_value.get("retries", DEFAULT_RETRIES),
+        backoff_factor=config_value.get("backoff_factor", DEFAULT_BACKOFF_FACTOR),
+    )
+
+
+def create_agent_graphql_client(
+    session: requests.Session,
+    url: str,
+    config_value: dict[str, Any],
+    scope: DagsterCloudInstanceScope = DagsterCloudInstanceScope.DEPLOYMENT,
+):
+    return DagsterCloudGraphQLClient(
+        url=url,
+        headers=get_agent_headers(config_value, scope=scope),
+        verify=config_value.get("verify", True),
+        timeout=config_value.get("timeout", DEFAULT_TIMEOUT),
+        cookies=config_value.get("cookies", {}),
+        # Requests library modifies proxies dictionary so create a copy
+        proxies=(
+            check.is_dict(config_value.get("proxies")).copy() if config_value.get("proxies") else {}
+        ),
+        session=session,
+        max_retries=config_value.get("retries", DEFAULT_RETRIES),
+        backoff_factor=config_value.get("backoff_factor", DEFAULT_BACKOFF_FACTOR),
+    )
+
+
+def _get_retry_after_sleep_time(headers):
+    # https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Retry-After
+
+    retry_after = headers.get("Retry-After")
+    if retry_after is None:
+        return None
+
+    if re.match(r"^\s*[0-9]+\s*$", retry_after):
+        seconds = int(retry_after)
+    else:
+        retry_date = parsedate_tz(retry_after)
+        if retry_date is None:
+            return None
+        retry_date = mktime_tz(retry_date)
+        seconds = retry_date - time.time()
+
+    return max(seconds, 0)
+
+
+@contextmanager
+def create_cloud_webserver_client(
+    url: str,
+    api_token: str,
+    retries=3,
+    deployment_name: str | None = None,
+    headers: dict[str, Any] | None = None,
+):
+    with create_graphql_requests_session(adapter_kwargs={}) as session:
+        yield DagsterCloudGraphQLClient(
+            session=session,
+            url=f"{url}/graphql",
+            headers={
+                **get_dagster_cloud_api_headers(
+                    api_token,
+                    scope=DagsterCloudInstanceScope.DEPLOYMENT,
+                    deployment_name=deployment_name,
+                ),
+                **(headers if headers else {}),
+            },
+            max_retries=retries,
+        )