Module panama.logging.file_handler

class PanamaFileHandler(logging.Handler):

    """Handler for file logs."""

    def __init__(self, path: str, task_name: str, remote_path: str, level: int = logging.INFO):
        """Initialize the json handler

        Args:
            path (str): path to the folder
            file_name (str) : name of the file where logging is written.
            remote_path (str): the path to the remote adls where the json will be sent at the end of a task.
            level (int, optional): minimum log level. Defaults to logging.INFO.
        """
        super().__init__(level=level)
        self.path = path
        self.remote_path = remote_path
        self.task_name = task_name
        self.start_time = datetime.astimezone(datetime.now(), tz=timezone("Europe/Rome"))
        self.file_name = f"{self.task_name.replace('_','')}_{self.start_time.strftime('%Y%m%d%H%M%S')}.json"
        self.set_local_path()
        self.file_pointer = None

    def _open_local_connection(self):
        """Method used to open connection in a local path"""
        self.file_pointer = open(self.local_path, "w+")

    def set_local_path(self, path_prefix=None):
        if path_prefix is None:
            if runtime() == "databricks":
                path_prefix = "/dbfs"
            else:
                path_prefix = os.getcwd()
        local_folder = os.path.join(path_prefix, self.path)
        if not os.path.exists(local_folder):
            os.makedirs(local_folder)
        # note: the file is not yet created.
        # the empty file is created when _open_local_connection is called
        self.local_path = os.path.join(local_folder, self.file_name)

    def emit(self, record: logging.LogRecord):
        """Method used to write logs to the target local json file."""
        if self.file_pointer is None:
            self._open_local_connection()

        # format record
        log_entry = self.format(record=record)
        self.file_pointer.write(log_entry + "\n")  # type: ignore
        self.file_pointer.flush()  # type: ignore
        os.fsync(self.file_pointer.fileno())  # type: ignore

    def move_to_remote(self):
        """Method used to move logs to remote storage."""
        spark = SparkSession.getActiveSession()
        dbutils = get_db_utils(spark)
        if self.file_pointer is not None:
            self.file_pointer.close()
            try:
                # handle non-databricks environment
                if runtime() == "databricks":
                    source = f"file:{self.local_path}"
                else:
                    source = f"{self.local_path}"
                dbutils.fs.mkdirs(self.remote_path)
                dbutils.fs.mv(source, self.remote_path, recurse=True)
            except Exception as e:
                if "java.io.FileNotFoundException" not in str(e):
                    print(self.local_path, "not found.")
                    raise e

class fileFilter(logging.Filter):
    def __init__(self, config_run: ConfigRun):
        self.config_run = config_run

    def _set_defaults(self, record):
        record.job_id = self.config_run.job_id
        record.job_run_id = self.config_run.job_run_id
        record.task_id = self.config_run.task_id

    def filter(self, record):
        self._set_defaults(record)
        return record.levelno < 99

class fileFormatter(logging.Formatter):
    def __init__(self):
        super().__init__(fmt="%(asctime)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S")

    def format(self, record):
        """Function to format the record as required by panama.

        Args:
            record (LogRecord): input record from logger.

        Returns:
            str: string to write to file.
        """
        super().format(record)

        # get trace if present
        if hasattr(record, "trace"):
            trace_content = record.trace  # type: ignore
        elif record.exc_text is not None:
            trace_content = record.exc_text
        else:
            trace_content = None

        # get extra content
        extra_content = getattr(record, "extra", dict())
        if len(extra_content) > 0:
            extra_content = json.dumps(extra_content)
        else:
            extra_content = None

        # populate message content
        msg = {
            "asctime": record.asctime,
            "levelname": record.levelname,
            "job_id": record.job_id,  # type: ignore
            "job_run_id": record.job_run_id,  # type: ignore
            "task_id": record.task_id,  # type: ignore
            "step_id": record.message,
        }

        # if trace is not none, add trace
        if trace_content is not None:
            msg["trace"] = trace_content
        # if extra content is not none, add extra content
        if extra_content is not None:
            msg["extras"] = extra_content

        return json.dumps(msg)