diff --git a/cloudquery/sdk/plugin/plugin.py b/cloudquery/sdk/plugin/plugin.py index 4765cf2..7fc76f6 100644 --- a/cloudquery/sdk/plugin/plugin.py +++ b/cloudquery/sdk/plugin/plugin.py @@ -38,6 +38,9 @@ def __init__(self, name: str, version: str) -> None: def init(self, spec: bytes) -> None: pass + def set_logger(self, logger) -> None: + pass + def name(self) -> str: return self._name diff --git a/cloudquery/sdk/scheduler/scheduler.py b/cloudquery/sdk/scheduler/scheduler.py index c4e594d..8f8d111 100644 --- a/cloudquery/sdk/scheduler/scheduler.py +++ b/cloudquery/sdk/scheduler/scheduler.py @@ -39,6 +39,8 @@ def __init__( self._max_depth = max_depth if logger is None: self._logger = structlog.get_logger() + else: + self._logger = logger if concurrency <= 0: raise ValueError("concurrency must be greater than 0") if max_depth <= 0: diff --git a/cloudquery/sdk/serve/plugin.py b/cloudquery/sdk/serve/plugin.py index d412a74..8df531a 100644 --- a/cloudquery/sdk/serve/plugin.py +++ b/cloudquery/sdk/serve/plugin.py @@ -4,6 +4,8 @@ import grpc import structlog import sys +import os + from cloudquery.discovery_v1 import discovery_pb2_grpc from cloudquery.plugin_v3 import plugin_pb2_grpc @@ -15,8 +17,50 @@ DOC_FORMATS = ["json", "markdown"] +_IS_WINDOWS = sys.platform == "win32" + +try: + import colorama +except ImportError: + colorama = None + +if _IS_WINDOWS: # pragma: no cover + # On Windows, use colors by default only if Colorama is installed. + _has_colors = colorama is not None +else: + # On other OSes, use colors by default. + _has_colors = True + + def get_logger(args): - log = structlog.get_logger(processors=[structlog.processors.JSONRenderer()]) + processors = [ + structlog.contextvars.merge_contextvars, + structlog.processors.add_log_level, + structlog.processors.StackInfoRenderer(), + structlog.dev.set_exc_info, + structlog.processors.TimeStamper(fmt="%Y-%m-%d %H:%M:%S", utc=False), + ] + if args.log_format == "text": + processors.append( + structlog.dev.ConsoleRenderer( + colors=os.environ.get("NO_COLOR", "") == "" + and ( + os.environ.get("FORCE_COLOR", "") != "" + or ( + _has_colors + and sys.stdout is not None + and hasattr(sys.stdout, "isatty") + and sys.stdout.isatty() + ) + ) + ) + ) + else: + processors.append(structlog.processors.JSONRenderer()) + + # if args.log_format == "json": + # processors.append(structlog.processors.JSONRenderer()) + log = structlog.get_logger(processors=processors) return log @@ -87,6 +131,7 @@ def run(self, args): def _serve(self, args): logger = get_logger(args) + self._plugin.set_logger(logger) self._server = grpc.server(futures.ThreadPoolExecutor(max_workers=10)) discovery_pb2_grpc.add_DiscoveryServicer_to_server( DiscoveryServicer([3]), self._server