Skip to content

Commit

Permalink
fix(cli): fix delete urn cli bug + stricter type annotations (datahub…
Browse files Browse the repository at this point in the history
  • Loading branch information
hsheth2 authored and cccs-Dustin committed Feb 1, 2023
1 parent 591a79e commit 344dd66
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 9 deletions.
1 change: 0 additions & 1 deletion metadata-ingestion/src/datahub/cli/delete_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -203,7 +203,6 @@ def delete(
aspect_name=aspect_name,
soft=soft,
dry_run=dry_run,
entity_type=entity_type,
start_time=start_time,
end_time=end_time,
cached_session_host=(session, host),
Expand Down
22 changes: 14 additions & 8 deletions metadata-ingestion/src/datahub/telemetry/telemetry.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
from typing import Any, Callable, Dict, Optional, TypeVar

from mixpanel import Consumer, Mixpanel
from typing_extensions import ParamSpec

import datahub as datahub_package
from datahub.cli.cli_utils import DATAHUB_ROOT_FOLDER, get_boolean_env_variable
Expand Down Expand Up @@ -222,20 +223,23 @@ def init_tracking(self) -> None:
def ping(
self,
event_name: str,
properties: Dict[str, Any] = {},
properties: Optional[Dict[str, Any]] = None,
server: Optional[DataHubGraph] = None,
) -> None:
"""
Send a single telemetry event.
Args:
event_name (str): name of the event to send.
properties (Optional[Dict[str, Any]]): metadata for the event
event_name: name of the event to send.
properties: metadata for the event
"""

if not self.enabled or self.mp is None:
return

if properties is None:
properties = {}

# send event
try:
logger.debug("Sending Telemetry")
Expand Down Expand Up @@ -266,10 +270,8 @@ def _server_props(self, server: Optional[DataHubGraph]) -> Dict[str, str]:

telemetry_instance = Telemetry()

T = TypeVar("T")


def suppress_telemetry() -> Any:
def suppress_telemetry() -> None:
"""disables telemetry for this invocation, doesn't affect persistent client settings"""
if telemetry_instance.enabled:
logger.debug("Disabling telemetry locally due to server config")
Expand All @@ -283,9 +285,13 @@ def get_full_class_name(obj):
return f"{module}.{obj.__class__.__name__}"


def with_telemetry(func: Callable[..., T]) -> Callable[..., T]:
_T = TypeVar("_T")
_P = ParamSpec("_P")


def with_telemetry(func: Callable[_P, _T]) -> Callable[_P, _T]:
@wraps(func)
def wrapper(*args: Any, **kwargs: Any) -> Any:
def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _T:
function = f"{func.__module__}.{func.__name__}"

telemetry_instance.init_tracking()
Expand Down

0 comments on commit 344dd66

Please sign in to comment.