lakehouse_engine.utils.engine_usage_stats

  1"""Utilities for recording the engine activity."""
  2
  3import ast
  4import json
  5import re
  6from datetime import datetime
  7from typing import Dict
  8from urllib.parse import urlparse
  9
 10from lakehouse_engine.core.definitions import CollectEngineUsage, EngineStats
 11from lakehouse_engine.core.exec_env import ExecEnv
 12from lakehouse_engine.utils.configs.config_utils import ConfigUtils
 13from lakehouse_engine.utils.logging_handler import LoggingHandler
 14from lakehouse_engine.utils.storage.file_storage_functions import FileStorageFunctions
 15
 16
 17class EngineUsageStats(object):
 18    """Engine Usage utilities class."""
 19
 20    _LOGGER = LoggingHandler(__name__).get_logger()
 21
 22    @classmethod
 23    def store_engine_usage(
 24        cls,
 25        acon: dict,
 26        func_name: str,
 27        collect_engine_usage: str = None,
 28        spark_confs: dict = None,
 29    ) -> None:
 30        """Collects and store Lakehouse Engine usage statistics.
 31
 32        These statistics include the acon and other relevant information, such as
 33        the lakehouse engine version and the functions/algorithms being used.
 34
 35        Args:
 36            acon: acon dictionary file.
 37            func_name: function name that called this log acon.
 38            collect_engine_usage: Lakehouse usage statistics collection strategy.
 39            spark_confs: optional dictionary with the spark confs to be used when
 40                collecting the engine usage.
 41        """
 42        try:
 43            if (
 44                collect_engine_usage
 45                in [
 46                    CollectEngineUsage.ENABLED.value,
 47                    CollectEngineUsage.PROD_ONLY.value,
 48                ]
 49                or ExecEnv.ENGINE_CONFIG.collect_engine_usage
 50                in CollectEngineUsage.ENABLED.value
 51            ):
 52                start_timestamp = datetime.now()
 53                timestamp_str = start_timestamp.strftime("%Y%m%d%H%M%S")
 54
 55                usage_stats: Dict = {"acon": ConfigUtils.remove_sensitive_info(acon)}
 56                EngineUsageStats.get_spark_conf_values(usage_stats, spark_confs)
 57
 58                engine_usage_path = None
 59                if usage_stats["environment"] == "prod":
 60                    engine_usage_path = ExecEnv.ENGINE_CONFIG.engine_usage_path
 61                elif collect_engine_usage != CollectEngineUsage.PROD_ONLY.value:
 62                    engine_usage_path = ExecEnv.ENGINE_CONFIG.engine_dev_usage_path
 63
 64                if engine_usage_path is not None:
 65                    usage_stats["function"] = func_name
 66                    usage_stats["engine_version"] = ConfigUtils.get_engine_version()
 67                    usage_stats["start_timestamp"] = start_timestamp
 68                    usage_stats["year"] = start_timestamp.year
 69                    usage_stats["month"] = start_timestamp.month
 70                    run_id_extracted = re.search(
 71                        "run-([1-9]\\w+)", usage_stats["run_id"]
 72                    )
 73                    usage_stats["run_id"] = (
 74                        run_id_extracted.group(1) if run_id_extracted else ""
 75                    )
 76
 77                    log_file_name = f"eng_usage_{func_name}_{timestamp_str}.json"
 78
 79                    usage_stats_str = json.dumps(usage_stats, default=str)
 80
 81                    url = urlparse(
 82                        f"{engine_usage_path}/{usage_stats['dp_name']}/"
 83                        f"{start_timestamp.year}/{start_timestamp.month}/"
 84                        f"{log_file_name}",
 85                        allow_fragments=False,
 86                    )
 87
 88                    try:
 89                        FileStorageFunctions.write_payload(
 90                            engine_usage_path, url, usage_stats_str
 91                        )
 92                        cls._LOGGER.info("Storing Lakehouse Engine usage statistics")
 93                    except FileNotFoundError as e:
 94                        cls._LOGGER.error(
 95                            f"Could not write engine stats into file: {e}."
 96                        )
 97        except Exception as e:
 98            cls._LOGGER.error(
 99                "Failed while collecting the lakehouse engine stats: "
100                f"Unexpected {e=}, {type(e)=}."
101            )
102
103    @classmethod
104    def get_spark_conf_values(cls, usage_stats: dict, spark_confs: dict) -> None:
105        """Get information from spark session configurations.
106
107        Args:
108            usage_stats: usage_stats dictionary file.
109            spark_confs: optional dictionary with the spark tags to be used when
110                collecting the engine usage.
111        """
112        spark_confs = (
113            EngineStats.DEF_SPARK_CONFS.value
114            if spark_confs is None
115            else EngineStats.DEF_SPARK_CONFS.value | spark_confs
116        )
117
118        for spark_conf_key, spark_conf_value in spark_confs.items():
119            # whenever the spark_conf_value has #, it means it is an array, so we need
120            # to split it and adequately process it
121            if "#" in spark_conf_value:
122                array_key = spark_conf_value.split("#")
123                array_values = ast.literal_eval(
124                    ExecEnv.SESSION.conf.get(array_key[0], "[]")
125                )
126                final_value = [
127                    key_val["value"]
128                    for key_val in array_values
129                    if key_val["key"] == array_key[1]
130                ]
131                usage_stats[spark_conf_key] = (
132                    final_value[0] if len(final_value) > 0 else ""
133                )
134            else:
135                usage_stats[spark_conf_key] = ExecEnv.SESSION.conf.get(
136                    spark_conf_value, ""
137                )