Log exception (#544)

* adds log-exception decorator

* flake8

src/ploomber/exceptions.py

@@ -28,6 +28,10 @@ class BaseException(ClickException):
     from chained exceptions
     """
 
+    def __init__(self, message, type_=None):
+        super().__init__(message)
+        self.type_ = type_
+
     def show(self, file: t.Optional[t.IO] = None) -> None:
         if file is None:
             file = get_text_stderr()

src/ploomber/telemetry/telemetry.py

@@ -39,6 +39,7 @@
 from pathlib import Path
 import sys
 import uuid
+from functools import wraps
 
 from ploomber.telemetry import validate_inputs
 from ploomber import __version__
@@ -417,3 +418,30 @@ def log_api(action,
                             properties=props)
 
         posthog.capture(distinct_id=uid, event=action, properties=props)
+
+
+def log_exception(action):
+    """Runs a function and logs exceptions, if any
+    """
+
+    def _log_exceptions(func):
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            try:
+                start = datetime.datetime.now()
+                return func(*args, **kwargs)
+            except Exception as e:
+                log_api(
+                    action=action,
+                    total_runtime=str(datetime.datetime.now() - start),
+                    metadata={
+                        # can we log None to posthog?
+                        'type': getattr(e, 'type_', ''),
+                        'exception': str(e),
+                    })
+                raise e
+
+        return wrapper
+
+    return _log_exceptions

tests/telemetry/test_telemetry.py

@@ -1,15 +1,19 @@
 import pathlib
-import click
 import sys
 from unittest.mock import Mock
 from pathlib import Path
+
+import click
 import pytest
+
 from ploomber.telemetry import telemetry
 from ploomber.telemetry.validate_inputs import str_param, opt_str_param
 from ploomber.cli import plot, install, build, interact, task, report, status
 import ploomber.dag.dag as dag_module
 from ploomber.tasks.tasks import PythonCallable
 from ploomber.products.file import File
+from ploomber.exceptions import BaseException
+
 from conftest import _write_sample_conda_env, _prepare_files
 
 
@@ -140,6 +144,7 @@ def test_pip_env(monkeypatch, inside_pip_env):
 # Ref: https://stackoverflow.com/questions/43878953/how-does-one-detect-if-
 # one-is-running-within-a-docker-container-within-python
 def test_docker_env(monkeypatch):
+
     def mock(input_path):
         return 'dockerenv' in str(input_path)
 
@@ -387,6 +392,7 @@ def test_parse_dag_products(monkeypatch):
 
 
 def test_parse_dag(monkeypatch, tmp_directory):
+
     def fn1(product):
         pass
 
@@ -474,3 +480,47 @@ def test_creates_config_directory(monkeypatch, tmp_nbs,
     assert Path('stats').is_dir()
     assert Path('stats', 'uid.yaml').is_file()
     assert Path('stats', 'config.yaml').is_file()
+
+
+def test_log_exception(monkeypatch):
+    mock = Mock()
+    mock_dt = Mock()
+    mock_dt.now.side_effect = [1, 2]
+    monkeypatch.setattr(telemetry, 'log_api', mock)
+    monkeypatch.setattr(telemetry.datetime, 'datetime', mock_dt)
+
+    @telemetry.log_exception('some-action')
+    def my_function():
+        raise ValueError('some error')
+
+    with pytest.raises(ValueError):
+        my_function()
+
+    mock.assert_called_once_with(action='some-action',
+                                 total_runtime='1',
+                                 metadata={
+                                     'type': '',
+                                     'exception': 'some error'
+                                 })
+
+
+def test_log_exception_logs_type(monkeypatch):
+    mock = Mock()
+    mock_dt = Mock()
+    mock_dt.now.side_effect = [1, 2]
+    monkeypatch.setattr(telemetry, 'log_api', mock)
+    monkeypatch.setattr(telemetry.datetime, 'datetime', mock_dt)
+
+    @telemetry.log_exception('some-action')
+    def my_function():
+        raise BaseException('some error', type_='some-type')
+
+    with pytest.raises(BaseException):
+        my_function()
+
+    mock.assert_called_once_with(action='some-action',
+                                 total_runtime='1',
+                                 metadata={
+                                     'type': 'some-type',
+                                     'exception': 'some error'
+                                 })