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

dagster_cloud_cli/core/pex_builder/gitlab_context.py

@@ -0,0 +1,30 @@
+import os
+
+from dagster_cloud_cli.core.pex_builder import git_context
+
+# Loads event details from within a gitlab pipeline
+
+
+class GitlabEvent:
+    def __init__(self, project_dir: str):
+        # https://docs.gitlab.com/ee/ci/variables/predefined_variables.html
+        self.commit_sha = os.environ["CI_COMMIT_SHA"]
+        self.branch_name = os.environ.get("CI_COMMIT_REF_NAME", "")
+        self.job_url = os.environ["CI_JOB_URL"]
+        self.project_name = os.environ["CI_PROJECT_NAME"]
+        self.project_url = os.environ["CI_PROJECT_URL"]
+
+        self.merge_request_iid: str | None = os.environ.get("CI_MERGE_REQUEST_IID")
+
+        self.branch_url = f"{self.project_url}/-/tree/{self.branch_name}"
+
+        if self.merge_request_iid:
+            self.merge_request_url = f"{self.project_url}/-/merge_requests/{self.merge_request_iid}"
+        else:
+            self.merge_request_url = None
+
+        self.git_metadata = git_context.get_git_commit_metadata(project_dir)
+
+
+def get_gitlab_event(project_dir) -> GitlabEvent:
+    return GitlabEvent(project_dir)

dagster_cloud_cli/core/pex_builder/parse_workspace.py

@@ -0,0 +1,39 @@
+import os
+from dataclasses import dataclass
+
+import yaml
+
+from dagster_cloud_cli import ui
+
+
+@dataclass
+class Location:
+    name: str
+    directory: str
+    build_folder: str
+    location_file: str
+
+
+def get_locations(dagster_cloud_yaml_file) -> list[Location]:
+    """Returns list of locations parsed from dagster_cloud.yaml."""
+    base_dir = os.path.abspath(os.path.dirname(dagster_cloud_yaml_file))
+
+    with open(dagster_cloud_yaml_file, encoding="utf-8") as yaml_file:
+        workspace_contents = yaml_file.read()
+        workspace_contents_yaml = yaml.safe_load(workspace_contents)
+
+        locations = []
+        for location in workspace_contents_yaml["locations"]:
+            location_dir = os.path.join(
+                base_dir, location.get("build", {"directory": "."}).get("directory")
+            )
+            locations.append(
+                Location(
+                    name=location["location_name"],
+                    directory=location_dir,
+                    build_folder=location_dir,
+                    location_file=os.path.abspath(dagster_cloud_yaml_file),
+                )
+            )
+        ui.print(f"Parsed {len(locations)} locations from {locations}")
+        return locations

dagster_cloud_cli/core/pex_builder/pex_registry.py

@@ -0,0 +1,148 @@
+import base64
+import json
+import os
+from tempfile import TemporaryDirectory
+from typing import Any
+
+import requests
+
+from dagster_cloud_cli import gql, ui
+
+GENERATE_PUT_URL_QUERY = """
+mutation GenerateServerlessPexUrlMutation($filenames: [String!]!) {
+    generateServerlessPexUrl(filenames: $filenames, method:PUT) {
+        url
+    }
+}
+"""
+
+GENERATE_GET_URL_QUERY = """
+mutation GenerateServerlessPexUrlMutation($filenames: [String!]!) {
+    generateServerlessPexUrl(filenames: $filenames, method:GET) {
+        url
+    }
+}
+"""
+
+
+def get_s3_urls_for_put(
+    dagster_cloud_url: str, dagster_cloud_api_token: str, filenames: list[str]
+) -> list[str] | None:
+    with gql.graphql_client_from_url(dagster_cloud_url, dagster_cloud_api_token) as client:
+        result = client.execute(
+            GENERATE_PUT_URL_QUERY,
+            variable_values={"filenames": filenames},
+            idempotent_mutation=True,
+        )
+
+        if result["data"]:
+            return [item["url"] for item in result["data"]["generateServerlessPexUrl"]]
+        else:
+            return None
+
+
+def get_s3_urls_for_get(
+    dagster_cloud_url: str, dagster_cloud_api_token: str, filenames: list[str]
+) -> list[str] | None:
+    with gql.graphql_client_from_url(dagster_cloud_url, dagster_cloud_api_token) as client:
+        result = client.execute(
+            GENERATE_GET_URL_QUERY,
+            variable_values={"filenames": filenames},
+            idempotent_mutation=True,
+        )
+
+        if result["data"]:
+            return [item["url"] for item in result["data"]["generateServerlessPexUrl"]]
+        else:
+            return None
+
+
+def requirements_hash_filename(requirements_hash: str, cache_tag: str | None):
+    # encode cache_tag as a filesystem safe string
+    if cache_tag:
+        cache_tag_suffix = "-" + base64.urlsafe_b64encode(cache_tag.encode("utf-8")).decode("utf-8")
+    else:
+        cache_tag_suffix = ""
+
+    return f"requirements-{requirements_hash}{cache_tag_suffix}.txt"
+
+
+def get_cached_deps_details(
+    dagster_cloud_url: str,
+    dagster_cloud_api_token: str,
+    requirements_hash: str,
+    cache_tag: str | None,
+) -> dict[str, Any] | None:
+    """Returns a metadata dict for the requirements_hash and cache_tag.
+
+    The dict contains:
+    'deps_pex_name': filename for the deps pex, eg 'deps-123234334.pex'
+    'dagster_version': dagster package version included in the deps, eg '1.0.14'
+    """
+    urls = get_s3_urls_for_get(
+        dagster_cloud_url,
+        dagster_cloud_api_token,
+        [requirements_hash_filename(requirements_hash, cache_tag)],
+    )
+    if not urls:
+        return None
+
+    url = urls[0]
+    if not url:
+        return None
+
+    result = requests.get(url)
+    if result.ok:
+        data = json.loads(result.content)
+        # Don't return partial information
+        if "deps_pex_name" in data and "dagster_version" in data:
+            return data
+
+    if result.status_code == 404:
+        ui.print("No dependencies in cache")
+    else:
+        ui.warn(f"Unable to read cached deps details {result}")
+    return None
+
+
+def set_cached_deps_details(
+    dagster_cloud_url: str,
+    dagster_cloud_api_token: str,
+    requirements_hash: str,
+    cache_tag: str | None,
+    deps_pex_name: str,
+    dagster_version: str,
+):
+    """Saves the deps_pex_name and dagster_version into the requirements hash file."""
+    filename = requirements_hash_filename(requirements_hash, cache_tag)
+    content = json.dumps({"deps_pex_name": deps_pex_name, "dagster_version": dagster_version})
+    with TemporaryDirectory() as tmp_dir:
+        filepath = os.path.join(tmp_dir, filename)
+        with open(filepath, "w", encoding="utf-8") as f:
+            f.write(content)
+
+        upload_files(dagster_cloud_url, dagster_cloud_api_token, [filepath])
+
+
+def upload_files(dagster_cloud_url: str, dagster_cloud_api_token: str, filepaths: list[str]):
+    filenames = [os.path.basename(filepath) for filepath in filepaths]
+    urls = get_s3_urls_for_put(dagster_cloud_url, dagster_cloud_api_token, filenames)
+    if not urls:
+        ui.error(f"Cannot upload files, did not get PUT urls for: {filenames}")
+        return
+
+    # we expect response list to be in the same order as the request
+    for _, filepath, url in zip(filenames, filepaths, urls):
+        if not url:
+            ui.print(f"No upload URL received for {filepath} - skipping")
+            continue
+
+        ui.print(f"Uploading {filepath} ...")
+        with open(filepath, "rb") as f:
+            response = requests.put(url, data=f)
+            if response.ok:
+                ui.print(f"Upload successful: {filepath}")
+            else:
+                ui.error(f"Upload failed for {filepath}: {response}")
+                ui.error(f"Upload URL: {url}")
+                ui.error(f"{response.content}")