Merge branch 'main' into move_catalog_schema

.changes/unreleased/Features-20240514-162052.yaml

@@ -0,0 +1,6 @@
+kind: Features
+body: Support adding callbacks to the event manager
+time: 2024-05-14T16:20:52.120336-07:00
+custom:
+  Author: QMalcolm
+  Issue: "131"

.changes/unreleased/Features-20240515-161106.yaml

@@ -0,0 +1,6 @@
+kind: Features
+body: Add a print event to support firing event to console in --quiet mode
+time: 2024-05-15T16:11:06.815526-07:00
+custom:
+  Author: ChenyuLInx
+  Issue: "8756"

dbt_common/__about__.py

@@ -1 +1 @@
-version = "1.0.4"
+version = "1.1.0"

dbt_common/events/base_types.py

@@ -6,14 +6,11 @@
 from google.protobuf.json_format import ParseDict, MessageToDict, MessageToJson
 from google.protobuf.message import Message
 from dbt_common.events.helpers import get_json_string_utcnow
-from typing import Optional
+from typing import Callable, Optional
 
 from dbt_common.invocation import get_invocation_id
 
-if sys.version_info >= (3, 8):
-    from typing import Protocol
-else:
-    from typing_extensions import Protocol
+from typing import Protocol
 
 
 # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
@@ -128,6 +125,9 @@ class EventMsg(Protocol):
     data: Message
 
 
+TCallback = Callable[[EventMsg], None]
+
+
 def msg_from_base_event(event: BaseEvent, level: Optional[EventLevel] = None):
     msg_class_name = f"{type(event).__name__}Msg"
     msg_cls = getattr(event.PROTO_TYPES_MODULE, msg_class_name)

dbt_common/events/event_manager.py

@@ -1,15 +1,15 @@
 import os
 import traceback
-from typing import Callable, List, Optional, Protocol, Tuple
+from typing import List, Optional, Protocol, Tuple
 
-from dbt_common.events.base_types import BaseEvent, EventLevel, msg_from_base_event, EventMsg
+from dbt_common.events.base_types import BaseEvent, EventLevel, msg_from_base_event, TCallback
 from dbt_common.events.logger import LoggerConfig, _Logger, _TextLogger, _JsonLogger, LineFormat
 
 
 class EventManager:
     def __init__(self) -> None:
         self.loggers: List[_Logger] = []
-        self.callbacks: List[Callable[[EventMsg], None]] = []
+        self.callbacks: List[TCallback] = []
 
     def fire_event(self, e: BaseEvent, level: Optional[EventLevel] = None) -> None:
         msg = msg_from_base_event(e, level=level)
@@ -37,13 +37,16 @@ def add_logger(self, config: LoggerConfig) -> None:
         )
         self.loggers.append(logger)
 
+    def add_callback(self, callback: TCallback) -> None:
+        self.callbacks.append(callback)
+
     def flush(self) -> None:
         for logger in self.loggers:
             logger.flush()
 
 
 class IEventManager(Protocol):
-    callbacks: List[Callable[[EventMsg], None]]
+    callbacks: List[TCallback]
     loggers: List[_Logger]
 
     def fire_event(self, e: BaseEvent, level: Optional[EventLevel] = None) -> None:
@@ -52,6 +55,9 @@ def fire_event(self, e: BaseEvent, level: Optional[EventLevel] = None) -> None:
     def add_logger(self, config: LoggerConfig) -> None:
         ...
 
+    def add_callback(self, callback: TCallback) -> None:
+        ...
+
 
 class TestEventManager(IEventManager):
     __test__ = False

dbt_common/events/event_manager_client.py

@@ -1,6 +1,7 @@
 # Since dbt-rpc does not do its own log setup, and since some events can
 # currently fire before logs can be configured by setup_event_logger(), we
 # create a default configuration with default settings and no file output.
+from dbt_common.events.base_types import TCallback
 from dbt_common.events.event_manager import IEventManager, EventManager
 
 _EVENT_MANAGER: IEventManager = EventManager()
@@ -16,6 +17,11 @@ def add_logger_to_manager(logger) -> None:
     _EVENT_MANAGER.add_logger(logger)
 
 
+def add_callback_to_manager(callback: TCallback) -> None:
+    global _EVENT_MANAGER
+    _EVENT_MANAGER.add_callback(callback)
+
+
 def ctx_set_event_manager(event_manager: IEventManager) -> None:
     global _EVENT_MANAGER
     _EVENT_MANAGER = event_manager

dbt_common/events/logger.py

@@ -13,6 +13,13 @@
 from dbt_common.events.format import timestamp_to_datetime_string
 from dbt_common.utils.encoding import ForgivingJSONEncoder
 
+PRINT_EVENT_NAME = "PrintEvent"
+
+
+def _is_print_event(msg: EventMsg) -> bool:
+    return msg.info.name == PRINT_EVENT_NAME
+
+
 # A Filter is a function which takes a BaseEvent and returns True if the event
 # should be logged, False otherwise.
 Filter = Callable[[EventMsg], bool]
@@ -120,7 +127,14 @@ def create_line(self, msg: EventMsg) -> str:
     def write_line(self, msg: EventMsg):
         line = self.create_line(msg)
         if self._python_logger is not None:
-            send_to_logger(self._python_logger, msg.info.level, line)
+            # We send PrintEvent to logger as error so it goes to stdout
+            # when --quiet flag is set.
+            # --quiet flag will filter out all events lower than ERROR.
+            if _is_print_event(msg):
+                level = "error"
+            else:
+                level = msg.info.level
+            send_to_logger(self._python_logger, level, line)
 
     def flush(self):
         if self._python_logger is not None:
@@ -138,8 +152,11 @@ def create_line(self, msg: EventMsg) -> str:
         return self.create_debug_line(msg) if self.use_debug_format else self.create_info_line(msg)
 
     def create_info_line(self, msg: EventMsg) -> str:
-        ts: str = datetime.utcnow().strftime("%H:%M:%S")
         scrubbed_msg: str = self.scrubber(msg.info.msg)  # type: ignore
+        if _is_print_event(msg):
+            # PrintEvent is a special case, we don't want to add a timestamp
+            return scrubbed_msg
+        ts: str = datetime.utcnow().strftime("%H:%M:%S")
         return f"{self._get_color_tag()}{ts}  {scrubbed_msg}"
 
     def create_debug_line(self, msg: EventMsg) -> str:

dbt_common/events/types.proto

@@ -119,3 +119,13 @@ message NoteMsg {
     EventInfo info = 1;
     Note data = 2;
 }
+
+// Z052
+message PrintEvent {
+    string msg = 1;
+}
+
+message PrintEventMsg {
+    EventInfo info = 1;
+    PrintEvent data = 2;
+}

dbt_common/events/types.py

@@ -125,3 +125,14 @@ def code(self) -> str:
 
     def message(self) -> str:
         return self.msg
+
+
+class PrintEvent(InfoLevel):
+    # Use this event to skip any formatting and just print a message
+    # This event will get to stdout even if the logger is set to ERROR
+    # This is to support commands that want --quiet option but also log something to stdout
+    def code(self) -> str:
+        return "Z052"
+
+    def message(self) -> str:
+        return self.msg

pyproject.toml

@@ -58,6 +58,7 @@ lint = [
 ]
 test = [
     "pytest>=7.3,<8.0",
+    "pytest-mock",
     "pytest-xdist>=3.2,<4.0",
     "pytest-cov>=4.1,<5.0",
     "hypothesis>=6.87,<7.0",

tests/unit/events/test_logger.py

@@ -0,0 +1,40 @@
+import json
+from dbt_common.events.logger import LoggerConfig, _TextLogger, _JsonLogger
+from dbt_common.events.base_types import EventLevel, msg_from_base_event
+from dbt_common.events.types import PrintEvent
+
+
+def test_create_print_line():
+    # No format, still fired even when error is the level
+    config = LoggerConfig(name="test_logger", level=EventLevel.ERROR)
+    logger = _TextLogger(config)
+    msg = msg_from_base_event(PrintEvent(msg="This is a print event"))
+    expected_line = "This is a print event"
+    actual_line = logger.create_line(msg)
+    assert actual_line == expected_line
+
+
+def test_create_print_json():
+    # JSON format still have event level being info
+    config = LoggerConfig(name="test_logger", level=EventLevel.ERROR)
+    logger = _JsonLogger(config)
+    msg = msg_from_base_event(PrintEvent(msg="This is a print event"))
+    expected_json = {
+        "data": {"msg": "This is a print event"},
+        "info": {
+            "category": "",
+            "code": "Z052",
+            "extra": {},
+            "level": "info",
+            "msg": "This is a print event",
+            "name": "PrintEvent",
+            "thread": "MainThread",
+        },
+    }
+    actual_line = logger.create_line(msg)
+    actual_json = json.loads(actual_line)
+    assert "info" in actual_json
+    actual_json["info"].pop("invocation_id")
+    actual_json["info"].pop("ts")
+    actual_json["info"].pop("pid")
+    assert actual_json == expected_json

tests/unit/test_event_manager.py

@@ -0,0 +1,11 @@
+from dbt_common.events.event_manager import EventManager
+from tests.unit.utils import EventCatcher
+
+
+class TestEventManager:
+    def test_add_callback(self) -> None:
+        event_manager = EventManager()
+        assert len(event_manager.callbacks) == 0
+
+        event_manager.add_callback(EventCatcher().catch)
+        assert len(event_manager.callbacks) == 1