From 17ecd8b121443a25fcedde28daa6138347a0147f Mon Sep 17 00:00:00 2001 From: sujitkp-blr Date: Thu, 27 Oct 2022 22:16:20 +0530 Subject: [PATCH] Adding dbt deployment environment instrumentation (#65) * Adding dbt deployment environment instrumentation * Adding profile parameters instrumentation --- dbt/adapters/hive/cloudera_tracking.py | 23 ++++++----------------- dbt/adapters/hive/connections.py | 11 ++++++++--- 2 files changed, 14 insertions(+), 20 deletions(-) diff --git a/dbt/adapters/hive/cloudera_tracking.py b/dbt/adapters/hive/cloudera_tracking.py index bc4fa57..eda4948 100644 --- a/dbt/adapters/hive/cloudera_tracking.py +++ b/dbt/adapters/hive/cloudera_tracking.py @@ -42,9 +42,8 @@ # Json object to store dbt profile(profile.yml) related information profile_info = {} -# Json object to store cml environment variables -cml_info = {} - +# Json object to store dbt deployment environment variables +dbt_deployment_env_info = {} def populate_platform_info(cred: Credentials, ver): """ @@ -69,22 +68,12 @@ def populate_platform_info(cred: Credentials, ver): # dbt adapter info e.g. impala-1.2.0 platform_info["dbt_adapter"] = f"{cred.type}-{ver.version}" - -def populate_cml_info(): +def populate_dbt_deployment_env_info(): """ - populate cml environment variables if available to be passed on for tracking + populate dbt deployment environment variables if available to be passed on for tracking """ default_value = "" # if environment variables doesn't exist add empty string as default - cml_info["ml_runtime_edition"] = os.environ.get('ML_RUNTIME_EDITION', default_value) - cml_info["ml_runtime_git_hash"] = os.environ.get('ML_RUNTIME_GIT_HASH', default_value) - cml_info["ml_runtime_kernel"] = os.environ.get('ML_RUNTIME_KERNEL', default_value) - cml_info["ml_runtime_editor"] = os.environ.get('ML_RUNTIME_EDITOR', default_value) - cml_info["ml_runtime_gbn"] = os.environ.get('ML_RUNTIME_GBN', default_value) - cml_info["ml_runtime_full_version"] = os.environ.get('ML_RUNTIME_FULL_VERSION', default_value) - cml_info["ml_runtime_description"] = os.environ.get('ML_RUNTIME_DESCRIPTION', default_value) - cml_info["ml_runtime_maintenance_version"] = os.environ.get('ML_RUNTIME_MAINTENANCE_VERSION', default_value) - cml_info["ml_runtime_metadata_version"] = os.environ.get('ML_RUNTIME_METADATA_VERSION', default_value) - + dbt_deployment_env_info["dbt_deployment_env"] = os.environ.get('DBT_DEPLOYMENT_ENV', default_value) def populate_unique_ids(cred: Credentials): host = str(cred.host).encode() @@ -201,7 +190,7 @@ def track_usage(tracking_payload): # inject other static payload to tracking_payload tracking_payload = _merge_keys(unique_ids, tracking_payload) tracking_payload = _merge_keys(platform_info, tracking_payload) - tracking_payload = _merge_keys(cml_info, tracking_payload) + tracking_payload = _merge_keys(dbt_deployment_env_info, tracking_payload) tracking_payload = _merge_keys(profile_info, tracking_payload) # form the tracking data diff --git a/dbt/adapters/hive/connections.py b/dbt/adapters/hive/connections.py index 3c5b2e8..4b030b1 100644 --- a/dbt/adapters/hive/connections.py +++ b/dbt/adapters/hive/connections.py @@ -22,7 +22,7 @@ import impala.dbapi from dbt.adapters.base import Credentials from dbt.adapters.sql import SQLConnectionManager -from dbt.contracts.connection import (AdapterResponse, Connection, +from dbt.contracts.connection import (AdapterResponse, AdapterRequiredConfig, Connection, ConnectionState) from dbt.events import AdapterLogger from dbt.events.functions import fire_event @@ -84,8 +84,8 @@ def __post_init__(self): tracker.usage_tracking = self.usage_tracking # get platform information for tracking tracker.populate_platform_info(self, ver) - # get cml information for tracking - tracker.populate_cml_info() + # get dbt deployment information for tracking + tracker.populate_dbt_deployment_env_info() # generate unique ids for tracking tracker.populate_unique_ids(self) @@ -168,6 +168,11 @@ def description(self): class HiveConnectionManager(SQLConnectionManager): TYPE = "hive" + def __init__(self, profile: AdapterRequiredConfig): + super().__init__(profile) + # generate profile related object for instrumentation. + tracker.generate_profile_info(self) + @classmethod def open(cls, connection): if connection.state == ConnectionState.OPEN: