@forward_logs
def task_1():
    logger.bind(key="value").info("This was logged in Task 1 (normal function)")

@task
@forward_logs
def task_2():
    logger.critical("This was logged in Task 2 (prefect task)")

@forward_logs
def task_3():
    logger.warning("This was logged in Task 3 (normal function)")

@task
@forward_logs
def task_4():
    logger.error("This was logged in Task 4 inside a subflow (prefect task)")

@task
@forward_logs
def task_5():
    <http://logger.info|logger.info>("This was logged in Task 5 inside a subflow (prefect task)")

@flow
def example_flow():
    subflow()
    task_1()
    task_2()
    task_3()

@flow
def subflow():
    task_4()
    task_5()

if __name__ == "__main__":
    serve(example_flow.to_deployment(name="example-flow"))

import functools

from loguru import logger
from prefect import flow, get_run_logger, serve, task

handler_ids = []


def forward_logs(func):
    """
    Loguru uses an anonymous logger, which prohibits us from using the
    PREFECT_LOGGING_EXTRA_LOGGERS variable to capture logs and forward them to the
    Prefect UI.

    Decorating functions at the flow-level will forward logs for all subflows, tasks and
    ordinary python functions called by that flow automatically to the Flow Run in the
    Prefect UI.

    Note that in order to have logs show up in the Task run view in the Prefect UI, that
    the Task itself must be decorated with this decorator.
    """

    # Reference:
    # <https://discourse.prefect.io/t/can-i-use-loguru-logs-in-prefect-flows/140>
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        global handler_ids

        try:
            # Passthrough in case we are outside of a task / flow context, e.g. `.fn`
            run_logger = get_run_logger()
        except Exception:
            return func(*args, **kwargs)

        for handler_id in handler_ids:
            try:
                logger.remove(handler_id)
            except ValueError:
                pass

        log_format = "{message}"

        handler_ids.append(
            logger.add(
                run_logger.debug,
                filter=lambda record: record["level"].name == "DEBUG",
                level="TRACE",
                format=log_format,
            )
        )

        handler_ids.append(
            logger.add(
                run_logger.warning,
                filter=lambda record: record["level"].name == "WARNING",
                level="TRACE",
                format=log_format,
            )
        )

        handler_ids.append(
            logger.add(
                run_logger.error,
                filter=lambda record: record["level"].name == "ERROR",
                level="TRACE",
                format=log_format,
            )
        )

        handler_ids.append(
            logger.add(
                run_logger.critical,
                filter=lambda record: record["level"].name == "CRITICAL",
                level="TRACE",
                format=log_format,
            )
        )

        handler_ids.append(
            logger.add(
                <http://run_logger.info|run_logger.info>,
                filter=lambda record: record["level"].name
                not in ["DEBUG", "WARNING", "ERROR", "CRITICAL"],
                level="TRACE",
                format=log_format,
            )
        )

        return func(*args, **kwargs)

    return wrapper