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

dagster_cloud_cli/commands/deployment/alert_policies/commands.py

@@ -0,0 +1,52 @@
+from pathlib import Path
+
+import yaml
+from typer import Option, Typer
+
+from dagster_cloud_cli import gql, ui
+from dagster_cloud_cli.config_utils import dagster_cloud_options
+
+DEFAULT_ALERT_POLICIES_YAML_FILENAME = "alert_policies.yaml"
+
+app = Typer(help="Interact with your alert policies.")
+
+
+@app.command(name="list")
+@dagster_cloud_options(allow_empty=True, requires_url=True)
+def list_command(
+    api_token: str,
+    url: str,
+    deployment: str | None,
+):
+    """List your alert policies, output in YAML format."""
+    with gql.graphql_client_from_url(url, api_token, deployment_name=deployment) as client:
+        alert_policies_response = gql.get_alert_policies(client)
+
+    ui.print_yaml(alert_policies_response)
+
+
+@app.command(name="sync")
+@dagster_cloud_options(allow_empty=True, requires_url=True)
+def sync_command(
+    api_token: str,
+    url: str,
+    deployment: str | None,
+    alert_policies_file: Path = Option(
+        DEFAULT_ALERT_POLICIES_YAML_FILENAME,
+        "--alert-policies",
+        "-a",
+        exists=True,
+        help="Path to alert policies file.",
+    ),
+):
+    """Sync your YAML configured alert policies to Dagster Cloud."""
+    with gql.graphql_client_from_url(url, api_token, deployment_name=deployment) as client:
+        with open(str(alert_policies_file), encoding="utf8") as f:
+            config = yaml.load(f.read(), Loader=yaml.SafeLoader)
+
+        try:
+            alert_policies = gql.reconcile_alert_policies(client, config)
+
+            ui.print(f"Synced alert policies: {', '.join(alert_policies)}")
+        except Exception as e:
+            raise ui.error(str(e))

dagster_cloud_cli/commands/integration/__init__.py

@@ -0,0 +1,6 @@
+from typer import Typer
+
+from dagster_cloud_cli.commands.integration.atlan import app as atlan_app
+
+app = Typer(help="Customize your integrations in Dagster Cloud.")
+app.add_typer(atlan_app, name="atlan", no_args_is_help=True)

dagster_cloud_cli/commands/integration/atlan/__init__.py

@@ -0,0 +1,94 @@
+import sys
+from typing import Annotated
+
+import validators
+from typer import Argument, Typer
+
+from dagster_cloud_cli import gql, ui
+from dagster_cloud_cli.config_utils import dagster_cloud_options
+
+app = Typer(help="Customize your Atlan integration.")
+
+
+@app.command(name="set-settings")
+@dagster_cloud_options(allow_empty=True, requires_url=True)
+def set_atlan_settings_command(
+    atlan_token: Annotated[
+        str,
+        Argument(
+            help="The token to use to access the Atlan API.",
+        ),
+    ],
+    atlan_domain: Annotated[
+        str,
+        Argument(
+            help="The domain of your Atlan tenant. Eg. your-organization.atlan.com.",
+        ),
+    ],
+    api_token: str,
+    url: str,
+):
+    """Upload your Atlan settings to enable the Dagster<>Atlan integration in Dagster Cloud."""
+    if not validators.domain(atlan_domain):
+        ui.print(
+            "Invalid domain. "
+            "Please provide the domain of your Altan tenant using the following format: your-organization.atlan.com."
+        )
+        sys.exit(1)
+
+    try:
+        with gql.graphql_client_from_url(url, api_token) as client:
+            organization, _ = gql.set_atlan_integration_settings(
+                client, token=atlan_token, domain=atlan_domain
+            )
+        ui.print(f"Atlan settings successfully set for organization: {organization}")
+
+    except Exception as e:
+        ui.print(f"Failed to set Atlan settings. Exception: {e}")
+        sys.exit(1)
+
+
+@app.command(name="delete-settings")
+@dagster_cloud_options(allow_empty=True, requires_url=True)
+def delete_atlan_settings_command(
+    api_token: str,
+    url: str,
+):
+    """Delete your Atlan settings to disable the Dagster<>Atlan integration in Dagster Cloud."""
+    try:
+        with gql.graphql_client_from_url(url, api_token) as client:
+            organization, _ = gql.delete_atlan_integration_settings(client)
+        ui.print(f"Atlan settings successfully deleted for organization: {organization}")
+    except Exception as e:
+        ui.print(f"Failed to delete Atlan settings. Exception: {e}")
+        sys.exit(1)
+
+
+@app.command(name="get-settings")
+@dagster_cloud_options(allow_empty=True, requires_url=True)
+def get_atlan_settings_command(
+    api_token: str,
+    url: str,
+):
+    """Get your current Atlan settings from Dagster Cloud."""
+    with gql.graphql_client_from_url(url, api_token) as client:
+        settings = gql.get_atlan_integration_settings(client)
+        ui.print(f"Token: {settings['token']}")
+        ui.print(f"Domain: {settings['domain']}")
+
+
+@app.command(name="preflight-check")
+@dagster_cloud_options(allow_empty=True, requires_url=True)
+def preflight_check_command(
+    api_token: str,
+    url: str,
+):
+    """Run a preflight check on your Atlan integration settings."""
+    with gql.graphql_client_from_url(url, api_token) as client:
+        result = gql.atlan_integration_preflight_check(client)
+        if result["success"]:
+            ui.print("Preflight check passed successfully!")
+        else:
+            ui.print("Preflight check failed:")
+            ui.print(f"  Error code: {result['error_code']}")
+            ui.print(f"  Error message: {result['error_message']}")

dagster_cloud_cli/commands/job/__init__.py

@@ -0,0 +1,117 @@
+import json
+import os
+import time
+from typing import Any
+from uuid import uuid4
+
+import typer
+from typer import Typer
+
+from dagster_cloud_cli import gql, ui
+from dagster_cloud_cli.config_utils import dagster_cloud_options
+
+app = Typer(help="Commands for working with Dagster Cloud jobs.")
+
+
+SINGLETON_REPOSITORY_NAME = "__repository__"
+
+
+@app.command(name="launch")
+@dagster_cloud_options(allow_empty=True, requires_url=True)
+def launch(
+    api_token: str,
+    url: str,
+    deployment: str | None,
+    location: str = typer.Option(..., "-l", "--location", help="Location name in the deployment."),
+    job: str = typer.Option(..., "-j", "--job", help="Job name to run."),
+    repository: str = typer.Option(
+        None,
+        "-r",
+        "--repository",
+        help=(
+            "Repository in the specified code location. Required if a repository is defined in the"
+            " specified code location."
+        ),
+    ),
+    tags: str = typer.Option(None, "--tags", help="JSON dict of tags to use for this job run."),
+    config: str = typer.Option(
+        None, "--config-json", help="JSON string of run config to use for this job run"
+    ),
+    asset_keys: list[str] | None = typer.Option(
+        None,
+        "--asset-key",
+        help=(
+            "Asset key to materialize. Can be specified multiple times to materialize multiple"
+            " assets."
+        ),
+    ),
+    wait: bool = typer.Option(
+        False,
+        "-w",
+        "--wait",
+        help=(
+            "Wait for the run to finish and print the result. This will block until the run is"
+            " finished."
+        ),
+    ),
+    interval: int = typer.Option(
+        30,
+        "-i",
+        "--interval",
+        help=(
+            "Interval in seconds to wait between checking the status of the run. Only used if"
+            " --wait is specified."
+        ),
+    ),
+):
+    """Launch a run for a job."""
+    loaded_tags: dict[str, Any] = json.loads(tags) if tags else {}
+    loaded_config: dict[str, Any] = json.loads(config) if config else {}
+
+    repository = repository or SINGLETON_REPOSITORY_NAME
+
+    headers: dict[str, str] = {
+        "Idempotency-Key": str(uuid4()),
+    }
+
+    with gql.graphql_client_from_url(
+        url,
+        api_token,
+        deployment_name=deployment,
+        headers=headers,
+        retries=int(os.getenv("DAGSTER_CLOUD_JOB_LAUNCH_RETRIES", "5")),
+    ) as client:
+        run_id = gql.launch_run(
+            client,
+            location,
+            repository,
+            job,
+            loaded_tags,
+            loaded_config,
+            asset_keys=asset_keys,
+        )
+    if wait:
+        ui.print(f"Run {run_id} launched, waiting for completion...")
+        status = None
+        with gql.graphql_client_from_url(url, api_token, deployment_name=deployment) as client:
+            while True:
+                time.sleep(interval)
+                try:
+                    status = gql.run_status(client, run_id)
+                except Exception as e:
+                    ui.error(f"Failed to get status for run {run_id}: {e}.")
+
+                if status in ["SUCCESS", "FAILURE", "CANCELED"]:
+                    break
+                else:
+                    ui.print(f"Run {run_id} is in progress (status: {status})...")
+
+        if status == "SUCCESS":
+            ui.print(f"Run {run_id} finished successfully.")
+        else:
+            ui.error(
+                f"Run {run_id} failed with status '{status}'. Check the Dagster Cloud UI for more details."
+            )
+
+    else:
+        ui.print(run_id)

dagster_cloud_cli/commands/metrics.py

@@ -0,0 +1,114 @@
+import os
+import time
+from contextvars import ContextVar
+from enum import Enum
+from functools import wraps
+
+from dagster_cloud_cli import gql
+from dagster_cloud_cli.types import CliEventTags, CliEventType
+from dagster_cloud_cli.ui import ExitWithMessage
+
+
+def get_source() -> CliEventTags.source:
+    default = CliEventTags.source.cli
+    unknown = CliEventTags.source.unknown
+    env_map = {
+        "BITBUCKET_BUILD_NUMBER": (
+            CliEventTags.source.bitbucket  # https://support.atlassian.com/bitbucket-cloud/docs/variables-and-secrets/
+        ),
+        "BUILDKITE": (
+            CliEventTags.source.buildkite  # https://buildkite.com/docs/pipelines/environment-variables
+        ),
+        "CIRCLECI": (
+            CliEventTags.source.circle_ci  # https://circleci.com/docs/variables/#built-in-environment-variables
+        ),
+        "CODEBUILD_BUILD_ID": (
+            CliEventTags.source.codebuild  # https://docs.aws.amazon.com/codebuild/latest/userguide/build-env-ref-env-vars.html
+        ),
+        "GITHUB_ACTION": (
+            CliEventTags.source.github  # https://docs.github.com/en/actions/learn-github-actions/variables#default-environment-variables
+        ),
+        "GITLAB_CI": (
+            CliEventTags.source.gitlab  # https://docs.gitlab.com/ee/ci/variables/predefined_variables.html
+        ),
+        "JENKINS_URL": (
+            CliEventTags.source.jenkins  # https://www.jenkins.io/doc/book/pipeline/jenkinsfile/#using-environment-variables
+        ),
+        "TRAVIS": (
+            CliEventTags.source.travis  # https://docs.travis-ci.com/user/environment-variables/#default-environment-variables
+        ),
+    }
+
+    sources = [value for key, value in env_map.items() if os.getenv(key)]
+    if len(sources) > 1:
+        return unknown
+    elif sources:
+        return sources[0]
+    else:
+        return default
+
+
+mark_cli_event_context_tags: ContextVar[list[str]] = ContextVar("mark_cli_event_context_tags")
+
+
+def instrument_add_tags(tags: list[Enum]):
+    """Add tags to the instrumentation within the current context.
+    If this is not wrapped in an @instrument decorator, this call has no effect.
+    """
+    new_tags = [str(tag.value) for tag in tags or []]
+    context_tags = mark_cli_event_context_tags.get([])
+    mark_cli_event_context_tags.set(context_tags + new_tags)
+
+
+def instrument(event_type: CliEventType, tags: list[Enum] | None = None):
+    str_tags = [str(tag.value) for tag in tags or []]
+    str_tags.append(str(get_source().value))
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            url = kwargs.get("url")
+            api_token = kwargs.get("api_token")
+            start_time = time.time()
+
+            with gql.graphql_client_from_url(
+                url=str(url), token=str(api_token), retries=0
+            ) as client:
+                context_token = mark_cli_event_context_tags.set([])
+                try:
+                    result = func(*args, **kwargs)
+                    gql.mark_cli_event(
+                        client=client,
+                        event_type=event_type,
+                        duration_seconds=time.time() - start_time,
+                        success=True,
+                        tags=str_tags + mark_cli_event_context_tags.get([]),
+                    )
+                    return result
+                except ExitWithMessage as ex:
+                    gql.mark_cli_event(
+                        client=client,
+                        event_type=event_type,
+                        duration_seconds=time.time() - start_time,
+                        success=False,
+                        tags=str_tags + mark_cli_event_context_tags.get([]),
+                        message=ex.message,
+                    )
+                    raise
+
+                except Exception:
+                    gql.mark_cli_event(
+                        client=client,
+                        event_type=event_type,
+                        duration_seconds=time.time() - start_time,
+                        success=False,
+                        tags=str_tags + mark_cli_event_context_tags.get([]),
+                        message="unexpected error",
+                    )
+                    raise
+                finally:
+                    mark_cli_event_context_tags.reset(context_token)
+
+        return wrapper
+
+    return decorator

dagster_cloud_cli/commands/organization/__init__.py

@@ -0,0 +1,134 @@
+from pathlib import Path
+from typing import Annotated
+
+import yaml
+from typer import Argument, Typer
+
+from dagster_cloud_cli import gql, ui
+from dagster_cloud_cli.commands.organization.saml import commands as saml_cli
+from dagster_cloud_cli.config_utils import dagster_cloud_options
+from dagster_cloud_cli.core.artifacts import download_artifact, upload_artifact
+from dagster_cloud_cli.core.headers.auth import DagsterCloudInstanceScope
+
+app = Typer(help="Customize your Dagster Cloud organization.")
+settings_app = Typer(help="Customize your organization settings.")
+
+app.add_typer(settings_app, name="settings", no_args_is_help=True)
+settings_app.add_typer(saml_cli.app, name="saml", no_args_is_help=True)
+
+# Legacy, to support the old `dagster-cloud settings` path
+# New command is `dagster-cloud organization
+legacy_settings_app = Typer(
+    help=(
+        "[Deprecated, in favor of dagster-cloud organization] Customize your dagster-cloud"
+        " settings."
+    ),
+    hidden=True,
+)
+legacy_settings_app.add_typer(saml_cli.app, name="saml")
+
+
+@settings_app.command(name="set-from-file")
+@dagster_cloud_options(allow_empty=True, requires_url=True)
+def set_from_filecommand(
+    organization: str,
+    api_token: str,
+    url: str,
+    file_path: Path = Argument(..., readable=True, metavar="SETTINGS_FILE_PATH"),
+):
+    """Set the Dagster Cloud organization settings from a YAML file."""
+    if not url and not organization:
+        raise ui.error("Must provide either organization name or URL.")
+    if not url:
+        url = gql.url_from_config(organization=organization)
+
+    with open(file_path, encoding="utf8") as f:
+        settings = {"settings": yaml.safe_load(f) or {}}
+    with gql.graphql_client_from_url(url, api_token) as client:
+        gql.set_organization_settings(client, settings)
+
+
+@settings_app.command(name="get")
+@dagster_cloud_options(allow_empty=True, requires_url=True)
+def get_command(
+    organization: str,
+    api_token: str,
+    url: str,
+):
+    """Get the Dagster Cloud organization settings."""
+    if not url and not organization:
+        raise ui.error("Must provide either organization name or URL.")
+    if not url:
+        url = gql.url_from_config(organization=organization)
+
+    with gql.graphql_client_from_url(url, api_token) as client:
+        settings = gql.get_organization_settings(client)
+        ui.print_yaml(settings)
+
+
+@app.command(name="upload-artifact")
+@dagster_cloud_options(allow_empty=True, requires_url=True)
+def upload_artifact_command(
+    key: Annotated[
+        str,
+        Argument(
+            help="The key for the artifact being uploaded.",
+        ),
+    ],
+    path: Annotated[
+        str,
+        Argument(
+            help="The path to the file to upload.",
+        ),
+    ],
+    organization: str,
+    api_token: str,
+    url: str,
+):
+    """Upload a organization scoped artifact."""
+    if not url and not organization:
+        raise ui.error("Must provide either organization name or URL.")
+    if not url:
+        url = gql.url_from_config(organization=organization)
+
+    upload_artifact(
+        url=url,
+        scope=DagsterCloudInstanceScope.ORGANIZATION,
+        api_token=api_token,
+        key=key,
+        path=path,
+    )
+
+
+@app.command(name="download-artifact")
+@dagster_cloud_options(allow_empty=True, requires_url=True)
+def download_artifact_command(
+    key: Annotated[
+        str,
+        Argument(
+            help="The key for the artifact to download.",
+        ),
+    ],
+    path: Annotated[
+        str,
+        Argument(
+            help="Path to the file that the contents should be written to.",
+        ),
+    ],
+    organization: str,
+    api_token: str,
+    url: str,
+):
+    """Download a organization scoped artifact."""
+    if not url and not organization:
+        raise ui.error("Must provide either organization name or URL.")
+    if not url:
+        url = gql.url_from_config(organization=organization)
+
+    download_artifact(
+        url=url,
+        scope=DagsterCloudInstanceScope.ORGANIZATION,
+        api_token=api_token,
+        key=key,
+        path=path,
+    )

dagster_cloud_cli/commands/organization/saml/commands.py

@@ -0,0 +1,85 @@
+from pathlib import Path
+
+import requests
+from typer import Argument, Option, Typer
+
+from dagster_cloud_cli import gql, ui
+from dagster_cloud_cli.config_utils import dagster_cloud_options
+from dagster_cloud_cli.core.errors import raise_http_error
+from dagster_cloud_cli.core.headers.auth import DagsterCloudInstanceScope
+from dagster_cloud_cli.core.headers.impl import get_dagster_cloud_api_headers
+
+app = Typer(help="Customize your SAML settings.")
+
+
+@app.command(name="upload-identity-provider-metadata")
+@dagster_cloud_options(allow_empty=True)
+def upload_identity_provider_metadata_command(
+    api_token: str,
+    organization: str,
+    url: str = Option(
+        None,
+        "--url",
+        help=(
+            "[DEPRECATED] Your Dagster Cloud url, in the form of"
+            " 'https://{ORGANIZATION_NAME}.dagster.cloud'."
+        ),
+        hidden=True,
+    ),
+    metadata_file: Path = Argument(..., readable=True, metavar="METADATA_FILE_PATH"),
+):
+    """Upload your identity provider's metadata to enable SAML Single sign-on (SSO)."""
+    if not url and not organization:
+        raise ui.error("Must provide either organization name or URL.")
+
+    if not url:
+        url = gql.url_from_config(organization=organization)
+
+    with open(metadata_file, "rb") as f:
+        response = requests.post(
+            url=f"{url}/upload_idp_metadata",
+            headers=get_dagster_cloud_api_headers(api_token, DagsterCloudInstanceScope.DEPLOYMENT),
+            files={"metadata.xml": f},
+        )
+
+    raise_http_error(response)
+
+    if not response.text == "SUCCESS":
+        raise ui.error(
+            "Upload failed; unexpected response. Verify that the correct URL is specified.\n"
+            f" {response.text}"
+        )
+
+    ui.print("The identity provider metadata was successfully uploaded.")
+
+
+@app.command(name="remove-identity-provider-metadata")
+@dagster_cloud_options(allow_empty=True)
+def remove_identity_provider_metadata_command(
+    api_token: str,
+    organization: str,
+    url: str = Option(
+        None,
+        "--url",
+        help=(
+            "[DEPRECATED] Your Dagster Cloud url, in the form of"
+            " 'https://{ORGANIZATION_NAME}.dagster.cloud'."
+        ),
+        hidden=True,
+    ),
+):
+    """Remove SAML Single sign-on (SSO) metadata."""
+    if not url and not organization:
+        raise ui.error("Must provide either organization name or URL.")
+
+    if not url:
+        url = gql.url_from_config(organization=organization)
+
+    response = requests.post(
+        url=f"{url}/remove_idp_metadata",
+        headers=get_dagster_cloud_api_headers(api_token, DagsterCloudInstanceScope.DEPLOYMENT),
+    )
+
+    raise_http_error(response)
+
+    ui.print("The identity provider metadata was successfully removed.")

dagster_cloud_cli/commands/run/__init__.py

@@ -0,0 +1,27 @@
+import typer
+from typer import Typer
+
+from dagster_cloud_cli import gql, ui
+from dagster_cloud_cli.config_utils import dagster_cloud_options
+
+app = Typer(help="Commands for working with Dagster Cloud runs.")
+
+
+@app.command(name="status")
+@dagster_cloud_options(allow_empty=True, requires_url=True)
+def status(
+    api_token: str,
+    url: str,
+    deployment: str | None,
+    run_id: str = typer.Argument(None),
+):
+    """Check the status of a run."""
+    # this should be replaced by using the `dagster` CLI against a dagster-cloud `dagster.yaml`
+    if not run_id:
+        raise ui.error("No run_id provided")
+
+    with gql.graphql_client_from_url(url, api_token, deployment_name=deployment) as client:
+        try:
+            ui.print(gql.run_status(client, run_id))
+        except Exception as e:
+            ui.error(f"Failed to get status for run {run_id}: {e}.")

dagster_cloud_cli/commands/serverless/Dockerfile

@@ -0,0 +1,38 @@
+# base image is templated in (default: python:3.11-slim)
+
+# copy both requirements.txt / setup.py / pyproject.toml to guarantee that at least one file will be copied over
+# for any pre-src COPY.  On some flavors of Docker (Linux), the COPY command will fail unless there
+# exists at least one file to copy over.
+COPY *dagster_cloud_pre_install.sh *setup.py *pyproject.toml *requirements.txt /
+RUN if [ -f "dagster_cloud_pre_install.sh" ]; then \
+        ./dagster_cloud_pre_install.sh; \
+    fi
+
+# This is to enable Docker caching - if setup.py or pyproject.toml exists, we install the
+# dependencies before copying all other files
+RUN if [ -f "setup.py" ]; then \
+        pip install .; \
+    elif [ -f "pyproject.toml" ]; then \
+        pip install .; \
+    fi
+
+WORKDIR /opt/dagster/app
+COPY . /opt/dagster/app
+
+# Install the rest of dependencies in case there is a requirements.txt
+# Note: this happens after the src COPY command into the working directory, to ensure that any
+# editable installs will still work
+RUN if [ -f "requirements.txt" ]; then \
+        pip install -r requirements.txt; \
+    fi
+
+# Run the custom dagster cloud post install script if available
+RUN if [ -f "dagster_cloud_post_install.sh" ]; then \
+        ./dagster_cloud_post_install.sh; \
+    fi
+
+# Make sure dagster-cloud is installed. Fail early here if not.
+RUN if ! dagster-cloud --version; then \
+        echo "Could not find the dagster-cloud package.  Make sure you include the dagster-cloud package in your setup.py, pyproject.toml, or requirements.txt file to ensure that it is installed as part of your workspace."; \
+        exit 1; \
+    fi