@datarecce/ui 0.1.30 → 0.1.32

package/recce-source/recce/event/__init__.py

@@ -0,0 +1,304 @@
+import hashlib
+import os
+import re
+import sys
+import threading
+import uuid
+from datetime import datetime, timezone
+from hashlib import sha256
+from typing import Dict
+
+import sentry_sdk
+
+from recce import get_runner, get_version, is_ci_env, is_recce_cloud_instance
+from recce import yaml as pyml
+from recce.event.collector import Collector
+from recce.git import current_branch, hosting_repo
+from recce.github import (
+    get_github_codespace_available_at,
+    get_github_codespace_info,
+    get_github_codespace_name,
+    is_github_codespace,
+)
+
+USER_HOME = os.path.expanduser("~")
+RECCE_USER_HOME = os.path.join(USER_HOME, ".recce")
+RECCE_USER_PROFILE = os.path.join(RECCE_USER_HOME, "profile.yml")
+RECCE_USER_EVENT_PATH = os.path.join(RECCE_USER_HOME, ".unsend_events.json")
+
+__version__ = get_version()
+_collector = Collector()
+user_profile_lock = threading.Lock()
+
+
+def init():
+    api_key = _get_api_key()
+    user_profile = load_user_profile()
+
+    # Amplitude init
+    _collector.set_api_key(api_key)
+    _collector.set_user_id(user_profile.get("user_id"))
+    _collector.set_unsend_events_file(RECCE_USER_EVENT_PATH)
+
+    # Sentry init
+    sentry_env = _get_sentry_env()
+    sentry_dns = _get_sentry_dns()
+    release_version = __version__ if sentry_env != "development" else None
+    sentry_sdk.init(
+        dsn=sentry_dns,
+        environment=sentry_env,
+        release=release_version,
+        # Set traces_sample_rate to 1.0 to capture 100%
+        # of transactions for performance monitoring.
+        # We recommend adjusting this value in production.
+        traces_sample_rate=1.0,
+    )
+    sentry_sdk.set_tag("recce.version", __version__)
+    sentry_sdk.set_tag("platform", sys.platform)
+    sentry_sdk.set_tag("is_ci_env", is_ci_env())
+    sentry_sdk.set_tag("is_github_codespace", is_github_codespace())
+    sentry_sdk.set_tag("is_recce_cloud_instance", is_recce_cloud_instance())
+    sentry_sdk.set_tag("system_timezone", get_system_timezone())
+
+
+def get_user_id():
+    return load_user_profile().get("user_id")
+
+
+def get_recce_api_token():
+    return load_user_profile().get("api_token")
+
+
+def update_recce_api_token(token):
+    return update_user_profile({"api_token": token})
+
+
+def is_anonymous_tracking():
+    return load_user_profile().get("anonymous_tracking", False)
+
+
+def _get_sentry_dns():
+    dns_file = os.path.normpath(os.path.join(os.path.dirname(__file__), "SENTRY_DNS"))
+    with open(dns_file, encoding="utf-8") as f:
+        dns = f.read().strip()
+        return dns
+
+
+def _get_sentry_env():
+    if ".dev" in __version__:
+        return "development"
+    elif re.match(r"^\d+\.\d+\.\d+\.\d{8}[a|b|rc]?.*$", __version__):
+        return "nightly"
+    elif "a" in __version__:
+        return "alpha"
+    elif "b" in __version__:
+        return "beta"
+    elif "rc" in __version__:
+        return "release-candidate"
+    return "production"
+
+
+def _get_api_key():
+    if os.getenv("RECCE_EVENT_API_KEY"):
+        # For local testing purpose
+        return os.getenv("RECCE_EVENT_API_KEY")
+
+    config_file = os.path.abspath(os.path.join(os.path.dirname(__file__), "CONFIG"))
+    try:
+        with open(config_file, encoding="utf-8") as fh:
+            config = pyml.load(fh)
+            return config.get("event_api_key")
+    except Exception:
+        return None
+
+
+def _generate_user_profile():
+    try:
+        os.makedirs(RECCE_USER_HOME, exist_ok=True)
+    except Exception:
+        # TODO: should show warning message but not raise exception
+        print("Please disable command tracking to continue.")
+        exit(1)
+    if is_github_codespace() is True:
+        salted_name = f"codespace-{get_github_codespace_name()}"
+        user_id = hashlib.sha256(salted_name.encode()).hexdigest()
+    else:
+        user_id = uuid.uuid4().hex
+    with open(RECCE_USER_PROFILE, "w+", encoding="utf-8") as f:
+        pyml.dump({"user_id": user_id, "anonymous_tracking": True}, f)
+    return dict(user_id=user_id, anonymous_tracking=True)
+
+
+def load_user_profile():
+    with user_profile_lock:
+        if not os.path.exists(RECCE_USER_PROFILE):
+            user_profile = _generate_user_profile()
+        else:
+            with open(RECCE_USER_PROFILE, "r", encoding="utf-8") as f:
+                user_profile = pyml.load(f)
+                if user_profile is None or user_profile.get("user_id") is None:
+                    user_profile = _generate_user_profile()
+
+        return user_profile
+
+
+def update_user_profile(update_values):
+    original = load_user_profile()
+    original.update(update_values)
+    with open(RECCE_USER_PROFILE, "w+", encoding="utf-8") as f:
+        pyml.dump(original, f)
+    return original
+
+
+def flush_events(command=None):
+    _collector.send_events()
+
+
+def should_log_event():
+    with open(RECCE_USER_PROFILE, "r", encoding="utf-8") as f:
+        user_profile = pyml.load(f)
+    # TODO: default anonymous_tracking to false if field is not present
+    tracking = user_profile.get("anonymous_tracking", False)
+    tracking = tracking and isinstance(tracking, bool)
+    if not tracking:
+        return False
+
+    if not _collector.is_ready():
+        return False
+
+    return True
+
+
+def log_event(prop, event_type, **kwargs):
+    if should_log_event() is False:
+        return
+
+    repo = hosting_repo()
+    if repo is not None:
+        prop["repository"] = sha256(repo.encode()).hexdigest()
+
+    branch = current_branch()
+    if branch is not None:
+        prop["branch"] = sha256(branch.encode()).hexdigest()
+
+    runner = get_runner()
+    if runner is not None:
+        prop["runner_type"] = runner
+
+        if runner == "github codespaces":
+            prop["codespaces_name"] = get_github_codespace_name()
+
+    payload = dict(
+        **prop,
+    )
+
+    _collector.log_event(payload, event_type)
+
+
+def log_api_event(endpoint_name, prop):
+    prop = dict(
+        **prop,
+        endpoint_name=endpoint_name,
+    )
+    log_event(prop, "api_event")
+    _collector.schedule_flush()
+
+
+def log_load_state(command="server", single_env=False):
+    from recce.models import CheckDAO
+
+    checks = 0
+    preset_checks = 0
+
+    for check in CheckDAO().list():
+        checks += 1
+        if check.is_preset:
+            preset_checks += 1
+
+    prop = dict(
+        command=command,
+        checks=checks,
+        preset_checks=preset_checks,
+    )
+
+    if command == "server":
+        prop["single_env"] = single_env
+
+    log_event(prop, "load_state")
+    if command == "server":
+        _collector.schedule_flush()
+
+
+def log_codespaces_events(command):
+    # Only log when the recce is running in GitHub Codespaces
+    codespace = get_github_codespace_info()
+    if codespace is None:
+        return
+
+    user_prop = dict(
+        location=codespace.get("location"),
+        is_prebuild=codespace.get("prebuild", False),
+    )
+
+    prop = dict(
+        machine=codespace.get("machine", {}).get("display_name"),
+        codespaces_name=get_github_codespace_name(),
+    )
+
+    # Codespace created event, send once
+    codespace_created_at = load_user_profile().get("codespace_created_at")
+    if codespace_created_at is None:
+        created_at = datetime.fromisoformat(codespace.get("created_at"))
+        prop["state"] = "created"
+        _collector.log_event(prop, "codespace_instance", event_triggered_at=created_at, user_properties=user_prop)
+        update_user_profile({"codespace_created_at": codespace.get("created_at")})
+
+    # Codespace available event, send multiple times as start/stop it
+    available_at = get_github_codespace_available_at(codespace)
+    if available_at and available_at.isoformat() != load_user_profile().get("codespace_available_at"):
+        prop["state"] = "available"
+        _collector.log_event(prop, "codespace_instance", event_triggered_at=available_at, user_properties=user_prop)
+        update_user_profile({"codespace_available_at": available_at.isoformat()})
+
+    # Codespace instance event should be flushed immediately
+    _collector.send_events()
+
+
+def log_single_env_event():
+    prop = dict(
+        action="launch_server",
+    )
+    log_event(prop, "[Experiment] single_environment")
+    _collector.schedule_flush()
+
+
+def log_performance(feature_name: str, metrics: Dict):
+    prop = metrics
+    log_event(prop, f"[Performance] {feature_name}")
+    _collector.schedule_flush()
+
+
+def log_connected_to_cloud():
+    log_event({"action": "connected_to_cloud"}, "Connect OSS to Cloud")
+    _collector.schedule_flush()
+
+
+def capture_exception(e):
+    user_id = load_user_profile().get("user_id")
+    if is_ci_env() is True:
+        user_id = f"{user_id}_CI"
+
+    sentry_sdk.set_tag("user_id", user_id)
+    sentry_sdk.capture_exception(e)
+
+
+def flush_exceptions():
+    sentry_sdk.flush()
+
+
+def set_exception_tag(key, value):
+    sentry_sdk.set_tag(key, value)
+
+
+def get_system_timezone():
+    return datetime.now(timezone.utc).astimezone().tzinfo

package/recce-source/recce/event/collector.py

@@ -0,0 +1,184 @@
+import json
+import os
+import platform
+import sys
+import time
+from contextlib import contextmanager
+from datetime import datetime, timezone
+from json import JSONDecodeError
+
+import portalocker
+import requests
+
+from recce import __version__, is_ci_env, is_recce_cloud_instance
+from recce.github import is_github_codespace
+
+
+class Collector:
+    def __init__(self):
+        self._api_endpoint = "https://api.amplitude.com/2/httpapi"
+        self._api_key = None
+        self._user_id = None
+
+        self._unsend_events_file = None
+        self._delete_threshold = 1000
+        self._upload_threshold = 10
+        self._is_ci: bool = is_ci_env()
+        self._is_github_codespace: bool = is_github_codespace()
+        self._is_recce_cloud_instance: bool = is_recce_cloud_instance()
+        self._flush_timer = None
+
+    def schedule_flush(self):
+        def _callback():
+            self._flush_timer = None
+            self.send_events()
+
+        # send async thread
+        import threading
+
+        if self._flush_timer:
+            try:
+                self._flush_timer.cancel()
+            except Exception:
+                # do nothing
+                pass
+        self._flush_timer = threading.Timer(1, _callback)
+        self._flush_timer.start()
+
+    def is_ready(self):
+        if self._api_key is None or self._user_id is None:
+            return False
+        return True
+
+    def set_api_key(self, api_key):
+        self._api_key = api_key
+
+    def set_user_id(self, user_id):
+        self._user_id = user_id
+
+    def set_unsend_events_file(self, unsend_events_file):
+        self._unsend_events_file = unsend_events_file
+        self._check_required_files()
+
+    def _log_event(
+        self,
+        user_id,
+        event_type,
+        created_at,
+        user_properties,
+        event_properties,
+    ):
+        event = dict(
+            user_id=user_id,
+            event_type=event_type,
+            ip="$remote",
+            time=int(time.mktime(created_at.timetuple())),
+            user_properties=user_properties,
+            event_properties=event_properties,
+            platform=sys.platform,
+            os_version=platform.platform(),
+            app_version=__version__,
+        )
+        self._store_to_file(event)
+        if self._is_full():
+            self.send_events()
+        self._cleanup_unsend_events()
+
+    def _get_user_id(self):
+        # when the recce is running in automation use cases
+        # replace the user id with project_id to avoid so many unique user id
+        user_id = self._user_id
+        if self._is_ci is True:
+            user_id = f"{self._user_id}_CI"
+        return user_id
+
+    def log_event(self, event_prop, event_type, event_triggered_at: datetime = None, user_properties=None):
+        # Use local timezone
+        if event_triggered_at is None:
+            created_at = datetime.now()
+        else:
+            # Convert to UTC timezone
+            created_at = event_triggered_at.astimezone(timezone.utc)
+        python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
+
+        # when the recce is running in automation use cases
+        # replace the user id with project_id to avoid so many unique user id
+        user_id = self._get_user_id()
+
+        default_user_properties = dict(
+            version=__version__,
+            python_version=python_version,
+            is_ci=self._is_ci,
+            is_github_codespace=self._is_github_codespace,
+            is_recce_cloud_instance=self._is_recce_cloud_instance,
+        )
+
+        if user_properties is not None:
+            default_user_properties.update(user_properties)
+
+        self._log_event(
+            user_id=user_id,
+            event_type=event_type,
+            created_at=created_at,
+            user_properties=default_user_properties,
+            event_properties=event_prop,
+        )
+
+    def _check_required_files(self):
+        user_home = os.path.dirname(self._unsend_events_file)
+        if not os.path.exists(user_home):
+            os.makedirs(user_home, exist_ok=True)
+        if not os.path.exists(self._unsend_events_file):
+            with portalocker.Lock(self._unsend_events_file, "w+", timeout=5) as f:
+                f.write(json.dumps({"unsend_events": []}))
+
+    def _is_full(self):
+        with portalocker.Lock(self._unsend_events_file, "r+", timeout=5) as f:
+            o = json.loads(f.read())
+            return len(o.get("unsend_events", [])) >= self._upload_threshold
+
+    @contextmanager
+    def load_json(self):
+        with portalocker.Lock(self._unsend_events_file, "r+", timeout=5) as f:
+            o = None
+            try:
+                o = json.loads(f.read())
+                yield o
+            except JSONDecodeError:
+                o = dict(unsend_events=[])
+                yield o
+            finally:
+                f.seek(0)
+                f.truncate()
+                if o is not None:
+                    f.write(json.dumps(o))
+
+    def send_events(self):
+        with self.load_json() as o:
+            payload = dict(
+                api_key=self._api_key,
+                events=o["unsend_events"],
+            )
+            o["unsend_events"] = []
+        try:
+            requests.post(self._api_endpoint, json=payload)
+        except Exception:
+            # TODO: handle exception when sending events
+            pass
+
+    def _store_to_file(self, event):
+        with self.load_json() as o:
+            events = o.get("unsend_events", None)
+            if events is None:
+                o["unsend_events"] = []
+
+            o["unsend_events"].append(event)
+
+    def _cleanup_unsend_events(self):
+        with self.load_json() as o:
+            events = o.get("unsend_events", None)
+            if events is None:
+                o["unsend_events"] = []
+
+            while len(o["unsend_events"]) > self._delete_threshold:
+                o["unsend_events"].pop(0)

package/recce-source/recce/event/track.py

@@ -0,0 +1,158 @@
+import logging
+import os
+import sys
+import time
+import traceback
+import typing as t
+from hashlib import sha256
+
+from click import Context
+from click.core import Command
+from rich.console import Console
+from rich.markup import escape
+
+from recce import event, get_runner
+from recce.core import load_context
+from recce.exceptions import RecceException
+from recce.git import current_branch, hosting_repo
+
+console = Console()
+
+_enable_traceback: bool = os.environ.get("RECCE_PRINT_TRACEBACK") == "1"
+logger = logging.getLogger("uvicorn")
+
+
+class TrackCommand(Command):
+    def __init__(
+        self,
+        name: t.Optional[str],
+        context_settings: t.Optional[t.Dict[str, t.Any]] = None,
+        callback: t.Optional[t.Callable[..., t.Any]] = None,
+        params: t.Any = None,
+        help: t.Optional[str] = None,
+        epilog: t.Optional[str] = None,
+        short_help: t.Optional[str] = None,
+        options_metavar: t.Optional[str] = "[OPTIONS]",
+        add_help_option: bool = True,
+        no_args_is_help: bool = False,
+        hidden: bool = False,
+        deprecated: bool = False,
+        beta: bool = False,
+    ) -> None:
+        super(TrackCommand, self).__init__(
+            name,
+            context_settings,
+            callback,
+            params,
+            help,
+            epilog,
+            short_help,
+            options_metavar,
+            add_help_option,
+            no_args_is_help,
+            hidden,
+            deprecated,
+        )
+
+    def _show_error_message(self, msg, params):
+        if params.get("debug"):
+            console.print_exception(show_locals=True)
+        else:
+            print(traceback.format_exc())
+            # console.print('[bold red]Error:[/bold red] ', end='')
+            # console.out(msg, highlight=False)
+
+    def _show_hint_message(self, hint):
+        console.print(f"[bold yellow]Hint[/bold yellow]:\n  {escape(hint)}")
+
+    def invoke(self, ctx: Context) -> t.Any:
+        status = False
+        start_time = time.time()
+        reason = "error"
+        event.set_exception_tag("command", ctx.command.name)
+        event.log_codespaces_events(ctx.command.name)
+
+        try:
+            ret = super(TrackCommand, self).invoke(ctx)
+            if ret is None or ret == 0:
+                status = True
+                reason = "ok"
+            else:
+                reason = "error"
+                sys.exit(ret)
+            return ret
+        except RecceException as e:
+            logger.debug(traceback.format_exc())
+            console.log("[Error] " + str(e))
+            reason = "error"
+            sys.exit(1)
+        except SystemExit as e:
+            reason = "error"
+            raise e
+        except KeyboardInterrupt as e:
+            reason = "aborted"
+            raise e
+        except Exception as e:
+            self._show_error_message(str(e), ctx.params)
+            event.capture_exception(e)
+            reason = "fatal"
+            event.flush_exceptions()
+            sys.exit(1)
+        finally:
+            end_time = time.time()
+            runner = get_runner()
+            repo = hosting_repo()
+            branch = current_branch()
+            command = ctx.command.name
+            duration = end_time - start_time
+            target_path = ctx.params.get("target_path", None)
+            target_base_path = ctx.params.get("target_base_path", None)
+            props = dict(
+                command=command,
+                status=status,
+                reason=reason,
+                duration=duration,
+                cloud=ctx.params.get("cloud", False),
+                review=ctx.params.get("review", False),
+                debug=ctx.params.get("debug", False),
+            )
+
+            if runner is not None:
+                props["runner_type"] = runner
+
+            if repo is not None:
+                props["repository"] = sha256(repo.encode()).hexdigest()
+
+            if branch is not None:
+                props["branch"] = sha256(branch.encode()).hexdigest()
+
+            if target_path is not None:
+                props["target_path"] = sha256(target_path.encode()).hexdigest()
+
+            if target_base_path is not None:
+                props["target_base_path"] = sha256(target_base_path.encode()).hexdigest()
+
+            try:
+                recce_context = load_context()
+            except Exception:
+                # it's not a ready-for-use project
+                recce_context = None
+
+            if recce_context is not None:
+                if recce_context.adapter_type == "dbt":
+                    props["adapter_type"] = "DBT"
+                    # Add dbt warehouse type only for dbt adapter
+                    try:
+                        from recce.adapter.dbt_adapter import DbtAdapter
+
+                        dbt_adapter: DbtAdapter = recce_context.adapter
+                        warehouse_type = dbt_adapter.adapter.type()
+                        props["warehouse_type"] = warehouse_type
+                    except Exception:
+                        # If we can't get the warehouse type, skip it
+                        pass
+                elif recce_context.adapter_type == "sqlmesh":
+                    props["adapter_type"] = "SQLMesh"
+
+            event.log_event(props, "command", params=ctx.params)
+            event.flush_events()

package/recce-source/recce/exceptions.py

@@ -0,0 +1,21 @@
+class RecceException(Exception):
+    def __init__(self, message, is_raise=True):
+        self.message = message
+        self.is_raise = is_raise
+        super().__init__(self.message)
+
+
+class RecceCancelException(RecceException):
+    def __init__(self):
+        super().__init__("Cancelled", False)
+
+
+class RecceConfigException(RecceException):
+    def __init__(self, message, cause=None):
+        super().__init__(message)
+        self.cause = cause
+
+    def __str__(self):
+        if self.cause:
+            return f"{super().__str__()}\n{str(self.cause)}"
+        return super().__str__()