dagster-cloud 1.8.2__py3-none-any.whl → 1.12.6__py3-none-any.whl

dagster_cloud/dagster_insights/snowflake/dbt_wrapper.py

@@ -1,7 +1,9 @@
-from typing import TYPE_CHECKING, Any, Iterable, Iterator, List, Optional, Tuple, Union
+from collections.abc import Iterable, Iterator
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 import yaml
 from dagster import (
+    AssetCheckEvaluation,
     AssetCheckResult,
     AssetExecutionContext,
     AssetKey,
@@ -11,8 +13,15 @@ from dagster import (
     Output,
 )
 
-from ..insights_utils import extract_asset_info_from_event, handle_raise_on_error
-from .snowflake_utils import OPAQUE_ID_SQL_SIGIL, build_opaque_id_metadata, marker_asset_key_for_job
+from dagster_cloud.dagster_insights.insights_utils import (
+    extract_asset_info_from_event,
+    handle_raise_on_error,
+)
+from dagster_cloud.dagster_insights.snowflake.snowflake_utils import (
+    OPAQUE_ID_SQL_SIGIL,
+    build_opaque_id_metadata,
+    marker_asset_key_for_job,
+)
 
 if TYPE_CHECKING:
     from dagster_dbt import DbtCliInvocation
@@ -23,11 +32,21 @@ def dbt_with_snowflake_insights(
     context: Union[OpExecutionContext, AssetExecutionContext],
     dbt_cli_invocation: "DbtCliInvocation",
     dagster_events: Optional[
-        Iterable[Union[Output, AssetMaterialization, AssetObservation, AssetCheckResult]]
+        Iterable[
+            Union[
+                Output,
+                AssetMaterialization,
+                AssetObservation,
+                AssetCheckResult,
+                AssetCheckEvaluation,
+            ]
+        ]
     ] = None,
     skip_config_check=False,
     record_observation_usage: bool = True,
-) -> Iterator[Union[Output, AssetMaterialization, AssetObservation, AssetCheckResult]]:
+) -> Iterator[
+    Union[Output, AssetMaterialization, AssetObservation, AssetCheckResult, AssetCheckEvaluation]
+]:
     """Wraps a dagster-dbt invocation to associate each Snowflake query with the produced
     asset materializations. This allows the cost of each query to be associated with the asset
     materialization that it produced.
@@ -38,7 +57,7 @@ def dbt_with_snowflake_insights(
     Args:
         context (AssetExecutionContext): The context of the asset that is being materialized.
         dbt_cli_invocation (DbtCliInvocation): The invocation of the dbt CLI to wrap.
-        dagster_events (Optional[Iterable[Union[Output, AssetObservation, AssetCheckResult]]]):
+        dagster_events (Optional[Iterable[Union[Output, AssetObservation, AssetCheckResult, AssetCheckEvaluation]]]):
             The events that were produced by the dbt CLI invocation. If not provided, it is assumed
             that the dbt CLI invocation has not yet been run, and it will be run and the events
             will be streamed.
@@ -89,10 +108,17 @@ def dbt_with_snowflake_insights(
     if dagster_events is None:
         dagster_events = dbt_cli_invocation.stream()
 
-    asset_and_partition_key_to_unique_id: List[Tuple[AssetKey, Optional[str], Any]] = []
+    asset_and_partition_key_to_unique_id: list[tuple[AssetKey, Optional[str], Any]] = []
     for dagster_event in dagster_events:
         if isinstance(
-            dagster_event, (AssetMaterialization, AssetObservation, Output, AssetCheckResult)
+            dagster_event,
+            (
+                AssetMaterialization,
+                AssetObservation,
+                Output,
+                AssetCheckResult,
+                AssetCheckEvaluation,
+            ),
         ):
             unique_id = dagster_event.metadata["unique_id"].value
             asset_key, partition = extract_asset_info_from_event(

dagster_cloud/dagster_insights/snowflake/definitions.py

@@ -1,14 +1,16 @@
 import warnings
+from collections.abc import Sequence
 from dataclasses import dataclass
 from datetime import date, datetime, timedelta, timezone
-from pprint import pprint
-from typing import TYPE_CHECKING, Optional, Sequence, Union
+from typing import TYPE_CHECKING, Optional, Union
 
 from dagster import (
     AssetExecutionContext,
+    AssetKey,
     AssetsDefinition,
     AssetSelection,
     HourlyPartitionsDefinition,
+    RunRequest,
     ScheduleDefinition,
     ScheduleEvaluationContext,
     TimeWindow,
@@ -18,9 +20,16 @@ from dagster import (
     fs_io_manager,
     schedule,
 )
+from dagster._core.definitions.unresolved_asset_job_definition import UnresolvedAssetJobDefinition
+from dagster._core.storage.tags import (
+    ASSET_PARTITION_RANGE_END_TAG,
+    ASSET_PARTITION_RANGE_START_TAG,
+)
 
-from ..metrics_utils import put_cost_information
-from .dagster_snowflake_insights import get_cost_data_for_hour
+from dagster_cloud.dagster_insights.metrics_utils import put_cost_information
+from dagster_cloud.dagster_insights.snowflake.dagster_snowflake_insights import (
+    get_cost_data_for_hour,
+)
 
 if TYPE_CHECKING:
     from dagster_snowflake import SnowflakeConnection
@@ -34,6 +43,22 @@ class SnowflakeInsightsDefinitions:
     schedule: ScheduleDefinition
 
 
+def _build_run_request_for_partition_key_range(
+    job: UnresolvedAssetJobDefinition,
+    asset_keys: Sequence[AssetKey],
+    partition_range_start: str,
+    partition_range_end: str,
+) -> RunRequest:
+    tags = {
+        ASSET_PARTITION_RANGE_START_TAG: partition_range_start,
+        ASSET_PARTITION_RANGE_END_TAG: partition_range_end,
+    }
+    partition_key = partition_range_start if partition_range_start == partition_range_end else None
+    return RunRequest(
+        job_name=job.name, asset_selection=asset_keys, partition_key=partition_key, tags=tags
+    )
+
+
 def create_snowflake_insights_asset_and_schedule(
     start_date: Union[datetime, date, str],
     name: Optional[str] = None,
@@ -97,7 +122,7 @@ def create_snowflake_insights_asset_and_schedule(
     def poll_snowflake_query_history_hour(
         context: AssetExecutionContext,
     ) -> None:
-        snowflake: "SnowflakeConnection" = getattr(context.resources, snowflake_resource_key)
+        snowflake: SnowflakeConnection = getattr(context.resources, snowflake_resource_key)
 
         start_hour = context.partition_time_window.start
         end_hour = context.partition_time_window.end
@@ -130,7 +155,7 @@ def create_snowflake_insights_asset_and_schedule(
         )
 
         if dry_run:
-            pprint(costs)
+            pass
         else:
             context.log.info(
                 f"Submitting cost information for {len(costs)} queries to Dagster Insights"
@@ -168,13 +193,14 @@ def create_snowflake_insights_asset_and_schedule(
             n_hours_ago = timestamp - timedelta(hours=schedule_batch_size_hrs)
             window = TimeWindow(start=n_hours_ago, end=timestamp)
 
-            partition_keys = partitions_def.get_partition_keys_in_time_window(window)
+            partition_key_range = partitions_def.get_partition_key_range_for_time_window(window)
 
-            for partition_key in partition_keys:
-                yield insights_job.run_request_for_partition(
-                    partition_key=partition_key,
-                    run_key=partition_key,
-                )
+            yield _build_run_request_for_partition_key_range(
+                insights_job,
+                [poll_snowflake_query_history_hour.key],
+                partition_key_range.start,
+                partition_key_range.end,
+            )
 
         insights_schedule = _insights_schedule
 

dagster_cloud/dagster_insights/snowflake/insights_snowflake_resource.py

@@ -1,19 +1,8 @@
 import warnings
+from collections.abc import Generator, Iterable, Iterator, Mapping, Sequence
 from contextlib import contextmanager
 from io import StringIO
-from typing import (
-    Any,
-    Generator,
-    Iterable,
-    Iterator,
-    Mapping,
-    Optional,
-    Sequence,
-    Tuple,
-    Type,
-    Union,
-    cast,
-)
+from typing import Any, Optional, Union, cast
 
 import snowflake.connector
 from dagster import (
@@ -23,7 +12,7 @@ from dagster import (
     OpExecutionContext,
     get_dagster_logger,
 )
-from dagster._annotations import experimental, public
+from dagster._annotations import beta, public
 from dagster._check import CheckError
 from dagster._core.errors import DagsterInvalidPropertyError
 from dagster._core.storage.event_log.sql_event_log import SqlDbConnection
@@ -31,13 +20,12 @@ from dagster_snowflake import SnowflakeConnection, SnowflakeResource
 from snowflake.connector.cursor import SnowflakeCursor
 
 from dagster_cloud.dagster_insights.insights_utils import get_current_context_and_asset_key
-
-from .snowflake_utils import meter_snowflake_query
+from dagster_cloud.dagster_insights.snowflake.snowflake_utils import meter_snowflake_query
 
 
-def get_current_context_and_asset_key_or_warn() -> (
-    Tuple[Union[OpExecutionContext, AssetExecutionContext, None], Optional[AssetKey]]
-):
+def get_current_context_and_asset_key_or_warn() -> tuple[
+    Union[OpExecutionContext, AssetExecutionContext, None], Optional[AssetKey]
+]:
     try:
         return get_current_context_and_asset_key()
     except (DagsterInvalidPropertyError, DagsterInvariantViolationError, CheckError):
@@ -80,7 +68,7 @@ class WrappedSnowflakeConnection(snowflake.connector.SnowflakeConnection):
         sql_text: str,
         remove_comments: bool = False,
         return_cursors: bool = True,
-        cursor_class: Type[SnowflakeCursor] = InsightsSnowflakeCursor,
+        cursor_class: type[SnowflakeCursor] = InsightsSnowflakeCursor,
         **kwargs,
     ) -> Iterable[SnowflakeCursor]:
         return super().execute_string(
@@ -95,14 +83,14 @@ class WrappedSnowflakeConnection(snowflake.connector.SnowflakeConnection):
         self,
         stream: StringIO,
         remove_comments: bool = False,
-        cursor_class: Type[SnowflakeCursor] = InsightsSnowflakeCursor,
+        cursor_class: type[SnowflakeCursor] = InsightsSnowflakeCursor,
         **kwargs,
     ) -> Generator[SnowflakeCursor, None, None]:
         return super().execute_stream(stream, remove_comments, cursor_class, **kwargs)  # type: ignore  # (bad stubs)
 
     def cursor(self, cursor_class=None) -> SnowflakeCursor:
         if cursor_class is None:
-            cursor = cast(InsightsSnowflakeCursor, super().cursor(InsightsSnowflakeCursor))
+            cursor = cast("InsightsSnowflakeCursor", super().cursor(InsightsSnowflakeCursor))
             cursor.set_asset_key(self._asset_key)
         else:
             cursor = super().cursor(cursor_class)
@@ -110,7 +98,7 @@ class WrappedSnowflakeConnection(snowflake.connector.SnowflakeConnection):
         return cursor
 
 
-@experimental
+@beta
 class InsightsSnowflakeResource(SnowflakeResource):
     """A wrapper around :py:class:`SnowflakeResource` which automatically tags
     Snowflake queries with comments which can be used to attribute Snowflake

dagster_cloud/definitions/job_selection.py

@@ -0,0 +1,36 @@
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from typing import AbstractSet  # noqa: UP035
+
+import dagster._check as check
+from dagster._record import record
+from dagster_shared.serdes import whitelist_for_serdes
+
+
+@whitelist_for_serdes
+@record
+class JobSelection(ABC):
+    code_location_name: str
+    repository_name: str
+
+    @abstractmethod
+    def resolve_job_names(self) -> AbstractSet[str]:
+        raise NotImplementedError()
+
+    @staticmethod
+    def names(
+        names: Sequence[str], code_location_name: str, repository_name: str
+    ) -> "NamesJobSelection":
+        check.invariant(len(names) == 1, "Only one job name is supported at this time")
+        return NamesJobSelection(
+            job_names=names, code_location_name=code_location_name, repository_name=repository_name
+        )
+
+
+@whitelist_for_serdes
+@record
+class NamesJobSelection(JobSelection):
+    job_names: Sequence[str]
+
+    def resolve_job_names(self) -> AbstractSet[str]:
+        return set(self.job_names)

dagster_cloud/execution/cloud_run_launcher/k8s.py

@@ -5,5 +5,5 @@ class CloudK8sRunLauncher(K8sRunLauncher):
     # Avoid call to `count_resume_run_attempts`, since resuming runs is not currently
     # supported in cloud and the method makes repeated event log calls.
     @property
-    def supports_run_worker_crash_recovery(self):
+    def supports_run_worker_crash_recovery(self):  # pyright: ignore[reportIncompatibleMethodOverride]
         return False

dagster_cloud/execution/cloud_run_launcher/process.py

@@ -6,11 +6,11 @@ from dagster._core.launcher import RunLauncher
 from dagster._core.launcher.base import LaunchRunContext
 from dagster._core.utils import parse_env_var
 from dagster._grpc.types import ExecuteRunArgs
-from dagster._serdes.ipc import open_ipc_subprocess
-from dagster_cloud_cli.core.workspace import get_instance_ref_for_user_code
+from dagster_shared.ipc import open_ipc_subprocess
 
 from dagster_cloud.execution.utils import TaskStatus
 from dagster_cloud.execution.utils.process import check_on_process, kill_process
+from dagster_cloud.workspace.user_code_launcher.utils import get_instance_ref_for_user_code
 
 PID_TAG = "process/pid"
 
@@ -59,7 +59,7 @@ class CloudProcessRunLauncher(RunLauncher):
                 for run_id in self._run_ids
                 if (
                     self._instance.get_run_by_id(run_id)
-                    and not self._instance.get_run_by_id(run_id).is_finished
+                    and not self._instance.get_run_by_id(run_id).is_finished  # pyright: ignore[reportOptionalMemberAccess]
                 )
             ]
 

dagster_cloud/execution/monitoring/__init__.py

@@ -2,22 +2,9 @@ import logging
 import os
 import sys
 import threading
+from collections.abc import Iterable, Mapping, Sequence
 from enum import Enum
-from typing import (
-    Any,
-    Dict,
-    Iterable,
-    List,
-    Mapping,
-    NamedTuple,
-    Optional,
-    Sequence,
-    Set,
-    Tuple,
-    TypedDict,
-    Union,
-    cast,
-)
+from typing import Any, NamedTuple, Optional, TypedDict, Union, cast
 
 import dagster._check as check
 import grpc
@@ -96,13 +83,13 @@ class CloudCodeServerHeartbeat(
         error: Optional[SerializableErrorInfo] = None,
         metadata: Optional[Mapping[str, Any]] = None,
     ):
-        return super(CloudCodeServerHeartbeat, cls).__new__(
+        return super().__new__(
             cls,
             location_name=check.str_param(location_name, "location_name"),
             server_status=check.inst_param(server_status, "server_status", CloudCodeServerStatus),
             error=check.opt_inst_param(error, "error", SerializableErrorInfo),
             metadata=cast(
-                CloudCodeServerHeartbeatMetadata, check.opt_mapping_param(metadata, "metadata")
+                "CloudCodeServerHeartbeatMetadata", check.opt_mapping_param(metadata, "metadata")
             ),
         )
 
@@ -113,7 +100,7 @@ class CloudCodeServerHeartbeat(
         return self._replace(
             error=None,
             metadata=cast(
-                CloudCodeServerHeartbeatMetadata,
+                "CloudCodeServerHeartbeatMetadata",
                 {},
             ),
         )
@@ -148,7 +135,7 @@ class CloudRunWorkerStatus(
         transient: Optional[bool] = None,
         run_worker_id: Optional[str] = None,
     ):
-        return super(CloudRunWorkerStatus, cls).__new__(
+        return super().__new__(
             cls,
             run_id=check.str_param(run_id, "run_id"),
             status_type=check.inst_param(status_type, "status_type", WorkerStatus),
@@ -180,7 +167,7 @@ class CloudRunWorkerStatuses(
     NamedTuple(
         "_CloudRunWorkerStatuses",
         [
-            ("statuses", List[CloudRunWorkerStatus]),
+            ("statuses", list[CloudRunWorkerStatus]),
             ("run_worker_monitoring_supported", bool),
             ("run_worker_monitoring_thread_alive", Optional[bool]),
         ],
@@ -188,11 +175,11 @@ class CloudRunWorkerStatuses(
 ):
     def __new__(
         cls,
-        statuses: Optional[List[CloudRunWorkerStatus]],
+        statuses: Optional[list[CloudRunWorkerStatus]],
         run_worker_monitoring_supported: bool,
         run_worker_monitoring_thread_alive: Optional[bool],
     ):
-        return super(CloudRunWorkerStatuses, cls).__new__(
+        return super().__new__(
             cls,
             statuses=check.opt_list_param(statuses, "statuses", of_type=CloudRunWorkerStatus),
             run_worker_monitoring_supported=check.bool_param(
@@ -229,13 +216,13 @@ def _is_grpc_unknown_error(error: Exception) -> bool:
 def get_cloud_run_worker_statuses(
     instance: DagsterCloudAgentInstance, deployment_names: Iterable[str], logger: logging.Logger
 ) -> Mapping[str, Sequence[CloudRunWorkerStatus]]:
-    statuses: Dict[str, Sequence[CloudRunWorkerStatus]] = {}
+    statuses: dict[str, Sequence[CloudRunWorkerStatus]] = {}
 
     # protected with a lock inside the method
     active_grpc_server_handles = instance.user_code_launcher.get_active_grpc_server_handles()
     active_grpc_server_handle_strings = [str(s) for s in active_grpc_server_handles]
 
-    active_non_isolated_run_ids_by_server_handle: Dict[
+    active_non_isolated_run_ids_by_server_handle: dict[
         str, Union[Sequence[str], _GetCurrentRunsError]
     ] = {}
     if instance.user_code_launcher.supports_get_current_runs_for_server_handle:
@@ -276,12 +263,19 @@ def get_cloud_run_worker_statuses(
             instance.ref_for_deployment(deployment_name)
         ) as scoped_instance:
             runs = scoped_instance.get_runs(RunsFilter(statuses=IN_PROGRESS_RUN_STATUSES))
-            statuses_for_deployment: List[CloudRunWorkerStatus] = []
+            statuses_for_deployment: list[CloudRunWorkerStatus] = []
             for run in runs:
                 if is_isolated_run(run):
                     launcher = scoped_instance.run_launcher
 
-                    run_worker_health = launcher.check_run_worker_health(run)
+                    try:
+                        run_worker_health = launcher.check_run_worker_health(run)
+                    except Exception:
+                        logger.exception(
+                            f"An exception occurred while checking run worker health for run '{run.run_id}'"
+                        )
+                        continue
+
                     run_worker_debug_info = None
 
                     if run_worker_health.status == WorkerStatus.FAILED:
@@ -385,8 +379,8 @@ def get_cloud_run_worker_statuses(
 
 def run_worker_monitoring_thread(
     instance: DagsterCloudAgentInstance,
-    deployments_to_check: Set[str],
-    statuses_dict: Dict[str, Sequence[CloudRunWorkerStatus]],
+    deployments_to_check: set[str],
+    statuses_dict: dict[str, Sequence[CloudRunWorkerStatus]],
     run_worker_monitoring_lock: threading.Lock,
     shutdown_event: threading.Event,
     monitoring_interval: Optional[int] = None,
@@ -409,8 +403,8 @@ def run_worker_monitoring_thread(
 
 def run_worker_monitoring_thread_iteration(
     instance: DagsterCloudAgentInstance,
-    deployments_to_check: Set[str],
-    statuses_dict: Dict[str, Sequence[CloudRunWorkerStatus]],
+    deployments_to_check: set[str],
+    statuses_dict: dict[str, Sequence[CloudRunWorkerStatus]],
     run_worker_monitoring_lock: threading.Lock,
     logger: logging.Logger,
 ) -> None:
@@ -432,11 +426,11 @@ def run_worker_monitoring_thread_iteration(
 
 def start_run_worker_monitoring_thread(
     instance: DagsterCloudAgentInstance,
-    deployments_to_check: Set[str],
-    statuses_dict: Dict[str, List[CloudRunWorkerStatus]],
+    deployments_to_check: set[str],
+    statuses_dict: dict[str, list[CloudRunWorkerStatus]],
     run_worker_monitoring_lock: threading.Lock,
     monitoring_interval: Optional[int] = None,
-) -> Tuple[threading.Thread, threading.Event]:
+) -> tuple[threading.Thread, threading.Event]:
     shutdown_event = threading.Event()
     thread = threading.Thread(
         target=run_worker_monitoring_thread,

dagster_cloud/execution/utils/process.py

@@ -1,10 +1,10 @@
 import os
 import sys
-from typing import Sequence
+from collections.abc import Sequence
 
-from dagster._serdes.ipc import interrupt_ipc_subprocess_pid, open_ipc_subprocess
+from dagster_shared.ipc import interrupt_ipc_subprocess_pid, open_ipc_subprocess
 
-from . import TaskStatus
+from dagster_cloud.execution.utils import TaskStatus
 
 
 def launch_process(args: Sequence[str]) -> int: