dvt-core 1.11.0b4__py3-none-any.whl

dvt/tracking.py

@@ -0,0 +1,529 @@
+import os
+import platform
+import traceback
+import uuid
+from contextlib import contextmanager
+from datetime import datetime
+from typing import Optional
+
+import pytz
+import requests
+from dvt import version as dbt_version
+from dvt.clients.yaml_helper import safe_load, yaml  # noqa:F401
+from dvt.events.types import (
+    DisableTracking,
+    FlushEvents,
+    FlushEventsFailure,
+    MainEncounteredError,
+    SendEventFailure,
+    SendingEvent,
+    TrackingInitializeFailure,
+)
+from packaging.version import Version
+from snowplow_tracker import Emitter, SelfDescribingJson, Subject, Tracker
+from snowplow_tracker import __version__ as snowplow_version  # type: ignore
+from snowplow_tracker import logger as sp_logger
+from snowplow_tracker.events import StructuredEvent
+
+from dbt.adapters.exceptions import FailedToConnectError
+from dbt_common.events.base_types import EventMsg
+from dbt_common.events.functions import fire_event, get_invocation_id, msg_to_dict
+from dbt_common.exceptions import NotImplementedError
+
+sp_logger.setLevel(100)
+
+COLLECTOR_URL = "fishtownanalytics.sinter-collect.com"
+COLLECTOR_PROTOCOL = "https"
+DBT_INVOCATION_ENV = "DBT_INVOCATION_ENV"
+
+ADAPTER_INFO_SPEC = "iglu:com.dbt/adapter_info/jsonschema/1-0-1"
+DEPRECATION_WARN_SPEC = "iglu:com.dbt/deprecation_warn/jsonschema/1-0-0"
+BEHAVIOR_CHANGE_WARN_SPEC = "iglu:com.dbt/behavior_change_warn/jsonschema/1-0-0"
+EXPERIMENTAL_PARSER = "iglu:com.dbt/experimental_parser/jsonschema/1-0-0"
+INVOCATION_ENV_SPEC = "iglu:com.dbt/invocation_env/jsonschema/1-0-0"
+INVOCATION_SPEC = "iglu:com.dbt/invocation/jsonschema/1-0-2"
+LOAD_ALL_TIMING_SPEC = "iglu:com.dbt/load_all_timing/jsonschema/1-0-3"
+PACKAGE_INSTALL_SPEC = "iglu:com.dbt/package_install/jsonschema/1-0-0"
+PARTIAL_PARSER = "iglu:com.dbt/partial_parser/jsonschema/1-0-1"
+PLATFORM_SPEC = "iglu:com.dbt/platform/jsonschema/1-0-0"
+PROJECT_ID_SPEC = "iglu:com.dbt/project_id/jsonschema/1-0-1"
+RESOURCE_COUNTS = "iglu:com.dbt/resource_counts/jsonschema/1-0-1"
+RPC_REQUEST_SPEC = "iglu:com.dbt/rpc_request/jsonschema/1-0-1"
+RUNNABLE_TIMING = "iglu:com.dbt/runnable/jsonschema/1-0-0"
+RUN_MODEL_SPEC = "iglu:com.dbt/run_model/jsonschema/1-1-0"
+PLUGIN_GET_NODES = "iglu:com.dbt/plugin_get_nodes/jsonschema/1-0-0"
+ARTIFACT_UPLOAD = "iglu:com.dbt/artifact_upload/jsonschema/1-0-0"
+
+SNOWPLOW_TRACKER_VERSION = Version(snowplow_version)
+
+# workaround in case real snowplow tracker is in the env
+# the argument was renamed in https://github.com/snowplow/snowplow-python-tracker/commit/39fd50a3aff98a5efdd5c5c7fb5518fe4761305b
+INIT_KW_ARGS = (
+    {"buffer_size": 30} if SNOWPLOW_TRACKER_VERSION < Version("0.13.0") else {"batch_size": 30}
+)
+
+
+class TimeoutEmitter(Emitter):
+    def __init__(self) -> None:
+        super().__init__(
+            COLLECTOR_URL,
+            protocol=COLLECTOR_PROTOCOL,
+            on_failure=self.handle_failure,
+            method="post",
+            # don't set this.
+            byte_limit=None,
+            **INIT_KW_ARGS,
+        )
+
+    @staticmethod
+    def handle_failure(num_ok, unsent):
+        # num_ok will always be 0, unsent will always be 1 entry long, because
+        # the buffer is length 1, so not much to talk about
+        fire_event(DisableTracking())
+        disable_tracking()
+
+    def _log_request(self, request, payload):
+        sp_logger.info(f"Sending {request} request to {self.endpoint}...")
+        sp_logger.debug(f"Payload: {payload}")
+
+    def _log_result(self, request, status_code):
+        msg = f"{request} request finished with status code: {status_code}"
+        if self.is_good_status_code(status_code):
+            sp_logger.info(msg)
+        else:
+            sp_logger.warning(msg)
+
+    def http_post(self, payload):
+        self._log_request("POST", payload)
+
+        r = requests.post(
+            self.endpoint,
+            data=payload,
+            headers={"content-type": "application/json; charset=utf-8"},
+            timeout=5.0,
+        )
+
+        self._log_result("GET", r.status_code)
+        return r
+
+    def http_get(self, payload):
+        self._log_request("GET", payload)
+
+        r = requests.get(self.endpoint, params=payload, timeout=5.0)
+
+        self._log_result("GET", r.status_code)
+        return r
+
+
+emitter = TimeoutEmitter()
+tracker = Tracker(
+    emitters=emitter,
+    namespace="cf",
+    app_id="dbt",
+)
+
+
+class User:
+    def __init__(self, cookie_dir) -> None:
+        self.do_not_track = True
+        self.cookie_dir = cookie_dir
+
+        self.id = None
+        self.invocation_id = get_invocation_id()
+        self.run_started_at = datetime.now(tz=pytz.utc)
+
+    def state(self):
+        return "do not track" if self.do_not_track else "tracking"
+
+    @property
+    def cookie_path(self):
+        return os.path.join(self.cookie_dir, ".user.yml")
+
+    def initialize(self):
+        self.do_not_track = False
+
+        cookie = self.get_cookie()
+        self.id = cookie.get("id")
+
+        subject = Subject()
+        subject.set_user_id(self.id)
+        tracker.set_subject(subject)
+
+    def disable_tracking(self):
+        self.do_not_track = True
+        self.id = None
+        self.cookie_dir = None
+        tracker.set_subject(None)
+
+    def set_cookie(self):
+        # If the user points dbt to a profile directory which exists AND
+        # contains a profiles.yml file, then we can set a cookie. If the
+        # specified folder does not exist, or if there is not a profiles.yml
+        # file in this folder, then an inconsistent cookie can be used. This
+        # will change in every dbt invocation until the user points to a
+        # profile dir file which contains a valid profiles.yml file.
+        #
+        # See: https://github.com/dbt-labs/dbt-core/issues/1645
+
+        user = {"id": str(uuid.uuid4())}
+
+        cookie_path = os.path.abspath(self.cookie_dir)
+        profiles_file = os.path.join(cookie_path, "profiles.yml")
+        if os.path.exists(cookie_path) and os.path.exists(profiles_file):
+            with open(self.cookie_path, "w") as fh:
+                yaml.dump(user, fh)
+
+        return user
+
+    def get_cookie(self):
+        if not os.path.isfile(self.cookie_path):
+            user = self.set_cookie()
+        else:
+            with open(self.cookie_path, "r") as fh:
+                try:
+                    user = safe_load(fh)
+                    if user is None:
+                        user = self.set_cookie()
+                except yaml.reader.ReaderError:
+                    user = self.set_cookie()
+        return user
+
+
+active_user: Optional[User] = None
+
+
+def get_platform_context():
+    data = {
+        "platform": platform.platform(),
+        "python": platform.python_version(),
+        "python_version": platform.python_implementation(),
+    }
+
+    return SelfDescribingJson(PLATFORM_SPEC, data)
+
+
+def get_dbt_env_context():
+    default = "manual"
+
+    dbt_invocation_env = os.getenv(DBT_INVOCATION_ENV, default)
+    if dbt_invocation_env == "":
+        dbt_invocation_env = default
+
+    data = {
+        "environment": dbt_invocation_env,
+    }
+
+    return SelfDescribingJson(INVOCATION_ENV_SPEC, data)
+
+
+def track(user, *args, **kwargs):
+    if user.do_not_track:
+        return
+
+    fire_event(SendingEvent(kwargs=str(kwargs)))
+    try:
+        tracker.track(StructuredEvent(*args, **kwargs))
+    except Exception:
+        fire_event(SendEventFailure())
+
+
+def track_project_id(options):
+    assert active_user is not None, "Cannot track project_id when active user is None"
+    context = [SelfDescribingJson(PROJECT_ID_SPEC, options)]
+
+    track(
+        active_user,
+        category="dbt",
+        action="project_id",
+        label=get_invocation_id(),
+        context=context,
+    )
+
+
+def track_adapter_info(options):
+    assert active_user is not None, "Cannot track adapter_info when active user is None"
+    context = [SelfDescribingJson(ADAPTER_INFO_SPEC, options)]
+
+    track(
+        active_user,
+        category="dbt",
+        action="adapter_info",
+        label=get_invocation_id(),
+        context=context,
+    )
+
+
+def track_invocation_start(invocation_context):
+    data = {"progress": "start", "result_type": None, "result": None}
+    data.update(invocation_context)
+    context = [
+        SelfDescribingJson(INVOCATION_SPEC, data),
+        get_platform_context(),
+        get_dbt_env_context(),
+    ]
+
+    track(active_user, category="dbt", action="invocation", label="start", context=context)
+
+
+def track_project_load(options):
+    context = [SelfDescribingJson(LOAD_ALL_TIMING_SPEC, options)]
+    assert active_user is not None, "Cannot track project loading time when active user is None"
+
+    track(
+        active_user,
+        category="dbt",
+        action="load_project",
+        label=get_invocation_id(),
+        context=context,
+    )
+
+
+def track_resource_counts(resource_counts):
+    context = [SelfDescribingJson(RESOURCE_COUNTS, resource_counts)]
+    assert active_user is not None, "Cannot track resource counts when active user is None"
+
+    track(
+        active_user,
+        category="dbt",
+        action="resource_counts",
+        label=get_invocation_id(),
+        context=context,
+    )
+
+
+def track_model_run(options):
+    context = [SelfDescribingJson(RUN_MODEL_SPEC, options)]
+    assert active_user is not None, "Cannot track model runs when active user is None"
+
+    track(
+        active_user, category="dbt", action="run_model", label=get_invocation_id(), context=context
+    )
+
+
+def track_rpc_request(options):
+    context = [SelfDescribingJson(RPC_REQUEST_SPEC, options)]
+    assert active_user is not None, "Cannot track rpc requests when active user is None"
+
+    track(
+        active_user,
+        category="dbt",
+        action="rpc_request",
+        label=get_invocation_id(),
+        context=context,
+    )
+
+
+def get_base_invocation_context():
+    assert (
+        active_user is not None
+    ), "initialize active user before calling get_base_invocation_context"
+    return {
+        "project_id": None,
+        "user_id": active_user.id,
+        "invocation_id": active_user.invocation_id,
+        "command": None,
+        "options": None,
+        "version": str(dbt_version.installed),
+        "run_type": "regular",
+        "adapter_type": None,
+        "adapter_unique_id": None,
+    }
+
+
+def track_package_install(command_name: str, project_hashed_name: Optional[str], options):
+    assert active_user is not None, "Cannot track package installs when active user is None"
+
+    invocation_data = get_base_invocation_context()
+
+    invocation_data.update({"project_id": project_hashed_name, "command": command_name})
+
+    context = [
+        SelfDescribingJson(INVOCATION_SPEC, invocation_data),
+        SelfDescribingJson(PACKAGE_INSTALL_SPEC, options),
+    ]
+
+    track(
+        active_user,
+        category="dbt",
+        action="package",
+        label=get_invocation_id(),
+        property_="install",
+        context=context,
+    )
+
+
+def track_deprecation_warn(options):
+
+    assert active_user is not None, "Cannot track deprecation warnings when active user is None"
+
+    context = [SelfDescribingJson(DEPRECATION_WARN_SPEC, options)]
+
+    track(
+        active_user,
+        category="dbt",
+        action="deprecation",
+        label=get_invocation_id(),
+        property_="warn",
+        context=context,
+    )
+
+
+def track_behavior_change_warn(msg: EventMsg) -> None:
+    if msg.info.name != "BehaviorChangeEvent" or active_user is None:
+        return
+
+    context = [SelfDescribingJson(BEHAVIOR_CHANGE_WARN_SPEC, msg_to_dict(msg))]
+    track(
+        active_user,
+        category="dbt",
+        action=msg.info.name,
+        label=get_invocation_id(),
+        context=context,
+    )
+
+
+def track_invocation_end(invocation_context, result_type=None):
+    data = {"progress": "end", "result_type": result_type, "result": None}
+    data.update(invocation_context)
+    context = [
+        SelfDescribingJson(INVOCATION_SPEC, data),
+        get_platform_context(),
+        get_dbt_env_context(),
+    ]
+
+    assert active_user is not None, "Cannot track invocation end when active user is None"
+
+    track(active_user, category="dbt", action="invocation", label="end", context=context)
+
+
+def track_invalid_invocation(args=None, result_type=None):
+    assert active_user is not None, "Cannot track invalid invocations when active user is None"
+    invocation_context = get_base_invocation_context()
+    invocation_context.update({"command": args.which})
+    data = {"progress": "invalid", "result_type": result_type, "result": None}
+    data.update(invocation_context)
+    context = [
+        SelfDescribingJson(INVOCATION_SPEC, data),
+        get_platform_context(),
+        get_dbt_env_context(),
+    ]
+    track(active_user, category="dbt", action="invocation", label="invalid", context=context)
+
+
+def track_experimental_parser_sample(options):
+    context = [SelfDescribingJson(EXPERIMENTAL_PARSER, options)]
+    assert (
+        active_user is not None
+    ), "Cannot track experimental parser info when active user is None"
+
+    track(
+        active_user,
+        category="dbt",
+        action="experimental_parser",
+        label=get_invocation_id(),
+        context=context,
+    )
+
+
+def track_partial_parser(options):
+    context = [SelfDescribingJson(PARTIAL_PARSER, options)]
+    assert active_user is not None, "Cannot track partial parser info when active user is None"
+
+    track(
+        active_user,
+        category="dbt",
+        action="partial_parser",
+        label=get_invocation_id(),
+        context=context,
+    )
+
+
+def track_plugin_get_nodes(options):
+    context = [SelfDescribingJson(PLUGIN_GET_NODES, options)]
+    assert active_user is not None, "Cannot track plugin node info when active user is None"
+
+    track(
+        active_user,
+        category="dbt",
+        action="plugin_get_nodes",
+        label=get_invocation_id(),
+        context=context,
+    )
+
+
+def track_runnable_timing(options):
+    context = [SelfDescribingJson(RUNNABLE_TIMING, options)]
+    assert active_user is not None, "Cannot track runnable info when active user is None"
+    track(
+        active_user,
+        category="dbt",
+        action="runnable_timing",
+        label=get_invocation_id(),
+        context=context,
+    )
+
+
+def track_artifact_upload(options):
+    context = [SelfDescribingJson(ARTIFACT_UPLOAD, options)]
+    assert active_user is not None, "Cannot track artifact upload when active user is None"
+
+    track(
+        active_user,
+        category="dbt",
+        action="artifact_upload",
+        label=get_invocation_id(),
+        context=context,
+    )
+
+
+def flush():
+    fire_event(FlushEvents())
+    try:
+        tracker.flush()
+    except Exception:
+        fire_event(FlushEventsFailure())
+
+
+def disable_tracking():
+    global active_user
+    if active_user is not None:
+        active_user.disable_tracking()
+    else:
+        active_user = User(None)
+
+
+def do_not_track():
+    global active_user
+    active_user = User(None)
+
+
+def initialize_from_flags(send_anonymous_usage_stats, profiles_dir):
+    global active_user
+    if send_anonymous_usage_stats:
+        active_user = User(profiles_dir)
+        try:
+            active_user.initialize()
+        except Exception:
+            fire_event(TrackingInitializeFailure(exc_info=traceback.format_exc()))
+            active_user = User(None)
+    else:
+        active_user = User(None)
+
+
+@contextmanager
+def track_run(run_command=None):
+    invocation_context = get_base_invocation_context()
+    invocation_context["command"] = run_command
+
+    track_invocation_start(invocation_context)
+    try:
+        yield
+        track_invocation_end(invocation_context, result_type="ok")
+    except (NotImplementedError, FailedToConnectError) as e:
+        fire_event(MainEncounteredError(exc=str(e)))
+        track_invocation_end(invocation_context, result_type="error")
+    except Exception:
+        track_invocation_end(invocation_context, result_type="error")
+        raise
+    finally:
+        flush()

dvt/utils/__init__.py

@@ -0,0 +1,3 @@
+# re-export utils
+from .utils import *  # noqa: F403
+from .utils import _coerce_decimal  # noqa: F401 somewhere in the codebase we use this

dvt/utils/artifact_upload.py

@@ -0,0 +1,151 @@
+import time
+import uuid
+import zipfile
+
+import dvt.tracking
+import requests
+from dvt._pydantic_shim import BaseSettings  # type: ignore
+from dvt.config.runtime import UnsetProfile, load_project
+from dvt.constants import MANIFEST_FILE_NAME, RUN_RESULTS_FILE_NAME
+from dvt.events.types import ArtifactUploadSkipped, ArtifactUploadSuccess
+from dvt.exceptions import DbtProjectError
+
+from dbt_common.events.functions import fire_event
+from dbt_common.exceptions import DbtBaseException as DbtException
+
+MAX_RETRIES = 3
+
+EXECUTION_ARTIFACTS = [MANIFEST_FILE_NAME, RUN_RESULTS_FILE_NAME]
+
+PRODUCED_ARTIFACTS_PATHS: set[str] = set()
+
+
+# artifact paths calling this will be uploaded to dbt Cloud
+def add_artifact_produced(artifact_path: str):
+    PRODUCED_ARTIFACTS_PATHS.add(artifact_path)
+
+
+class ArtifactUploadConfig(BaseSettings):
+    tenant_hostname: str
+    DBT_CLOUD_TOKEN: str
+    DBT_CLOUD_ACCOUNT_ID: str
+    DBT_CLOUD_ENVIRONMENT_ID: str
+
+    def get_ingest_url(self):
+        return f"https://{self.tenant_hostname}/api/private/accounts/{self.DBT_CLOUD_ACCOUNT_ID}/environments/{self.DBT_CLOUD_ENVIRONMENT_ID}/ingests/"
+
+    def get_complete_url(self, ingest_id):
+        return f"{self.get_ingest_url()}{ingest_id}/"
+
+    def get_headers(self, invocation_id=None):
+        if invocation_id is None:
+            invocation_id = str(uuid.uuid4())
+        return {
+            "Accept": "application/json",
+            "X-Invocation-Id": invocation_id,
+            "Authorization": f"Token {self.DBT_CLOUD_TOKEN}",
+        }
+
+
+def _retry_with_backoff(operation_name, func, max_retries=MAX_RETRIES, retry_codes=None):
+    """Execute a function with exponential backoff retry logic.
+
+    Args:
+        operation_name: Name of the operation for error messages
+        func: Function to execute that returns (success, result)
+        max_retries: Maximum number of retry attempts
+
+    Returns:
+        The result from the function if successful
+
+    Raises:
+        DbtException: If all retry attempts fail
+    """
+    if retry_codes is None:
+        retry_codes = [500, 502, 503, 504]
+    retry_delay = 1
+    for attempt in range(max_retries + 1):
+        try:
+            success, result = func()
+            if success:
+                return result
+
+            if result.status_code not in retry_codes:
+                raise DbtException(f"Error {operation_name}: {result}")
+            if attempt == max_retries:  # Last attempt
+                raise DbtException(f"Error {operation_name}: {result}")
+        except requests.RequestException as e:
+            if attempt == max_retries:  # Last attempt
+                raise DbtException(f"Error {operation_name}: {str(e)}")
+
+        time.sleep(retry_delay)
+        retry_delay *= 2  # exponential backoff
+
+
+def upload_artifacts(project_dir, target_path, command):
+    # Check if there are artifacts to upload for this command
+    if not PRODUCED_ARTIFACTS_PATHS:
+        fire_event(ArtifactUploadSkipped(msg="No artifacts to upload for current command"))
+        return
+
+    # read configurations
+    try:
+        project = load_project(
+            project_dir, version_check=False, profile=UnsetProfile(), cli_vars=None
+        )
+        if not project.dbt_cloud or "tenant_hostname" not in project.dbt_cloud:
+            raise DbtProjectError("dbt_cloud.tenant_hostname not found in dbt_project.yml")
+        tenant_hostname = project.dbt_cloud["tenant_hostname"]
+        if not tenant_hostname:
+            raise DbtProjectError("dbt_cloud.tenant_hostname is empty in dbt_project.yml")
+    except Exception as e:
+        raise DbtProjectError(
+            f"Error reading dbt_cloud.tenant_hostname from dbt_project.yml: {str(e)}"
+        )
+
+    config = ArtifactUploadConfig(tenant_hostname=tenant_hostname)
+
+    if not target_path:
+        target_path = "target"
+
+    # Create zip file with artifacts
+    zip_file_name = "target.zip"
+    with zipfile.ZipFile(zip_file_name, "w") as z:
+        for artifact_path in PRODUCED_ARTIFACTS_PATHS:
+            z.write(artifact_path, artifact_path.split("/")[-1])
+
+    # Step 1: Create ingest request with retry
+    def create_ingest():
+        response = requests.post(url=config.get_ingest_url(), headers=config.get_headers())
+        return response.status_code == 200, response
+
+    response = _retry_with_backoff("creating ingest request", create_ingest)
+    response_data = response.json()
+    ingest_id = response_data["data"]["id"]
+    upload_url = response_data["data"]["upload_url"]
+
+    # Step 2: Upload the zip file to the provided URL with retry
+    with open(zip_file_name, "rb") as f:
+        file_data = f.read()
+
+        def upload_file():
+            upload_response = requests.put(url=upload_url, data=file_data)
+            return upload_response.status_code in (200, 204), upload_response
+
+        _retry_with_backoff("uploading artifacts", upload_file)
+
+    # Step 3: Mark the ingest as successful with retry
+    def complete_ingest():
+        complete_response = requests.patch(
+            url=config.get_complete_url(ingest_id),
+            headers=config.get_headers(),
+            json={"upload_status": "SUCCESS"},
+        )
+        return complete_response.status_code == 204, complete_response
+
+    _retry_with_backoff("completing ingest", complete_ingest)
+
+    fire_event(ArtifactUploadSuccess(msg=f"command {command} completed successfully"))
+    if dbt.tracking.active_user is not None:
+        dbt.tracking.track_artifact_upload({"command": command})
+    PRODUCED_ARTIFACTS_PATHS.clear()