snowflake-cli 3.0.2__py3-none-any.whl → 3.2.0__py3-none-any.whl

snowflake/cli/api/exceptions.py

@@ -188,3 +188,44 @@ class IncompatibleParametersError(UsageError):
         super().__init__(
             f"Parameters {comma_separated_options} and {options_with_quotes[-1]} are incompatible and cannot be used simultaneously."
         )
+
+
+class UnmetParametersError(UsageError):
+    def __init__(self, options: list[str]):
+        options_with_quotes = [f"'{option}'" for option in options]
+        comma_separated_options = ", ".join(options_with_quotes[:-1])
+        super().__init__(
+            f"Parameters {comma_separated_options} and {options_with_quotes[-1]} must be used simultaneously."
+        )
+
+
+class NoWarehouseSelectedInSessionError(ClickException):
+    def __init__(self, msg: str):
+        super().__init__(
+            "Received the following error message while executing SQL statement:\n"
+            f"'{msg}'\n"
+            "Please provide a warehouse for the active session role in your project definition file, config.toml file, or via command line."
+        )
+
+
+class DoesNotExistOrUnauthorizedError(ClickException):
+    def __init__(self, msg: str):
+        super().__init__(
+            "Received the following error message while executing SQL statement:\n"
+            f"'{msg}'\n"
+            "Please check the name of the resource you are trying to query or the permissions of the role you are using to run the query."
+        )
+
+
+class CouldNotUseObjectError(ClickException):
+    def __init__(self, object_type: ObjectType, name: str):
+        super().__init__(
+            f"Could not use {object_type} {name}. Object does not exist, or operation cannot be performed."
+        )
+
+
+class ShowSpecificObjectMultipleRowsError(RuntimeError):
+    def __init__(self, show_obj_query: str):
+        super().__init__(
+            f"Received multiple rows from result of SQL statement: {show_obj_query}. Usage of 'show_specific_object' may not be properly scoped."
+        )

snowflake/cli/api/identifiers.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import re
+from pathlib import Path
 
 from click import ClickException
 from snowflake.cli.api.exceptions import FQNInconsistencyError, FQNNameError
@@ -121,8 +122,15 @@ class FQN:
         name = stage
         if stage.startswith("@"):
             name = stage[1:]
+        if stage.startswith("~"):
+            return cls(name="~", database=None, schema=None)
         return cls.from_string(name)
 
+    @classmethod
+    def from_stage_path(cls, stage_path: str) -> "FQN":
+        stage = Path(stage_path).parts[0]
+        return cls.from_stage(stage)
+
     @classmethod
     def from_identifier_model_v1(cls, model: ObjectIdentifierBaseModel) -> "FQN":
         """Create an instance from object model."""

snowflake/cli/api/metrics.py

@@ -11,8 +11,25 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations
 
-from typing import Dict, Optional
+import time
+import uuid
+from contextlib import contextmanager
+from dataclasses import dataclass, field, replace
+from heapq import nsmallest
+from typing import ClassVar, Dict, Iterator, List, Optional, Set
+
+
+class CLIMetricsInvalidUsageError(RuntimeError):
+    """
+    Indicative of bug in the code where a call to CLIMetrics was made erroneously
+
+    We do not want metrics errors to break the execution of commands,
+    so only raise this error in the event that an invariant was broken during setup
+    """
+
+    pass
 
 
 class _TypePrefix:
@@ -52,20 +69,135 @@ class CLICounterField:
         f"{_TypePrefix.FEATURES}.{_DomainPrefix.APP}.post_deploy_scripts"
     )
     PACKAGE_SCRIPTS = f"{_TypePrefix.FEATURES}.{_DomainPrefix.APP}.package_scripts"
+    EVENT_SHARING = f"{_TypePrefix.FEATURES}.{_DomainPrefix.APP}.event_sharing"
+    EVENT_SHARING_WARNING = (
+        f"{_TypePrefix.FEATURES}.{_DomainPrefix.APP}.event_sharing_warning"
+    )
+    EVENT_SHARING_ERROR = (
+        f"{_TypePrefix.FEATURES}.{_DomainPrefix.APP}.event_sharing_error"
+    )
+
+
+@dataclass
+class CLIMetricsSpan:
+    """
+    class for holding metrics span data and encapsulating related operations
+    """
+
+    # keys for dict representation
+    ID_KEY: ClassVar[str] = "id"
+    NAME_KEY: ClassVar[str] = "name"
+    PARENT_KEY: ClassVar[str] = "parent"
+    PARENT_ID_KEY: ClassVar[str] = "parent_id"
+    START_TIME_KEY: ClassVar[str] = "start_time"
+    EXECUTION_TIME_KEY: ClassVar[str] = "execution_time"
+    ERROR_KEY: ClassVar[str] = "error"
+    # total number of spans started under this span, inclusive of itself and its children's children (recursively)
+    SPAN_COUNT_IN_SUBTREE_KEY: ClassVar[str] = "span_count_in_subtree"
+    # the number of spans in the path between the current span and the topmost parent span, inclusive of both
+    SPAN_DEPTH_KEY: ClassVar[str] = "span_depth"
+    # denotes whether direct children were trimmed from telemetry payload
+    TRIMMED_KEY: ClassVar[str] = "trimmed"
+
+    # constructor vars
+    name: str
+    start_time: float  # relative to when the command first started executing
+    parent: Optional[CLIMetricsSpan] = None
+
+    # vars for reporting
+    span_id: str = field(init=False, default_factory=lambda: uuid.uuid4().hex)
+    execution_time: Optional[float] = field(init=False, default=None)
+    error: Optional[BaseException] = field(init=False, default=None)
+    span_depth: int = field(init=False, default=1)
+    span_count_in_subtree: int = field(init=False, default=1)
+
+    # vars for postprocessing
+    # spans started directly under this one
+    children: Set[CLIMetricsSpan] = field(init=False, default_factory=set)
+
+    # private state
+    # start time of the step from the monotonic clock in order to calculate execution time
+    _monotonic_start: float = field(
+        init=False, default_factory=lambda: time.monotonic()
+    )
+
+    def __hash__(self) -> int:
+        return hash(self.span_id)
+
+    def __post_init__(self):
+        if not self.name:
+            raise CLIMetricsInvalidUsageError("span name must not be empty")
+
+        if self.parent:
+            self.parent.add_child(self)
+            self.span_depth = self.parent.span_depth + 1
+
+    def increment_subtree_node_count(self) -> None:
+        self.span_count_in_subtree += 1
+
+        if self.parent:
+            self.parent.increment_subtree_node_count()
+
+    def add_child(self, child: CLIMetricsSpan) -> None:
+        self.children.add(child)
+        self.increment_subtree_node_count()
+
+    def finish(self, error: Optional[BaseException] = None) -> None:
+        """
+        Sets the execution time and (optionally) error raised for the span
+
+        If already called, this method is a no-op
+        """
+        if self.execution_time is not None:
+            return
+
+        if error:
+            self.error = error
 
+        self.execution_time = time.monotonic() - self._monotonic_start
 
+    def to_dict(self) -> Dict:
+        """
+        Custom dict conversion function to be used for reporting telemetry
+        """
+
+        return {
+            self.ID_KEY: self.span_id,
+            self.NAME_KEY: self.name,
+            self.START_TIME_KEY: self.start_time,
+            self.PARENT_KEY: self.parent.name if self.parent is not None else None,
+            self.PARENT_ID_KEY: (
+                self.parent.span_id if self.parent is not None else None
+            ),
+            self.EXECUTION_TIME_KEY: self.execution_time,
+            self.ERROR_KEY: type(self.error).__name__ if self.error else None,
+            self.SPAN_COUNT_IN_SUBTREE_KEY: self.span_count_in_subtree,
+            self.SPAN_DEPTH_KEY: self.span_depth,
+        }
+
+
+@dataclass
 class CLIMetrics:
     """
     Class to track various metrics across the execution of a command
     """
 
-    def __init__(self):
-        self._counters: Dict[str, int] = {}
+    # limits for reporting purposes
+    SPAN_DEPTH_LIMIT: ClassVar[int] = 5
+    SPAN_TOTAL_LIMIT: ClassVar[int] = 100
+
+    _counters: Dict[str, int] = field(init=False, default_factory=dict)
+    # stack of in progress spans as command is executing
+    _in_progress_spans: List[CLIMetricsSpan] = field(init=False, default_factory=list)
+    # list of finished steps for telemetry to process
+    _completed_spans: List[CLIMetricsSpan] = field(init=False, default_factory=list)
+    # monotonic clock time of when this class was initialized to approximate when the command first started executing
+    _monotonic_start: float = field(
+        init=False, default_factory=lambda: time.monotonic(), compare=False
+    )
 
-    def __eq__(self, other):
-        if isinstance(other, CLIMetrics):
-            return self._counters == other._counters
-        return False
+    def clone(self) -> CLIMetrics:
+        return replace(self)
 
     def get_counter(self, name: str) -> Optional[int]:
         return self._counters.get(name)
@@ -86,7 +218,91 @@ class CLIMetrics:
         else:
             self._counters[name] += value
 
+    @property
+    def current_span(self) -> Optional[CLIMetricsSpan]:
+        return self._in_progress_spans[-1] if len(self._in_progress_spans) > 0 else None
+
+    @contextmanager
+    def start_span(self, name: str) -> Iterator[CLIMetricsSpan]:
+        """
+        Starts a new span that tracks various metrics throughout its execution
+
+        Assumes that parent spans contain the entirety of their child spans
+        If not provided, parent spans are automatically populated as the most recently executed spans
+
+        Spans are not emitted in telemetry if depth/total limits are exceeded
+
+        :raises CliMetricsInvalidUsageError: if the step name is empty
+        """
+        new_span = CLIMetricsSpan(
+            name=name,
+            start_time=time.monotonic() - self._monotonic_start,
+            parent=self.current_span,
+        )
+
+        self._in_progress_spans.append(new_span)
+
+        try:
+            yield new_span
+        except BaseException as err:
+            new_span.finish(error=err)
+            raise
+        else:
+            new_span.finish()
+        finally:
+            self._completed_spans.append(new_span)
+            self._in_progress_spans.remove(new_span)
+
     @property
     def counters(self) -> Dict[str, int]:
         # return a copy of the original dict to avoid mutating the original
         return self._counters.copy()
+
+    @property
+    def num_spans_past_depth_limit(self) -> int:
+        return len(
+            [
+                span
+                for span in self._completed_spans
+                if span.span_depth > self.SPAN_DEPTH_LIMIT
+            ]
+        )
+
+    @property
+    def num_spans_past_total_limit(self) -> int:
+        return max(0, len(self._completed_spans) - self.SPAN_TOTAL_LIMIT)
+
+    @property
+    def completed_spans(self) -> List[Dict]:
+        """
+        Returns the completed spans tracked throughout a command, sorted by start time, for reporting telemetry
+
+        Ensures that the spans we send are within the configured limits and marks
+        certain spans as trimmed if their children would bypass the limits we set
+        """
+        # take spans breadth-first within the depth and total limits
+        # since we care more about the big picture than granularity
+        spans_to_report = set(
+            nsmallest(
+                n=self.SPAN_TOTAL_LIMIT,
+                iterable=(
+                    span
+                    for span in self._completed_spans
+                    if span.span_depth <= self.SPAN_DEPTH_LIMIT
+                ),
+                key=lambda span: (span.span_depth, span.start_time),
+            )
+        )
+
+        # sort by start time to make reading the payload easier
+        sorted_spans_to_report = sorted(
+            spans_to_report, key=lambda span: span.start_time
+        )
+
+        return [
+            {
+                **span.to_dict(),
+                CLIMetricsSpan.TRIMMED_KEY: not span.children <= spans_to_report,
+            }
+            for span in sorted_spans_to_report
+        ]

snowflake/cli/api/output/types.py

@@ -37,7 +37,7 @@ class ObjectResult(CommandResult):
 
 
 class CollectionResult(CommandResult):
-    def __init__(self, elements: t.Iterable[t.Dict]):
+    def __init__(self, elements: t.Iterable[t.Dict] | t.Generator[t.Dict, None, None]):
         self._elements = elements
 
     @property