sentry-sdk 0.7.5__py2.py3-none-any.whl → 2.46.0__py2.py3-none-any.whl

sentry_sdk/integrations/asyncio.py

@@ -0,0 +1,144 @@
+import sys
+import functools
+
+import sentry_sdk
+from sentry_sdk.consts import OP
+from sentry_sdk.integrations import Integration, DidNotEnable
+from sentry_sdk.utils import event_from_exception, logger, reraise
+
+try:
+    import asyncio
+    from asyncio.tasks import Task
+except ImportError:
+    raise DidNotEnable("asyncio not available")
+
+from typing import cast, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any, Callable, TypeVar
+    from collections.abc import Coroutine
+
+    from sentry_sdk._types import ExcInfo
+
+    T = TypeVar("T", bound=Callable[..., Any])
+
+
+def get_name(coro):
+    # type: (Any) -> str
+    return (
+        getattr(coro, "__qualname__", None)
+        or getattr(coro, "__name__", None)
+        or "coroutine without __name__"
+    )
+
+
+def _wrap_coroutine(wrapped):
+    # type: (Coroutine[Any, Any, Any]) -> Callable[[T], T]
+    # Only __name__ and __qualname__ are copied from function to coroutine in CPython
+    return functools.partial(
+        functools.update_wrapper,
+        wrapped=wrapped,  # type: ignore
+        assigned=("__name__", "__qualname__"),
+        updated=(),
+    )
+
+
+def patch_asyncio():
+    # type: () -> None
+    orig_task_factory = None
+    try:
+        loop = asyncio.get_running_loop()
+        orig_task_factory = loop.get_task_factory()
+
+        def _sentry_task_factory(loop, coro, **kwargs):
+            # type: (asyncio.AbstractEventLoop, Coroutine[Any, Any, Any], Any) -> asyncio.Future[Any]
+
+            @_wrap_coroutine(coro)
+            async def _task_with_sentry_span_creation():
+                # type: () -> Any
+                result = None
+
+                with sentry_sdk.isolation_scope():
+                    with sentry_sdk.start_span(
+                        op=OP.FUNCTION,
+                        name=get_name(coro),
+                        origin=AsyncioIntegration.origin,
+                    ):
+                        try:
+                            result = await coro
+                        except StopAsyncIteration as e:
+                            raise e from None
+                        except Exception:
+                            reraise(*_capture_exception())
+
+                return result
+
+            task = None
+
+            # Trying to use user set task factory (if there is one)
+            if orig_task_factory:
+                task = orig_task_factory(
+                    loop, _task_with_sentry_span_creation(), **kwargs
+                )
+
+            if task is None:
+                # The default task factory in `asyncio` does not have its own function
+                # but is just a couple of lines in `asyncio.base_events.create_task()`
+                # Those lines are copied here.
+
+                # WARNING:
+                # If the default behavior of the task creation in asyncio changes,
+                # this will break!
+                task = Task(_task_with_sentry_span_creation(), loop=loop, **kwargs)
+                if task._source_traceback:  # type: ignore
+                    del task._source_traceback[-1]  # type: ignore
+
+            # Set the task name to include the original coroutine's name
+            try:
+                cast("asyncio.Task[Any]", task).set_name(
+                    f"{get_name(coro)} (Sentry-wrapped)"
+                )
+            except AttributeError:
+                # set_name might not be available in all Python versions
+                pass
+
+            return task
+
+        loop.set_task_factory(_sentry_task_factory)  # type: ignore
+
+    except RuntimeError:
+        # When there is no running loop, we have nothing to patch.
+        logger.warning(
+            "There is no running asyncio loop so there is nothing Sentry can patch. "
+            "Please make sure you call sentry_sdk.init() within a running "
+            "asyncio loop for the AsyncioIntegration to work. "
+            "See https://docs.sentry.io/platforms/python/integrations/asyncio/"
+        )
+
+
+def _capture_exception():
+    # type: () -> ExcInfo
+    exc_info = sys.exc_info()
+
+    client = sentry_sdk.get_client()
+
+    integration = client.get_integration(AsyncioIntegration)
+    if integration is not None:
+        event, hint = event_from_exception(
+            exc_info,
+            client_options=client.options,
+            mechanism={"type": "asyncio", "handled": False},
+        )
+        sentry_sdk.capture_event(event, hint=hint)
+
+    return exc_info
+
+
+class AsyncioIntegration(Integration):
+    identifier = "asyncio"
+    origin = f"auto.function.{identifier}"
+
+    @staticmethod
+    def setup_once():
+        # type: () -> None
+        patch_asyncio()

sentry_sdk/integrations/asyncpg.py

@@ -0,0 +1,208 @@
+from __future__ import annotations
+import contextlib
+from typing import Any, TypeVar, Callable, Awaitable, Iterator
+
+import sentry_sdk
+from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.integrations import _check_minimum_version, Integration, DidNotEnable
+from sentry_sdk.tracing import Span
+from sentry_sdk.tracing_utils import add_query_source, record_sql_queries
+from sentry_sdk.utils import (
+    ensure_integration_enabled,
+    parse_version,
+    capture_internal_exceptions,
+)
+
+try:
+    import asyncpg  # type: ignore[import-not-found]
+    from asyncpg.cursor import BaseCursor  # type: ignore
+
+except ImportError:
+    raise DidNotEnable("asyncpg not installed.")
+
+
+class AsyncPGIntegration(Integration):
+    identifier = "asyncpg"
+    origin = f"auto.db.{identifier}"
+    _record_params = False
+
+    def __init__(self, *, record_params: bool = False):
+        AsyncPGIntegration._record_params = record_params
+
+    @staticmethod
+    def setup_once() -> None:
+        # asyncpg.__version__ is a string containing the semantic version in the form of "<major>.<minor>.<patch>"
+        asyncpg_version = parse_version(asyncpg.__version__)
+        _check_minimum_version(AsyncPGIntegration, asyncpg_version)
+
+        asyncpg.Connection.execute = _wrap_execute(
+            asyncpg.Connection.execute,
+        )
+
+        asyncpg.Connection._execute = _wrap_connection_method(
+            asyncpg.Connection._execute
+        )
+        asyncpg.Connection._executemany = _wrap_connection_method(
+            asyncpg.Connection._executemany, executemany=True
+        )
+        asyncpg.Connection.cursor = _wrap_cursor_creation(asyncpg.Connection.cursor)
+        asyncpg.Connection.prepare = _wrap_connection_method(asyncpg.Connection.prepare)
+        asyncpg.connect_utils._connect_addr = _wrap_connect_addr(
+            asyncpg.connect_utils._connect_addr
+        )
+
+
+T = TypeVar("T")
+
+
+def _wrap_execute(f: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
+    async def _inner(*args: Any, **kwargs: Any) -> T:
+        if sentry_sdk.get_client().get_integration(AsyncPGIntegration) is None:
+            return await f(*args, **kwargs)
+
+        # Avoid recording calls to _execute twice.
+        # Calls to Connection.execute with args also call
+        # Connection._execute, which is recorded separately
+        # args[0] = the connection object, args[1] is the query
+        if len(args) > 2:
+            return await f(*args, **kwargs)
+
+        query = args[1]
+        with record_sql_queries(
+            cursor=None,
+            query=query,
+            params_list=None,
+            paramstyle=None,
+            executemany=False,
+            span_origin=AsyncPGIntegration.origin,
+        ) as span:
+            res = await f(*args, **kwargs)
+
+        with capture_internal_exceptions():
+            add_query_source(span)
+
+        return res
+
+    return _inner
+
+
+SubCursor = TypeVar("SubCursor", bound=BaseCursor)
+
+
+@contextlib.contextmanager
+def _record(
+    cursor: SubCursor | None,
+    query: str,
+    params_list: tuple[Any, ...] | None,
+    *,
+    executemany: bool = False,
+) -> Iterator[Span]:
+    integration = sentry_sdk.get_client().get_integration(AsyncPGIntegration)
+    if integration is not None and not integration._record_params:
+        params_list = None
+
+    param_style = "pyformat" if params_list else None
+
+    with record_sql_queries(
+        cursor=cursor,
+        query=query,
+        params_list=params_list,
+        paramstyle=param_style,
+        executemany=executemany,
+        record_cursor_repr=cursor is not None,
+        span_origin=AsyncPGIntegration.origin,
+    ) as span:
+        yield span
+
+
+def _wrap_connection_method(
+    f: Callable[..., Awaitable[T]], *, executemany: bool = False
+) -> Callable[..., Awaitable[T]]:
+    async def _inner(*args: Any, **kwargs: Any) -> T:
+        if sentry_sdk.get_client().get_integration(AsyncPGIntegration) is None:
+            return await f(*args, **kwargs)
+        query = args[1]
+        params_list = args[2] if len(args) > 2 else None
+        with _record(None, query, params_list, executemany=executemany) as span:
+            _set_db_data(span, args[0])
+            res = await f(*args, **kwargs)
+
+        return res
+
+    return _inner
+
+
+def _wrap_cursor_creation(f: Callable[..., T]) -> Callable[..., T]:
+    @ensure_integration_enabled(AsyncPGIntegration, f)
+    def _inner(*args: Any, **kwargs: Any) -> T:  # noqa: N807
+        query = args[1]
+        params_list = args[2] if len(args) > 2 else None
+
+        with _record(
+            None,
+            query,
+            params_list,
+            executemany=False,
+        ) as span:
+            _set_db_data(span, args[0])
+            res = f(*args, **kwargs)
+            span.set_data("db.cursor", res)
+
+        return res
+
+    return _inner
+
+
+def _wrap_connect_addr(f: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
+    async def _inner(*args: Any, **kwargs: Any) -> T:
+        if sentry_sdk.get_client().get_integration(AsyncPGIntegration) is None:
+            return await f(*args, **kwargs)
+
+        user = kwargs["params"].user
+        database = kwargs["params"].database
+
+        with sentry_sdk.start_span(
+            op=OP.DB,
+            name="connect",
+            origin=AsyncPGIntegration.origin,
+        ) as span:
+            span.set_data(SPANDATA.DB_SYSTEM, "postgresql")
+            addr = kwargs.get("addr")
+            if addr:
+                try:
+                    span.set_data(SPANDATA.SERVER_ADDRESS, addr[0])
+                    span.set_data(SPANDATA.SERVER_PORT, addr[1])
+                except IndexError:
+                    pass
+            span.set_data(SPANDATA.DB_NAME, database)
+            span.set_data(SPANDATA.DB_USER, user)
+
+            with capture_internal_exceptions():
+                sentry_sdk.add_breadcrumb(
+                    message="connect", category="query", data=span._data
+                )
+            res = await f(*args, **kwargs)
+
+        return res
+
+    return _inner
+
+
+def _set_db_data(span: Span, conn: Any) -> None:
+    span.set_data(SPANDATA.DB_SYSTEM, "postgresql")
+
+    addr = conn._addr
+    if addr:
+        try:
+            span.set_data(SPANDATA.SERVER_ADDRESS, addr[0])
+            span.set_data(SPANDATA.SERVER_PORT, addr[1])
+        except IndexError:
+            pass
+
+    database = conn._params.database
+    if database:
+        span.set_data(SPANDATA.DB_NAME, database)
+
+    user = conn._params.user
+    if user:
+        span.set_data(SPANDATA.DB_USER, user)

sentry_sdk/integrations/atexit.py

@@ -1,28 +1,29 @@
-from __future__ import absolute_import
-
 import os
 import sys
 import atexit
 
-from sentry_sdk.hub import Hub
+import sentry_sdk
 from sentry_sdk.utils import logger
 from sentry_sdk.integrations import Integration
+from typing import TYPE_CHECKING
 
-if False:
+if TYPE_CHECKING:
     from typing import Any
     from typing import Optional
 
 
 def default_callback(pending, timeout):
+    # type: (int, int) -> None
     """This is the default shutdown callback that is set on the options.
     It prints out a message to stderr that informs the user that some events
     are still pending and the process is waiting for them to flush out.
     """
 
     def echo(msg):
+        # type: (str) -> None
         sys.stderr.write(msg + "\n")
 
-    echo("Sentry is attempting to send %i pending error messages" % pending)
+    echo("Sentry is attempting to send %i pending events" % pending)
     echo("Waiting up to %s seconds" % timeout)
     echo("Press Ctrl-%s to quit" % (os.name == "nt" and "Break" or "C"))
     sys.stderr.flush()
@@ -42,9 +43,15 @@ class AtexitIntegration(Integration):
         # type: () -> None
         @atexit.register
         def _shutdown():
+            # type: () -> None
+            client = sentry_sdk.get_client()
+            integration = client.get_integration(AtexitIntegration)
+
+            if integration is None:
+                return
+
             logger.debug("atexit: got shutdown signal")
-            hub = Hub.main
-            integration = hub.get_integration(AtexitIntegration)
-            if integration is not None:
-                logger.debug("atexit: shutting down client")
-                hub.client.close(callback=integration.callback)
+            logger.debug("atexit: shutting down client")
+            sentry_sdk.get_isolation_scope().end_session()
+
+            client.close(callback=integration.callback)