8000 [Hackweek] Add explain plan to db spans. (#2315) · GitLukeW/sentry-python@2faf03d · GitHub
[go: up one dir, main page]
Skip to content

Commit 2faf03d

Browse files
antonpirkersentrivana
antonpirker
and
sentrivana
authored
[Hackweek] Add explain plan to db spans. (getsentry#2315)
This is a proof of concept of adding the explain plan to db spans. The explain plan will be added to the span in the `db.explain_plan` data item. There is a cache to make sure that the explain plan for each db query is only executed ever X seconds and there is also a max number of elements that are cached. To make sure we do not put to much strain on CPU or memory. 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" } } ``` Now you have a explain in the `span.data.db.explain_plan` in your database spans. --------- Co-authored-by: Ivana Kellyerova <ivana.kellyerova@sentry.io>
1 parent a0d0c3d commit 2faf03d

File tree

8 files changed

+182
-1
lines changed

8 files changed

+182
-1
lines changed

‎scripts/build_aws_lambda_layer.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -76,9 +76,10 @@ def zip(self):
7676

7777
shutil.copy(
7878
os.path.join(self.base_dir, self.out_zip_filename),
79-
os.path.abspath(DIST_PATH)
79+
os.path.abspath(DIST_PATH),
8080
)
8181

82+
8283
def build_packaged_zip():
8384
with tempfile.TemporaryDirectory() as base_dir:
8485
layer_builder = LayerBuilder(base_dir)

‎sentry_sdk/consts.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
Experiments = TypedDict(
3636
"Experiments",
3737
{
38+
"attach_explain_plans": dict[str, Any],
3839
"max_spans": Optional[int],
3940
"record_sql_params": Optional[bool],
4041
# TODO: Remove these 2 profiling related experiments

‎sentry_sdk/db/__init__.py

Whitespace-only changes.

‎sentry_sdk/db/explain_plan/__init__.py

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
import datetime
2+
3+
from sentry_sdk.consts import TYPE_CHECKING
4+
5+
if TYPE_CHECKING:
6+
from typing import Any
7+
8+
9+
EXPLAIN_CACHE = {}
10+
EXPLAIN_CACHE_SIZE = 50
11+
EXPLAIN_CACHE_TIMEOUT_SECONDS = 60 * 60 * 24
12+
13+
14+
def cache_statement(statement, options):
15+
# type: (str, dict[str, Any]) -> None
16+
global EXPLAIN_CACHE
17+
18+
now = datetime.datetime.utcnow()
19+
explain_cache_timeout_seconds = options.get(
20+
"explain_cache_timeout_seconds", EXPLAIN_CACHE_TIMEOUT_SECONDS
21+
)
22+
expiration_time = now + datetime.timedelta(seconds=explain_cache_timeout_seconds)
23+
24+
EXPLAIN_CACHE[hash(statement)] = expiration_time
25+
26+
27+
def remove_expired_cache_items():
28+
# type: () -> None
29+
"""
30+
Remove expired cache items from the cache.
31+
"""
32+
global EXPLAIN_CACHE
33+
34+
now = datetime.datetime.utcnow()
35+
36+
for key, expiration_time in EXPLAIN_CACHE.items():
37+
expiration_in_the_past = expiration_time < now
38+
if expiration_in_the_past:
39+
del EXPLAIN_CACHE[key]
40+
41+
42+
def should_run_explain_plan(statement, options):
43+
# type: (str, dict[str, Any]) -> bool
44+
"""
45+
Check cache if the explain plan for the given statement should be run.
46+
"""
47+
global EXPLAIN_CACHE
48+
49+
remove_expired_cache_items()
50+
51+
key = hash(statement)
52+
if key in EXPLAIN_CACHE:
53+
return False
54+
55+
explain_cache_size = options.get("explain_cache_size", EXPLAIN_CACHE_SIZE)
56+
cache_is_full = len(EXPLAIN_CACHE.keys()) >= explain_cache_size
57+
if cache_is_full:
58+
return False
59+
60+
return True

‎sentry_sdk/db/explain_plan/django.py

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
from sentry_sdk.consts import TYPE_CHECKING
2+
from sentry_sdk.db.explain_plan import cache_statement, should_run_explain_plan
3+
4+
if TYPE_CHECKING:
5+
from typing import Any
6+
from typing import Callable
7+
8+
from sentry_sdk.tracing import Span
9+
10+
11+
def attach_explain_plan_to_span(
12+
span, connection, statement, parameters, mogrify, options
13+
):
14+
# type: (Span, Any, str, Any, Callable[[str, Any], bytes], dict[str, Any]) -> None
15+
"""
16+
Run EXPLAIN or EXPLAIN ANALYZE on the given statement and attach the explain plan to the span data.
17+
18+
Usage:
19+
```
20+
sentry_sdk.init(
21+
dsn="...",
22+
_experiments={
23+
"attach_explain_plans": {
24+
"explain_cache_size": 1000, # Run explain plan for the 1000 most run queries
25+
"explain_cache_timeout_seconds": 60 * 60 * 24, # Run the explain plan for each statement only every 24 hours
26+
"use_explain_analyze": True, # Run "explain analyze" instead of only "explain"
27+
}
28+
}
29+
```
30+
"""
31+
if not statement.strip().upper().startswith("SELECT"):
32+
return
33+
34+
if not should_run_explain_plan(statement, options):
35+
return
36+
37+
analyze = "ANALYZE" if options.get("use_explain_analyze", False) else ""
38+
explain_statement = ("EXPLAIN %s " % analyze) + mogrify(
39+
statement, parameters
40+
).decode("utf-8")
41+
42+
with connection.cursor() as cursor:
43+
cursor.execute(explain_statement)
44+
explain_plan = [row for row in cursor.fetchall()]
45+
46+
span.set_data("db.explain_plan", explain_plan)
47+
cache_statement(statement, options)

‎sentry_sdk/db/explain_plan/sqlalchemy.py

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
from __future__ import absolute_import
2+
3+
from sentry_sdk.consts import TYPE_CHECKING
4+
from sentry_sdk.db.explain_plan import cache_statement, should_run_explain_plan
5+
from sentry_sdk.integrations import DidNotEnable
6+
7+
try:
8+
from sqlalchemy.sql import text # type: ignore
9+
except ImportError:
10+
raise DidNotEnable("SQLAlchemy not installed.")
11+
12+
if TYPE_CHECKING:
13+
from typing import Any
14+
15+
from sentry_sdk.tracing import Span
16+
17+
18+
def attach_explain_plan_to_span(span, connection, statement, parameters, options):
19+
# type: (Span, Any, str, Any, dict[str, Any]) -> None
20+
"""
21+
Run EXPLAIN or EXPLAIN ANALYZE on the given statement and attach the explain plan to the span data.
22+
23+
Usage:
24+
```
25+
sentry_sdk.init(
26+
dsn="...",
27+
_experiments={
28+
"attach_explain_plans": {
29+
"explain_cache_size": 1000, # Run explain plan for the 1000 most run queries
30+
"explain_cache_timeout_seconds": 60 * 60 * 24, # Run the explain plan for each statement only every 24 hours
31+
"use_explain_analyze": True, # Run "explain analyze" instead of only "explain"
32+
}
33+
}
34+
```
35+
"""
36+
if not statement.strip().upper().startswith("SELECT"):
37+
return
38+
39+
if not should_run_explain_plan(statement, options):
40+
return
41+
42+
analyze = "ANALYZE" if options.get("use_explain_analyze", False) else ""
43+
explain_statement = (("EXPLAIN %s " % analyze) + statement) % parameters
44+
45+
result = connection.execute(text(explain_statement))
46+
explain_plan = [row for row in result]
47+
48+
span.set_data("db.explain_plan", explain_plan)
49+
cache_statement(statement, options)

‎sentry_sdk/integrations/django/__init__.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from sentry_sdk._compat import string_types, text_type
1010
from sentry_sdk._types import TYPE_CHECKING
1111
from sentry_sdk.consts import OP, SPANDATA
12+
from sentry_sdk.db.explain_plan.django import attach_explain_plan_to_span
1213
from sentry_sdk.hub import Hub, _should_send_default_pii
1314
from sentry_sdk.scope import add_global_event_processor
1415
from sentry_sdk.serializer import add_global_repr_processor
@@ -613,6 +614,17 @@ def execute(self, sql, params=None):
613614
hub, self.cursor, sql, params, paramstyle="format", executemany=False
614615
) as span:
615616
_set_db_data(span, self)
617+
if hub.client:
618+
options = hub.client.options["_experiments"].get("attach_explain_plans")
619+
if options is not None:
620+
attach_explain_plan_to_span(
621+
span,
622+
self.cursor.connection,
623+
sql,
624+
params,
625+
self.mogrify,
626+
options,
627+
)
616628
return real_execute(self, sql, params)
617629

618630
def executemany(self, sql, param_list):

‎sentry_sdk/integrations/sqlalchemy.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from sentry_sdk._compat import text_type
44
from sentry_sdk._types import TYPE_CHECKING
55
from sentry_sdk.consts import SPANDATA
6+
from sentry_sdk.db.explain_plan.sqlalchemy import attach_explain_plan_to_span
67
from sentry_sdk.hub import Hub
78
from sentry_sdk.integrations import Integration, DidNotEnable
89
from sentry_sdk.tracing_utils import record_sql_queries
@@ -68,6 +69,16 @@ def _before_cursor_execute(
6869

6970
if span is not None:
7071
_set_db_data(span, conn)
72+
if hub.client:
73+
options = hub.client.options["_experiments"].get("attach_explain_plans")
74+
if options is not None:
75+
attach_explain_plan_to_span(
76+
span,
77+
conn,
78+
statement,
79+
parameters,
80+
options,
81+
)
7182
context._sentry_sql_span = span
7283

7384

0 commit comments

Comments
 (0)
0