getsentry · antonpirker · Oct 2, 2023 · Aug 21, 2023 · Aug 21, 2023 · Aug 21, 2023
@@ -76,9 +76,10 @@ def zip(self):
 
         shutil.copy(
             os.path.join(self.base_dir, self.out_zip_filename),
-            os.path.abspath(DIST_PATH)
+            os.path.abspath(DIST_PATH),
         )
 
+
 def build_packaged_zip():
     with tempfile.TemporaryDirectory() as base_dir:
         layer_builder = LayerBuilder(base_dir)

@@ -35,6 +35,7 @@
     Experiments = TypedDict(
         "Experiments",
         {
+            "attach_explain_plans": dict[str, Any],
             "max_spans": Optional[int],
             "record_sql_params": Optional[bool],
             # TODO: Remove these 2 profiling related experiments

@@ -0,0 +1,60 @@
+import datetime
+
+from sentry_sdk.consts import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any
+
+
+EXPLAIN_CACHE = {}
+EXPLAIN_CACHE_SIZE = 50
+EXPLAIN_CACHE_TIMEOUT_SECONDS = 60 * 60 * 24
+
+
+def cache_statement(statement, options):
+    # type: (str, dict[str, Any]) -> None
+    global EXPLAIN_CACHE
+
+    now = datetime.datetime.utcnow()
+    explain_cache_timeout_seconds = options.get(
+        "explain_cache_timeout_seconds", EXPLAIN_CACHE_TIMEOUT_SECONDS
+    )
+    expiration_time = now + datetime.timedelta(seconds=explain_cache_timeout_seconds)
+
+    EXPLAIN_CACHE[hash(statement)] = expiration_time
+
+
+def remove_expired_cache_items():
+    # type: () -> None
+    """
+    Remove expired cache items from the cache.
+    """
+    global EXPLAIN_CACHE
+
+    now = datetime.datetime.utcnow()
+
+    for key, expiration_time in EXPLAIN_CACHE.items():
+        expiration_in_the_past = expiration_time < now
+        if expiration_in_the_past:
+            del EXPLAIN_CACHE[key]
+
+
+def should_run_explain_plan(statement, options):
+    # type: (str, dict[str, Any]) -> bool
+    """
+    Check cache if the explain plan for the given statement should be run.
+    """
+    global EXPLAIN_CACHE
+
+    remove_expired_cache_items()
+
+    key = hash(statement)
+    if key in EXPLAIN_CACHE:
+        return False
+
+    explain_cache_size = options.get("explain_cache_size", EXPLAIN_CACHE_SIZE)
+    cache_is_full = len(EXPLAIN_CACHE.keys()) >= explain_cache_size
+    if cache_is_full:
+        return False
+
+    return True
@@ -0,0 +1,47 @@
+from sentry_sdk.consts import TYPE_CHECKING
+from sentry_sdk.db.explain_plan import cache_statement, should_run_explain_plan
+
+if TYPE_CHECKING:
+    from typing import Any
+    from typing import Callable
+
+    from sentry_sdk.tracing import Span
+
+
+def attach_explain_plan_to_span(
+    span, connection, statement, parameters, mogrify, options
+):
+    # type: (Span, Any, str, Any, Callable[[str, Any], bytes], dict[str, Any]) -> None
+    """
+    Run EXPLAIN or EXPLAIN ANALYZE on the given statement and attach the explain plan to the span data.
+
+    Usage:
+    ```
+    sentry_sdk.init(
+        dsn="...",
+        _experiments={
+            "attach_explain_plans": {
+                "explain_cache_size": 1000,  # Run explain plan for the 1000 most run queries
+                "explain_cache_timeout_seconds": 60 * 60 * 24,  # Run the explain plan for each statement only every 24 hours
+                "use_explain_analyze": True,  # Run "explain analyze" instead of only "explain"
+            }
+        }
+    ```
+    """
+    if not statement.strip().upper().startswith("SELECT"):
+        return
+<
8000
/span>
+    if not should_run_explain_plan(statement, options):
+        return
+
+    analyze = "ANALYZE" if options.get("use_explain_analyze", False) else ""
+    explain_statement = ("EXPLAIN %s " % analyze) + mogrify(
+        statement, parameters
+    ).decode("utf-8")
+
+    with connection.cursor() as cursor:
+        cursor.execute(explain_statement)
+        explain_plan = [row for row in cursor.fetchall()]
+
+        span.set_data("db.explain_plan", explain_plan)
+        cache_statement(statement, options)
@@ -0,0 +1,49 @@
+from __future__ import absolute_import
+
+from sentry_sdk.consts import TYPE_CHECKING
+from sentry_sdk.db.explain_plan import cache_statement, should_run_explain_plan
+from sentry_sdk.integrations import DidNotEnable
+
+try:
+    from sqlalchemy.sql import text  # type: ignore
+except ImportError:
+    raise DidNotEnable("SQLAlchemy not installed.")
+
+if TYPE_CHECKING:
+    from typing import Any
+
+    from sentry_sdk.tracing import Span
+
+
+def attach_explain_plan_to_span(span, connection, statement, parameters, options):
+    # type: (Span, Any, str, Any, dict[str, Any]) -> None
+    """
+    Run EXPLAIN or EXPLAIN ANALYZE on the given statement and attach the explain plan to the span data.
+
+    Usage:
+    ```
+    sentry_sdk.init(
+        dsn="...",
+        _experiments={
+            "attach_explain_plans": {
+                "explain_cache_size": 1000,  # Run explain plan for the 1000 most run queries
+                "explain_cache_timeout_seconds": 60 * 60 * 24,  # Run the explain plan for each statement only every 24 hours
+                "use_explain_analyze": True,  # Run "explain analyze" instead of only "explain"
+            }
+        }
+    ```
+    """
+    if not statement.strip().upper().startswith("SELECT"):
+        return
+
+    if not should_run_explain_plan(statement, options):
+        return
+
+    analyze = "ANALYZE" if options.get("use_explain_analyze", False) else ""
+    explain_statement = (("EXPLAIN %s " % analyze) + statement) % parameters
+
+    result = connection.execute(text(explain_statement))
+    explain_plan = [row for row in result]
+
+    span.set_data("db.explain_plan", explain_plan)
+    cache_statement(statement, options)
@@ -9,6 +9,7 @@
 from sentry_sdk._compat import string_types, text_type
 from sentry_sdk._types import TYPE_CHECKING
 from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.db.explain_plan.django import attach_explain_plan_to_span
 from sentry_sdk.hub import Hub, _should_send_default_pii
 from sentry_sdk.scope import add_global_event_processor
 from sentry_sdk.serializer import add_global_repr_processor
@@ -613,6 +614,17 @@ def execute(self, sql, params=None):
             hub, self.cursor, sql, params, paramstyle="format", executemany=False
         ) as span:
             _set_db_data(span, self)
+            if hub.client:
+                options = hub.client.options["_experiments"].get("attach_explain_plans")
+                if options is not None:
+                    attach_explain_plan_to_span(
+                        span,
+                        self.cursor.connection,
+                        sql,
+                        params,
+                        self.mogrify,
+                        options,
+                    )
             return real_execute(self, sql, params)
 
     def executemany(self, sql, param_list):

@@ -3,6 +3,7 @@
 from sentry_sdk._compat import text_type
 from sentry_sdk._types import TYPE_CHECKING
 from sentry_sdk.consts import SPANDATA
+from sentry_sdk.db.explain_plan.sqlalchemy import attach_explain_plan_to_span
 from sentry_sdk.hub import Hub
 from sentry_sdk.integrations import Integration, DidNotEnable
 from sentry_sdk.tracing_utils import record_sql_queries
@@ -68,6 +69,16 @@ def _before_cursor_execute(
 
     if span is not None:
         _set_db_data(span, conn)
+        if hub.client:
+            options = hub.client.options["_experiments"].get("attach_explain_plans")
+            if options is not None:
+                attach_explain_plan_to_span(
+                    span,
+                    conn,
+                    statement,
+                    parameters,
+                    options,
+                )
         context._sentry_sql_span = span