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

dagster_cloud_cli/commands/workspace/__init__.py

@@ -0,0 +1,358 @@
+import time
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any
+
+import yaml
+from dagster_shared import check
+from dagster_shared.serdes.serdes import deserialize_value
+from typer import Argument, Option, Typer
+
+from dagster_cloud_cli import gql, ui
+from dagster_cloud_cli.config_utils import (
+    DEFAULT_LOCATION_LOAD_TIMEOUT,
+    DEPLOYMENT_CLI_OPTIONS,
+    dagster_cloud_options,
+    get_location_document,
+)
+from dagster_cloud_cli.core.graphql_client import DagsterCloudGraphQLClient
+from dagster_cloud_cli.core.workspace import CodeLocationDeployData
+from dagster_cloud_cli.utils import add_options
+
+DEFAULT_LOCATIONS_YAML_FILENAME = "locations.yaml"
+
+app = Typer(help="Manage your Dagster Cloud workspace.")
+
+
+def _add_or_update_location(
+    client: DagsterCloudGraphQLClient,
+    location_document: dict[str, Any],
+    location_load_timeout: int,
+    agent_heartbeat_timeout: int,
+    url: str | None = None,
+) -> None:
+    try:
+        gql.add_or_update_code_location(client, location_document)
+        name = location_document["location_name"]
+        ui.print(f"Added or updated location {name}.")
+        wait_for_load(
+            client,
+            [name],
+            location_load_timeout=location_load_timeout,
+            agent_heartbeat_timeout=agent_heartbeat_timeout,
+            url=url,
+        )
+    except Exception as e:
+        raise ui.error(str(e))
+
+
+@app.command(name="add-location", short_help="Add or update a code location image.")
+@dagster_cloud_options(allow_empty=True, requires_url=True)
+@add_options(DEPLOYMENT_CLI_OPTIONS)
+def add_command(
+    api_token: str,
+    url: str,
+    deployment: str | None,
+    location_load_timeout: int,
+    agent_heartbeat_timeout: int,
+    location: str = Argument(None, help="Code location name."),
+    **kwargs,
+):
+    """Add or update the image for a code location in the deployment."""
+    with gql.graphql_client_from_url(url, api_token, deployment_name=deployment) as client:
+        location_document = get_location_document(location, kwargs)
+        _add_or_update_location(
+            client,
+            location_document,
+            location_load_timeout=location_load_timeout,
+            agent_heartbeat_timeout=agent_heartbeat_timeout,
+            url=url,
+        )
+
+
+def list_locations(location_names: list[str]) -> str:
+    if len(location_names) == 0:
+        return ""
+    elif len(location_names) == 1:
+        return location_names[0]
+    else:
+        return f"{', '.join(location_names[:-1])}, and {location_names[-1]}"
+
+
+@app.command(name="update-location", short_help="Update a code location image.")
+@dagster_cloud_options(allow_empty=True, requires_url=True)
+@add_options(DEPLOYMENT_CLI_OPTIONS)
+def update_command(
+    api_token: str,
+    url: str,
+    deployment: str | None,
+    location_load_timeout: int,
+    agent_heartbeat_timeout: int,
+    location: str = Argument(None, help="Code location name."),
+    **kwargs,
+):
+    """Update the image for a code location in a deployment."""
+    with gql.graphql_client_from_url(url, api_token, deployment_name=deployment) as client:
+        location_document = get_location_document(location, kwargs)
+        _add_or_update_location(
+            client,
+            location_document,
+            location_load_timeout=location_load_timeout,
+            agent_heartbeat_timeout=agent_heartbeat_timeout,
+            url=url,
+        )
+
+
+def _format_error(load_error: Mapping[str, Any]):
+    result = [
+        load_error["message"],
+        "".join(load_error["stack"]),
+    ]
+
+    for chain_link in load_error["errorChain"]:
+        result.append(
+            "The above exception was caused by the following exception:"
+            if chain_link["isExplicitLink"]
+            else "The above exception occurred during handling of the following exception:"
+        )
+
+        result.extend([chain_link["error"]["message"], "".join(chain_link["error"]["stack"])])
+
+    return "\n".join(result)
+
+
+def wait_for_load(
+    client,
+    locations,
+    location_load_timeout=DEFAULT_LOCATION_LOAD_TIMEOUT,
+    agent_heartbeat_timeout: int | None = DEFAULT_LOCATION_LOAD_TIMEOUT,
+    url: str | None = None,
+):
+    start_time = time.time()
+    if url:
+        ui.print(f"View the status of your locations at {url}/locations.\n")
+    ui.print(f"Waiting for agent to sync changes to {list_locations(locations)}...")
+
+    if not location_load_timeout and not agent_heartbeat_timeout:
+        return
+
+    has_agent_heartbeat = False
+    iterations = 0
+    while True:
+        if not has_agent_heartbeat:
+            if agent_heartbeat_timeout and time.time() - start_time > agent_heartbeat_timeout:
+                raise ui.error(
+                    "No Dagster Cloud agent is actively heartbeating. Make sure that you have a"
+                    " Dagster Cloud agent running."
+                )
+            try:
+                agents = gql.fetch_agent_status(client)
+            except Exception as e:
+                raise ui.error("Unable to query agent status: " + str(e))
+
+            has_agent_heartbeat = any(a["status"] == "RUNNING" for a in agents)
+
+        if time.time() - start_time > location_load_timeout:
+            raise ui.error("Timed out waiting for location data to update.")
+
+        try:
+            nodes = gql.fetch_code_locations(client)
+        except Exception as e:
+            raise ui.error(str(e))
+
+        nodes_by_location = {node["name"]: node for node in nodes}
+
+        if all(
+            location in nodes_by_location
+            and nodes_by_location[location].get("loadStatus") == "LOADED"
+            for location in locations
+        ):
+            error_locations = [
+                location
+                for location in locations
+                if "locationOrLoadError" in nodes_by_location[location]
+                and nodes_by_location[location]["locationOrLoadError"]["__typename"]
+                == "PythonError"
+            ]
+
+            if error_locations:
+                error_string = (
+                    "Some locations failed to load after being synced by the agent:\n"
+                    + "\n".join(
+                        [
+                            f"Error loading {error_location}:\n"
+                            f"{_format_error(nodes_by_location[error_location]['locationOrLoadError'])}"
+                            for error_location in error_locations
+                        ]
+                    )
+                )
+                raise ui.error(error_string)
+            else:
+                ui.print(
+                    f"Agent synced changes to {list_locations(locations)}. Changes should now be"
+                    " visible in Dagster Cloud."
+                )
+                break
+
+        time.sleep(3)
+        iterations += 1
+
+        if iterations % 3 == 0:
+            ui.print(
+                f"Still waiting for agent to sync changes to {list_locations(locations)}. This can"
+                " take a few minutes."
+            )
+
+
+@app.command(
+    name="delete-location",
+)
+@dagster_cloud_options(allow_empty=True, requires_url=True)
+def delete_command(
+    api_token: str,
+    url: str,
+    deployment: str | None,
+    location: str = Argument(..., help="Code location name."),
+):
+    """Delete a code location from the deployment."""
+    with gql.graphql_client_from_url(url, api_token, deployment_name=deployment) as client:
+        try:
+            gql.delete_code_location(client, location)
+            ui.print(f"Deleted location {location}.")
+        except Exception as e:
+            raise ui.error(str(e))
+
+
+@app.command(
+    name="list",
+)
+@dagster_cloud_options(allow_empty=True, requires_url=True)
+def list_command(
+    url: str,
+    deployment: str | None,
+    api_token: str,
+):
+    """List code locations in the deployment."""
+    with gql.graphql_client_from_url(url, api_token, deployment_name=deployment) as client:
+        execute_list_command(client)
+
+
+def execute_list_command(client):
+    list_res = gql.fetch_workspace_entries(client)
+
+    ui.print("Listing locations...")
+
+    for location in list_res:
+        metadata = deserialize_value(
+            location["serializedDeploymentMetadata"], CodeLocationDeployData
+        )
+
+        location_desc = [location["locationName"]]
+        if metadata.python_file:
+            location_desc.append(f"File: {metadata.python_file}")
+        if metadata.package_name:
+            location_desc.append(f"Package: {metadata.package_name}")
+        if metadata.image:
+            location_desc.append(f"Image: {metadata.image}")
+
+        ui.print("\t".join(location_desc))
+
+
+@app.command(name="pull")
+@dagster_cloud_options(allow_empty=True, requires_url=True)
+def pull_command(
+    url: str,
+    api_token: str,
+    deployment: str | None,
+):
+    """Retrieve code location definitions as a workspace.yaml file."""
+    with gql.graphql_client_from_url(url, api_token, deployment_name=deployment) as client:
+        document = gql.fetch_locations_as_document(client)
+        ui.print_yaml(document or {})
+
+
+@app.command(
+    name="sync",
+    short_help="Sync deployment code locations with a workspace.yaml file.",
+)
+@dagster_cloud_options(allow_empty=True, requires_url=True)
+def sync_command(
+    url: str,
+    api_token: str,
+    deployment: str | None,
+    location_load_timeout: int,
+    agent_heartbeat_timeout: int,
+    workspace: Path = Option(
+        "workspace.yaml",
+        "--workspace",
+        "-w",
+        exists=True,
+        help="Path to workspace file.",
+    ),
+):
+    """Sync the workspace with the contents of a workspace.yaml file."""
+    with gql.graphql_client_from_url(url, api_token, deployment_name=deployment) as client:
+        execute_sync_command(
+            client,
+            workspace,
+            location_load_timeout=location_load_timeout,
+            agent_heartbeat_timeout=agent_heartbeat_timeout,
+            url=url,
+        )
+
+
+def format_workspace_config(workspace_config) -> dict[str, Any]:
+    """Ensures the input workspace config is in the modern format, migrating an input
+    in the old format if need be.
+    """
+    check.dict_param(workspace_config, "workspace_config")
+
+    if isinstance(workspace_config.get("locations"), dict):
+        # Convert legacy formatted locations to modern format
+        updated_locations = []
+        for name, location in workspace_config["locations"].items():
+            new_location = {
+                k: v
+                for k, v in location.items()
+                if k
+                not in ("python_file", "package_name", "module_name", "autoload_defs_module_name")
+            }
+            new_location["code_source"] = {}
+            if "python_file" in location:
+                new_location["code_source"]["python_file"] = location["python_file"]
+            if "package_name" in location:
+                new_location["code_source"]["package_name"] = location["package_name"]
+            if "module_name" in location:
+                new_location["code_source"]["module_name"] = location["module_name"]
+            if "autoload_defs_module_name" in location:
+                new_location["code_source"]["autoload_defs_module_name"] = location[
+                    "autoload_defs_module_name"
+                ]
+
+            new_location["location_name"] = name
+            updated_locations.append(new_location)
+        return {"locations": updated_locations}
+    else:
+        check.is_list(workspace_config.get("locations"))
+        return workspace_config
+
+
+def execute_sync_command(
+    client, workspace, location_load_timeout, agent_heartbeat_timeout, url: str | None = None
+):
+    with open(str(workspace), encoding="utf8") as f:
+        config = yaml.load(f.read(), Loader=yaml.SafeLoader)
+        processed_config = format_workspace_config(config)
+
+    try:
+        locations = gql.reconcile_code_locations(client, processed_config)
+        ui.print(f"Synced locations: {', '.join(locations)}")
+        wait_for_load(
+            client,
+            locations,
+            location_load_timeout=location_load_timeout,
+            agent_heartbeat_timeout=agent_heartbeat_timeout,
+            url=url,
+        )
+    except Exception as e:
+        raise ui.error(str(e))

dagster_cloud_cli/config/__init__.py

@@ -0,0 +1,6 @@
+from dagster_cloud_cli.config.jinja_template_loader import (
+    JinjaTemplateLoader as JinjaTemplateLoader,
+)
+from dagster_cloud_cli.config.merger import (
+    DagsterCloudConfigDefaultsMerger as DagsterCloudConfigDefaultsMerger,
+)

dagster_cloud_cli/config/jinja_template_loader.py

@@ -0,0 +1,45 @@
+from collections.abc import Mapping
+from os import getenv
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+
+class JinjaTemplateLoader:
+    """Reads and render a file with jinja2 templating support and provide additional templating functionality.
+
+    Ref: https://jinja.palletsprojects.com/en/stable/
+
+    In addition to the built-in global functions and filters, this class provides the following:
+
+    Custom jinja2 global function(s):
+    - get_env(name: str, default: str = "") -> str
+      Get environment variable if present or return a default value.
+
+    Custom jinja2 filter(s):
+    - to_yaml(input: Any, **kwargs) -> str
+      Renders a value into its yaml representation.
+    """
+
+    @staticmethod
+    def to_yaml(value: Any, **kwargs) -> str:
+        return yaml.dump(value, **kwargs)
+
+    @staticmethod
+    def get_env(name: str, default: str = "") -> str:
+        return getenv(name, default)
+
+    def render(self, filepath: str, context: Mapping[str, Any]) -> str:
+        # defer for import performance
+        from jinja2 import Environment, FileSystemLoader
+
+        path = Path(filepath)
+        template_path = path.resolve().parent
+        filename = path.name
+
+        env = Environment(loader=FileSystemLoader(template_path))
+        env.filters["to_yaml"] = self.to_yaml
+        env.globals["get_env"] = self.get_env
+        template = env.get_template(filename)
+        return template.render(context)

dagster_cloud_cli/config/list_merge_strategies.py

@@ -0,0 +1,77 @@
+from collections.abc import Callable
+from typing import Any
+
+from dagster_shared import check
+from dagster_shared.utils.hash import hash_collection
+
+ListMergeFunction = Callable[[list[Any], list[Any]], list[Any]]
+
+
+def replace(from_list: list[Any], _: list[Any]) -> list[Any]:
+    return from_list
+
+
+def append(from_list: list[Any], onto_list: list[Any]) -> list[Any]:
+    return onto_list + from_list
+
+
+def deduplicate(from_list: list[Any], onto_list: list[Any]) -> list[Any]:
+    return _deduplicate_list(onto_list + from_list)
+
+
+def get_list_merger_for_key_value_separated_strings(separator: str) -> ListMergeFunction:
+    def func(from_list: list[Any], onto_list: list[Any]) -> list[Any]:
+        fl = check.list_param(from_list, "from_list", of_type=str)
+        ol = check.list_param(onto_list, "onto_list", of_type=str)
+        return _merge_key_value_separated_strings(fl, ol, separator)
+
+    return func
+
+
+def get_list_merger_for_identifiable_dicts(name_key: str) -> ListMergeFunction:
+    def func(from_list: list[Any], onto_list: list[Any]) -> list[Any]:
+        fl = check.list_param(from_list, "from_list", of_type=dict)
+        ol = check.list_param(onto_list, "onto_list", of_type=dict)
+        return _merge_object_list_by_identity(name_key, fl, ol)
+
+    return func
+
+
+def _deduplicate_list(lst: list[Any]) -> list[Any]:
+    result = []
+    hashes = set()
+    for value in lst:
+        value_hash = hash_collection(value) if isinstance(value, (list, dict)) else hash(value)
+        if value_hash not in hashes:
+            hashes.add(value_hash)
+            result.append(value)
+    return result
+
+
+def _merge_key_value_separated_strings(
+    from_list: list[str], onto_list: list[str], separator: str
+) -> list[str]:
+    from_items = [item.split(separator)[0] for item in from_list]
+    original_items = []
+
+    for item in onto_list:
+        parts = item.split(separator)
+        if len(parts) > 2:
+            raise ValueError(f"Invalid key-value pair: {item}")
+
+        if parts[0] not in from_items:
+            original_items.append(item)
+
+    return original_items + from_list
+
+
+def _merge_object_list_by_identity(
+    name_key: str, from_list: list[dict[str, Any]], onto_list: list[dict[str, Any]]
+) -> list[dict[str, Any]]:
+    onto_items = {item[name_key]: item for item in onto_list}
+    from_items = {item[name_key]: item for item in from_list}
+
+    merged_items = {
+        key: item for key, item in onto_items.items() if key not in from_items
+    } | from_items
+    return list(merged_items.values())

dagster_cloud_cli/config/loader.py

@@ -0,0 +1,23 @@
+from typing import Any
+
+from dagster_cloud_cli.config.jinja_template_loader import JinjaTemplateLoader
+from dagster_cloud_cli.config.merger import DagsterCloudConfigDefaultsMerger
+from dagster_cloud_cli.config.models import ProcessedDagsterCloudConfig, load_dagster_cloud_yaml
+
+
+class DagsterCloudLocationsConfigLoader:
+    """Load a Dagster Cloud config from a template, with context values, merge in defaults to code location
+    and return a ProcessedDagsterCloudConfig object.
+    """
+
+    @staticmethod
+    def load_from_template(
+        template_file: str, context: dict[str, Any]
+    ) -> ProcessedDagsterCloudConfig:
+        template_loader = JinjaTemplateLoader()
+        template = template_loader.render(template_file, context)
+
+        dagster_cloud_config = load_dagster_cloud_yaml(template)
+
+        config_merger = DagsterCloudConfigDefaultsMerger()
+        return config_merger.process(dagster_cloud_config)

dagster_cloud_cli/config/merger.py

@@ -0,0 +1,92 @@
+from copy import deepcopy
+from typing import Any
+
+from dagster_shared import check
+from dagster_shared.utils import remove_none_recursively
+
+from dagster_cloud_cli.config.list_merge_strategies import (
+    deduplicate,
+    get_list_merger_for_identifiable_dicts,
+    get_list_merger_for_key_value_separated_strings,
+    replace,
+)
+from dagster_cloud_cli.config.models import DagsterCloudYaml, Location, ProcessedDagsterCloudConfig
+
+
+class DagsterCloudConfigDefaultsMerger:
+    """Transform the DagsterCloudYaml provided configuration input and return the processed config resulting
+    from merging the defaults with each location.
+    """
+
+    def __init__(self):
+        equal_sign_list_merger = get_list_merger_for_key_value_separated_strings("=")
+        named_object_list_merger = get_list_merger_for_identifiable_dicts("name")
+        keyed_object_list_merger = get_list_merger_for_identifiable_dicts("key")
+
+        self._default_list_strategy = deduplicate
+        self._list_strategy_exceptions = {
+            "container_context.docker.env_vars": equal_sign_list_merger,
+            "container_context.ecs.env_vars": equal_sign_list_merger,
+            "container_context.ecs.run_sidecar_containers": named_object_list_merger,
+            "container_context.ecs.secrets": named_object_list_merger,
+            "container_context.ecs.server_sidecar_containers": named_object_list_merger,
+            "container_context.ecs.volumes": named_object_list_merger,
+            "container_context.ecs.run_ecs_tags": keyed_object_list_merger,
+            "container_context.ecs.server_ecs_tags": keyed_object_list_merger,
+            "container_context.k8s.env_vars": equal_sign_list_merger,
+            "container_context.k8s.image_pull_secrets": named_object_list_merger,
+            "container_context.k8s.run_k8s_config.container_config.args": replace,
+            "container_context.k8s.run_k8s_config.container_config.command": replace,
+            "container_context.k8s.server_k8s_config.container_config.args": replace,
+            "container_context.k8s.server_k8s_config.container_config.command": replace,
+            "container_context.k8s.volumes": named_object_list_merger,
+            "container_context.k8s.volume_mounts": named_object_list_merger,
+        }
+
+    def process(self, source_config: DagsterCloudYaml) -> ProcessedDagsterCloudConfig:
+        if source_config.defaults is None:
+            return ProcessedDagsterCloudConfig(locations=source_config.locations)
+
+        source_dict = source_config.model_dump()
+        return self.process_dict(source_dict)
+
+    def process_dict(self, source_dict: dict[str, Any]) -> ProcessedDagsterCloudConfig:
+        processed_config = ProcessedDagsterCloudConfig(locations=[])
+
+        defaults = remove_none_recursively(source_dict.get("defaults", {}))
+        for location in source_dict.get("locations", []):
+            location_data = remove_none_recursively(location)
+            if not defaults:
+                processed_config.locations.append(Location.model_validate(location_data))
+                continue
+
+            transformed_location = self.merge(defaults, location_data)
+            processed_config.locations.append(Location.model_validate(transformed_location))
+
+        return processed_config
+
+    def merge(self, onto_dict: dict[str, Any], from_dict: dict[str, Any]):
+        dest = deepcopy(onto_dict)
+        return self._merge_dict("", dest, from_dict)
+
+    def _merge_dict(self, parent_key: str, onto_dict: dict[str, Any], from_dict: dict[str, Any]):
+        for key, value in from_dict.items():
+            key_path = ".".join([parent_key, key]) if parent_key else key
+            if key in onto_dict and onto_dict[key] is not None:
+                if isinstance(value, dict):
+                    check.invariant(isinstance(onto_dict[key], dict))
+                    onto_dict[key] = self._merge_dict(key_path, onto_dict[key], value)
+                elif isinstance(value, list):
+                    check.invariant(isinstance(onto_dict[key], list))
+                    onto_dict[key] = self._merge_list(key_path, value, onto_dict[key])
+                else:
+                    onto_dict[key] = value
+            else:
+                onto_dict[key] = value
+        return onto_dict
+
+    def _merge_list(self, key_path: str, from_list: list[Any], onto_list: list[Any]) -> list[Any]:
+        if key_path in self._list_strategy_exceptions:
+            return self._list_strategy_exceptions[key_path](from_list, onto_list)
+
+        return self._default_list_strategy(from_list, onto_list)

dagster_cloud_cli/config/models.py

@@ -0,0 +1,76 @@
+"""Use a pydantic definition to validate dagster_cloud.yaml."""
+
+from typing import Any
+
+import yaml
+from pydantic import BaseModel, Field, field_validator, model_validator
+
+
+class CodeSource(BaseModel, extra="forbid"):
+    package_name: str | None = None
+    module_name: str | None = None
+    python_file: str | None = None
+    autoload_defs_module_name: str | None = None
+
+    @model_validator(mode="before")
+    def exactly_one_source_defined(cls, values: dict[str, str | None]) -> dict[str, str | None]:
+        defined = [key for key, value in values.items() if value]
+        if len(defined) > 1:
+            raise ValueError(
+                "only one of the following fields should be defined: " + ", ".join(defined)
+            )
+        elif not defined:
+            raise ValueError(
+                "one of package_name, module_name, python_file, or autoload_defs_module_name must be specified"
+            )
+        return values
+
+
+class Build(BaseModel, extra="forbid"):
+    directory: str | None = None
+    registry: str | None = None
+
+
+class Location(BaseModel, extra="forbid"):
+    location_name: str
+    code_source: CodeSource | None = None
+    build: Build | None = None
+    working_directory: str | None = None
+    image: str | None = None
+    executable_path: str | None = None
+    attribute: str | None = None
+    container_context: dict[str, Any] | None = None
+    agent_queue: str | None = None
+
+
+class LocationDefaults(BaseModel, extra="forbid"):
+    build: Build | None = None
+    working_directory: str | None = None
+    image: str | None = None
+    executable_path: str | None = None
+    container_context: dict[str, Any] | None = None
+    agent_queue: str | None = None
+
+
+class DagsterCloudYaml(BaseModel, extra="forbid"):
+    defaults: LocationDefaults | None = Field(
+        description="Default values for code locations", default=None
+    )
+    locations: list[Location] = Field(description="List of code locations")
+
+    @field_validator("locations")
+    def no_duplicate_names(cls, v: list[Location]) -> list[Location]:
+        names = set()
+        for location in v:
+            if location.location_name in names:
+                raise ValueError(f"duplicate location name: {location.location_name}")
+            names.add(location.location_name)
+        return v
+
+
+class ProcessedDagsterCloudConfig(BaseModel, extra="forbid"):
+    locations: list[Location] = Field(description="List of code locations")
+
+
+def load_dagster_cloud_yaml(text) -> DagsterCloudYaml:
+    return DagsterCloudYaml.model_validate(yaml.safe_load(text))